[multistage] initial commit to support all scalar/transform/filter ma…

…tch (#9707)

* initial commit to support all scalar/transform/filter match

* fix test issue

* also fix the rest of the list

* typo

* address comments and fix tests

* add more tests

Co-authored-by: Rong Rong <[email protected]>

pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java

@@ -96,24 +96,31 @@ public static void init() {
    */
   public static void registerFunction(Method method, boolean nullableParameters) {
     registerFunction(method.getName(), method, nullableParameters);
-
-    // Calcite ScalarFunctionImpl doesn't allow customized named functions. TODO: fix me.
-    if (method.getAnnotation(Deprecated.class) == null) {
-      FUNCTION_MAP.put(method.getName(), ScalarFunctionImpl.create(method));
-    }
   }
 
   /**
    * Registers a method with the given function name.
    */
   public static void registerFunction(String functionName, Method method, boolean nullableParameters) {
+    registerFunctionInfoMap(functionName, method, nullableParameters);
+    registerCalciteNamedFunctionMap(functionName, method, nullableParameters);
+  }
+
+  private static void registerFunctionInfoMap(String functionName, Method method, boolean nullableParameters) {
     FunctionInfo functionInfo = new FunctionInfo(method, method.getDeclaringClass(), nullableParameters);
     String canonicalName = canonicalize(functionName);
     Map<Integer, FunctionInfo> functionInfoMap = FUNCTION_INFO_MAP.computeIfAbsent(canonicalName, k -> new HashMap<>());
-    Preconditions.checkState(functionInfoMap.put(method.getParameterCount(), functionInfo) == null,
+    FunctionInfo existFunctionInfo = functionInfoMap.put(method.getParameterCount(), functionInfo);
+    Preconditions.checkState(existFunctionInfo == null || existFunctionInfo.getMethod() == functionInfo.getMethod(),
         "Function: %s with %s parameters is already registered", functionName, method.getParameterCount());
   }
 
+  private static void registerCalciteNamedFunctionMap(String functionName, Method method, boolean nullableParameters) {
+    if (method.getAnnotation(Deprecated.class) == null) {
+      FUNCTION_MAP.put(functionName, ScalarFunctionImpl.create(method));
+    }
+  }
+
   public static Map<String, List<Function>> getRegisteredCalciteFunctionMap() {
     return FUNCTION_MAP.map();
   }
@@ -147,4 +154,42 @@ public static FunctionInfo getFunctionInfo(String functionName, int numParameter
   private static String canonicalize(String functionName) {
     return StringUtils.remove(functionName, '_').toLowerCase();
   }
+
+  /**
+   * Placeholders for scalar function, they register and represents the signature for transform and filter predicate
+   * so that v2 engine can understand and plan them correctly.
+   */
+  private static class PlaceholderScalarFunctions {
+
+    @ScalarFunction(names = {"jsonExtractScalar", "json_extract_scalar"})
+    public static Object jsonExtractScalar(String jsonFieldName, String jsonPath, String resultsType) {
+      throw new UnsupportedOperationException("Placeholder scalar function, should not reach here");
+    }
+
+    @ScalarFunction(names = {"jsonExtractScalar", "json_extract_scalar"})
+    public static Object jsonExtractScalar(String jsonFieldName, String jsonPath, String resultsType,
+        Object defaultValue) {
+      throw new UnsupportedOperationException("Placeholder scalar function, should not reach here");
+    }
+
+    @ScalarFunction(names = {"jsonExtractKey", "json_extract_key"})
+    public static String jsonExtractKey(String jsonFieldName, String jsonPath) {
+      throw new UnsupportedOperationException("Placeholder scalar function, should not reach here");
+    }
+
+    @ScalarFunction(names = {"textContains", "text_contains"})
+    public static boolean textContains(String text, String pattern) {
+      throw new UnsupportedOperationException("Placeholder scalar function, should not reach here");
+    }
+
+    @ScalarFunction(names = {"textMatch", "text_match"})
+    public static boolean textMatch(String text, String pattern) {
+      throw new UnsupportedOperationException("Placeholder scalar function, should not reach here");
+    }
+
+    @ScalarFunction(names = {"jsonMatch", "json_match"})
+    public static boolean jsonMatch(String text, String pattern) {
+      throw new UnsupportedOperationException("Placeholder scalar function, should not reach here");
+    }
+  }
 }

pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java

@@ -25,7 +25,7 @@
 
 
 public enum TransformFunctionType {
-  // Aggregation functions for single-valued columns
+  // arithmetic functions for single-valued columns
   ADD("add", "plus"),
   SUB("sub", "minus"),
   MULT("mult", "times"),
@@ -48,6 +48,7 @@ public enum TransformFunctionType {
   LEAST("least"),
   GREATEST("greatest"),
 
+  // predicate functions
   EQUALS("equals"),
   NOT_EQUALS("not_equals"),
   GREATER_THAN("greater_than"),
@@ -68,10 +69,17 @@ public enum TransformFunctionType {
   OR("or"),
   NOT("not"),   // NOT operator doesn't cover the transform for NOT IN and NOT LIKE
 
-  CAST("cast"),
+  // CASE WHEN function parsed as 'CASE_WHEN'
   CASE("case"),
+
+  // date type conversion functions
+  CAST("cast"),
+
+  // string functions
   JSONEXTRACTSCALAR("jsonExtractScalar"),
   JSONEXTRACTKEY("jsonExtractKey"),
+
+  // date time functions
   TIMECONVERT("timeConvert"),
   DATETIMECONVERT("dateTimeConvert"),
   DATETRUNC("dateTrunc"),
@@ -87,6 +95,10 @@ public enum TransformFunctionType {
   MINUTE("minute"),
   SECOND("second"),
   MILLISECOND("millisecond"),
+
+  EXTRACT("extract"),
+
+  // array functions
   // The only column accepted by "cardinality" function is multi-value array, thus putting "cardinality" as alias.
   // TODO: once we support other types of multiset, we should make CARDINALITY its own function
   ARRAYLENGTH("arrayLength", "cardinality"),
@@ -96,12 +108,12 @@ public enum TransformFunctionType {
   ARRAYSUM("arraySum"),
   VALUEIN("valueIn"),
   MAPVALUE("mapValue"),
+
+  // special functions
   INIDSET("inIdSet"),
   LOOKUP("lookUp"),
   GROOVY("groovy"),
 
-  EXTRACT("extract"),
-
   // Regexp functions
   REGEXP_EXTRACT("regexpExtract"),
 

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

@@ -30,27 +30,27 @@ public class ArithmeticFunctions {
   private ArithmeticFunctions() {
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"add", "plus"})
   public static double plus(double a, double b) {
     return a + b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"sub", "minus"})
   public static double minus(double a, double b) {
     return a - b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"mult", "times"})
   public static double times(double a, double b) {
     return a * b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"div", "divide"})
   public static double divide(double a, double b) {
     return a / b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"div", "divide"})
   public static double divide(double a, double b, double defaultValue) {
     return (b == 0) ? defaultValue : a / b;
   }
@@ -135,14 +135,14 @@ public static double power(double a, double exponent) {
 
   // Big Decimal Implementation has been used here to avoid overflows
   // when multiplying by Math.pow(10, scale) for rounding
-  @ScalarFunction
+  @ScalarFunction(names = {"roundDecimal", "round_decimal"})
   public static double roundDecimal(double a, int scale) {
     return BigDecimal.valueOf(a).setScale(scale, RoundingMode.HALF_UP).doubleValue();
   }
 
   // TODO: The function should ideally be named 'round'
   // but it is not possible because of existing DateTimeFunction with same name.
-  @ScalarFunction
+  @ScalarFunction(names = {"roundDecimal", "round_decimal"})
   public static double roundDecimal(double a) {
     return Math.round(a);
   }

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

@@ -28,27 +28,27 @@ public class ComparisonFunctions {
   private ComparisonFunctions() {
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"greater_than", "greaterThan"})
   public static boolean greaterThan(double a, double b) {
     return a > b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"greater_than_or_equal", "greaterThanOrEqual"})
   public static boolean greaterThanOrEqual(double a, double b) {
     return a >= b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"less_than", "lessThan"})
   public static boolean lessThan(double a, double b) {
     return a < b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"less_than_or_equal", "lessThanOrEqual"})
   public static boolean lessThanOrEqual(double a, double b) {
     return a <= b;
   }
 
-  @ScalarFunction
+  @ScalarFunction(names = {"not_equals", "notEquals"})
   public static boolean notEquals(double a, double b) {
     return Math.abs(a - b) >= DOUBLE_COMPARISON_TOLERANCE;
   }