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/NeuralHighlighterManagerTests.java src/test/java/org/opensearch/neuralsearch/highlight/NeuralHighlighterEngineTests.java

@@ -27,21 +27,21 @@
 /**
  * Tests for the NeuralHighlighterEngine class
  */
-public class NeuralHighlighterManagerTests extends OpenSearchTestCase {
+public class NeuralHighlighterEngineTests extends OpenSearchTestCase {
 
     private static final String TEST_FIELD = "test_field";
     private static final String MODEL_ID = "test_model_id";
     private static final String TEST_CONTENT = "This is a test content. For highlighting purposes. With multiple sentences.";
     private static final String TEST_QUERY = "test content";
 
-    private NeuralHighlighterEngine manager;
+    private NeuralHighlighterEngine highlighterEngine;
     private MLCommonsClientAccessor mlCommonsClientAccessor;
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
         mlCommonsClientAccessor = mock(MLCommonsClientAccessor.class);
-        manager = new NeuralHighlighterEngine(mlCommonsClientAccessor);
+        highlighterEngine = new NeuralHighlighterEngine(mlCommonsClientAccessor);
 
         // Setup default mock behavior
         setupDefaultMockBehavior();
@@ -75,22 +75,22 @@ public void testGetModelId() {
         Map<String, Object> options = new HashMap<>();
         options.put("model_id", MODEL_ID);
 
-        String modelId = manager.getModelId(options);
+        String modelId = highlighterEngine.getModelId(options);
         assertEquals("Should extract model ID correctly", MODEL_ID, modelId);
     }
 
     public void testGetModelIdMissing() {
         Map<String, Object> options = new HashMap<>();
 
-        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> manager.getModelId(options));
+        IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> highlighterEngine.getModelId(options));
         assertNotNull(exception);
         assertTrue(exception.getMessage().contains("Missing required option: model_id"));
     }
 
     public void testExtractOriginalQuery() {
         // Test with TermQuery
         TermQuery termQuery = new TermQuery(new Term(TEST_FIELD, "term"));
-        String queryText = manager.extractOriginalQuery(termQuery, TEST_FIELD);
+        String queryText = highlighterEngine.extractOriginalQuery(termQuery, TEST_FIELD);
         assertEquals("Should extract term text", "term", queryText);
 
         // Test with BooleanQuery
@@ -99,19 +99,19 @@ public void testExtractOriginalQuery() {
         builder.add(new TermQuery(new Term(TEST_FIELD, "term2")), BooleanClause.Occur.MUST);
         BooleanQuery booleanQuery = builder.build();
 
-        queryText = manager.extractOriginalQuery(booleanQuery, TEST_FIELD);
+        queryText = highlighterEngine.extractOriginalQuery(booleanQuery, TEST_FIELD);
         assertEquals("Should extract combined terms", "term1 term2", queryText);
 
         // Test with NeuralKNNQuery
         NeuralKNNQuery neuralQuery = mock(NeuralKNNQuery.class);
         when(neuralQuery.getOriginalQueryText()).thenReturn("neural query");
 
-        queryText = manager.extractOriginalQuery(neuralQuery, TEST_FIELD);
+        queryText = highlighterEngine.extractOriginalQuery(neuralQuery, TEST_FIELD);
         assertEquals("Should extract neural query text", "neural query", queryText);
     }
 
     public void testGetHighlightedSentences() {
-        String result = manager.getHighlightedSentences(MODEL_ID, TEST_QUERY, TEST_CONTENT);
+        String result = highlighterEngine.getHighlightedSentences(MODEL_ID, TEST_QUERY, TEST_CONTENT);
 
         assertNotNull("Should return highlighted text", result);
         assertTrue("Should contain highlighting tags", result.contains("<em>") && result.contains("</em>"));
@@ -139,7 +139,7 @@ public void testApplyHighlighting() {
         highlights.add(resultMap);
 
         String text = "This is a test string";
-        String result = manager.applyHighlighting(text, highlights);
+        String result = highlighterEngine.applyHighlighting(text, highlights);
 
         assertEquals("Should apply highlights correctly", "<em>This</em> is <em>a te</em>st string", result);
     }
@@ -165,7 +165,7 @@ public void testApplyHighlightingWithOverlaps() {
         highlights.add(resultMap);
 
         String text = "This is a test string";
-        String result = manager.applyHighlighting(text, highlights);
+        String result = highlighterEngine.applyHighlighting(text, highlights);
 
         // Should merge the overlapping highlights
         assertEquals("Should merge overlapping highlights", "<em>This is a </em>test string", result);
@@ -205,7 +205,7 @@ public void testApplyHighlightingWithInvalidPositions() {
         highlights.add(resultMap);
 
         String text = "This is a test string";
-        String result = manager.applyHighlighting(text, highlights);
+        String result = highlighterEngine.applyHighlighting(text, highlights);
 
         // Should only apply the valid highlight
         assertEquals("Should only apply valid highlights", "<em>This</em> is a test string", result);