From fa1b203335734e980acab6a5fc68123aeaf747ac Mon Sep 17 00:00:00 2001
From: Bryan Cutler <cutlerb@gmail.com>
Date: Tue, 20 Aug 2024 14:51:51 -0700
Subject: [PATCH] feat: Add support for UUID comparison functions

This adds UUID comparison functions that were previously unsupported.
Functions added are <, >, <=, >=. Equal was already supported. The
ordering is done lexicographically and conforms to IETF [RFC 4122].
The ordering also matches Presto Java after #[23847].

[RFC 4122]: https://datatracker.ietf.org/doc/html/rfc4122.html
[23847]: https://github.com/prestodb/presto/pull/23847
---
 velox/functions/prestosql/UuidFunctions.h     |  23 ++++
 .../prestosql/tests/UuidFunctionsTest.cpp     | 114 ++++++++++++++++++
 2 files changed, 137 insertions(+)

diff --git a/velox/functions/prestosql/UuidFunctions.h b/velox/functions/prestosql/UuidFunctions.h
index 9e4b319a7fff..4da630e986f2 100644
--- a/velox/functions/prestosql/UuidFunctions.h
+++ b/velox/functions/prestosql/UuidFunctions.h
@@ -20,10 +20,29 @@
 
 #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::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) { \
+      result = (uuidCompExpr);                                                 \
+    }                                                                          \
+  };
+
+VELOX_GEN_BINARY_EXPR_UUID(LtFunction, (uint128_t)lhs < (uint128_t)rhs);
+VELOX_GEN_BINARY_EXPR_UUID(GtFunction, (uint128_t)lhs > (uint128_t)rhs);
+VELOX_GEN_BINARY_EXPR_UUID(LteFunction, (uint128_t)lhs <= (uint128_t)rhs);
+VELOX_GEN_BINARY_EXPR_UUID(GteFunction, (uint128_t)lhs >= (uint128_t)rhs);
+
+#undef VELOX_GEN_BINARY_EXPR_UUID
+
 template <typename T>
 struct UuidFunction {
   VELOX_DEFINE_FUNCTION_TYPES(T);
@@ -42,6 +61,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
index 4fc636d966f3..6afbe2b82747 100644
--- a/velox/functions/prestosql/tests/UuidFunctionsTest.cpp
+++ b/velox/functions/prestosql/tests/UuidFunctionsTest.cpp
@@ -108,5 +108,119 @@ 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(
+      true,
+      uuidEval(
+          "33355449-2c7d-43d7-967a-f53cd23215ad",
+          "<",
+          "ffffffff-ffff-ffff-ffff-ffffffffffff"));
+  ASSERT_EQ(
+      false,
+      uuidEval(
+          "33355449-2c7d-43d7-967a-f53cd23215ad",
+          "<",
+          "00000000-0000-0000-0000-000000000000"));
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "f768f36d-4f09-4da7-a298-3564d8f3c986",
+          ">",
+          "00000000-0000-0000-0000-000000000000"));
+  ASSERT_EQ(
+      false,
+      uuidEval(
+          "f768f36d-4f09-4da7-a298-3564d8f3c986",
+          ">",
+          "ffffffff-ffff-ffff-ffff-ffffffffffff"));
+
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "33355449-2c7d-43d7-967a-f53cd23215ad",
+          "<=",
+          "33355449-2c7d-43d7-967a-f53cd23215ad"));
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "33355449-2c7d-43d7-967a-f53cd23215ad",
+          "<=",
+          "ffffffff-ffff-ffff-ffff-ffffffffffff"));
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "33355449-2c7d-43d7-967a-f53cd23215ad",
+          ">=",
+          "33355449-2c7d-43d7-967a-f53cd23215ad"));
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "ffffffff-ffff-ffff-ffff-ffffffffffff",
+          ">=",
+          "33355449-2c7d-43d7-967a-f53cd23215ad"));
+
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "f768f36d-4f09-4da7-a298-3564d8f3c986",
+          "==",
+          "f768f36d-4f09-4da7-a298-3564d8f3c986"));
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "eed9f812-4b0c-472f-8a10-4ae7bff79a47",
+          "!=",
+          "f768f36d-4f09-4da7-a298-3564d8f3c986"));
+
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "11000000-0000-0022-0000-000000000000",
+          "<",
+          "22000000-0000-0011-0000-000000000000"));
+  ASSERT_EQ(
+      true,
+      uuidEval(
+          "00000000-0000-0000-2200-000000000011",
+          ">",
+          "00000000-0000-0000-1100-000000000022"));
+  ASSERT_EQ(
+      false,
+      uuidEval(
+          "00000000-0000-0000-0000-000000000011",
+          ">",
+          "22000000-0000-0000-0000-000000000000"));
+  ASSERT_EQ(
+      false,
+      uuidEval(
+          "11000000-0000-0000-0000-000000000000",
+          "<",
+          "00000000-0000-0000-0000-000000000022"));
+
+  std::string lhs = "12342345-3456-4567-5678-678978908901";
+  std::string rhs = "23451234-4567-3456-6789-567889017890";
+  ASSERT_EQ(true, uuidEval(lhs, "<", rhs));
+
+  for (vector_size_t i = 0; i < lhs.size(); i++) {
+    if (lhs[i] == '-') {
+      continue;
+    }
+    lhs[i] = '0';
+    rhs[i] = '0';
+    bool expected = boost::lexical_cast<boost::uuids::uuid>(lhs) <
+        boost::lexical_cast<boost::uuids::uuid>(rhs);
+    ASSERT_EQ(expected, uuidEval(lhs, "<", rhs));
+  }
+}
+
 } // namespace
 } // namespace facebook::velox::functions::prestosql