Address code reivews

Signed-off-by: Andy Kwok <[email protected]>
Bit-Quill · Jan 28, 2025 · 3d79741 · 3d79741
1 parent 66e17ce
commit 3d79741
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 44 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java b/core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java
@@ -130,36 +130,72 @@ private static FunctionResolver score(BuiltinFunctionName score) {
   public static class OpenSearchFunction extends FunctionExpression {
     private final FunctionName functionName;
     private final List<Expression> arguments;
-    private final ExprType returnType;
 
     @Getter @Setter private boolean isScoreTracked;
 
-    public OpenSearchFunction(
-        FunctionName functionName, List<Expression> arguments, ExprType returnType) {
-      super(functionName, arguments);
-      this.functionName = functionName;
-      this.arguments = arguments;
-      this.returnType = returnType;
-      this.isScoreTracked = false;
-    }
-
     /**
      * Required argument constructor.
      *
      * @param functionName name of the function
      * @param arguments a list of expressions
      */
     public OpenSearchFunction(FunctionName functionName, List<Expression> arguments) {
-      this(functionName, arguments, BOOLEAN);
+      super(functionName, arguments);
+      this.functionName = functionName;
+      this.arguments = arguments;
+      this.isScoreTracked = false;
     }
 
     @Override
     public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
       throw new UnsupportedOperationException(
-          String.format(
-              "OpenSearch defined function [%s] is only supported in WHERE clause, HAVING clause"
-                  + " and Eval operation.",
-              functionName));
+              String.format(
+                      "OpenSearch defined function [%s] is only supported in WHERE and HAVING clause.",
+                      functionName));
+    }
+
+    @Override
+    public ExprType type() {
+      return BOOLEAN;
+    }
+
+    @Override
+    public String toString() {
+      List<String> args =
+              arguments.stream()
+                      .map(
+                              arg ->
+                                      String.format(
+                                              "%s=%s",
+                                              ((NamedArgumentExpression) arg).getArgName(),
+                                              ((NamedArgumentExpression) arg).getValue().toString()))
+                      .collect(Collectors.toList());
+      return String.format("%s(%s)", functionName, String.join(", ", args));
+    }
+  }
+
+  /**
+   * Static class to identify functional Expression which specifically designed for OpenSearch storage runtime
+   */
+  public static class OpenSearchExecutableFunction extends FunctionExpression {
+    private final FunctionName functionName;
+    private final List<Expression> arguments;
+    private final ExprType returnType;
+
+    public OpenSearchExecutableFunction(
+            FunctionName functionName, List<Expression> arguments, ExprType returnType) {
+      super(functionName, arguments);
+      this.functionName = functionName;
+      this.arguments = arguments;
+      this.returnType = returnType;
+    }
+
+    @Override
+    public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+      throw new UnsupportedOperationException(
+              String.format(
+                      "OpenSearch defined function [%s] is only supported in Eval operation.",
+                      functionName));
     }
 
     @Override
@@ -170,15 +206,33 @@ public ExprType type() {
     @Override
     public String toString() {
       List<String> args =
-          arguments.stream()
-              .map(
-                  arg ->
-                      String.format(
-                          "%s=%s",
-                          ((NamedArgumentExpression) arg).getArgName(),
-                          ((NamedArgumentExpression) arg).getValue().toString()))
-              .collect(Collectors.toList());
+              arguments.stream()
+                      .map(
+                              arg ->
+                                      String.format(
+                                              "%s=%s",
+                                              ((NamedArgumentExpression) arg).getArgName(),
+                                              ((NamedArgumentExpression) arg).getValue().toString()))
+                      .collect(Collectors.toList());
       return String.format("%s(%s)", functionName, String.join(", ", args));
     }
+
+    /**
+     * Util method to generate probe implementation with given list of argument types, with marker
+     * class `OpenSearchFunction` to annotate this is an OpenSearch specific expression.
+     *
+     * @param returnType return type.
+     * @return Binary Function Implementation.
+     */
+    public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
+    openSearchImpl(ExprType returnType, List<ExprType> args) {
+      return functionName -> {
+        FunctionSignature functionSignature = new FunctionSignature(functionName, args);
+        FunctionBuilder functionBuilder =
+                (functionProperties, arguments) ->
+                        new OpenSearchFunctions.OpenSearchExecutableFunction(functionName, arguments, returnType);
+        return Pair.of(functionSignature, functionBuilder);
+      };
+    }
   }
 }
