Fixed wrapped bool queries for latest core, fixed failing tests

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

Author: martin-gaievski

qa/rolling-upgrade/src/test/java/org/opensearch/neuralsearch/bwc/rolling/BatchIngestionIT.java

@@ -11,6 +11,7 @@
 import java.nio.file.Path;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static org.opensearch.neuralsearch.util.BatchIngestionUtils.prepareDataForBulkIngestion;
 import static org.opensearch.neuralsearch.util.TestUtils.NODES_BWC_CLUSTER;
@@ -20,6 +21,7 @@ public class BatchIngestionIT extends AbstractRollingUpgradeTestCase {
     private static final String SPARSE_PIPELINE = "BatchIngestionIT_sparse_pipeline_rolling";
     private static final String TEXT_FIELD_NAME = "passage_text";
     private static final String EMBEDDING_FIELD_NAME = "passage_embedding";
+    private static final Set<MLModelState> READY_FOR_INFERENCE_STATES = Set.of(MLModelState.LOADED, MLModelState.DEPLOYED);
 
     public void testBatchIngestion_SparseEncodingProcessor_E2EFlow() throws Exception {
         waitForClusterHealthGreen(NODES_BWC_CLUSTER, 90);
@@ -98,7 +100,7 @@ private void waitForModelToLoad(String modelId) throws Exception {
 
         for (int attempt = 0; attempt < maxAttempts; attempt++) {
             MLModelState state = getModelState(modelId);
-            if (state == MLModelState.LOADED) {
+            if (READY_FOR_INFERENCE_STATES.contains(state)) {
                 logger.info("Model {} is now loaded after {} attempts", modelId, attempt + 1);
                 return;
             }

src/main/java/org/opensearch/neuralsearch/query/HybridQuery.java

@@ -14,6 +14,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.Callable;
+import java.util.stream.Collectors;
 
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
@@ -42,6 +43,26 @@ public final class HybridQuery extends Query implements Iterable<Query> {
      * @param filterQueries list of filters that will be applied to each sub query. Each filter from the list is added as bool "filter" clause. If this is null sub queries will be executed as is
      */
     public HybridQuery(final Collection<Query> subQueries, final List<Query> filterQueries, final HybridQueryContext hybridQueryContext) {
+        this(
+            subQueries,
+            hybridQueryContext,
+            filterQueries == null
+                ? null
+                : filterQueries.stream().map(query -> new BooleanClause(query, BooleanClause.Occur.FILTER)).collect(Collectors.toList())
+        );
+    }
+
+    /**
+     * Create new instance of hybrid query object based on collection of sub queries and boolean clauses that are used as filters for each sub-query
+     * @param subQueries
+     * @param hybridQueryContext
+     * @param booleanClauses
+     */
+    public HybridQuery(
+        final Collection<Query> subQueries,
+        final HybridQueryContext hybridQueryContext,
+        final List<BooleanClause> booleanClauses
+    ) {
         Objects.requireNonNull(subQueries, "collection of queries must not be null");
         if (subQueries.isEmpty()) {
             throw new IllegalArgumentException("collection of queries must not be empty");
@@ -50,14 +71,14 @@ public HybridQuery(final Collection<Query> subQueries, final List<Query> filterQ
         if (Objects.nonNull(paginationDepth) && paginationDepth == 0) {
             throw new IllegalArgumentException("pagination_depth must not be zero");
         }
-        if (Objects.isNull(filterQueries) || filterQueries.isEmpty()) {
+        if (Objects.isNull(booleanClauses) || booleanClauses.isEmpty()) {
             this.subQueries = new ArrayList<>(subQueries);
         } else {
             List<Query> modifiedSubQueries = new ArrayList<>();
             for (Query subQuery : subQueries) {
                 BooleanQuery.Builder builder = new BooleanQuery.Builder();
                 builder.add(subQuery, BooleanClause.Occur.MUST);
-                filterQueries.forEach(filterQuery -> builder.add(filterQuery, BooleanClause.Occur.FILTER));
+                booleanClauses.forEach(filterQuery -> builder.add(booleanClauses));
                 modifiedSubQueries.add(builder.build());
             }
             this.subQueries = modifiedSubQueries;

src/main/java/org/opensearch/neuralsearch/query/NeuralKNNQuery.java

@@ -5,7 +5,6 @@
 package org.opensearch.neuralsearch.query;
 
 import lombok.Getter;
-import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.QueryVisitor;
@@ -45,8 +44,8 @@ public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float bo
     }
 
     @Override
-    public Query rewrite(IndexReader reader) throws IOException {
-        Query rewritten = knnQuery.rewrite(reader);
+    public Query rewrite(IndexSearcher indexSearcher) throws IOException {
+        Query rewritten = knnQuery.rewrite(indexSearcher);
         if (rewritten == knnQuery) {
             return this;
         }

src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java

@@ -79,14 +79,11 @@ protected Query extractHybridQuery(final SearchContext searchContext, final Quer
         if (isHybridQueryWrappedInBooleanQuery(searchContext, query)) {
             List<BooleanClause> booleanClauses = ((BooleanQuery) query).clauses();
             if (!(booleanClauses.get(0).query() instanceof HybridQuery)) {
-                throw new IllegalStateException("cannot process hybrid query due to incorrect structure of top level query");
+                throw new IllegalArgumentException("hybrid query must be a top level query and cannot be wrapped into other queries");
             }
-            HybridQuery hybridQuery = (HybridQuery) booleanClauses.stream().findFirst().get().query();
-            List<Query> filterQueries = booleanClauses.stream()
-                .filter(clause -> BooleanClause.Occur.FILTER == clause.occur())
-                .map(BooleanClause::query)
-                .collect(Collectors.toList());
-            HybridQuery hybridQueryWithFilter = new HybridQuery(hybridQuery.getSubQueries(), filterQueries, hybridQuery.getQueryContext());
+            HybridQuery hybridQuery = (HybridQuery) booleanClauses.get(0).query();
+            List<BooleanClause> filterQueries = booleanClauses.stream().skip(1).collect(Collectors.toList());
+            HybridQuery hybridQueryWithFilter = new HybridQuery(hybridQuery.getSubQueries(), hybridQuery.getQueryContext(), filterQueries);
             return hybridQueryWithFilter;
         }
         return query;

src/main/java/org/opensearch/neuralsearch/util/HybridQueryUtil.java

@@ -24,33 +24,9 @@ public class HybridQueryUtil {
      * This method validates whether the query object is an instance of hybrid query
      */
     public static boolean isHybridQuery(final Query query, final SearchContext searchContext) {
-        if (query instanceof HybridQuery) {
+        if (query instanceof HybridQuery
+            || (Objects.nonNull(searchContext.parsedQuery()) && searchContext.parsedQuery().query() instanceof HybridQuery)) {
             return true;
-        } else if (isWrappedHybridQuery(query)) {
-            /* Checking if this is a hybrid query that is wrapped into a Bool query by core Opensearch code
-            https://github.com/opensearch-project/OpenSearch/blob/main/server/src/main/java/org/opensearch/search/DefaultSearchContext.java#L367-L370.
-            main reason for that is performance optimization, at time of writing we are ok with loosing on performance if that's unblocks
-            hybrid query for indexes with nested field types.
-            in such case we consider query a valid hybrid query. Later in the code we will extract it and execute as a main query for
-            this search request.
-            below is sample structure of such query:
-
-            Boolean {
-                should: {
-                       hybrid: {
-                           sub_query1 {}
-                           sub_query2 {}
-                       }
-                }
-                filter: {
-                       exists: {
-                           field: "_primary_term"
-                       }
-                }
-            }
-            */
-            // we have already checked if query in instance of Boolean in higher level else if condition
-            return hasNestedFieldOrNestedDocs(query, searchContext) || hasAliasFilter(query, searchContext);
         }
         return false;
     }

src/test/java/org/opensearch/neuralsearch/query/HybridQuerySortIT.java

@@ -546,20 +546,18 @@ public void testExplainAndSort_whenIndexWithMultipleShards_thenSuccessful() {
             hit1DetailsForHit1DetailsForHit1DetailsForHit1.get("description")
         );
         assertTrue((double) hit1DetailsForHit1DetailsForHit1DetailsForHit1.get("value") > 0.0f);
-        assertEquals(3, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").size());
+        assertEquals(2, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").size());
 
-        assertEquals("boost", getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("description"));
-        assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("value") > 0.0f);
         assertEquals(
             "idf, computed as log(1 + (N - n + 0.5) / (n + 0.5)) from:",
-            getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("description")
+            getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("description")
         );
-        assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("value") > 0.0f);
+        assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(0).get("value") > 0.0f);
         assertEquals(
             "tf, computed as freq / (freq + k1 * (1 - b + b * dl / avgdl)) from:",
-            getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(2).get("description")
+            getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("description")
         );
-        assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(2).get("value") > 0.0f);
+        assertTrue((double) getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit1, "details").get(1).get("value") > 0.0f);
 
         // hit 4
         Map<String, Object> searchHit4 = nestedHits.get(3);
@@ -588,7 +586,7 @@ public void testExplainAndSort_whenIndexWithMultipleShards_thenSuccessful() {
             hit1DetailsForHit1DetailsForHit1DetailsForHit4.get("description")
         );
         assertTrue((double) hit1DetailsForHit1DetailsForHit1DetailsForHit4.get("value") > 0.0f);
-        assertEquals(3, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit4, "details").size());
+        assertEquals(2, getListOfValues(hit1DetailsForHit1DetailsForHit1DetailsForHit4, "details").size());
 
         // hit 6
         Map<String, Object> searchHit6 = nestedHits.get(5);

src/test/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcherTests.java

@@ -57,6 +57,7 @@
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.TextFieldMapper;
 import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.ParsedQuery;
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.index.query.DisMaxQueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
@@ -702,8 +703,8 @@ public void testWrappedHybridQuery_whenHybridWrappedIntoBoolAndIncorrectStructur
 
         when(searchContext.query()).thenReturn(query);
 
-        IllegalStateException exception = expectThrows(
-            IllegalStateException.class,
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
             () -> hybridQueryPhaseSearcher.searchWith(
                 searchContext,
                 contextIndexSearcher,
@@ -716,7 +717,7 @@ public void testWrappedHybridQuery_whenHybridWrappedIntoBoolAndIncorrectStructur
 
         org.hamcrest.MatcherAssert.assertThat(
             exception.getMessage(),
-            containsString("cannot process hybrid query due to incorrect structure of top level query")
+            containsString("hybrid query must be a top level query and cannot be wrapped into other queries")
         );
 
         releaseResources(directory, w, reader);
@@ -818,6 +819,8 @@ public void testWrappedHybridQuery_whenHybridWrappedIntoBoolBecauseOfNested_then
         Query query = builder.build();
 
         when(searchContext.query()).thenReturn(query);
+        Query hybridQuery = queryBuilder.toQuery(mockQueryShardContext);
+        when(searchContext.parsedQuery()).thenReturn(new ParsedQuery(hybridQuery));
 
         CollectorManager<? extends Collector, ReduceableSearchResult> collectorManager = HybridCollectorManager
             .createHybridCollectorManager(searchContext);
@@ -1110,6 +1113,9 @@ public void testAliasWithFilter_whenHybridWrappedIntoBoolBecauseOfIndexAlias_the
         when(searchContext.query()).thenReturn(query);
         when(searchContext.aliasFilter()).thenReturn(termFilter);
 
+        Query hybridQuery = queryBuilder.toQuery(mockQueryShardContext);
+        when(searchContext.parsedQuery()).thenReturn(new ParsedQuery(hybridQuery));
+
         CollectorManager<? extends Collector, ReduceableSearchResult> collectorManager = HybridCollectorManager
             .createHybridCollectorManager(searchContext);
         Map<Class<?>, CollectorManager<? extends Collector, ReduceableSearchResult>> queryCollectorManagers = new HashMap<>();