Support polymorphic scalar comparison functions in the multi-stage query engine

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

@@ -312,7 +312,7 @@ private static SqlOperandTypeChecker getOperandTypeChecker(Method method) {
     private static SqlTypeFamily getSqlTypeFamily(Class<?> clazz) {
       // NOTE: Pinot allows some non-standard type conversions such as Timestamp <-> long, boolean <-> int etc. Do not
       //       restrict the type family for now. We only restrict the type family for String so that cast can be added.
-      //       Explicit cast is required to correctly convert boolean and Timestamp to String. Without explicit case,
+      //       Explicit cast is required to correctly convert boolean and Timestamp to String. Without explicit cast,
       //       BOOLEAN and TIMESTAMP type will be converted with their internal stored format which is INT and LONG
       //       respectively. E.g. true will be converted to "1", timestamp will be converted to long value string.
       // TODO: Revisit this.

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

@@ -40,7 +40,9 @@ public interface PinotScalarFunction {
    * doesn't need to be registered (e.g. standard SqlFunction).
    */
   @Nullable
-  PinotSqlFunction toPinotSqlFunction();
+  default PinotSqlFunction toPinotSqlFunction() {
+    return null;
+  }
 
   /**
    * Returns the {@link FunctionInfo} for the given argument types, or {@code null} if there is no matching.

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

@@ -22,42 +22,9 @@
 
 public class ComparisonFunctions {
 
-  private static final double DOUBLE_COMPARISON_TOLERANCE = 1e-7d;
-
   private ComparisonFunctions() {
   }
 
-  @ScalarFunction
-  public static boolean greaterThan(double a, double b) {
-    return a > b;
-  }
-
-  @ScalarFunction
-  public static boolean greaterThanOrEqual(double a, double b) {
-    return a >= b;
-  }
-
-  @ScalarFunction
-  public static boolean lessThan(double a, double b) {
-    return a < b;
-  }
-
-  @ScalarFunction
-  public static boolean lessThanOrEqual(double a, double b) {
-    return a <= b;
-  }
-
-  @ScalarFunction
-  public static boolean notEquals(double a, double b) {
-    return Math.abs(a - b) >= DOUBLE_COMPARISON_TOLERANCE;
-  }
-
-  @ScalarFunction
-  public static boolean equals(double a, double b) {
-    // To avoid approximation errors
-    return Math.abs(a - b) < DOUBLE_COMPARISON_TOLERANCE;
-  }
-
   @ScalarFunction
   public static boolean between(double val, double a, double b) {
     return val >= a && val <= b;

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

@@ -0,0 +1,176 @@
+/**
+ * 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.scalar.comparison;
+
+import java.math.BigDecimal;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import javax.annotation.Nullable;
+import org.apache.pinot.common.function.FunctionInfo;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.spi.annotations.ScalarFunction;
+
+/**
+ * Polymorphic equals (=) scalar function implementation
+ */
+@ScalarFunction
+public class EqualsScalarFunction extends PolymorphicComparisonScalarFunction {
+
+  private static final Map<ColumnDataType, FunctionInfo> TYPE_FUNCTION_INFO_MAP = new HashMap<>();
+  private static final FunctionInfo DOUBLE_EQUALS_WITH_TOLERANCE;
+
+  static {
+    try {
+       DOUBLE_EQUALS_WITH_TOLERANCE = new FunctionInfo(
+           EqualsScalarFunction.class.getMethod("doubleEqualsWithTolerance", double.class, double.class),
+           EqualsScalarFunction.class, false);
+
+      // Set nullable parameters to false for each function because the return value should be null if any argument
+      // is null
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.INT, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("intEquals", int.class, int.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.LONG, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("longEquals", long.class, long.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.FLOAT, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("floatEquals", float.class, float.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.DOUBLE, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("doubleEquals", double.class, double.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.BIG_DECIMAL, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("bigDecimalEquals", BigDecimal.class, BigDecimal.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.STRING, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("stringEquals", String.class, String.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.BYTES, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("bytesEquals", byte[].class, byte[].class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.OBJECT, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("objectEquals", Object.class, Object.class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.INT_ARRAY, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("intArrayEquals", int[].class, int[].class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.LONG_ARRAY, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("longArrayEquals", long[].class, long[].class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.FLOAT_ARRAY, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("floatArrayEquals", float[].class, float[].class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.DOUBLE_ARRAY, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("doubleArrayEquals", double[].class, double[].class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.STRING_ARRAY, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("stringArrayEquals", String[].class, String[].class),
+          EqualsScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.BYTES_ARRAY, new FunctionInfo(
+          EqualsScalarFunction.class.getMethod("bytesArrayEquals", byte[][].class, byte[][].class),
+          EqualsScalarFunction.class, false));
+    } catch (NoSuchMethodException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @Override
+  protected FunctionInfo functionInfoForType(ColumnDataType argumentType) {
+    return TYPE_FUNCTION_INFO_MAP.get(argumentType);
+  }
+
+  @Nullable
+  @Override
+  public FunctionInfo getFunctionInfo(int numArguments) {
+    if (numArguments != 2) {
+      return null;
+    }
+
+    // For backward compatibility
+    return DOUBLE_EQUALS_WITH_TOLERANCE;
+  }
+
+  @Override
+  public String getName() {
+    return "equals";
+  }
+
+  public static boolean intEquals(int a, int b) {
+    return a == b;
+  }
+
+  public static boolean longEquals(long a, long b) {
+    return a == b;
+  }
+
+  public static boolean floatEquals(float a, float b) {
+    return a == b;
+  }
+
+  public static boolean doubleEquals(double a, double b) {
+    return a == b;
+  }
+
+  public static boolean doubleEqualsWithTolerance(double a, double b) {
+    // To avoid approximation errors
+    return Math.abs(a - b) < DOUBLE_COMPARISON_TOLERANCE;
+  }
+
+  public static boolean bigDecimalEquals(BigDecimal a, BigDecimal b) {
+    return a.compareTo(b) == 0;
+  }
+
+  public static boolean stringEquals(String a, String b) {
+    return a.equals(b);
+  }
+
+  public static boolean bytesEquals(byte[] a, byte[] b) {
+    return Arrays.equals(a, b);
+  }
+
+  public static boolean objectEquals(Object a, Object b) {
+    return Objects.equals(a, b);
+  }
+
+  public static boolean intArrayEquals(int[] a, int[] b) {
+    return Arrays.equals(a, b);
+  }
+
+  public static boolean longArrayEquals(long[] a, long[] b) {
+    return Arrays.equals(a, b);
+  }
+
+  public static boolean floatArrayEquals(float[] a, float[] b) {
+    return Arrays.equals(a, b);
+  }
+
+  public static boolean doubleArrayEquals(double[] a, double[] b) {
+    return Arrays.equals(a, b);
+  }
+
+  public static boolean stringArrayEquals(String[] a, String[] b) {
+    return Arrays.equals(a, b);
+  }
+
+  public static boolean bytesArrayEquals(byte[][] a, byte[][]b) {
+    return Arrays.deepEquals(a, b);
+  }
+}

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

@@ -0,0 +1,98 @@
+/**
+ * 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.scalar.comparison;
+
+import java.math.BigDecimal;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.common.function.FunctionInfo;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.spi.annotations.ScalarFunction;
+
+
+/**
+ * Polymorphic greaterThanOrEqual (>=) scalar function implementation
+ */
+@ScalarFunction
+public class GreaterThanOrEqualScalarFunction extends PolymorphicComparisonScalarFunction {
+
+  private static final Map<ColumnDataType, FunctionInfo> TYPE_FUNCTION_INFO_MAP = new HashMap<>();
+
+  static {
+    try {
+      // Set nullable parameters to false for each function because the return value should be null if any argument
+      // is null
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.INT, new FunctionInfo(
+          GreaterThanOrEqualScalarFunction.class.getMethod("intGreaterThanOrEqual", int.class, int.class),
+          GreaterThanOrEqualScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.LONG, new FunctionInfo(
+          GreaterThanOrEqualScalarFunction.class.getMethod("longGreaterThanOrEqual", long.class, long.class),
+          GreaterThanOrEqualScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.FLOAT, new FunctionInfo(
+          GreaterThanOrEqualScalarFunction.class.getMethod("floatGreaterThanOrEqual", float.class, float.class),
+          GreaterThanOrEqualScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.DOUBLE, new FunctionInfo(
+          GreaterThanOrEqualScalarFunction.class.getMethod("doubleGreaterThanOrEqual", double.class, double.class),
+          GreaterThanOrEqualScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.BIG_DECIMAL, new FunctionInfo(
+          GreaterThanOrEqualScalarFunction.class.getMethod("bigDecimalGreaterThanOrEqual",
+              BigDecimal.class, BigDecimal.class),
+          GreaterThanOrEqualScalarFunction.class, false));
+      TYPE_FUNCTION_INFO_MAP.put(ColumnDataType.STRING, new FunctionInfo(
+          GreaterThanOrEqualScalarFunction.class.getMethod("stringGreaterThanOrEqual", String.class, String.class),
+          GreaterThanOrEqualScalarFunction.class, false));
+    } catch (NoSuchMethodException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @Override
+  protected FunctionInfo functionInfoForType(ColumnDataType argumentType) {
+    return TYPE_FUNCTION_INFO_MAP.get(argumentType);
+  }
+
+  @Override
+  public String getName() {
+    return "greaterThanOrEqual";
+  }
+
+  public static boolean intGreaterThanOrEqual(int a, int b) {
+    return a >= b;
+  }
+
+  public static boolean longGreaterThanOrEqual(long a, long b) {
+    return a >= b;
+  }
+
+  public static boolean floatGreaterThanOrEqual(float a, float b) {
+    return a >= b;
+  }
+
+  public static boolean doubleGreaterThanOrEqual(double a, double b) {
+    return a >= b;
+  }
+
+  public static boolean bigDecimalGreaterThanOrEqual(BigDecimal a, BigDecimal b) {
+    return a.compareTo(b) >= 0;
+  }
+
+  public static boolean stringGreaterThanOrEqual(String a, String b) {
+    return a.compareTo(b) >= 0;
+  }
+}