Support relevance query functions pushdown implementation in Calcite

core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java

@@ -11,6 +11,8 @@
 import static org.opensearch.sql.ast.expression.SpanUnit.NONE;
 import static org.opensearch.sql.ast.expression.SpanUnit.UNKNOWN;
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
+import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.MULTI_FIELDS_RELEVANCE_FUNCTION_SET;
+import static org.opensearch.sql.calcite.utils.UserDefinedFunctionUtils.SINGLE_FIELD_RELEVANCE_FUNCTION_SET;
 
 import java.math.BigDecimal;
 import java.util.ArrayList;
@@ -59,6 +61,7 @@
 import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.SpanUnit;
+import org.opensearch.sql.ast.expression.UnresolvedArgument;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.expression.When;
 import org.opensearch.sql.ast.expression.WindowFunction;
@@ -500,29 +503,56 @@ private List<RexNode> castArgument(
   public RexNode visitFunction(Function node, CalcitePlanContext context) {
     List<UnresolvedExpression> args = node.getFuncArgs();
     List<RexNode> arguments = new ArrayList<>();
-    for (UnresolvedExpression arg : args) {
-      if (arg instanceof LambdaFunction) {
-        CalcitePlanContext lambdaContext =
-            prepareLambdaContext(
-                context, (LambdaFunction) arg, arguments, node.getFuncName(), null);
-        RexNode lambdaNode = analyze(arg, lambdaContext);
-        if (node.getFuncName().equalsIgnoreCase("reduce")) { // analyze again with calculate type
-          lambdaContext =
+
+    if (SINGLE_FIELD_RELEVANCE_FUNCTION_SET.contains(node.getFuncName())
+        || MULTI_FIELDS_RELEVANCE_FUNCTION_SET.contains(node.getFuncName())) {
+      List<RexNode> varArgRexNodeList = new ArrayList<>();
+      for (UnresolvedExpression arg : args) {
+        String unresolvedArgName = ((UnresolvedArgument) arg).getArgName();
+        // Relevance function first two arguments are always "field", "query" or "fields", "query"
+        if ("field".equals(unresolvedArgName)
+            || "query".equals(unresolvedArgName)
+            || "fields".equals(unresolvedArgName)) {
+          arguments.add(analyze(arg, context));
+        } else {
+          varArgRexNodeList.add(
+              context.rexBuilder.makeLiteral(
+                  unresolvedArgName,
+                  context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR),
+                  true));
+          varArgRexNodeList.add(analyze(arg, context));
+        }
+      }
+      if (!varArgRexNodeList.isEmpty()) {
+        arguments.add(
+            context.rexBuilder.makeCall(
+                SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR, varArgRexNodeList));
+      }
+    } else {
+      for (UnresolvedExpression arg : args) {
+        if (arg instanceof LambdaFunction) {
+          CalcitePlanContext lambdaContext =
               prepareLambdaContext(
-                  context,
-                  (LambdaFunction) arg,
-                  arguments,
-                  node.getFuncName(),
-                  lambdaNode.getType());
-          lambdaNode = analyze(arg, lambdaContext);
+                  context, (LambdaFunction) arg, arguments, node.getFuncName(), null);
+          RexNode lambdaNode = analyze(arg, lambdaContext);
+          if (node.getFuncName().equalsIgnoreCase("reduce")) { // analyze again with calculate type
+            lambdaContext =
+                prepareLambdaContext(
+                    context,
+                    (LambdaFunction) arg,
+                    arguments,
+                    node.getFuncName(),
+                    lambdaNode.getType());
+            lambdaNode = analyze(arg, lambdaContext);
+          }
+          arguments.add(lambdaNode);
+        } else {
+          arguments.add(analyze(arg, context));
         }
-        arguments.add(lambdaNode);
-      } else {
-        arguments.add(analyze(arg, context));
       }
-    }
 
-    arguments = castArgument(arguments, node.getFuncName(), context.rexBuilder);
+      arguments = castArgument(arguments, node.getFuncName(), context.rexBuilder);
+    }
 
     RexNode resolvedNode =
         PPLFuncImpTable.INSTANCE.resolve(
@@ -669,6 +699,22 @@ public RexNode visitWhen(When node, CalcitePlanContext context) {
 
   @Override
   public RexNode visitRelevanceFieldList(RelevanceFieldList node, CalcitePlanContext context) {
-    throw new CalciteUnsupportedException("Relevance fields expression is unsupported in Calcite");
+    List<RexNode> varArgRexNodeList = new ArrayList<>();
+    node.getFieldList()
+        .forEach(
+            (k, v) -> {
+              varArgRexNodeList.add(
+                  context.rexBuilder.makeLiteral(
+                      k,
+                      context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.VARCHAR),
+                      true));
+              varArgRexNodeList.add(
+                  context.rexBuilder.makeLiteral(
+                      v,
+                      context.rexBuilder.getTypeFactory().createSqlType(SqlTypeName.DOUBLE),
+                      true));
+            });
+    return context.rexBuilder.makeCall(
+        SqlStdOperatorTable.MAP_VALUE_CONSTRUCTOR, varArgRexNodeList);
   }
 }

core/src/main/java/org/opensearch/sql/calcite/utils/UserDefinedFunctionUtils.java

@@ -10,11 +10,13 @@
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.*;
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT.*;
 
+import com.google.common.collect.ImmutableSet;
 import java.time.Instant;
 import java.time.ZoneId;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.TimeZone;
 import javax.annotation.Nullable;
 import org.apache.calcite.DataContext;
@@ -70,6 +72,10 @@ public class UserDefinedFunctionUtils {
       TYPE_FACTORY.createMapType(
           TYPE_FACTORY.createSqlType(SqlTypeName.VARCHAR),
           createArrayType(TYPE_FACTORY, TYPE_FACTORY.createSqlType(SqlTypeName.VARCHAR), false));
+  public static Set<String> SINGLE_FIELD_RELEVANCE_FUNCTION_SET =
+      ImmutableSet.of("match", "match_phrase", "match_bool_prefix", "match_phrase_prefix");
+  public static Set<String> MULTI_FIELDS_RELEVANCE_FUNCTION_SET =
+      ImmutableSet.of("simple_query_string", "query_string", "multi_match");
 
   public static RelBuilder.AggCall TransferUserDefinedAggFunction(
       Class<? extends UserDefinedAggFunction> UDAF,

core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java

@@ -47,6 +47,7 @@
 import org.opensearch.sql.expression.function.jsonUDF.JsonSetFunctionImpl;
 import org.opensearch.sql.expression.function.udf.CryptographicFunction;
 import org.opensearch.sql.expression.function.udf.GrokFunction;
+import org.opensearch.sql.expression.function.udf.RelevanceQueryFunction;
 import org.opensearch.sql.expression.function.udf.SpanFunction;
 import org.opensearch.sql.expression.function.udf.condition.EarliestFunction;
 import org.opensearch.sql.expression.function.udf.condition.LatestFunction;
@@ -327,6 +328,17 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
   public static final SqlOperator TRANSFORM = new TransformFunctionImpl().toUDF("transform");
   public static final SqlOperator REDUCE = new ReduceFunctionImpl().toUDF("reduce");
 
+  public static final SqlOperator MATCH = new RelevanceQueryFunction().toUDF("match");
+  public static final SqlOperator MATCH_PHRASE = new RelevanceQueryFunction().toUDF("match_phrase");
+  public static final SqlOperator MATCH_BOOL_PREFIX =
+      new RelevanceQueryFunction().toUDF("match_bool_prefix");
+  public static final SqlOperator MATCH_PHRASE_PREFIX =
+      new RelevanceQueryFunction().toUDF("match_phrase_prefix");
+  public static final SqlOperator SIMPLE_QUERY_STRING =
+      new RelevanceQueryFunction().toUDF("simple_query_string");
+  public static final SqlOperator QUERY_STRING = new RelevanceQueryFunction().toUDF("query_string");
+  public static final SqlOperator MULTI_MATCH = new RelevanceQueryFunction().toUDF("multi_match");
+
   /**
    * Invoking an implementor registered in {@link RexImpTable}, need to use reflection since they're
    * all private Use method directly in {@link BuiltInMethod} if possible, most operators'

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -119,6 +119,10 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.LTRIM;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAKEDATE;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAKETIME;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_BOOL_PREFIX;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_PHRASE;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_PHRASE_PREFIX;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAX;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MD5;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MICROSECOND;
@@ -133,6 +137,7 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MONTHNAME;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MONTH_OF_YEAR;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MULTIPLY;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MULTI_MATCH;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.NOT;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.NOTEQUAL;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.NOW;
@@ -146,6 +151,7 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.POW;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.POWER;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.QUARTER;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.QUERY_STRING;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.RADIANS;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.RAND;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.REDUCE;
@@ -161,6 +167,7 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.SHA1;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.SHA2;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.SIGN;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.SIMPLE_QUERY_STRING;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.SIN;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.SPAN;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.SQRT;
@@ -693,6 +700,13 @@ void populate() {
       registerOperator(SHA2, PPLBuiltinOperators.SHA2);
       registerOperator(CIDRMATCH, PPLBuiltinOperators.CIDRMATCH);
       registerOperator(INTERNAL_GROK, PPLBuiltinOperators.GROK);
+      registerOperator(MATCH, PPLBuiltinOperators.MATCH);
+      registerOperator(MATCH_PHRASE, PPLBuiltinOperators.MATCH_PHRASE);
+      registerOperator(MATCH_BOOL_PREFIX, PPLBuiltinOperators.MATCH_BOOL_PREFIX);
+      registerOperator(MATCH_PHRASE_PREFIX, PPLBuiltinOperators.MATCH_PHRASE_PREFIX);
+      registerOperator(SIMPLE_QUERY_STRING, PPLBuiltinOperators.SIMPLE_QUERY_STRING);
+      registerOperator(QUERY_STRING, PPLBuiltinOperators.QUERY_STRING);
+      registerOperator(MULTI_MATCH, PPLBuiltinOperators.MULTI_MATCH);
 
       // Register PPL Datetime UDF operator
       registerOperator(TIMESTAMP, PPLBuiltinOperators.TIMESTAMP);

core/src/main/java/org/opensearch/sql/expression/function/udf/RelevanceQueryFunction.java

@@ -0,0 +1,55 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.udf;
+
+import java.util.List;
+import org.apache.calcite.adapter.enumerable.NotNullImplementor;
+import org.apache.calcite.adapter.enumerable.NullPolicy;
+import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
+import org.apache.calcite.linq4j.tree.Expression;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.sql.type.CompositeOperandTypeChecker;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.opensearch.sql.expression.function.ImplementorUDF;
+import org.opensearch.sql.expression.function.UDFOperandMetadata;
+
+public class RelevanceQueryFunction extends ImplementorUDF {
+
+  public RelevanceQueryFunction() {
+    super(new RelevanceQueryImplementor(), NullPolicy.ANY);
+  }
+
+  @Override
+  public SqlReturnTypeInference getReturnTypeInference() {
+    return ReturnTypes.BOOLEAN;
+  }
+
+  @Override
+  public UDFOperandMetadata getOperandMetadata() {
+    return UDFOperandMetadata.wrap(
+        (CompositeOperandTypeChecker)
+            OperandTypes.STRING_STRING
+                .or(
+                    OperandTypes.family(
+                        SqlTypeFamily.STRING, SqlTypeFamily.STRING, SqlTypeFamily.MAP))
+                .or(OperandTypes.family(SqlTypeFamily.MAP, SqlTypeFamily.STRING))
+                .or(
+                    OperandTypes.family(
+                        SqlTypeFamily.MAP, SqlTypeFamily.STRING, SqlTypeFamily.MAP)));
+  }
+
+  public static class RelevanceQueryImplementor implements NotNullImplementor {
+    @Override
+    public Expression implement(
+        RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
+      throw new UnsupportedOperationException(
+          "Relevance search query functions are only supported when they are pushed down");
+    }
+  }
+}

integ-test/src/test/java/org/opensearch/sql/calcite/pushdown/CalciteNoPushdownIT.java

@@ -47,9 +47,6 @@
   CalciteJsonFunctionsIT.class,
   CalciteLegacyAPICompatibilityIT.class,
   CalciteLikeQueryIT.class,
-  CalciteMatchBoolPrefixIT.class,
-  CalciteMatchIT.class,
-  CalciteMatchPhraseIT.class,
   CalciteMathematicalFunctionIT.class,
   CalciteNewAddedCommandsIT.class,
   CalciteNowLikeFunctionIT.class,
@@ -87,15 +84,12 @@
   CalcitePPLTrendlineIT.class,
   CalcitePrometheusDataSourceCommandsIT.class,
   CalciteQueryAnalysisIT.class,
-  CalciteQueryStringIT.class,
   CalciteRareCommandIT.class,
-  CalciteRelevanceFunctionIT.class,
   CalciteRenameCommandIT.class,
   CalciteResourceMonitorIT.class,
   CalciteSearchCommandIT.class,
   CalciteSettingsIT.class,
   CalciteShowDataSourcesCommandIT.class,
-  CalciteSimpleQueryStringIT.class,
   CalciteSortCommandIT.class,
   CalciteStatsCommandIT.class,
   CalciteSystemFunctionIT.class,

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteMatchBoolPrefixIT.java

@@ -12,7 +12,6 @@ public class CalciteMatchBoolPrefixIT extends MatchBoolPrefixIT {
   public void init() throws Exception {
     super.init();
     enableCalcite();
-    // TODO: "https://github.com/opensearch-project/sql/issues/3462"
-    // disallowCalciteFallback();
+    disallowCalciteFallback();
   }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteMatchIT.java

@@ -12,7 +12,6 @@ public class CalciteMatchIT extends MatchIT {
   public void init() throws Exception {
     super.init();
     enableCalcite();
-    // TODO: "https://github.com/opensearch-project/sql/issues/3462"
-    // disallowCalciteFallback();
+    disallowCalciteFallback();
   }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteMatchPhraseIT.java

@@ -12,7 +12,6 @@ public class CalciteMatchPhraseIT extends MatchPhraseIT {
   public void init() throws Exception {
     super.init();
     enableCalcite();
-    // TODO: "https://github.com/opensearch-project/sql/issues/3462"
-    // disallowCalciteFallback();
+    disallowCalciteFallback();
   }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteMatchPhrasePrefixIT.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.remote;
+
+import org.opensearch.sql.ppl.MatchPhrasePrefixIT;
+
+public class CalciteMatchPhrasePrefixIT extends MatchPhrasePrefixIT {
+  @Override
+  public void init() throws Exception {
+    super.init();
+    enableCalcite();
+    disallowCalciteFallback();
+  }
+}

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteMultiMatchIT.java

@@ -0,0 +1,17 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.remote;
+
+import org.opensearch.sql.ppl.MultiMatchIT;
+
+public class CalciteMultiMatchIT extends MultiMatchIT {
+  @Override
+  public void init() throws Exception {
+    super.init();
+    enableCalcite();
+    disallowCalciteFallback();
+  }
+}

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteQueryStringIT.java

@@ -12,7 +12,6 @@ public class CalciteQueryStringIT extends QueryStringIT {
   public void init() throws Exception {
     super.init();
     enableCalcite();
-    // TODO: "https://github.com/opensearch-project/sql/issues/3462"
-    // disallowCalciteFallback();
+    disallowCalciteFallback();
   }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteRelevanceFunctionIT.java

@@ -12,7 +12,6 @@ public class CalciteRelevanceFunctionIT extends RelevanceFunctionIT {
   public void init() throws Exception {
     super.init();
     enableCalcite();
-    // TODO: "https://github.com/opensearch-project/sql/issues/3462"
-    // disallowCalciteFallback();
+    disallowCalciteFallback();
   }
 }

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteSimpleQueryStringIT.java

@@ -12,8 +12,6 @@ public class CalciteSimpleQueryStringIT extends SimpleQueryStringIT {
   public void init() throws Exception {
     super.init();
     enableCalcite();
-    // Search Functions are not supported
-    // TODO: "https://github.com/opensearch-project/sql/issues/3462"
-    // disallowCalciteFallback();
+    disallowCalciteFallback();
   }
 }