Test var arg function resolve such as concat

Signed-off-by: Chen Dai <[email protected]>
dai-chen · Jan 24, 2023 · 2f4b960 · 2f4b960
1 parent 6a90fe6
commit 2f4b960
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 4 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionRepository.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionRepository.java
@@ -7,6 +7,7 @@
 
 import static org.opensearch.sql.ast.expression.Cast.getCastFunctionName;
 import static org.opensearch.sql.ast.expression.Cast.isCastFunction;
+import static org.opensearch.sql.data.type.ExprCoreType.ARRAY;
 
 import com.google.common.annotations.VisibleForTesting;
 import java.util.ArrayList;
@@ -179,7 +180,9 @@ private FunctionBuilder getFunctionBuilder(
     List<ExprType> sourceTypes = functionSignature.getParamTypeList();
     List<ExprType> targetTypes = resolvedSignature.getKey().getParamTypeList();
     FunctionBuilder funcBuilder = resolvedSignature.getValue();
-    if (isCastFunction(functionName) || sourceTypes.equals(targetTypes)) {
+    if (isCastFunction(functionName)
+        || isVarArgFunction(targetTypes)
+        || sourceTypes.equals(targetTypes)) {
       return funcBuilder;
     }
     return castArguments(sourceTypes,
@@ -230,4 +233,8 @@ private Function<FunctionProperties, Expression> cast(Expression arg, ExprType t
     return functionProperties -> (Expression) compile(functionProperties,
         castFunctionName, List.of(arg));
   }
+
+  private boolean isVarArgFunction(List<ExprType> argTypes) {
+    return argTypes.size() == 1 && argTypes.get(0) == ARRAY;
+  }
 }
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/FunctionSignature.java b/core/src/main/java/org/opensearch/sql/expression/function/FunctionSignature.java
@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.expression.function;
 
+import static org.opensearch.sql.data.type.ExprCoreType.ARRAY;
+
 import java.util.List;
 import java.util.stream.Collectors;
 import lombok.EqualsAndHashCode;
@@ -37,6 +39,11 @@ public int match(FunctionSignature functionSignature) {
     List<ExprType> functionTypeList = functionSignature.getParamTypeList();
     if (!functionName.equals(functionSignature.getFunctionName())
         || paramTypeList.size() != functionTypeList.size()) {
+
+      // TODO: improve to support regular and array type mixed, ex. func(int,string,array)
+      if (functionTypeList.size() == 1 && functionTypeList.get(0) == ARRAY) {
+        return EXACTLY_MATCH;
+      }
       return NOT_MATCH;
     }
 

diff --git a/core/src/main/java/org/opensearch/sql/expression/text/TextFunction.java b/core/src/main/java/org/opensearch/sql/expression/text/TextFunction.java
@@ -6,20 +6,30 @@
 
 package org.opensearch.sql.expression.text;
 
+import static org.opensearch.sql.data.type.ExprCoreType.ARRAY;
 import static org.opensearch.sql.data.type.ExprCoreType.INTEGER;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 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 java.util.Collections;
+import java.util.stream.Collectors;
 import lombok.experimental.UtilityClass;
+import org.apache.commons.lang3.tuple.Pair;
+import org.opensearch.sql.data.model.ExprCollectionValue;
 import org.opensearch.sql.data.model.ExprIntegerValue;
 import org.opensearch.sql.data.model.ExprStringValue;
 import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.FunctionExpression;
+import org.opensearch.sql.expression.env.Environment;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
 import org.opensearch.sql.expression.function.DefaultFunctionResolver;
 import org.opensearch.sql.expression.function.FunctionName;
+import org.opensearch.sql.expression.function.FunctionSignature;
 import org.opensearch.sql.expression.function.SerializableBiFunction;
 import org.opensearch.sql.expression.function.SerializableTriFunction;
 
@@ -148,9 +158,26 @@ private DefaultFunctionResolver upper() {
    * (STRING, STRING) -> STRING
    */
   private DefaultFunctionResolver concat() {
-    return define(BuiltinFunctionName.CONCAT.getName(),
-        impl(nullMissingHandling((str1, str2) ->
-            new ExprStringValue(str1.stringValue() + str2.stringValue())), STRING, STRING, STRING));
+    FunctionName concatFuncName = BuiltinFunctionName.CONCAT.getName();
+    return define(concatFuncName, funcName ->
+      Pair.of(
+          new FunctionSignature(concatFuncName, Collections.singletonList(ARRAY)),
+          (funcProp, args) -> new FunctionExpression(funcName, args) {
+            @Override
+            public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+                  return new ExprStringValue(
+                      args.stream()
+                          .map(arg -> arg.valueOf(valueEnv))
+                          .map(argVal -> String.valueOf(argVal.value()))
+                          .collect(Collectors.joining()));
+            }
+
+            @Override
+            public ExprType type() {
+              return STRING;
+            }
+          }
+      ));
   }
 
   /**