diff --git a/core/src/main/java/org/opensearch/sql/expression/ip/IPFunctions.java b/core/src/main/java/org/opensearch/sql/expression/ip/IPFunctions.java
@@ -11,6 +11,7 @@
 import static org.opensearch.sql.expression.function.FunctionDSL.define;
 import static org.opensearch.sql.expression.function.FunctionDSL.impl;
 import static org.opensearch.sql.expression.function.FunctionDSL.nullMissingHandling;
+import static org.opensearch.sql.expression.function.OpenSearchFunctions.OpenSearchExecutableFunction.openSearchImpl;
 
 import inet.ipaddr.IPAddress;
 import java.util.Arrays;
@@ -79,22 +80,4 @@ private DefaultFunctionResolver geoIp() {
         openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING)),
         openSearchImpl(BOOLEAN, Arrays.asList(STRING, STRING, STRING)));
   }
-
-  /**
-   * Util method to generate probe implementation with given list of argument types, with marker
-   * class `OpenSearchFunction` to annotate this is an OpenSearch specific expression.
-   *
-   * @param returnType return type.
-   * @return Binary Function Implementation.
-   */
-  public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
-      openSearchImpl(ExprType returnType, List<ExprType> args) {
-    return functionName -> {
-      FunctionSignature functionSignature = new FunctionSignature(functionName, args);
-      FunctionBuilder functionBuilder =
-          (functionProperties, arguments) ->
-              new OpenSearchFunctions.OpenSearchFunction(functionName, arguments, returnType);
-      return Pair.of(functionSignature, functionBuilder);
-    };
-  }
 }
diff --git a/.../src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchEvalOperator.java b/.../src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchEvalOperator.java
@@ -50,7 +50,7 @@ protected Map<String, ExprValue> eval(Environment<Expression, ExprValue> env) {
     Map<String, ExprValue> evalResultMap = new LinkedHashMap<>();
     for (Pair<ReferenceExpression, Expression> pair : this.getExpressionList()) {
       ExprValue value;
-      if (pair.getValue() instanceof OpenSearchFunctions.OpenSearchFunction openSearchExpr) {
+      if (pair.getValue() instanceof OpenSearchFunctions.OpenSearchExecutableFunction openSearchExpr) {
         value = OpenSearchEvalProcessor.process(openSearchExpr, env, nodeClient);
       } else {
         value = pair.getValue().valueOf(env);

diff --git a/...src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchEvalProcessor.java b/...src/main/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchEvalProcessor.java
@@ -37,7 +37,7 @@ public class OpenSearchEvalProcessor {
    * @return evaluation result.
    */
   public static ExprValue process(
-      OpenSearchFunctions.OpenSearchFunction funcExpression,
+      OpenSearchFunctions.OpenSearchExecutableFunction funcExpression,
       Environment<Expression, ExprValue> env,
       NodeClient nodeClient) {
 

diff --git a/.../test/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchEvalOperatorTest.java b/.../test/java/org/opensearch/sql/opensearch/planner/physical/OpenSearchEvalOperatorTest.java
@@ -120,7 +120,7 @@ public void testEvalOperatorOnGeoIpExpression() {
         List.of(
             ImmutablePair.of(
                 new ReferenceExpression("ipEnrichmentResult", OpenSearchTextType.of()),
-                new OpenSearchFunctions.OpenSearchFunction(
+                new OpenSearchFunctions.OpenSearchExecutableFunction(
                     BuiltinFunctionName.GEOIP.getName(),
                     List.of(
                         DSL.literal("my-datasource"),