From b26bfd7afaa50a58ecbe0ecf208015ad324c7414 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 20 Aug 2024 21:48:27 -0700 Subject: [PATCH] Fix neural search build & increment OS version to 2.17 (#869) * Upgrade BWC version and fix neural search build (#868) Signed-off-by: Varun Jain * Incrementing opensearch version to 2.17.0-SNAPSHOT Signed-off-by: Varun Jain * Incorporate fix from https://github.com/opensearch-project/OpenSearch/pull/15238 from OpenSearch Core Signed-off-by: Varun Jain * Address Naveen's comment Signed-off-by: Varun Jain --------- Signed-off-by: Varun Jain --- ...backwards_compatibility_tests_workflow.yml | 8 +-- TRIAGING.md | 22 +++--- build.gradle | 2 +- gradle.properties | 4 +- .../ml/MLCommonsClientAccessor.java | 4 +- ...ValidateDependentPluginInstallationIT.java | 3 +- .../query/HybridQueryBuilderTests.java | 69 +++++++++++++++++-- .../neuralsearch/query/HybridQueryTests.java | 54 ++++++++++++++- .../query/HybridQueryPhaseSearcherTests.java | 19 ++++- 9 files changed, 154 insertions(+), 31 deletions(-) diff --git a/.github/workflows/backwards_compatibility_tests_workflow.yml b/.github/workflows/backwards_compatibility_tests_workflow.yml index 54f846448..4d552a892 100644 --- a/.github/workflows/backwards_compatibility_tests_workflow.yml +++ b/.github/workflows/backwards_compatibility_tests_workflow.yml @@ -15,8 +15,8 @@ jobs: matrix: java: [ 11, 17, 21 ] os: [ubuntu-latest,windows-latest] - bwc_version : ["2.9.0","2.10.0","2.11.0","2.12.0","2.13.0","2.14.0","2.15.0"] - opensearch_version : [ "2.16.0-SNAPSHOT" ] + bwc_version : ["2.9.0","2.10.0","2.11.0","2.12.0","2.13.0","2.14.0","2.15.0","2.16.0"] + opensearch_version : [ "2.17.0-SNAPSHOT" ] name: NeuralSearch Restart-Upgrade BWC Tests runs-on: ${{ matrix.os }} @@ -42,8 +42,8 @@ jobs: matrix: java: [ 11, 17, 21 ] os: [ubuntu-latest,windows-latest] - bwc_version: [ "2.11.0","2.12.0","2.13.0","2.14.0","2.15.0" ] - opensearch_version: [ "2.16.0-SNAPSHOT" ] + bwc_version: [ "2.11.0","2.12.0","2.13.0","2.14.0","2.15.0", "2.16.0" ] + opensearch_version: [ "2.17.0-SNAPSHOT" ] name: NeuralSearch Rolling-Upgrade BWC Tests runs-on: ${{ matrix.os }} diff --git a/TRIAGING.md b/TRIAGING.md index a666486b0..fb757051a 100644 --- a/TRIAGING.md +++ b/TRIAGING.md @@ -1,8 +1,8 @@ -The maintainers of the k-NN/neural-search Repo's seek to promote an inclusive and engaged community of contributors. In -order to facilitate this, bi-weekly triage meetings are open-to-all and attendance is encouraged for anyone who hopes to -contribute, discuss an issue, or learn more about the project. To learn more about contributing to the +The maintainers of the k-NN/neural-search Repo's seek to promote an inclusive and engaged community of contributors. In +order to facilitate this, bi-weekly triage meetings are open-to-all and attendance is encouraged for anyone who hopes to +contribute, discuss an issue, or learn more about the project. To learn more about contributing to the k-NN/neural-search Repo visit the [Contributing](./CONTRIBUTING.md) documentation. ### Do I need to attend for my issue to be addressed/triaged? @@ -11,19 +11,19 @@ Attendance is not required for your issue to be triaged or addressed. All new is ### What happens if my issue does not get covered this time? -Each meeting we seek to address all new issues. However, should we run out of time before your issue is discussed, you +Each meeting we seek to address all new issues. However, should we run out of time before your issue is discussed, you are always welcome to attend the next meeting or to follow up on the issue post itself. ### How do I join the Backlog & Triage meeting? -Meetings are hosted regularly at 5 PM Pacific Time on Tuesdays bi-weekly and can be joined via the links posted on the -[OpenSearch Meetup Group](https://www.meetup.com/opensearch/events/) list of events. The event will be titled +Meetings are hosted regularly at 5 PM Pacific Time on Tuesdays bi-weekly and can be joined via the links posted on the +[OpenSearch Meetup Group](https://www.meetup.com/opensearch/events/) list of events. The event will be titled `Development Backlog & Triage Meeting - k-NN/neural-search`. -After joining the Chime meeting, you can enable your video / voice to join the discussion. If you do not have a webcam +After joining the Chime meeting, you can enable your video / voice to join the discussion. If you do not have a webcam or microphone available, you can still join in via the text chat. -If you have an issue you'd like to bring forth please consider getting a link to the issue so it can be presented to +If you have an issue you'd like to bring forth please consider getting a link to the issue so it can be presented to everyone in the meeting. ### Is there an agenda for each week? @@ -48,13 +48,13 @@ No, all are welcome and encouraged to attend. Attending the Backlog & Triage mee ### What if I have an issue that is almost a duplicate, should I open a new one to be triaged? -You can always open an issue including one that you think may be a duplicate. However, in cases where you believe there -is an important distinction to be made between an existing issue and your newly created one, you are encouraged to +You can always open an issue including one that you think may be a duplicate. However, in cases where you believe there +is an important distinction to be made between an existing issue and your newly created one, you are encouraged to attend the triaging meeting to explain. ### What if I have follow-up questions on an issue? -If you have an existing issue you would like to discuss, you can always comment on the issue itself. Alternatively, you +If you have an existing issue you would like to discuss, you can always comment on the issue itself. Alternatively, you are welcome to come to the triage meeting to discuss. ### Is this meeting a good place to get help setting up k-NN/neural-search features on my OpenSearch instance? diff --git a/build.gradle b/build.gradle index f81089ef0..626862372 100644 --- a/build.gradle +++ b/build.gradle @@ -14,7 +14,7 @@ import java.util.concurrent.Callable buildscript { ext { - opensearch_version = System.getProperty("opensearch.version", "2.16.0-SNAPSHOT") + opensearch_version = System.getProperty("opensearch.version", "2.17.0-SNAPSHOT") buildVersionQualifier = System.getProperty("build.version_qualifier", "") isSnapshot = "true" == System.getProperty("build.snapshot", "true") version_tokens = opensearch_version.tokenize('-') diff --git a/gradle.properties b/gradle.properties index 1deb663d8..d275429d7 100644 --- a/gradle.properties +++ b/gradle.properties @@ -7,8 +7,8 @@ # https://github.com/opensearch-project/OpenSearch/blob/main/libs/core/src/main/java/org/opensearch/Version.java . # Wired compatibility of OpenSearch works like 3.x version is compatible with 2.(latest-major) version. # Therefore, to run rolling-upgrade BWC Test on local machine the BWC version here should be set 2.(latest-major). -systemProp.bwc.version=2.16.0-SNAPSHOT -systemProp.bwc.bundle.version=2.15.0 +systemProp.bwc.version=2.17.0-SNAPSHOT +systemProp.bwc.bundle.version=2.16.0 # For fixing Spotless check with Java 17 org.gradle.jvmargs=--add-exports jdk.compiler/com.sun.tools.javac.api=ALL-UNNAMED \ diff --git a/src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java b/src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java index 03c849a50..9b35aa68e 100644 --- a/src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java +++ b/src/main/java/org/opensearch/neuralsearch/ml/MLCommonsClientAccessor.java @@ -15,8 +15,8 @@ import java.util.Objects; import java.util.stream.Collectors; -import org.apache.logging.log4j.util.Strings; import org.opensearch.core.action.ActionListener; +import org.opensearch.core.common.Strings; import org.opensearch.core.common.util.CollectionUtils; import org.opensearch.ml.client.MachineLearningNodeClient; import org.opensearch.ml.common.FunctionName; @@ -237,7 +237,7 @@ private List> buildVectorFromResponse(MLOutput mlOutput) { final List tensorsList = tensors.getMlModelTensors(); for (final ModelTensor tensor : tensorsList) { if (Objects.isNull(tensor.getData())) { - if (Objects.nonNull(tensor.getDataAsMap()) && Strings.isNotBlank((String) tensor.getDataAsMap().get("message"))) { + if (Objects.nonNull(tensor.getDataAsMap()) && Strings.hasText((String) tensor.getDataAsMap().get("message"))) { String errorFromModel = (String) tensor.getDataAsMap().get("message"); throw new IllegalStateException( String.format(Locale.ROOT, "%s: %s", EXCEPTION_MESSAGE_PREFIX_MODEL_PREDICT_FAILED, errorFromModel) diff --git a/src/test/java/org/opensearch/neuralsearch/ValidateDependentPluginInstallationIT.java b/src/test/java/org/opensearch/neuralsearch/ValidateDependentPluginInstallationIT.java index ba9521303..e03a1d3fd 100644 --- a/src/test/java/org/opensearch/neuralsearch/ValidateDependentPluginInstallationIT.java +++ b/src/test/java/org/opensearch/neuralsearch/ValidateDependentPluginInstallationIT.java @@ -75,7 +75,8 @@ private void createBasicKnnIndex() throws IOException { .endObject() .toString(); mapping = mapping.substring(1, mapping.length() - 1); - createIndex(KNN_INDEX_NAME, Settings.EMPTY, mapping); + Settings settings = Settings.builder().put("index.knn", true).build(); + createIndex(KNN_INDEX_NAME, settings, mapping); } private Set getAllInstalledPlugins() throws IOException { diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index d3a8cf3fa..2a6fa49a3 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -17,18 +17,27 @@ import static org.opensearch.neuralsearch.query.NeuralQueryBuilder.MODEL_ID_FIELD; import static org.opensearch.neuralsearch.query.NeuralQueryBuilder.QUERY_TEXT_FIELD; +import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.ArrayList; +import java.util.Optional; +import java.util.Set; import java.util.function.Supplier; +import java.util.stream.Collectors; import org.apache.lucene.search.MatchNoDocsQuery; import org.apache.lucene.search.Query; import org.apache.lucene.search.TermQuery; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; import org.opensearch.Version; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; import org.opensearch.common.xcontent.XContentFactory; import org.opensearch.core.ParseField; import org.opensearch.core.common.ParsingException; @@ -47,8 +56,13 @@ import org.opensearch.index.query.QueryBuilders; import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorDataType; +import org.opensearch.knn.index.engine.KNNEngine; +import org.opensearch.knn.index.engine.KNNMethodContext; +import org.opensearch.knn.index.engine.MethodComponentContext; +import org.opensearch.knn.index.mapper.KNNMappingConfig; import org.opensearch.knn.index.mapper.KNNVectorFieldType; import org.opensearch.knn.index.query.KNNQuery; import org.opensearch.knn.index.query.KNNQueryBuilder; @@ -69,6 +83,26 @@ public class HybridQueryBuilderTests extends OpenSearchQueryTestCase { static final float BOOST = 1.8f; static final Supplier TEST_VECTOR_SUPPLIER = () -> new float[4]; static final QueryBuilder TEST_FILTER = new MatchAllQueryBuilder(); + @Mock + private ClusterService clusterService; + private AutoCloseable openMocks; + + @Override + public void setUp() throws Exception { + super.setUp(); + openMocks = MockitoAnnotations.openMocks(this); + // This is required to make sure that before every test we are initializing the KNNSettings. Not doing this + // leads to failures of unit tests cases when a unit test is run separately. Try running this test: + // ./gradlew ':test' --tests "org.opensearch.knn.training.TrainingJobTests.testRun_success" and see it fails + // but if run along with other tests this test passes. + initKNNSettings(); + } + + @Override + public void tearDown() throws Exception { + super.tearDown(); + openMocks.close(); + } @SneakyThrows public void testDoToQuery_whenNoSubqueries_thenBuildSuccessfully() { @@ -86,11 +120,14 @@ public void testDoToQuery_whenOneSubquery_thenBuildSuccessfully() { Index dummyIndex = new Index("dummy", "dummy"); QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + KNNMethodContext knnMethodContext = new KNNMethodContext(KNNEngine.FAISS, SpaceType.L2, MethodComponentContext.EMPTY); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); + when(mockKNNMappingConfig.getKnnMethodContext()).thenReturn(Optional.of(knnMethodContext)); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockKNNVectorField.getVectorDataType()).thenReturn(VectorDataType.FLOAT); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); - when(mockKNNVectorField.getSpaceType()).thenReturn(SpaceType.L2); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) .queryText(QUERY_TEXT) @@ -116,10 +153,13 @@ public void testDoToQuery_whenMultipleSubqueries_thenBuildSuccessfully() { Index dummyIndex = new Index("dummy", "dummy"); QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + KNNMethodContext knnMethodContext = new KNNMethodContext(KNNEngine.FAISS, SpaceType.L2, MethodComponentContext.EMPTY); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); + when(mockKNNMappingConfig.getKnnMethodContext()).thenReturn(Optional.of(knnMethodContext)); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockKNNVectorField.getVectorDataType()).thenReturn(VectorDataType.FLOAT); - when(mockKNNVectorField.getSpaceType()).thenReturn(SpaceType.L2); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) @@ -367,8 +407,10 @@ public void testToXContent_whenIncomingJsonIsCorrect_thenSuccessful() { Index dummyIndex = new Index("dummy", "dummy"); QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder().fieldName(VECTOR_FIELD_NAME) @@ -584,9 +626,11 @@ public void testRewrite_whenMultipleSubQueries_thenReturnBuilderForEachSubQuery( QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); Index dummyIndex = new Index("dummy", "dummy"); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) createMapperService().fieldType(TEXT_FIELD_NAME); @@ -737,4 +781,17 @@ private void setUpClusterService() { ClusterService clusterService = NeuralSearchClusterTestUtils.mockClusterService(Version.CURRENT); NeuralSearchClusterUtil.instance().initialize(clusterService); } + + private void initKNNSettings() { + Set> defaultClusterSettings = new HashSet<>(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + defaultClusterSettings.addAll( + KNNSettings.state() + .getSettings() + .stream() + .filter(s -> s.getProperties().contains(Setting.Property.NodeScope)) + .collect(Collectors.toList()) + ); + when(clusterService.getClusterSettings()).thenReturn(new ClusterSettings(Settings.EMPTY, defaultClusterSettings)); + KNNSettings.state().setClusterService(clusterService); + } } diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java index 10727246c..f821e7ddf 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryTests.java @@ -13,8 +13,11 @@ import java.io.IOException; import java.util.Arrays; +import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Optional; +import java.util.Set; import java.util.stream.Collectors; import org.apache.lucene.document.FieldType; @@ -33,13 +36,24 @@ import org.apache.lucene.store.Directory; import org.apache.lucene.tests.analysis.MockAnalyzer; import org.apache.lucene.tests.search.QueryUtils; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; import org.opensearch.core.index.Index; import org.opensearch.index.mapper.TextFieldMapper; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.index.query.QueryShardContext; +import org.opensearch.knn.index.KNNSettings; import org.opensearch.knn.index.SpaceType; import org.opensearch.knn.index.VectorDataType; +import org.opensearch.knn.index.engine.KNNEngine; +import org.opensearch.knn.index.engine.KNNMethodContext; +import org.opensearch.knn.index.engine.MethodComponentContext; +import org.opensearch.knn.index.mapper.KNNMappingConfig; import org.opensearch.knn.index.mapper.KNNVectorFieldType; import org.opensearch.knn.index.query.KNNQueryBuilder; @@ -54,6 +68,26 @@ public class HybridQueryTests extends OpenSearchQueryTestCase { static final String TERM_ANOTHER_QUERY_TEXT = "anotherkeyword"; static final float[] VECTOR_QUERY = new float[] { 1.0f, 2.0f, 2.1f, 0.6f }; static final int K = 2; + @Mock + protected ClusterService clusterService; + private AutoCloseable openMocks; + + @Override + public void setUp() throws Exception { + super.setUp(); + openMocks = MockitoAnnotations.openMocks(this); + // This is required to make sure that before every test we are initializing the KNNSettings. Not doing this + // leads to failures of unit tests cases when a unit test is run separately. Try running this test: + // ./gradlew ':test' --tests "org.opensearch.knn.training.TrainingJobTests.testRun_success" and see it fails + // but if run along with other tests this test passes. + initKNNSettings(); + } + + @Override + public void tearDown() throws Exception { + super.tearDown(); + openMocks.close(); + } @SneakyThrows public void testQueryBasics_whenMultipleDifferentQueries_thenSuccessful() { @@ -116,10 +150,13 @@ public void testRewrite_whenRewriteQuery_thenSuccessful() { Index dummyIndex = new Index("dummy", "dummy"); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + KNNMethodContext mockKNNMethodContext = new KNNMethodContext(KNNEngine.FAISS, SpaceType.L2, MethodComponentContext.EMPTY); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); + when(mockKNNMappingConfig.getKnnMethodContext()).thenReturn(Optional.ofNullable(mockKNNMethodContext)); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); - when(mockKNNVectorField.getSpaceType()).thenReturn(SpaceType.L2); when(mockKNNVectorField.getVectorDataType()).thenReturn(VectorDataType.FLOAT); KNNQueryBuilder knnQueryBuilder = new KNNQueryBuilder(VECTOR_FIELD_NAME, VECTOR_QUERY, K); Query knnQuery = knnQueryBuilder.toQuery(mockQueryShardContext); @@ -316,4 +353,17 @@ public void testFilter_whenSubQueriesWithFilterPassed_thenSuccessful() { } assertEquals(2, countOfQueries); } + + private void initKNNSettings() { + Set> defaultClusterSettings = new HashSet<>(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + defaultClusterSettings.addAll( + KNNSettings.state() + .getSettings() + .stream() + .filter(s -> s.getProperties().contains(Setting.Property.NodeScope)) + .collect(Collectors.toList()) + ); + when(clusterService.getClusterSettings()).thenReturn(new ClusterSettings(Settings.EMPTY, defaultClusterSettings)); + KNNSettings.state().setClusterService(clusterService); + } } diff --git a/src/test/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcherTests.java b/src/test/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcherTests.java index bfa519306..0a88324fb 100644 --- a/src/test/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcherTests.java +++ b/src/test/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcherTests.java @@ -62,6 +62,8 @@ import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.index.remote.RemoteStoreEnums; import org.opensearch.index.shard.IndexShard; +import org.opensearch.index.shard.SearchOperationListener; +import org.opensearch.knn.index.mapper.KNNMappingConfig; import org.opensearch.knn.index.mapper.KNNVectorFieldType; import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.OpenSearchQueryTestCase; @@ -93,8 +95,10 @@ public void testQueryType_whenQueryIsHybrid_thenCallHybridDocCollector() { HybridQueryPhaseSearcher hybridQueryPhaseSearcher = spy(new HybridQueryPhaseSearcher()); QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); MapperService mapperService = createMapperService(); TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) mapperService.fieldType(TEXT_FIELD_NAME); @@ -138,6 +142,7 @@ public void testQueryType_whenQueryIsHybrid_thenCallHybridDocCollector() { when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService); @@ -209,6 +214,7 @@ public void testQueryType_whenQueryIsNotHybrid_thenDoNotCallHybridDocCollector() when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.queryResult()).thenReturn(new QuerySearchResult()); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); @@ -280,6 +286,7 @@ public void testQueryResult_whenOneSubQueryWithHits_thenHybridResultsAreSet() { when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); QuerySearchResult querySearchResult = new QuerySearchResult(); when(searchContext.queryResult()).thenReturn(querySearchResult); @@ -366,6 +373,7 @@ public void testQueryResult_whenMultipleTextSubQueriesWithSomeHits_thenHybridRes when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService); @@ -452,6 +460,7 @@ public void testWrappedHybridQuery_whenHybridWrappedIntoBool_thenFail() { when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService); @@ -676,6 +685,7 @@ public void testWrappedHybridQuery_whenHybridWrappedIntoBoolBecauseOfNested_then when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService); @@ -766,6 +776,7 @@ public void testBoolQuery_whenTooManyNestedLevels_thenSuccess() { when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService); @@ -816,8 +827,10 @@ public void testAggregations_whenMetricAggregation_thenSuccessful() { HybridQueryPhaseSearcher hybridQueryPhaseSearcher = spy(new HybridQueryPhaseSearcher()); QueryShardContext mockQueryShardContext = mock(QueryShardContext.class); KNNVectorFieldType mockKNNVectorField = mock(KNNVectorFieldType.class); + KNNMappingConfig mockKNNMappingConfig = mock(KNNMappingConfig.class); + when(mockKNNVectorField.getKnnMappingConfig()).thenReturn(mockKNNMappingConfig); when(mockQueryShardContext.index()).thenReturn(dummyIndex); - when(mockKNNVectorField.getDimension()).thenReturn(4); + when(mockKNNVectorField.getKnnMappingConfig().getDimension()).thenReturn(4); when(mockQueryShardContext.fieldMapper(eq(VECTOR_FIELD_NAME))).thenReturn(mockKNNVectorField); MapperService mapperService = createMapperService(); TextFieldMapper.TextFieldType fieldType = (TextFieldMapper.TextFieldType) mapperService.fieldType(TEXT_FIELD_NAME); @@ -861,6 +874,7 @@ public void testAggregations_whenMetricAggregation_thenSuccessful() { when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService); @@ -954,6 +968,7 @@ public void testAliasWithFilter_whenHybridWrappedIntoBoolBecauseOfIndexAlias_the when(searchContext.searcher()).thenReturn(contextIndexSearcher); IndexShard indexShard = mock(IndexShard.class); when(indexShard.shardId()).thenReturn(new ShardId("test", "test", 0)); + when(indexShard.getSearchOperationListener()).thenReturn(mock(SearchOperationListener.class)); when(searchContext.indexShard()).thenReturn(indexShard); when(searchContext.bucketCollectorProcessor()).thenReturn(SearchContext.NO_OP_BUCKET_COLLECTOR_PROCESSOR); when(searchContext.mapperService()).thenReturn(mapperService);