[feature](function) support variadic template type in SQL function

fe/fe-common/src/main/java/org/apache/doris/catalog/StructType.java

@@ -31,8 +31,12 @@
 import com.google.gson.annotations.SerializedName;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
  * Describes a STRUCT type. STRUCT types have a list of named struct fields.
@@ -58,6 +62,19 @@ public StructType(ArrayList<StructField> fields) {
         }
     }
 
+    public StructType(List<Type> types) {
+        Preconditions.checkNotNull(types);
+        ArrayList<StructField> newFields = new ArrayList<>();
+        for (Type type : types) {
+            newFields.add(new StructField(type));
+        }
+        this.fields = newFields;
+    }
+
+    public StructType(Type... types) {
+        this(Arrays.asList(types));
+    }
+
     public StructType() {
         this.fields = Lists.newArrayList();
     }
@@ -171,6 +188,91 @@ public boolean matchesType(Type t) {
         return true;
     }
 
+    @Override
+    public boolean hasTemplateType() {
+        for (StructField field : fields) {
+            if (field.type.hasTemplateType()) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    @Override
+    public Type specializeTemplateType(Type specificType, Map<String, Type> specializedTypeMap,
+            boolean useSpecializedType) throws TypeException {
+        StructType specificStructType = null;
+        if (specificType instanceof StructType) {
+            specificStructType = (StructType) specificType;
+        } else if (!useSpecializedType) {
+            throw new TypeException(specificType + " is not StructType");
+        }
+
+        List<Type> newTypes = Lists.newArrayList();
+        for (int i = 0; i < fields.size(); i++) {
+            if (fields.get(i).type.hasTemplateType()) {
+                newTypes.add(fields.get(i).type.specializeTemplateType(
+                        specificStructType != null ? specificStructType.fields.get(i).type : specificType,
+                        specializedTypeMap, useSpecializedType));
+            }
+        }
+
+        Type newStructType = new StructType(newTypes);
+        if (Type.canCastTo(specificType, newStructType)
+                || (useSpecializedType && !(specificType instanceof StructType))) {
+            return newStructType;
+        } else {
+            throw new TypeException(specificType + " can not cast to specialize type " + newStructType);
+        }
+    }
+
+    @Override
+    public boolean needExpandTemplateType() {
+        Preconditions.checkNotNull(fields);
+        return fields.get(fields.size() - 1).type.needExpandTemplateType();
+    }
+
+    /**
+     * A struct variadic template is like `STRUCT<TYPES>` or `STRUCT<T, TYPES>`...
+     * So that we only need to expand the last field in struct variadic template.
+     */
+    @Override
+    public void collectTemplateExpandSize(Type[] args, Map<String, Integer> expandSizeMap) throws TypeException {
+        Preconditions.checkState(needExpandTemplateType());
+        if (args == null || args.length == 0) {
+            throw new TypeException("can not expand template type in struct since input args is empty.");
+        }
+        if (!(args[0] instanceof StructType)) {
+            throw new TypeException(args[0] + " is not StructType");
+        }
+        StructType structType = (StructType) args[0];
+        if (structType.fields.size() < fields.size()) {
+            throw new TypeException("the field size of input struct type " + structType
+                    + " is less than struct template " + this);
+        }
+        Type[] types = structType.fields.subList(fields.size() - 1, structType.fields.size())
+                .stream().map(field -> field.type).toArray(Type[]::new);
+        // only the last field is expandable
+        fields.get(fields.size() - 1).type.collectTemplateExpandSize(types, expandSizeMap);
+    }
+
+    /**
+     * A struct variadic template is like `STRUCT<TYPES>` or `STRUCT<T, TYPES>`...
+     * This method is used to expand a variadic template.
+     * e.g. Expand `STRUCT<T, TYPES>` to `STRUCT<T, TYPES_1, TYPES_2, ..., TYPES_SIZE>`
+     */
+    @Override
+    public List<Type> expandVariadicTemplateType(Map<String, Integer> expandSizeMap) {
+        Type type = fields.get(fields.size() - 1).type;
+        if (type.needExpandTemplateType()) {
+            List<Type> types = fields.subList(0, fields.size() - 1).stream().map(field -> field.type)
+                    .collect(Collectors.toList());
+            types.addAll(type.expandVariadicTemplateType(expandSizeMap));
+            return Lists.newArrayList(new StructType(types));
+        }
+        return Lists.newArrayList(this);
+    }
+
     @Override
     public boolean equals(Object other) {
         if (!(other instanceof StructType)) {

fe/fe-common/src/main/java/org/apache/doris/catalog/TemplateType.java

@@ -20,9 +20,12 @@
 import org.apache.doris.thrift.TColumnType;
 import org.apache.doris.thrift.TTypeDesc;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
+import com.google.common.collect.Lists;
 import com.google.gson.annotations.SerializedName;
 
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -34,8 +37,16 @@ public class TemplateType extends Type {
     @SerializedName(value = "name")
     private final String name;
 
-    public TemplateType(String name) {
+    @SerializedName(value = "isVariadic")
+    private final boolean isVariadic;
+
+    public TemplateType(String name, boolean isVariadic) {
         this.name = name;
+        this.isVariadic = isVariadic;
+    }
+
+    public TemplateType(String name) {
+        this(name, false);
     }
 
     @Override
@@ -49,7 +60,7 @@ public boolean equals(Object other) {
             return false;
         }
         TemplateType o = (TemplateType) other;
-        return o.name.equals(name);
+        return o.name.equals(name) && o.isVariadic == isVariadic;
     }
 
     @Override
@@ -63,6 +74,11 @@ public boolean hasTemplateType() {
         return true;
     }
 
+    @Override
+    public boolean needExpandTemplateType() {
+        return isVariadic;
+    }
+
     @Override
     public Type specializeTemplateType(Type specificType, Map<String, Type> specializedTypeMap,
                                        boolean useSpecializedType) throws TypeException {
@@ -94,6 +110,31 @@ public Type specializeTemplateType(Type specificType, Map<String, Type> speciali
         return specializedTypeMap.get(name);
     }
 
+    @Override
+    public void collectTemplateExpandSize(Type[] args, Map<String, Integer> expandSizeMap)
+            throws TypeException {
+        Preconditions.checkState(isVariadic);
+        expandSizeMap.computeIfAbsent(name, k -> args.length);
+        if (expandSizeMap.get(name) != args.length) {
+            throw new TypeException(
+                    String.format("can not expand variadic template type %s to %s size since it's "
+                            + "already expand as %s size", name, args.length, expandSizeMap.get(name)));
+        }
+    }
+
+    @Override
+    public List<Type> expandVariadicTemplateType(Map<String, Integer> expandSizeMap) {
+        if (needExpandTemplateType() && expandSizeMap.containsKey(name)) {
+            List<Type> types = Lists.newArrayList();
+            int size = expandSizeMap.get(name);
+            for (int index = 0; index < size; index++) {
+                types.add(new TemplateType(name + "_" + index));
+            }
+            return types;
+        }
+        return Lists.newArrayList(this);
+    }
+
     @Override
     public String toSql(int depth) {
         return name;

fe/fe-common/src/main/java/org/apache/doris/catalog/Type.java

@@ -524,7 +524,12 @@ public boolean hasTemplateType() {
         return false;
     }
 
-    // return a new type without template type, by specialize tempalte type in this type
+    // only used for struct type and variadic template type
+    public boolean needExpandTemplateType() {
+        return false;
+    }
+
+    // return a new type without template type, by specialize template type in this type
     public Type specializeTemplateType(Type specificType, Map<String, Type> specializedTypeMap,
                                        boolean useSpecializedType) throws TypeException {
         if (hasTemplateType()) {
@@ -535,6 +540,22 @@ public Type specializeTemplateType(Type specificType, Map<String, Type> speciali
         }
     }
 
+    /**
+     * Only used for struct type and variadic template type,
+     * collect variadic template's expand size based on the input arguments
+     */
+    public void collectTemplateExpandSize(Type[] args, Map<String, Integer> expandSizeMap)
+            throws TypeException {
+    }
+
+    /**
+     * Only used for struct type and variadic template type,
+     * Do expand variadic template type
+     */
+    public List<Type> expandVariadicTemplateType(Map<String, Integer> expandSizeMap) {
+        return Lists.newArrayList(this);
+    }
+
     /**
      * Returns true if Impala supports this type in the metdata. It does not mean we
      * can manipulate data of this type. For tables that contain columns with these

fe/fe-core/src/main/java/org/apache/doris/catalog/Function.java

@@ -40,6 +40,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -847,6 +848,28 @@ public boolean hasTemplateArg() {
         return false;
     }
 
+    public boolean hasVariadicTemplateArg() {
+        for (Type t : getArgs()) {
+            if (t.needExpandTemplateType()) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    // collect expand size of variadic template
+    public void collectTemplateExpandSize(Type[] args, Map<String, Integer> expandSizeMap) throws TypeException {
+        for (int i = argTypes.length - 1; i >= 0; i--) {
+            if (argTypes[i].hasTemplateType()) {
+                if (argTypes[i].needExpandTemplateType()) {
+                    argTypes[i].collectTemplateExpandSize(
+                            Arrays.copyOfRange(args, i, args.length), expandSizeMap);
+                }
+            }
+        }
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) {

fe/fe-core/src/main/java/org/apache/doris/catalog/FunctionSet.java

@@ -1227,9 +1227,14 @@ public Function getFunction(Function desc, Function.CompareMode mode, boolean is
 
         List<Function> normalFunctions = Lists.newArrayList();
         List<Function> templateFunctions = Lists.newArrayList();
+        List<Function> variadicTemplateFunctions = Lists.newArrayList();
         for (Function fn : fns) {
             if (fn.hasTemplateArg()) {
-                templateFunctions.add(fn);
+                if (!fn.hasVariadicTemplateArg()) {
+                    templateFunctions.add(fn);
+                } else {
+                    variadicTemplateFunctions.add(fn);
+                }
             } else {
                 normalFunctions.add(fn);
             }
@@ -1244,12 +1249,29 @@ public Function getFunction(Function desc, Function.CompareMode mode, boolean is
         // then specialize template functions and try them
         List<Function> specializedTemplateFunctions = Lists.newArrayList();
         for (Function f : templateFunctions) {
-            f = FunctionSet.specializeTemplateFunction(f, desc);
+            f = FunctionSet.specializeTemplateFunction(f, desc, false);
             if (f != null) {
                 specializedTemplateFunctions.add(f);
             }
         }
-        return getFunction(desc, mode, specializedTemplateFunctions);
+
+        // try template function second
+        fn = getFunction(desc, mode, specializedTemplateFunctions);
+        if (fn != null) {
+            return fn;
+        }
+
+        // then specialize variadic template function and try them
+        List<Function> specializedVariadicTemplateFunctions = Lists.newArrayList();
+        for (Function f : variadicTemplateFunctions) {
+            f = FunctionSet.specializeTemplateFunction(f, desc, true);
+            if (f != null) {
+                specializedVariadicTemplateFunctions.add(f);
+            }
+        }
+
+        // try variadic template function
+        return getFunction(desc, mode, specializedVariadicTemplateFunctions);
     }
 
     private Function getFunction(Function desc, Function.CompareMode mode, List<Function> fns) {
@@ -1292,19 +1314,43 @@ private Function getFunction(Function desc, Function.CompareMode mode, List<Func
         return null;
     }
 
-    public static Function specializeTemplateFunction(Function templateFunction, Function requestFunction) {
+    public static Function specializeTemplateFunction(Function templateFunction, Function requestFunction, boolean isVariadic) {
         try {
             boolean hasTemplateType = false;
             LOG.debug("templateFunction signature: " + templateFunction.signatureString()
                         + "  return: " + templateFunction.getReturnType());
             LOG.debug("requestFunction signature: " + requestFunction.signatureString()
                         + "  return: " + requestFunction.getReturnType());
+            List<Type> newArgTypes = Lists.newArrayList();
+            List<Type> newRetType = Lists.newArrayList();
+            if (isVariadic) {
+                Map<String, Integer> expandSizeMap = Maps.newHashMap();
+                templateFunction.collectTemplateExpandSize(requestFunction.getArgs(), expandSizeMap);
+                // expand the variadic template in arg types
+                for (Type argType : templateFunction.getArgs()) {
+                    if (argType.needExpandTemplateType()) {
+                        newArgTypes.addAll(argType.expandVariadicTemplateType(expandSizeMap));
+                    } else {
+                        newArgTypes.add(argType);
+                    }
+                }
+
+                // expand the variadic template in ret type
+                if (templateFunction.getReturnType().needExpandTemplateType()) {
+                    newRetType.addAll(templateFunction.getReturnType().expandVariadicTemplateType(expandSizeMap));
+                    Preconditions.checkState(newRetType.size() == 1);
+                } else {
+                    newRetType.add(templateFunction.getReturnType());
+                }
+            } else {
+                newArgTypes.addAll(Lists.newArrayList(templateFunction.getArgs()));
+                newRetType.add(templateFunction.getReturnType());
+            }
             Function specializedFunction = templateFunction;
             if (templateFunction instanceof ScalarFunction) {
                 ScalarFunction f = (ScalarFunction) templateFunction;
-                specializedFunction = new ScalarFunction(f.getFunctionName(), Lists.newArrayList(f.getArgs()),
-                                            f.getReturnType(), f.hasVarArgs(), f.getSymbolName(), f.getBinaryType(),
-                                            f.isUserVisible(), f.isVectorized(), f.getNullableMode());
+                specializedFunction = new ScalarFunction(f.getFunctionName(), newArgTypes, newRetType.get(0), f.hasVarArgs(),
+                        f.getSymbolName(), f.getBinaryType(), f.isUserVisible(), f.isVectorized(), f.getNullableMode());
             } else {
                 // TODO(xk)
             }