Merge pull request #105 from Bit-Quill/dev-relevancy-improvements-rb1

Refactor RelevanceQuery -- add SingleFieldQuery and MultiFieldQuery
base classes.

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MatchBoolPrefixQuery.java

@@ -14,7 +14,7 @@
  * Initializes MatchBoolPrefixQueryBuilder from a FunctionExpression.
  */
 public class MatchBoolPrefixQuery
-    extends RelevanceQuery<MatchBoolPrefixQueryBuilder> {
+    extends SingleFieldQuery<MatchBoolPrefixQueryBuilder> {
   /**
    * Constructor for MatchBoolPrefixQuery to configure RelevanceQuery
    * with support of optional parameters.
@@ -41,7 +41,12 @@ public MatchBoolPrefixQuery() {
    * @return  Object of executed query
    */
   @Override
-  protected MatchBoolPrefixQueryBuilder createQueryBuilder(String field, String query) {
+  protected MatchBoolPrefixQueryBuilder createBuilder(String field, String query) {
     return QueryBuilders.matchBoolPrefixQuery(field, query);
   }
+
+  @Override
+  protected String getQueryName() {
+    return MatchBoolPrefixQueryBuilder.NAME;
+  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MatchPhrasePrefixQuery.java

@@ -12,7 +12,7 @@
 /**
     * Lucene query that builds a match_phrase_prefix query.
     */
-public class MatchPhrasePrefixQuery  extends RelevanceQuery<MatchPhrasePrefixQueryBuilder> {
+public class MatchPhrasePrefixQuery  extends SingleFieldQuery<MatchPhrasePrefixQueryBuilder> {
   /**
    *  Default constructor for MatchPhrasePrefixQuery configures how RelevanceQuery.build() handles
    * named arguments.
@@ -29,7 +29,12 @@ public MatchPhrasePrefixQuery() {
   }
 
   @Override
-  protected MatchPhrasePrefixQueryBuilder createQueryBuilder(String field, String query) {
+  protected MatchPhrasePrefixQueryBuilder createBuilder(String field, String query) {
     return QueryBuilders.matchPhrasePrefixQuery(field, query);
   }
+
+  @Override
+  protected String getQueryName() {
+    return MatchPhrasePrefixQueryBuilder.NAME;
+  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MatchPhraseQuery.java

@@ -23,7 +23,7 @@
 /**
  * Lucene query that builds a match_phrase query.
  */
-public class MatchPhraseQuery extends RelevanceQuery<MatchPhraseQueryBuilder> {
+public class MatchPhraseQuery extends SingleFieldQuery<MatchPhraseQueryBuilder> {
   /**
    *  Default constructor for MatchPhraseQuery configures how RelevanceQuery.build() handles
    * named arguments.
@@ -39,7 +39,12 @@ public MatchPhraseQuery() {
   }
 
   @Override
-  protected MatchPhraseQueryBuilder createQueryBuilder(String field, String query) {
+  protected MatchPhraseQueryBuilder createBuilder(String field, String query) {
     return QueryBuilders.matchPhraseQuery(field, query);
   }
+
+  @Override
+  protected String getQueryName() {
+    return MatchPhraseQueryBuilder.NAME;
+  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MatchQuery.java

@@ -6,15 +6,14 @@
 package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;
 
 import com.google.common.collect.ImmutableMap;
-import java.util.Map;
 import org.opensearch.index.query.MatchQueryBuilder;
 import org.opensearch.index.query.Operator;
 import org.opensearch.index.query.QueryBuilders;
 
 /**
  * Initializes MatchQueryBuilder from a FunctionExpression.
  */
-public class MatchQuery extends RelevanceQuery<MatchQueryBuilder> {
+public class MatchQuery extends SingleFieldQuery<MatchQueryBuilder> {
   /**
    *  Default constructor for MatchQuery configures how RelevanceQuery.build() handles
    * named arguments.
@@ -40,7 +39,12 @@ public MatchQuery() {
   }
 
   @Override
-  protected MatchQueryBuilder createQueryBuilder(String field, String query) {
+  protected MatchQueryBuilder createBuilder(String field, String query) {
     return QueryBuilders.matchQuery(field, query);
   }
+
+  @Override
+  protected String getQueryName() {
+    return MatchQueryBuilder.NAME;
+  }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MultiFieldQuery.java

@@ -0,0 +1,37 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;
+
+import com.google.common.collect.ImmutableMap;
+import java.util.Map;
+import org.opensearch.index.query.QueryBuilder;
+import org.opensearch.sql.expression.NamedArgumentExpression;
+
+/**
+ *  Base class to represent relevance queries that search multiple fields.
+ * @param <T> The builder class for the OpenSearch query.
+ */
+abstract class MultiFieldQuery<T extends QueryBuilder> extends RelevanceQuery<T> {
+
+  public MultiFieldQuery(Map<String, QueryBuilderStep<T>> queryBuildActions) {
+    super(queryBuildActions);
+  }
+
+  @Override
+  public T createQueryBuilder(NamedArgumentExpression fields, NamedArgumentExpression queryExpr) {
+    var fieldsAndWeights = fields
+        .getValue()
+        .valueOf(null)
+        .tupleValue()
+        .entrySet()
+        .stream()
+        .collect(ImmutableMap.toImmutableMap(e -> e.getKey(), e -> e.getValue().floatValue()));
+    var query = queryExpr.getValue().valueOf(null).stringValue();
+    return createBuilder(fieldsAndWeights, query);
+  }
+
+  protected abstract  T createBuilder(ImmutableMap<String, Float> fields, String query);
+}

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MultiMatchQuery.java

@@ -6,19 +6,11 @@
 package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;
 
 import com.google.common.collect.ImmutableMap;
-import java.util.Iterator;
-import java.util.Map;
-import java.util.Objects;
 import org.opensearch.index.query.MultiMatchQueryBuilder;
 import org.opensearch.index.query.Operator;
-import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
-import org.opensearch.sql.exception.SemanticCheckException;
-import org.opensearch.sql.expression.Expression;
-import org.opensearch.sql.expression.FunctionExpression;
-import org.opensearch.sql.expression.NamedArgumentExpression;
 
-public class MultiMatchQuery extends RelevanceQuery<MultiMatchQueryBuilder> {
+public class MultiMatchQuery extends MultiFieldQuery<MultiMatchQueryBuilder> {
   /**
    *  Default constructor for MultiMatch configures how RelevanceQuery.build() handles
    * named arguments.
@@ -47,42 +39,12 @@ public MultiMatchQuery() {
   }
 
   @Override
-  public QueryBuilder build(FunctionExpression func) {
-    if (func.getArguments().size() < 2) {
-      throw new SemanticCheckException("'multi_match' must have at least two arguments");
-    }
-    Iterator<Expression> iterator = func.getArguments().iterator();
-    var fields = (NamedArgumentExpression) iterator.next();
-    var query = (NamedArgumentExpression) iterator.next();
-    // Fields is a map already, but we need to convert types.
-    var fieldsAndWeights = fields
-        .getValue()
-        .valueOf(null)
-        .tupleValue()
-        .entrySet()
-        .stream()
-        .collect(ImmutableMap.toImmutableMap(Map.Entry::getKey, e -> e.getValue().floatValue()));
-
-    MultiMatchQueryBuilder queryBuilder = createQueryBuilder(null,
-            query.getValue().valueOf(null).stringValue())
-        .fields(fieldsAndWeights);
-    while (iterator.hasNext()) {
-      NamedArgumentExpression arg = (NamedArgumentExpression) iterator.next();
-      if (!queryBuildActions.containsKey(arg.getArgName())) {
-        throw new SemanticCheckException(
-            String.format("Parameter %s is invalid for %s function.",
-                arg.getArgName(), queryBuilder.getWriteableName()));
-      }
-      (Objects.requireNonNull(
-          queryBuildActions
-              .get(arg.getArgName())))
-          .apply(queryBuilder, arg.getValue().valueOf(null));
-    }
-    return queryBuilder;
+  protected MultiMatchQueryBuilder createBuilder(ImmutableMap<String, Float> fields, String query) {
+    return QueryBuilders.multiMatchQuery(query).fields(fields);
   }
 
   @Override
-  protected MultiMatchQueryBuilder createQueryBuilder(String field, String query) {
-    return QueryBuilders.multiMatchQuery(query);
+  protected String getQueryName() {
+    return MultiMatchQueryBuilder.NAME;
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/QueryStringQuery.java

@@ -23,7 +23,7 @@
 /**
  * Class for Lucene query that builds the query_string query.
  */
-public class QueryStringQuery extends RelevanceQuery<QueryStringQueryBuilder> {
+public class QueryStringQuery extends MultiFieldQuery<QueryStringQueryBuilder> {
   /**
    *  Default constructor for QueryString configures how RelevanceQuery.build() handles
    * named arguments.
@@ -66,54 +66,22 @@ public QueryStringQuery() {
         .build());
   }
 
-  /**
-   * Override base build function for multi-field query support.
-   * @param func  function : 'query_string' function
-   * @return : QueryBuilder for query_string query
-   */
-  @Override
-  public QueryBuilder build(FunctionExpression func) {
-    Iterator<Expression> iterator = func.getArguments().iterator();
-    if (func.getArguments().size() < 2) {
-      throw new SemanticCheckException("'query_string' must have at least two arguments");
-    }
-    NamedArgumentExpression fields = (NamedArgumentExpression) iterator.next();
-    NamedArgumentExpression query = (NamedArgumentExpression) iterator.next();
-    // Fields is a map already, but we need to convert types.
-    var fieldsAndWeights = fields
-        .getValue()
-        .valueOf(null)
-        .tupleValue()
-        .entrySet()
-        .stream()
-        .collect(ImmutableMap.toImmutableMap(e -> e.getKey(), e -> e.getValue().floatValue()));
-
-    QueryStringQueryBuilder queryBuilder = createQueryBuilder(null,
-        query.getValue().valueOf(null).stringValue())
-        .fields(fieldsAndWeights);
-    while (iterator.hasNext()) {
-      NamedArgumentExpression arg = (NamedArgumentExpression) iterator.next();
-      if (!queryBuildActions.containsKey(arg.getArgName())) {
-        throw new SemanticCheckException(
-            String.format("Parameter %s is invalid for %s function.",
-                arg.getArgName(), queryBuilder.getWriteableName()));
-      }
-      (Objects.requireNonNull(
-          queryBuildActions
-              .get(arg.getArgName())))
-          .apply(queryBuilder, arg.getValue().valueOf(null));
-    }
-    return queryBuilder;
-  }
 
   /**
    * Builds QueryBuilder with query value and other default parameter values set.
-   * @param field : Field value in query_string query
+   *
+   * @param fields : A map of field names and their boost values
    * @param query : Query value for query_string query
    * @return : Builder for query_string query
    */
   @Override
-  protected QueryStringQueryBuilder createQueryBuilder(String field, String query) {
-    return QueryBuilders.queryStringQuery(query);
+  protected QueryStringQueryBuilder createBuilder(ImmutableMap<String, Float> fields,
+                                                  String query) {
+    return QueryBuilders.queryStringQuery(query).fields(fields);
+  }
+
+  @Override
+  protected String getQueryName() {
+    return QueryStringQueryBuilder.NAME;
   }
 }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/RelevanceQuery.java

@@ -12,6 +12,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.BiFunction;
+import lombok.RequiredArgsConstructor;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.data.model.ExprValue;
@@ -24,26 +25,20 @@
 /**
  * Base class for query abstraction that builds a relevance query from function expression.
  */
+@RequiredArgsConstructor
 public abstract class RelevanceQuery<T extends QueryBuilder> extends LuceneQuery {
-  protected Map<String, QueryBuilderStep<T>> queryBuildActions;
-
-  protected RelevanceQuery(Map<String, QueryBuilderStep<T>> actionMap) {
-    queryBuildActions = actionMap;
-  }
+  private final Map<String, QueryBuilderStep<T>> queryBuildActions;
 
   @Override
   public QueryBuilder build(FunctionExpression func) {
     List<Expression> arguments = func.getArguments();
     if (arguments.size() < 2) {
-      String queryName = createQueryBuilder("dummy_field", "").getWriteableName();
       throw new SyntaxCheckException(
-          String.format("%s requires at least two parameters", queryName));
+          String.format("%s requires at least two parameters", getQueryName()));
     }
     NamedArgumentExpression field = (NamedArgumentExpression) arguments.get(0);
     NamedArgumentExpression query = (NamedArgumentExpression) arguments.get(1);
-    T queryBuilder = createQueryBuilder(
-        field.getValue().valueOf(null).stringValue(),
-        query.getValue().valueOf(null).stringValue());
+    T queryBuilder = createQueryBuilder(field, query);
 
     Iterator<Expression> iterator = arguments.listIterator(2);
     Set<String> visitedParms = new HashSet();
@@ -69,15 +64,18 @@ public QueryBuilder build(FunctionExpression func) {
     return queryBuilder;
   }
 
-  protected abstract T createQueryBuilder(String field, String query);
+  protected abstract T createQueryBuilder(NamedArgumentExpression field,
+                                          NamedArgumentExpression query);
+
+  protected abstract String getQueryName();
 
   /**
    * Convenience interface for a function that updates a QueryBuilder
    * based on ExprValue.
+   *
    * @param <T> Concrete query builder
    */
-  public interface QueryBuilderStep<T extends QueryBuilder> extends
+  protected interface QueryBuilderStep<T extends QueryBuilder> extends
       BiFunction<T, ExprValue, T> {
-
   }
 }