Add main classes for Query and basic unit tests

src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.neuralsearch.plugin;
 
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -23,6 +24,7 @@
 import org.opensearch.neuralsearch.ml.MLCommonsClientAccessor;
 import org.opensearch.neuralsearch.processor.TextEmbeddingProcessor;
 import org.opensearch.neuralsearch.processor.factory.TextEmbeddingProcessorFactory;
+import org.opensearch.neuralsearch.query.HybridQueryBuilder;
 import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
 import org.opensearch.plugins.ActionPlugin;
 import org.opensearch.plugins.ExtensiblePlugin;
@@ -60,8 +62,9 @@ public Collection<Object> createComponents(
     }
 
     public List<QuerySpec<?>> getQueries() {
-        return Collections.singletonList(
-            new QuerySpec<>(NeuralQueryBuilder.NAME, NeuralQueryBuilder::new, NeuralQueryBuilder::fromXContent)
+        return Arrays.asList(
+            new QuerySpec<>(NeuralQueryBuilder.NAME, NeuralQueryBuilder::new, NeuralQueryBuilder::fromXContent),
+            new QuerySpec<>(HybridQueryBuilder.NAME, HybridQueryBuilder::new, HybridQueryBuilder::fromXContent)
         );
     }
 

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

@@ -0,0 +1,160 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.neuralsearch.query;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Objects;
+
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchNoDocsQuery;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.QueryVisitor;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Weight;
+
+/**
+ * Implementation fo Query interface for type "hybrid". It allows execution of multiple sub-queries and collect individual
+ * scores for each sub-query.
+ */
+public final class HybridQuery extends Query implements Iterable<Query> {
+
+    private final List<Query> subQueries;
+
+    public HybridQuery(Collection<Query> subQueries) {
+        Objects.requireNonNull(subQueries, "Collection of queries must not be null");
+        if (subQueries.isEmpty()) {
+            throw new IllegalArgumentException("Collection of queries must not be empty");
+        }
+        this.subQueries = new ArrayList<>(subQueries);
+    }
+
+    /**
+     * Returns an iterator over sub-queries that are parts of this hybrid query
+     * @return iterator
+     */
+    @Override
+    public Iterator<Query> iterator() {
+        return getSubQueries().iterator();
+    }
+
+    /**
+     * Prints a query to a string, with field assumed to be the default field and omitted.
+     * @param field default field
+     * @return string representation of hybrid query
+     */
+    @Override
+    public String toString(String field) {
+        StringBuilder buffer = new StringBuilder();
+        buffer.append("(");
+        Iterator<Query> it = subQueries.iterator();
+        for (int i = 0; it.hasNext(); i++) {
+            Query subquery = it.next();
+            if (subquery instanceof BooleanQuery) { // wrap sub-boolean in parents
+                buffer.append("(");
+                buffer.append(subquery.toString(field));
+                buffer.append(")");
+            } else buffer.append(subquery.toString(field));
+            if (i != subQueries.size() - 1) buffer.append(" | ");
+        }
+        buffer.append(")");
+        return buffer.toString();
+    }
+
+    /**
+     * Re-writes queries into primitive queries. Callers are expected to call rewrite multiple times if necessary,
+     * until the rewritten query is the same as the original query.
+     * @param reader
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public Query rewrite(IndexReader reader) throws IOException {
+        if (subQueries.isEmpty()) {
+            return new MatchNoDocsQuery("empty HybridQuery");
+        }
+
+        boolean actuallyRewritten = false;
+        List<Query> rewrittenSubQueries = new ArrayList<>();
+        for (Query subQuery : subQueries) {
+            Query rewrittenSub = subQuery.rewrite(reader);
+            /* we keep rewrite sub-query unless it's not equal to itself, it may take multiple levels of recursive calls
+               queries need to be rewritten from high-level clauses into lower-level clauses because low-level clauses
+               perform better. For hybrid query we need to track progress of re-write for all sub-queries */
+            actuallyRewritten |= rewrittenSub != subQuery;
+            rewrittenSubQueries.add(rewrittenSub);
+        }
+
+        if (actuallyRewritten) {
+            return new HybridQuery(rewrittenSubQueries);
+        }
+
+        return super.rewrite(reader);
+    }
+
+    /**
+     * Recurse through the query tree, visiting all child queries and execute provided visitor. Part of multiple
+     * standard workflows, e.g. IndexSearcher.rewrite
+     * @param queryVisitor a QueryVisitor to be called by each query in the tree
+     */
+    @Override
+    public void visit(QueryVisitor queryVisitor) {
+        QueryVisitor v = queryVisitor.getSubVisitor(BooleanClause.Occur.SHOULD, this);
+        for (Query q : subQueries) {
+            q.visit(v);
+        }
+    }
+
+    /**
+     * Override and implement query instance equivalence properly in a subclass. This is required so that QueryCache works properly.
+     * @param other query object that when compare with this query object
+     * @return
+     */
+    @Override
+    public boolean equals(Object other) {
+        return sameClassAs(other) && equalsTo(getClass().cast(other));
+    }
+
+    private boolean equalsTo(HybridQuery other) {
+        return Objects.equals(subQueries, other.subQueries);
+    }
+
+    /**
+     * Override and implement query hash code properly in a subclass. This is required so that QueryCache works properly.
+     * @return hash code of this object
+     */
+    @Override
+    public int hashCode() {
+        int h = classHash();
+        h = 31 * h + Objects.hashCode(subQueries);
+        return h;
+    }
+
+    public Collection<Query> getSubQueries() {
+        return Collections.unmodifiableCollection(subQueries);
+    }
+
+    /**
+     * Create the Weight used to score this query
+     *
+     * @param searcher
+     * @param scoreMode How the produced scorers will be consumed.
+     * @param boost The boost that is propagated by the parent queries.
+     * @return
+     * @throws IOException
+     */
+    @Override
+    public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+        return new HybridQueryWeight(this, searcher, scoreMode, boost);
+    }
+}