Support hybrid query with neural highlighter

Signed-off-by: Junqiu Lei <junqiu@amazon.com>

Author: junqiu-lei

src/main/java/org/opensearch/neuralsearch/highlight/extractor/BooleanQueryTextExtractor.java

@@ -24,9 +24,8 @@ public BooleanQueryTextExtractor(QueryTextExtractorRegistry registry) {
 
     @Override
     public String extractQueryText(Query query, String fieldName) {
-        if (!(query instanceof BooleanQuery booleanQuery)) {
-            return "";
-        }
+        validateQueryType(query, BooleanQuery.class);
+        BooleanQuery booleanQuery = (BooleanQuery) query;
 
         StringBuilder sb = new StringBuilder();
 
@@ -38,15 +37,14 @@ public String extractQueryText(Query query, String fieldName) {
 
             try {
                 String clauseText = registry.extractQueryText(clause.query(), fieldName);
-                if (!clauseText.isEmpty()) {
-                    if (sb.length() > 0) {
+                if (clauseText.isEmpty() == false) {
+                    if (sb.isEmpty() == false) {
                         sb.append(" ");
                     }
                     sb.append(clauseText);
                 }
             } catch (IllegalArgumentException e) {
-                // If a clause has empty query text, just skip it
-                log.debug("Skipping clause with empty query text: {}", clause.query());
+                log.warn(String.format("Failed to extract text from clause %s: %s", clause, e.getMessage()), e);
             }
         }
 

src/main/java/org/opensearch/neuralsearch/highlight/extractor/HybridQueryTextExtractor.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.highlight.extractor;
+
+import org.apache.lucene.search.Query;
+import org.opensearch.neuralsearch.query.HybridQuery;
+
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+/**
+ * Extractor for hybrid queries that combines text from all sub-queries
+ */
+public class HybridQueryTextExtractor implements QueryTextExtractor {
+
+    private final QueryTextExtractorRegistry registry;
+
+    public HybridQueryTextExtractor(QueryTextExtractorRegistry registry) {
+        this.registry = registry;
+    }
+
+    @Override
+    public String extractQueryText(Query query, String fieldName) {
+        validateQueryType(query, HybridQuery.class);
+        HybridQuery hybridQuery = (HybridQuery) query;
+
+        // Create a set to avoid duplicates while maintaining order
+        Set<String> queryTexts = new LinkedHashSet<>();
+
+        // Extract text from each sub-query
+        for (Query subQuery : hybridQuery.getSubQueries()) {
+            String extractedText = registry.extractQueryText(subQuery, fieldName);
+            if (!extractedText.isEmpty()) {
+                queryTexts.add(extractedText);
+            }
+        }
+
+        // Join with spaces
+        return String.join(" ", queryTexts).trim();
+    }
+}

src/main/java/org/opensearch/neuralsearch/highlight/extractor/NeuralQueryTextExtractor.java

@@ -14,9 +14,8 @@ public class NeuralQueryTextExtractor implements QueryTextExtractor {
 
     @Override
     public String extractQueryText(Query query, String fieldName) {
-        if (query instanceof NeuralKNNQuery neuralQuery) {
-            return neuralQuery.getOriginalQueryText();
-        }
-        return "";
+        validateQueryType(query, NeuralKNNQuery.class);
+        NeuralKNNQuery neuralQuery = (NeuralKNNQuery) query;
+        return neuralQuery.getOriginalQueryText();
     }
 }

src/main/java/org/opensearch/neuralsearch/highlight/extractor/QueryTextExtractor.java

@@ -10,6 +10,21 @@
  * Interface for extracting query text from different query types
  */
 public interface QueryTextExtractor {
+    /**
+     * Validates if the query is of the expected type
+     *
+     * @param query The query to validate
+     * @param expectedType The expected query type
+     * @throws IllegalArgumentException if the query is not of the expected type
+     */
+    default void validateQueryType(Query query, Class<? extends Query> expectedType) {
+        if (expectedType.isInstance(query) == false) {
+            throw new IllegalArgumentException(
+                String.format("Expected %s but got %s", expectedType.getSimpleName(), query.getClass().getSimpleName())
+            );
+        }
+    }
+
     /**
      * Extracts text from a query for highlighting
      *

src/main/java/org/opensearch/neuralsearch/highlight/extractor/QueryTextExtractorRegistry.java

@@ -8,6 +8,7 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
 import org.opensearch.neuralsearch.query.NeuralKNNQuery;
+import org.opensearch.neuralsearch.query.HybridQuery;
 
 import lombok.extern.log4j.Log4j2;
 
@@ -35,6 +36,7 @@ public QueryTextExtractorRegistry() {
     private void initialize() {
         register(NeuralKNNQuery.class, new NeuralQueryTextExtractor());
         register(TermQuery.class, new TermQueryTextExtractor());
+        register(HybridQuery.class, new HybridQueryTextExtractor(this));
 
         // BooleanQueryTextExtractor needs a reference to this registry
         // so we need to register it after creating the registry instance

src/main/java/org/opensearch/neuralsearch/highlight/extractor/TermQueryTextExtractor.java

@@ -15,9 +15,8 @@ public class TermQueryTextExtractor implements QueryTextExtractor {
 
     @Override
     public String extractQueryText(Query query, String fieldName) {
-        if (!(query instanceof TermQuery termQuery)) {
-            return "";
-        }
+        validateQueryType(query, TermQuery.class);
+        TermQuery termQuery = (TermQuery) query;
 
         Term term = termQuery.getTerm();
         // Only include terms from the field we're highlighting

src/test/java/org/opensearch/neuralsearch/highlight/QueryTextExtractorTests.java

@@ -7,14 +7,21 @@
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
 import org.opensearch.neuralsearch.highlight.extractor.BooleanQueryTextExtractor;
 import org.opensearch.neuralsearch.highlight.extractor.NeuralQueryTextExtractor;
 import org.opensearch.neuralsearch.highlight.extractor.QueryTextExtractorRegistry;
 import org.opensearch.neuralsearch.highlight.extractor.TermQueryTextExtractor;
+import org.opensearch.neuralsearch.highlight.extractor.HybridQueryTextExtractor;
 import org.opensearch.neuralsearch.query.NeuralKNNQuery;
+import org.opensearch.neuralsearch.query.HybridQuery;
+import org.opensearch.neuralsearch.query.HybridQueryContext;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -48,8 +55,14 @@ public void testTermQueryExtractor() {
 
         // Test with non-TermQuery
         BooleanQuery booleanQuery = new BooleanQuery.Builder().build();
-        result = extractor.extractQueryText(booleanQuery, "content");
-        assertEquals("Should return empty string for non-TermQuery", "", result);
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> extractor.extractQueryText(booleanQuery, "content")
+        );
+        assertTrue(
+            "Should throw IllegalArgumentException with correct message",
+            exception.getMessage().contains("Expected TermQuery but got BooleanQuery")
+        );
     }
 
     /**
@@ -67,8 +80,14 @@ public void testNeuralQueryExtractor() {
 
         // Test with non-NeuralKNNQuery
         TermQuery termQuery = new TermQuery(new Term("content", "term"));
-        result = extractor.extractQueryText(termQuery, "content");
-        assertEquals("Should return empty string for non-NeuralKNNQuery", "", result);
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> extractor.extractQueryText(termQuery, "content")
+        );
+        assertTrue(
+            "Should throw IllegalArgumentException with correct message",
+            exception.getMessage().contains("Expected NeuralKNNQuery but got TermQuery")
+        );
     }
 
     /**
@@ -91,8 +110,14 @@ public void testBooleanQueryExtractor() {
 
         // Test with non-BooleanQuery
         TermQuery termQuery = new TermQuery(new Term("content", "term"));
-        result = extractor.extractQueryText(termQuery, "content");
-        assertEquals("Should return empty string for non-BooleanQuery", "", result);
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> extractor.extractQueryText(termQuery, "content")
+        );
+        assertTrue(
+            "Should throw IllegalArgumentException with correct message",
+            exception.getMessage().contains("Expected BooleanQuery but got TermQuery")
+        );
 
         // Test with empty clauses
         BooleanQuery emptyQuery = new BooleanQuery.Builder().build();
@@ -160,4 +185,111 @@ public void visit(org.apache.lucene.search.QueryVisitor visitor) {
         result = registry.extractQueryText(customQuery, "content");
         assertEquals("Should use custom extractor", "custom-extracted", result);
     }
+
+    /**
+     * Tests the HybridQueryTextExtractor
+     */
+    public void testHybridQueryExtractor() {
+        // Create a hybrid query with multiple sub-queries
+        List<Query> subQueries = new ArrayList<>();
+
+        // Add a term query
+        TermQuery termQuery = new TermQuery(new Term("content", "machine"));
+        subQueries.add(termQuery);
+
+        // Add a boolean query (match query)
+        BooleanQuery.Builder boolBuilder = new BooleanQuery.Builder();
+        boolBuilder.add(new TermQuery(new Term("content", "learning")), BooleanClause.Occur.MUST);
+        subQueries.add(boolBuilder.build());
+
+        // Add a neural query
+        NeuralKNNQuery neuralQuery = mock(NeuralKNNQuery.class);
+        when(neuralQuery.getOriginalQueryText()).thenReturn("AI systems that can learn");
+        subQueries.add(neuralQuery);
+
+        // Create the hybrid query
+        HybridQuery hybridQuery = new HybridQuery(subQueries, HybridQueryContext.builder().build());
+
+        // Test extraction
+        String result = registry.extractQueryText(hybridQuery, "content");
+        assertEquals("Should combine all query texts correctly", "machine learning AI systems that can learn", result);
+
+        // Test with non-HybridQuery
+        TermQuery nonHybridQuery = new TermQuery(new Term("content", "term"));
+        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> {
+            HybridQueryTextExtractor extractor = new HybridQueryTextExtractor(registry);
+            extractor.extractQueryText(nonHybridQuery, "content");
+        });
+        assertTrue(
+            "Should throw IllegalArgumentException with correct message",
+            exception.getMessage().contains("Expected HybridQuery but got TermQuery")
+        );
+    }
+
+    /**
+     * Tests the HybridQueryTextExtractor with empty or invalid sub-queries
+     */
+    public void testHybridQueryExtractorWithEmptyQueries() {
+        // Create a hybrid query with no valid text
+        List<Query> subQueries = new ArrayList<>();
+
+        // Add a term query with non-matching field
+        TermQuery termQuery = new TermQuery(new Term("title", "machine"));
+        subQueries.add(termQuery);
+
+        // Create the hybrid query
+        HybridQuery hybridQuery = new HybridQuery(subQueries, HybridQueryContext.builder().build());
+
+        // Test extraction
+        String result = registry.extractQueryText(hybridQuery, "content");
+        assertEquals("Should return empty string for no valid text", "", result);
+    }
+
+    /**
+     * Tests the HybridQueryTextExtractor with duplicate texts
+     */
+    public void testHybridQueryExtractorWithDuplicates() {
+        List<Query> subQueries = new ArrayList<>();
+
+        // Add two term queries with the same text
+        subQueries.add(new TermQuery(new Term("content", "duplicate")));
+        subQueries.add(new TermQuery(new Term("content", "duplicate")));
+
+        // Add a neural query with overlapping text
+        NeuralKNNQuery neuralQuery = mock(NeuralKNNQuery.class);
+        when(neuralQuery.getOriginalQueryText()).thenReturn("duplicate text");
+        subQueries.add(neuralQuery);
+
+        // Create the hybrid query
+        HybridQuery hybridQuery = new HybridQuery(subQueries, HybridQueryContext.builder().build());
+
+        // Test extraction
+        String result = registry.extractQueryText(hybridQuery, "content");
+        assertEquals("Should deduplicate terms", "duplicate duplicate text", result);
+    }
+
+    /**
+     * Tests the HybridQueryTextExtractor with nested queries
+     */
+    public void testHybridQueryExtractorWithNestedQueries() {
+        List<Query> subQueries = new ArrayList<>();
+
+        // Create a boolean query with multiple terms
+        BooleanQuery.Builder boolBuilder = new BooleanQuery.Builder();
+        boolBuilder.add(new TermQuery(new Term("content", "nested")), BooleanClause.Occur.MUST);
+        boolBuilder.add(new TermQuery(new Term("content", "terms")), BooleanClause.Occur.MUST);
+        subQueries.add(boolBuilder.build());
+
+        // Add a neural query
+        NeuralKNNQuery neuralQuery = mock(NeuralKNNQuery.class);
+        when(neuralQuery.getOriginalQueryText()).thenReturn("neural text");
+        subQueries.add(neuralQuery);
+
+        // Create the hybrid query
+        HybridQuery hybridQuery = new HybridQuery(subQueries, HybridQueryContext.builder().build());
+
+        // Test extraction
+        String result = registry.extractQueryText(hybridQuery, "content");
+        assertEquals("Should handle nested queries correctly", "nested terms neural text", result);
+    }
 }