add spark comparison functions

velox/functions/sparksql/CMakeLists.txt

@@ -16,6 +16,7 @@ add_library(
   ArraySort.cpp
   Bitwise.cpp
   CompareFunctionsNullSafe.cpp
+  Comparisons.cpp
   Decimal.cpp
   DecimalArithmetic.cpp
   Hash.cpp

velox/functions/sparksql/Comparisons.cpp

@@ -0,0 +1,146 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed 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.
+ */
+#include "velox/functions/sparksql/LeastGreatest.h"
+
+#include "velox/expression/EvalCtx.h"
+#include "velox/expression/Expr.h"
+#include "velox/functions/sparksql/Comparisons.h"
+#include "velox/type/Type.h"
+
+namespace facebook::velox::functions::sparksql {
+namespace {
+
+template <typename Cmp, TypeKind kind>
+class ComparisonFunction final : public exec::VectorFunction {
+  using T = typename TypeTraits<kind>::NativeType;
+
+  bool isDefaultNullBehavior() const override {
+    return true;
+  }
+
+  bool supportsFlatNoNullsFastPath() const override {
+    return true;
+  }
+
+  void apply(
+      const SelectivityVector& rows,
+      std::vector<VectorPtr>& args,
+      const TypePtr& outputType,
+      exec::EvalCtx& context,
+      VectorPtr& result) const override {
+    exec::DecodedArgs decodedArgs(rows, args, context);
+    DecodedVector* decoded0 = decodedArgs.at(0);
+    DecodedVector* decoded1 = decodedArgs.at(1);
+    context.ensureWritable(rows, BOOLEAN(), result);
+    auto* flatResult = result->asFlatVector<bool>();
+    flatResult->mutableRawValues<uint64_t>();
+    const Cmp cmp;
+    if (decoded0->isIdentityMapping() && decoded1->isIdentityMapping()) {
+      auto decoded0Values = *args[0]->as<FlatVector<T>>();
+      auto decoded1Values = *args[1]->as<FlatVector<T>>();
+      rows.applyToSelected([&](vector_size_t i) {
+        flatResult->set(
+            i, cmp(decoded0Values.valueAt(i), decoded1Values.valueAt(i)));
+      });
+    } else if (decoded0->isIdentityMapping() && decoded1->isConstantMapping()) {
+      auto decoded0Values = *args[0]->as<FlatVector<T>>();
+      auto constantValue = decoded1->valueAt<T>(0);
+      rows.applyToSelected([&](vector_size_t i) {
+        flatResult->set(i, cmp(decoded0Values.valueAt(i), constantValue));
+      });
+    } else if (decoded0->isConstantMapping() && decoded1->isIdentityMapping()) {
+      auto constantValue = decoded0->valueAt<T>(0);
+      auto decoded1Values = *args[1]->as<FlatVector<T>>();
+      rows.applyToSelected([&](vector_size_t i) {
+        flatResult->set(i, cmp(constantValue, decoded1Values.valueAt(i)));
+      });
+    } else {
+      rows.applyToSelected([&](vector_size_t i) {
+        flatResult->set(
+            i, cmp(decoded0->valueAt<T>(i), decoded1->valueAt<T>(i)));
+      });
+    }
+  }
+};
+
+template <template <typename> class Cmp>
+std::shared_ptr<exec::VectorFunction> makeImpl(
+    const std::string& functionName,
+    const std::vector<exec::VectorFunctionArg>& args) {
+  VELOX_CHECK_EQ(args.size(), 2);
+  for (size_t i = 1; i < args.size(); i++) {
+    VELOX_CHECK(*args[i].type == *args[0].type);
+  }
+  switch (args[0].type->kind()) {
+#define SCALAR_CASE(kind)                            \
+  case TypeKind::kind:                               \
+    return std::make_shared<ComparisonFunction<      \
+        Cmp<TypeTraits<TypeKind::kind>::NativeType>, \
+        TypeKind::kind>>();
+    SCALAR_CASE(BOOLEAN)
+    SCALAR_CASE(TINYINT)
+    SCALAR_CASE(SMALLINT)
+    SCALAR_CASE(INTEGER)
+    SCALAR_CASE(BIGINT)
+    SCALAR_CASE(REAL)
+    SCALAR_CASE(DOUBLE)
+    SCALAR_CASE(VARCHAR)
+    SCALAR_CASE(VARBINARY)
+    SCALAR_CASE(TIMESTAMP)
+    SCALAR_CASE(DATE)
+    SCALAR_CASE(SHORT_DECIMAL)
+    SCALAR_CASE(LONG_DECIMAL)
+#undef SCALAR_CASE
+    default:
+      VELOX_NYI(
+          "{} does not support arguments of type {}",
+          functionName,
+          args[0].type->kind());
+  }
+}
+
+} // namespace
+
+std::shared_ptr<exec::VectorFunction> makeEqualTo(
+    const std::string& functionName,
+    const std::vector<exec::VectorFunctionArg>& args) {
+  return makeImpl<Equal>(functionName, args);
+}
+
+std::shared_ptr<exec::VectorFunction> makeLessThan(
+    const std::string& functionName,
+    const std::vector<exec::VectorFunctionArg>& args) {
+  return makeImpl<Less>(functionName, args);
+}
+
+std::shared_ptr<exec::VectorFunction> makeGreaterThan(
+    const std::string& functionName,
+    const std::vector<exec::VectorFunctionArg>& args) {
+  return makeImpl<Greater>(functionName, args);
+}
+
+std::shared_ptr<exec::VectorFunction> makeLessThanOrEqual(
+    const std::string& functionName,
+    const std::vector<exec::VectorFunctionArg>& args) {
+  return makeImpl<LessOrEqual>(functionName, args);
+}
+
+std::shared_ptr<exec::VectorFunction> makeGreaterThanOrEqual(
+    const std::string& functionName,
+    const std::vector<exec::VectorFunctionArg>& args) {
+  return makeImpl<GreaterOrEqual>(functionName, args);
+}
+} // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/Comparisons.h

@@ -13,6 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#pragma once
+
+#include <memory>
+
 #include "velox/expression/VectorFunction.h"
 
 namespace facebook::velox::functions::sparksql {
@@ -59,4 +63,66 @@ struct Equal {
   }
 };
 
+template <typename T>
+struct LessOrEqual {
+  constexpr bool operator()(const T& a, const T& b) const {
+    Less<T> less;
+    Equal<T> equal;
+    return less(a, b) || equal(a, b);
+  }
+};
+
+template <typename T>
+struct GreaterOrEqual : private Less<T> {
+  constexpr bool operator()(const T& a, const T& b) const {
+    Less<T> less;
+    Equal<T> equal;
+    return less(b, a) || equal(a, b);
+  }
+};
+
+/// Supported Types:
+/// TINYINT
+/// SMALLINT
+/// INTEGER
+/// BIGINT
+/// REAL
+/// DOUBLE
+/// BOOLEAN
+/// VARCHAR
+/// TIMESTAMP
+
+/// Special cases:
+/// NaN in Spark is handled differently from standard floating point semantics.
+/// It is considered larger than any other numeric values.
+
+std::shared_ptr<exec::VectorFunction> makeEqualTo(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs);
+
+std::shared_ptr<exec::VectorFunction> makeLessThan(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs);
+
+std::shared_ptr<exec::VectorFunction> makeGreaterThan(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs);
+
+std::shared_ptr<exec::VectorFunction> makeLessThanOrEqual(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs);
+
+std::shared_ptr<exec::VectorFunction> makeGreaterThanOrEqual(
+    const std::string& name,
+    const std::vector<exec::VectorFunctionArg>& inputArgs);
+
+inline std::vector<std::shared_ptr<exec::FunctionSignature>>
+comparisonSignatures() {
+  return {exec::FunctionSignatureBuilder()
+              .typeVariable("T")
+              .returnType("boolean")
+              .argumentType("T")
+              .argumentType("T")
+              .build()};
+}
 } // namespace facebook::velox::functions::sparksql

