Add UUID comparison functions and comparison unit test

facebookincubator · Oct 7, 2024 · f8ee222 · f8ee222
1 parent 08dd2d4
commit f8ee222
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 0 deletions.
diff --git a/velox/functions/prestosql/UuidFunctions.h b/velox/functions/prestosql/UuidFunctions.h
@@ -20,10 +20,49 @@
 
 #include "velox/functions/Macros.h"
 #include "velox/functions/Registerer.h"
+#include "velox/functions/prestosql/Comparisons.h"
 #include "velox/functions/prestosql/types/UuidType.h"
 
+namespace facebook::velox::detail {
+
+FOLLY_ALWAYS_INLINE static int compare64(uint64_t x, uint64_t y)
+{
+  return x < y ? -1 : (x == y ? 0 : 1);
+}
+
+static int compare128(const int128_t& lhs, const int128_t& rhs) {
+  int comp = compare64(HugeInt::upper(lhs), HugeInt::upper(rhs));
+  if (comp != 0) {
+    return comp;
+  }
+  return compare64(HugeInt::lower(lhs), HugeInt::lower(rhs));
+}
+
+} // namespace facebook::velox::detail
+
 namespace facebook::velox::functions {
 
+#define VELOX_GEN_BINARY_EXPR_UUID(Name, uuidCompExpr)                        \
+  template <typename T>                                                       \
+  struct Name##Uuid {                                                         \
+    VELOX_DEFINE_FUNCTION_TYPES(T);                                           \
+                                                                              \
+    FOLLY_ALWAYS_INLINE void call(                                            \
+        bool& result,                                                         \
+        const arg_type<Uuid>& lhs,                                            \
+        const arg_type<Uuid>& rhs) {                                          \
+      int compareResult = detail::compare128(lhs, rhs);                       \
+      result = (uuidCompExpr);                                                \
+    }                                                                         \
+  };
+
+VELOX_GEN_BINARY_EXPR_UUID(LtFunction, compareResult < 0);
+VELOX_GEN_BINARY_EXPR_UUID(GtFunction, compareResult > 0);
+VELOX_GEN_BINARY_EXPR_UUID(LteFunction, compareResult <= 0);
+VELOX_GEN_BINARY_EXPR_UUID(GteFunction, compareResult >= 0);
+
+#undef VELOX_GEN_BINARY_EXPR_UUID
+
 template <typename T>
 struct UuidFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
@@ -42,6 +81,10 @@ struct UuidFunction {
 inline void registerUuidFunctions(const std::string& prefix) {
   registerUuidType();
   registerFunction<UuidFunction, Uuid>({prefix + "uuid"});
+  registerFunction<LtFunctionUuid, bool, Uuid, Uuid>({prefix + "lt"});
+  registerFunction<GtFunctionUuid, bool, Uuid, Uuid>({prefix + "gt"});
+  registerFunction<LteFunctionUuid, bool, Uuid, Uuid>({prefix + "lte"});
+  registerFunction<GteFunctionUuid, bool, Uuid, Uuid>({prefix + "gte"});
 }
 
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/tests/UuidFunctionsTest.cpp b/velox/functions/prestosql/tests/UuidFunctionsTest.cpp
@@ -108,5 +108,29 @@ TEST_F(UuidFunctionsTest, unsupportedCast) {
       evaluate("cast(123 as uuid())", input), "Cannot cast BIGINT to UUID.");
 }
 
+TEST_F(UuidFunctionsTest, comparisons) {
+  const auto uuidEval = [&](const std::optional<std::string>& lhs, const std::string& operation, const std::optional<std::string>& rhs) {
+    return evaluateOnce<bool>(fmt::format("cast(c0 as uuid) {} cast(c1 as uuid)", operation), lhs, rhs);
+  };
+
+  ASSERT_EQ(uuidEval("33355449-2c7d-43d7-967a-f53cd23215ad", "<", "ffffffff-ffff-ffff-ffff-ffffffffffff"), true);
+  ASSERT_EQ(uuidEval("33355449-2c7d-43d7-967a-f53cd23215ad", "<", "00000000-0000-0000-0000-000000000000"), false);
+  ASSERT_EQ(uuidEval("f768f36d-4f09-4da7-a298-3564d8f3c986", ">", "00000000-0000-0000-0000-000000000000"), true);
+  ASSERT_EQ(uuidEval("f768f36d-4f09-4da7-a298-3564d8f3c986", ">", "ffffffff-ffff-ffff-ffff-ffffffffffff"), false);
+
+  ASSERT_EQ(uuidEval("33355449-2c7d-43d7-967a-f53cd23215ad", "<=", "33355449-2c7d-43d7-967a-f53cd23215ad"), true);
+  ASSERT_EQ(uuidEval("33355449-2c7d-43d7-967a-f53cd23215ad", "<=", "ffffffff-ffff-ffff-ffff-ffffffffffff"), true);
+  ASSERT_EQ(uuidEval("33355449-2c7d-43d7-967a-f53cd23215ad", ">=", "33355449-2c7d-43d7-967a-f53cd23215ad"), true);
+  ASSERT_EQ(uuidEval("ffffffff-ffff-ffff-ffff-ffffffffffff", ">=", "33355449-2c7d-43d7-967a-f53cd23215ad"), true);
+
+  ASSERT_EQ(uuidEval("f768f36d-4f09-4da7-a298-3564d8f3c986", "==", "f768f36d-4f09-4da7-a298-3564d8f3c986"), true);
+  ASSERT_EQ(uuidEval("eed9f812-4b0c-472f-8a10-4ae7bff79a47", "!=", "f768f36d-4f09-4da7-a298-3564d8f3c986"), true);
+
+  ASSERT_EQ(uuidEval("11000000-0000-0022-0000-000000000000", "<", "22000000-0000-0011-0000-000000000000"), true);
+  ASSERT_EQ(uuidEval("00000000-0000-0000-2200-000000000011", ">", "00000000-0000-0000-1100-000000000022"), true);
+  ASSERT_EQ(uuidEval("00000000-0000-0000-0000-000000000011", ">", "22000000-0000-0000-0000-000000000000"), false);
+  ASSERT_EQ(uuidEval("11000000-0000-0000-0000-000000000000", "<", "00000000-0000-0000-0000-000000000022"), false);
+}
+
 } // namespace
 } // namespace facebook::velox::functions::prestosql