[multistage][feature] FunctionRegistry unification

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/ArrayFunctions.java

@@ -18,13 +18,24 @@
  */
 package org.apache.pinot.common.function.scalar;
 
+import com.google.common.base.Preconditions;
 import it.unimi.dsi.fastutil.ints.IntArrayList;
 import it.unimi.dsi.fastutil.ints.IntLinkedOpenHashSet;
 import it.unimi.dsi.fastutil.ints.IntSet;
 import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
 import it.unimi.dsi.fastutil.objects.ObjectSet;
 import java.math.BigDecimal;
 import java.util.Arrays;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.sql.type.NonNullableAccessors;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
+import org.apache.calcite.sql.type.SameOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.calcite.sql.type.SqlTypeFamily;
+import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.commons.lang3.ArrayUtils;
 import org.apache.pinot.spi.annotations.ScalarFunction;
 import org.apache.pinot.spi.utils.CommonConstants.NullValuePlaceHolder;
@@ -228,61 +239,98 @@ public static String arrayElementAtString(String[] arr, int idx) {
     return idx > 0 && idx <= arr.length ? arr[idx - 1] : NullValuePlaceHolder.STRING;
   }
 
-  @ScalarFunction(names = {"array", "arrayValueConstructor"}, isVarArg = true)
-  public static Object arrayValueConstructor(Object... arr) {
-    if (arr.length == 0) {
-      return arr;
+  @ScalarFunction(names = {"arrayElementAt", "array_element_at"})
+  public static class ArrayElementAt {
+    public static final SqlReturnTypeInference RETURN_TYPE_INFERENCE = (opBinding) -> {
+      final RelDataTypeFactory typeFactory = opBinding.getTypeFactory();
+      final RelDataType operandType = opBinding.getOperandType(0);
+      Preconditions.checkState(operandType.getSqlTypeName() == SqlTypeName.ARRAY);
+      return typeFactory.createTypeWithNullability(NonNullableAccessors.getComponentTypeOrThrow(operandType), true);
+    };
+    public static final SqlOperandTypeChecker OPERAND_TYPE_CHECKER =
+        OperandTypes.family(SqlTypeFamily.ARRAY, SqlTypeFamily.INTEGER);
+
+    public static int eval(int[] arr, int idx) {
+      return idx > 0 && idx <= arr.length ? arr[idx - 1] : NullValuePlaceHolder.INT;
     }
-    Class<?> clazz = arr[0].getClass();
-    if (clazz == Integer.class) {
-      int[] intArr = new int[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        intArr[i] = (Integer) arr[i];
-      }
-      return intArr;
+
+    public static long eval(long[] arr, int idx) {
+      return idx > 0 && idx <= arr.length ? arr[idx - 1] : NullValuePlaceHolder.LONG;
     }
-    if (clazz == Long.class) {
-      long[] longArr = new long[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        longArr[i] = (Long) arr[i];
-      }
-      return longArr;
+
+    public static float eval(float[] arr, int idx) {
+      return idx > 0 && idx <= arr.length ? arr[idx - 1] : NullValuePlaceHolder.FLOAT;
     }
-    if (clazz == Float.class) {
-      float[] floatArr = new float[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        floatArr[i] = (Float) arr[i];
-      }
-      return floatArr;
+
+    public static double eval(double[] arr, int idx) {
+      return idx > 0 && idx <= arr.length ? arr[idx - 1] : NullValuePlaceHolder.DOUBLE;
     }
-    if (clazz == Double.class) {
-      double[] doubleArr = new double[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        doubleArr[i] = (Double) arr[i];
-      }
-      return doubleArr;
+
+    public static String eval(String[] arr, int idx) {
+      return idx > 0 && idx <= arr.length ? arr[idx - 1] : NullValuePlaceHolder.STRING;
     }
-    if (clazz == Boolean.class) {
-      boolean[] boolArr = new boolean[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        boolArr[i] = (Boolean) arr[i];
+  }
+
+  @ScalarFunction(names = {"array", "arrayValueConstructor"}, isVarArg = true)
+  public static class ArrayValueConstructor {
+    public static final SqlReturnTypeInference RETURN_TYPE_INFERENCE = ReturnTypes.TO_ARRAY;
+    public static final SqlOperandTypeChecker OPERAND_TYPE_CHECKER = new SameOperandTypeChecker(-1);
+
+    public static Object eval(Object... arr) {
+      if (arr.length == 0) {
+        return arr;
       }
-      return boolArr;
-    }
-    if (clazz == BigDecimal.class) {
-      BigDecimal[] bigDecimalArr = new BigDecimal[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        bigDecimalArr[i] = (BigDecimal) arr[i];
+      Class<?> clazz = arr[0].getClass();
+      if (clazz == Integer.class) {
+        int[] intArr = new int[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          intArr[i] = (Integer) arr[i];
+        }
+        return intArr;
       }
-      return bigDecimalArr;
-    }
-    if (clazz == String.class) {
-      String[] strArr = new String[arr.length];
-      for (int i = 0; i < arr.length; i++) {
-        strArr[i] = (String) arr[i];
+      if (clazz == Long.class) {
+        long[] longArr = new long[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          longArr[i] = (Long) arr[i];
+        }
+        return longArr;
+      }
+      if (clazz == Float.class) {
+        float[] floatArr = new float[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          floatArr[i] = (Float) arr[i];
+        }
+        return floatArr;
       }
-      return strArr;
+      if (clazz == Double.class) {
+        double[] doubleArr = new double[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          doubleArr[i] = (Double) arr[i];
+        }
+        return doubleArr;
+      }
+      if (clazz == Boolean.class) {
+        boolean[] boolArr = new boolean[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          boolArr[i] = (Boolean) arr[i];
+        }
+        return boolArr;
+      }
+      if (clazz == BigDecimal.class) {
+        BigDecimal[] bigDecimalArr = new BigDecimal[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          bigDecimalArr[i] = (BigDecimal) arr[i];
+        }
+        return bigDecimalArr;
+      }
+      if (clazz == String.class) {
+        String[] strArr = new String[arr.length];
+        for (int i = 0; i < arr.length; i++) {
+          strArr[i] = (String) arr[i];
+        }
+        return strArr;
+      }
+      return arr;
     }
-    return arr;
   }
 }

pinot-common/src/main/java/org/apache/pinot/common/function/scalar/StringFunctions.java

@@ -140,7 +140,7 @@ public static String substring(String input, int beginIndex, int length) {
    * @param seperator
    * @return The two input strings joined by the seperator
    */
-  @ScalarFunction(names = "concat_ws")
+  @ScalarFunction(names = {"concatWS", "concat_ws"})
   public static String concatws(String seperator, String input1, String input2) {
     return concat(input1, input2, seperator);
   }

pinot-common/src/main/java/org/apache/pinot/common/function/schema/PinotFunction.java

@@ -0,0 +1,33 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.function.schema;
+
+import org.apache.calcite.schema.Function;
+import org.apache.calcite.sql.type.SqlOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+
+
+/**
+ * Function Schema used to resolve function signature by {@link org.apache.pinot.common.function.FunctionRegistry} and
+ * {@link org.apache.pinot.common.function.sql.PinotCalciteCatalogReader}.
+ */
+public interface PinotFunction extends Function {
+  SqlOperandTypeChecker getOperandTypeChecker();
+  SqlReturnTypeInference getReturnTypeInference();
+}

pinot-common/src/main/java/org/apache/pinot/common/function/schema/PinotScalarFunction.java

@@ -0,0 +1,95 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.common.function.schema;
+
+import java.lang.reflect.Method;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.schema.ScalarFunction;
+import org.apache.calcite.schema.impl.ReflectiveFunctionBase;
+import org.apache.calcite.sql.type.SqlOperandTypeChecker;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.apache.pinot.common.function.FunctionInfo;
+
+
+/**
+ * Pinot specific implementation of the {@link ScalarFunction}.
+ *
+ * @see "{@link org.apache.calcite.schema.impl.ScalarFunctionImpl}"
+ */
+public class PinotScalarFunction extends ReflectiveFunctionBase implements PinotFunction, ScalarFunction {
+  private final FunctionInfo _functionInfo;
+  private final String _name;
+  private final Method _method;
+  private final SqlOperandTypeChecker _sqlOperandTypeChecker;
+  private final SqlReturnTypeInference _sqlReturnTypeInference;
+  private final boolean _isNullableParameter;
+  private final boolean _isVarArgs;
+
+  public PinotScalarFunction(String name, Method method, boolean isNullableParameter, boolean isVarArg) {
+    this(name, method, isNullableParameter, isVarArg, null, null);
+  }
+
+  public PinotScalarFunction(String name, Method method, boolean isNullableParameter, boolean isVarArgs,
+    SqlOperandTypeChecker sqlOperandTypeChecker, SqlReturnTypeInference sqlReturnTypeInference) {
+    super(method);
+    _name = name;
+    _method = method;
+    _isNullableParameter = isNullableParameter;
+    _isVarArgs = isVarArgs;
+    _functionInfo = new FunctionInfo(method, method.getDeclaringClass(), isNullableParameter);
+    _sqlOperandTypeChecker = sqlOperandTypeChecker;
+    _sqlReturnTypeInference = sqlReturnTypeInference;
+  }
+
+  @Override
+  public RelDataType getReturnType(RelDataTypeFactory typeFactory) {
+    return typeFactory.createJavaType(method.getReturnType());
+  }
+
+  public String getName() {
+    return _name;
+  }
+
+  public Method getMethod() {
+    return _method;
+  }
+
+  public FunctionInfo getFunctionInfo() {
+    return _functionInfo;
+  }
+
+  @Override
+  public SqlOperandTypeChecker getOperandTypeChecker() {
+    return _sqlOperandTypeChecker;
+  }
+
+  @Override
+  public SqlReturnTypeInference getReturnTypeInference() {
+    return _sqlReturnTypeInference;
+  }
+
+  public boolean isNullableParameter() {
+    return _isNullableParameter;
+  }
+
+  public boolean isVarArgs() {
+    return _isVarArgs;
+  }
+}