velox/functions/sparksql/Register.cpp

@@ -24,6 +24,7 @@
 #include "velox/functions/sparksql/ArraySort.h"
 #include "velox/functions/sparksql/Bitwise.h"
 #include "velox/functions/sparksql/CompareFunctionsNullSafe.h"
+#include "velox/functions/sparksql/Comparisons.h"
 #include "velox/functions/sparksql/DateTime.h"
 #include "velox/functions/sparksql/DateTimeFunctions.h"
 #include "velox/functions/sparksql/Decimal.h"
@@ -126,7 +127,19 @@ void registerFunctions(const std::string& prefix) {
 
   // Register 'in' functions.
   registerIn(prefix);
-
+  // Register compare functions
+  exec::registerStatefulVectorFunction(
+      prefix + "equalto", comparisonSignatures(), makeEqualTo);
+  exec::registerStatefulVectorFunction(
+      prefix + "lessthan", comparisonSignatures(), makeLessThan);
+  exec::registerStatefulVectorFunction(
+      prefix + "greaterthan", comparisonSignatures(), makeGreaterThan);
+  exec::registerStatefulVectorFunction(
+      prefix + "lessthanorequal", comparisonSignatures(), makeLessThanOrEqual);
+  exec::registerStatefulVectorFunction(
+      prefix + "greaterthanorequal",
+      comparisonSignatures(),
+      makeGreaterThanOrEqual);
   // Compare nullsafe functions
   exec::registerStatefulVectorFunction(
       prefix + "equalnullsafe", equalNullSafeSignatures(), makeEqualNullSafe);

velox/functions/sparksql/RegisterCompare.cpp

@@ -21,13 +21,6 @@
 namespace facebook::velox::functions::sparksql {
 
 void registerCompareFunctions(const std::string& prefix) {
-  registerBinaryScalar<EqFunction, bool>({prefix + "equalto"});
-  registerBinaryScalar<NeqFunction, bool>({prefix + "notequalto"});
-  registerBinaryScalar<LtFunction, bool>({prefix + "lessthan"});
-  registerBinaryScalar<GtFunction, bool>({prefix + "greaterthan"});
-  registerBinaryScalar<LteFunction, bool>({prefix + "lessthanorequal"});
-  registerBinaryScalar<GteFunction, bool>({prefix + "greaterthanorequal"});
-
   registerFunction<BetweenFunction, bool, int8_t, int8_t, int8_t>(
       {prefix + "between"});
   registerFunction<BetweenFunction, bool, int16_t, int16_t, int16_t>(

velox/functions/sparksql/tests/CMakeLists.txt

@@ -17,7 +17,7 @@ add_executable(
   ArithmeticTest.cpp
   ArraySortTest.cpp
   BitwiseTest.cpp
-  CompareNullSafeTests.cpp
+  CompareTests.cpp
   DateTimeFunctionsTest.cpp
   DecimalArithmeticTest.cpp
   ElementAtTest.cpp