From debafba6ef490438ed541cf45b9fa9e354c39638 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 16 Jan 2024 16:32:15 -0800 Subject: [PATCH 1/5] Enable support for default model id on HybridQueryBuilder Signed-off-by: Varun Jain --- .../query/HybridQueryBuilder.java | 15 ++++++++++++++ .../NeuralQueryEnricherProcessorIT.java | 20 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java index b0104639b..bf7469359 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java @@ -13,6 +13,7 @@ import java.util.stream.Collectors; import org.apache.commons.lang.builder.EqualsBuilder; +import org.apache.lucene.search.BooleanClause.Occur; import org.apache.lucene.search.Query; import org.opensearch.common.lucene.search.Queries; import org.opensearch.core.ParseField; @@ -27,6 +28,7 @@ import org.opensearch.index.query.QueryShardContext; import org.opensearch.index.query.QueryShardException; import org.opensearch.index.query.Rewriteable; +import org.opensearch.index.query.QueryBuilderVisitor; import lombok.Getter; import lombok.NoArgsConstructor; @@ -295,4 +297,17 @@ private Collection toQueries(Collection queryBuilders, Quer }).filter(Objects::nonNull).collect(Collectors.toList()); return queries; } + + /** + * visit method to parse the HybridQueryBuilder by a visitor + * @return + */ + @Override + public void visit(QueryBuilderVisitor visitor) { + visitor.accept(this); + QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST); + for (QueryBuilder subQueryBuilder : queries) { + subQueryBuilder.visit(subVisitor); + } + } } diff --git a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java index 3c8adf5b3..b37db7c82 100644 --- a/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java +++ b/src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java @@ -15,6 +15,7 @@ import org.junit.Before; import org.opensearch.common.settings.Settings; import org.opensearch.neuralsearch.BaseNeuralSearchIT; +import org.opensearch.neuralsearch.query.HybridQueryBuilder; import org.opensearch.neuralsearch.query.NeuralQueryBuilder; import com.google.common.primitives.Floats; @@ -62,6 +63,25 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() { } + @SneakyThrows + public void testNeuralQueryEnricherProcessor_whenHybridQueryBuilderAndNoModelIdPassed_thenSuccess() { + initializeIndexIfNotExist(); + String modelId = getDeployedModelId(); + createSearchRequestProcessor(modelId, search_pipeline); + createPipelineProcessor(modelId, ingest_pipeline); + updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline)); + NeuralQueryBuilder neuralQueryBuilder = new NeuralQueryBuilder(); + neuralQueryBuilder.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1); + neuralQueryBuilder.queryText("Hello World"); + neuralQueryBuilder.k(1); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(neuralQueryBuilder); + Map response = search(index, hybridQueryBuilder, 2); + + assertFalse(response.isEmpty()); + + } + @SneakyThrows private void initializeIndexIfNotExist() { if (index.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(index)) { From a3ef808453bea1917da061b6debb0ac2a95a86ad Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 16 Jan 2024 16:56:40 -0800 Subject: [PATCH 2/5] Adding tests and updating changelog.md Signed-off-by: Varun Jain --- CHANGELOG.md | 1 + .../neuralsearch/query/HybridQueryBuilder.java | 1 - .../neuralsearch/query/HybridQueryBuilderTests.java | 9 +++++++++ .../java/org/opensearch/neuralsearch/TestUtils.java | 6 ++++++ 4 files changed, 16 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ede506f2b..922016575 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ### Bug Fixes - Fixing multiple issues reported in #497 ([#524](https://github.com/opensearch-project/neural-search/pull/524)) - Fix Flaky test reported in #433 ([#533](https://github.com/opensearch-project/neural-search/pull/533)) +- Enable support for default model id on HybridQueryBuilder ([#541](https://github.com/opensearch-project/neural-search/pull/541)) ### Infrastructure - BWC tests for Neural Search ([#515](https://github.com/opensearch-project/neural-search/pull/515)) - Github action to run integ tests in secure opensearch cluster ([#535](https://github.com/opensearch-project/neural-search/pull/535)) diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java index bf7469359..a7da508bf 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java @@ -300,7 +300,6 @@ private Collection toQueries(Collection queryBuilders, Quer /** * visit method to parse the HybridQueryBuilder by a visitor - * @return */ @Override public void visit(QueryBuilderVisitor visitor) { diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index 7bae96780..8f062c425 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -20,6 +20,7 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.ArrayList; import java.util.function.Supplier; import org.apache.lucene.search.MatchNoDocsQuery; @@ -52,6 +53,7 @@ import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils; import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; +import static org.opensearch.neuralsearch.TestUtils.createTestVisitor; import com.carrotsearch.randomizedtesting.RandomizedTest; @@ -708,6 +710,13 @@ public void testBoost_whenDefaultBoostSet_thenBuildSuccessfully() { assertNotNull(hybridQueryBuilder); } + public void testVisit() { + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add(new NeuralQueryBuilder()).add(new NeuralSparseQueryBuilder()); + List visitedQueries = new ArrayList<>(); + hybridQueryBuilder.visit(createTestVisitor(visitedQueries)); + assertEquals(3, visitedQueries.size()); + } + private Map getInnerMap(Object innerObject, String queryName, String fieldName) { if (!(innerObject instanceof Map)) { fail("field name does not map to nested object"); diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java b/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java index 99e396395..205352d3f 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java @@ -8,6 +8,8 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; import static org.opensearch.test.OpenSearchTestCase.randomFloat; import java.util.ArrayList; @@ -310,4 +312,8 @@ public static String getModelId(Map pipeline, String processor) return (String) textEmbeddingProcessor.get("model_id"); } + + public static QueryBuilderVisitor createTestVisitor(List visitedQueries) { + return QueryBuilderVisitor.NO_OP_VISITOR; + } } From a523130edcab04ba0cf6ecdc072b0e26db793474 Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 16 Jan 2024 17:06:26 -0800 Subject: [PATCH 3/5] Optimizing code Signed-off-by: Varun Jain --- .../query/HybridQueryBuilderTests.java | 14 +++++++------- .../org/opensearch/neuralsearch/TestUtils.java | 6 ------ 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index 8f062c425..6b73f2709 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -9,6 +9,12 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.opensearch.core.xcontent.ToXContent.EMPTY_PARAMS; +import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.QueryShardContext; +import org.opensearch.index.query.TermQueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; import static org.opensearch.index.query.AbstractQueryBuilder.BOOST_FIELD; import static org.opensearch.index.query.AbstractQueryBuilder.DEFAULT_BOOST; import static org.opensearch.knn.index.query.KNNQueryBuilder.FILTER_FIELD; @@ -42,18 +48,12 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.mapper.TextFieldMapper; -import org.opensearch.index.query.MatchAllQueryBuilder; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.QueryBuilders; -import org.opensearch.index.query.QueryShardContext; -import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.knn.index.VectorDataType; import org.opensearch.knn.index.mapper.KNNVectorFieldMapper; import org.opensearch.knn.index.query.KNNQuery; import org.opensearch.knn.index.query.KNNQueryBuilder; import org.opensearch.neuralsearch.util.NeuralSearchClusterTestUtils; import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil; -import static org.opensearch.neuralsearch.TestUtils.createTestVisitor; import com.carrotsearch.randomizedtesting.RandomizedTest; @@ -713,7 +713,7 @@ public void testBoost_whenDefaultBoostSet_thenBuildSuccessfully() { public void testVisit() { HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add(new NeuralQueryBuilder()).add(new NeuralSparseQueryBuilder()); List visitedQueries = new ArrayList<>(); - hybridQueryBuilder.visit(createTestVisitor(visitedQueries)); + hybridQueryBuilder.visit(QueryBuilderVisitor.NO_OP_VISITOR); assertEquals(3, visitedQueries.size()); } diff --git a/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java b/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java index 205352d3f..99e396395 100644 --- a/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java +++ b/src/testFixtures/java/org/opensearch/neuralsearch/TestUtils.java @@ -8,8 +8,6 @@ import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.QueryBuilderVisitor; import static org.opensearch.test.OpenSearchTestCase.randomFloat; import java.util.ArrayList; @@ -312,8 +310,4 @@ public static String getModelId(Map pipeline, String processor) return (String) textEmbeddingProcessor.get("model_id"); } - - public static QueryBuilderVisitor createTestVisitor(List visitedQueries) { - return QueryBuilderVisitor.NO_OP_VISITOR; - } } From 0e63d85d07724010d7f0a198014ca51b479fdbeb Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Tue, 16 Jan 2024 17:24:03 -0800 Subject: [PATCH 4/5] modyfing the tests4 Signed-off-by: Varun Jain --- .../query/HybridQueryBuilderTests.java | 13 +++++---- .../query/OpenSearchQueryTestCase.java | 27 +++++++++++++++---- 2 files changed, 28 insertions(+), 12 deletions(-) diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index 6b73f2709..72ea8fa7b 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -9,12 +9,6 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.opensearch.core.xcontent.ToXContent.EMPTY_PARAMS; -import org.opensearch.index.query.MatchAllQueryBuilder; -import org.opensearch.index.query.QueryBuilder; -import org.opensearch.index.query.QueryBuilders; -import org.opensearch.index.query.QueryShardContext; -import org.opensearch.index.query.TermQueryBuilder; -import org.opensearch.index.query.QueryBuilderVisitor; import static org.opensearch.index.query.AbstractQueryBuilder.BOOST_FIELD; import static org.opensearch.index.query.AbstractQueryBuilder.DEFAULT_BOOST; import static org.opensearch.knn.index.query.KNNQueryBuilder.FILTER_FIELD; @@ -48,6 +42,11 @@ import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.mapper.TextFieldMapper; +import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.QueryShardContext; +import org.opensearch.index.query.TermQueryBuilder; import org.opensearch.knn.index.VectorDataType; import org.opensearch.knn.index.mapper.KNNVectorFieldMapper; import org.opensearch.knn.index.query.KNNQuery; @@ -713,7 +712,7 @@ public void testBoost_whenDefaultBoostSet_thenBuildSuccessfully() { public void testVisit() { HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add(new NeuralQueryBuilder()).add(new NeuralSparseQueryBuilder()); List visitedQueries = new ArrayList<>(); - hybridQueryBuilder.visit(QueryBuilderVisitor.NO_OP_VISITOR); + hybridQueryBuilder.visit(createTestVisitor(visitedQueries)); assertEquals(3, visitedQueries.size()); } diff --git a/src/test/java/org/opensearch/neuralsearch/query/OpenSearchQueryTestCase.java b/src/test/java/org/opensearch/neuralsearch/query/OpenSearchQueryTestCase.java index 97e1f7c5d..a1e8210e6 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/OpenSearchQueryTestCase.java +++ b/src/test/java/org/opensearch/neuralsearch/query/OpenSearchQueryTestCase.java @@ -7,6 +7,14 @@ import static java.util.Collections.emptyMap; import static java.util.Collections.singletonMap; import static java.util.stream.Collectors.toList; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.Explanation; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.Weight; +import org.apache.lucene.search.BooleanClause; +import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; import static org.opensearch.neuralsearch.settings.NeuralSearchSettings.NEURAL_SEARCH_HYBRID_SEARCH_DISABLED; import java.io.IOException; @@ -20,11 +28,6 @@ import org.apache.lucene.document.FieldType; import org.apache.lucene.document.TextField; import org.apache.lucene.index.LeafReaderContext; -import org.apache.lucene.search.DocIdSetIterator; -import org.apache.lucene.search.Explanation; -import org.apache.lucene.search.Query; -import org.apache.lucene.search.Scorer; -import org.apache.lucene.search.Weight; import org.opensearch.Version; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.common.CheckedConsumer; @@ -230,4 +233,18 @@ public float getMaxScore(int upTo) { protected static void initFeatureFlags() { System.setProperty(NEURAL_SEARCH_HYBRID_SEARCH_DISABLED.getKey(), "true"); } + + protected static QueryBuilderVisitor createTestVisitor(List visitedQueries) { + return new QueryBuilderVisitor() { + @Override + public void accept(QueryBuilder qb) { + visitedQueries.add(qb); + } + + @Override + public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { + return this; + } + }; + } } From b8d504a0b695ad750b2b5a86dde00266e68e3f0f Mon Sep 17 00:00:00 2001 From: Varun Jain Date: Wed, 17 Jan 2024 13:03:05 -0800 Subject: [PATCH 5/5] Addressing Heemin comment Signed-off-by: Varun Jain --- .../org/opensearch/neuralsearch/query/HybridQueryBuilder.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java index a7da508bf..aa4242c2e 100644 --- a/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java +++ b/src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java @@ -304,6 +304,8 @@ private Collection toQueries(Collection queryBuilders, Quer @Override public void visit(QueryBuilderVisitor visitor) { visitor.accept(this); + // getChildVisitor of NeuralSearchQueryVisitor return this. + // therefore any argument can be passed. Here we have used Occcur.MUST as an argument. QueryBuilderVisitor subVisitor = visitor.getChildVisitor(Occur.MUST); for (QueryBuilder subQueryBuilder : queries) { subQueryBuilder.visit(subVisitor);