Support comparison for ipprefix [3/n] (#11546)

Summary: We need to support comparison for IPPrefix, which has underlying type of RowVector<HUGEINT, TINYINT>. Since the first 16 bytes of IPPrefix is pretty much the same as IPAddress, we can just convert the underlying type of IPPrefix to RowVector<IPADDRESS, TINYINT>. Since IPADDRESS already has its custom comparison defined, and rowVector by default compares left to right, this takes of comparison for us for free. Differential Revision: D65521072
facebookincubator · Dec 8, 2024 · 8726e15 · 8726e15
1 parent 929affe
commit 8726e15
Show file tree

Hide file tree

Showing 4 changed files with 292 additions and 1 deletion.
diff --git a/velox/expression/fuzzer/ExpressionFuzzer.cpp b/velox/expression/fuzzer/ExpressionFuzzer.cpp
@@ -664,6 +664,7 @@ bool ExpressionFuzzer::isSupportedSignature(
   if (usesTypeName(signature, "opaque") ||
       usesTypeName(signature, "timestamp with time zone") ||
       usesTypeName(signature, "interval day to second") ||
+      usesTypeName(signature, "ipprefix") ||
       (!options_.enableDecimalType && usesTypeName(signature, "decimal")) ||
       (!options_.enableComplexTypes && useComplexType) ||
       (options_.enableComplexTypes && usesTypeName(signature, "unknown"))) {

diff --git a/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp
@@ -17,6 +17,7 @@
 #include "velox/functions/lib/RegistrationHelpers.h"
 #include "velox/functions/prestosql/Comparisons.h"
 #include "velox/functions/prestosql/types/IPAddressType.h"
+#include "velox/functions/prestosql/types/IPPrefixType.h"
 #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
 #include "velox/type/Type.h"
 
@@ -32,6 +33,7 @@ void registerNonSimdizableScalar(const std::vector<std::string>& aliases) {
   registerFunction<T, TReturn, TimestampWithTimezone, TimestampWithTimezone>(
       aliases);
   registerFunction<T, TReturn, IPAddress, IPAddress>(aliases);
+  registerFunction<T, TReturn, IPPrefix, IPPrefix>(aliases);
 }
 } // namespace
 
@@ -40,6 +42,7 @@ void registerComparisonFunctions(const std::string& prefix) {
   // independent of DateTimeFunctions
   registerTimestampWithTimeZoneType();
   registerIPAddressType();
+  registerIPPrefixType();
 
   registerNonSimdizableScalar<EqFunction, bool>({prefix + "eq"});
   VELOX_REGISTER_VECTOR_FUNCTION(udf_simd_comparison_eq, prefix + "eq");
@@ -121,6 +124,8 @@ void registerComparisonFunctions(const std::string& prefix) {
       TimestampWithTimezone>({prefix + "between"});
   registerFunction<BetweenFunction, bool, IPAddress, IPAddress, IPAddress>(
       {prefix + "between"});
+  registerFunction<BetweenFunction, bool, IPPrefix, IPPrefix, IPPrefix>(
+      {prefix + "between"});
 }
 
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/tests/ComparisonsTest.cpp b/velox/functions/prestosql/tests/ComparisonsTest.cpp
@@ -20,6 +20,7 @@
 #include "velox/functions/lib/RegistrationHelpers.h"
 #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
 #include "velox/functions/prestosql/types/IPAddressType.h"
+#include "velox/functions/prestosql/types/IPPrefixType.h"
 #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h"
 #include "velox/type/tests/utils/CustomTypesForTesting.h"
 #include "velox/type/tz/TimeZoneMap.h"
@@ -1183,6 +1184,290 @@ TEST_F(ComparisonsTest, TimestampWithTimezone) {
            false}));
 }
 
+TEST_F(ComparisonsTest, IPPrefixType) {
+  auto ipprefix = [](const std::string& ipprefixString) {
+    auto tryIpPrefix = ipaddress::tryParseIpPrefixString(ipprefixString);
+    return variant::row(
+        {tryIpPrefix.value().first, tryIpPrefix.value().second});
+  };
+
+  // Comparison Operator Tests
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")}});
+    auto expected = makeFlatVector<bool>(
+        {false, false, true, true, true, false, false, true, true});
+    auto result = evaluate("c0 > c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")}});
+    auto expected = makeFlatVector<bool>(
+        {false, true, false, false, false, false, true, false, false});
+    auto result = evaluate("c0 < c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")}});
+    auto expected = makeFlatVector<bool>(
+        {true, false, true, true, true, true, false, true, true});
+    auto result = evaluate("c0 >= c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")}});
+    auto expected = makeFlatVector<bool>(
+        {true, true, false, false, false, true, true, false, false});
+    auto result = evaluate("c0 <= c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/24")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")},
+         {ipprefix("1.2.3.5/24")}});
+    auto expected = makeFlatVector<bool>(
+        {true, false, false, false, false, true, false, false, false, true});
+    auto result = evaluate("c0 = c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")}});
+    auto expected = makeFlatVector<bool>(
+        {false, true, true, true, true, false, true, true, true});
+    auto result = evaluate("c0 <> c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  // Distinct from test
+  {
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")}});
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("1.2.0.0/24")},
+         {ipprefix("::1/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")}});
+    auto expected = makeFlatVector<bool>(
+        {false, true, true, true, true, false, true, true, true});
+    auto result =
+        evaluate("c0 is distinct from c1", makeRowVector({left, right}));
+    test::assertEqualVectors(result, expected);
+  }
+
+  // Inbetween test
+  {
+    auto inbetween = makeArrayOfRowVector(
+        IPPREFIX(),
+        {{ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.4/32")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("1.2.3.5/32")},
+         {ipprefix("1.2.0.0/25")},
+         {ipprefix("::1/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80")},
+         {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+         {ipprefix("::1/128")},
+         {ipprefix("::2222/128")}
+
+        });
+    auto left = makeArrayOfRowVector(
+        IPPREFIX(),
+        {
+            {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+            {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+            {ipprefix("1.2.3.4/32")},
+            {ipprefix("1.2.3.4/32")},
+            {ipprefix("1.2.0.0/24")},
+            {ipprefix("::1/128")},
+            {ipprefix("1.2.3.4/32")},
+            {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+            {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")},
+            {ipprefix("::1/128")},
+            {ipprefix("::1/128")},
+        });
+    auto right = makeArrayOfRowVector(
+        IPPREFIX(),
+        {
+            {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+            {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8329/128")},
+            {ipprefix("1.2.3.4/32")},
+            {ipprefix("1.2.3.4/32")},
+            {ipprefix("1.2.0.0/24")},
+            {ipprefix("::1/128")},
+            {ipprefix("1.2.3.4/32")},
+            {ipprefix("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64")},
+            {ipprefix("2001:0db8:0000:0000:0000:ff00:0042:8321/128")},
+            {ipprefix("::1234/128")},
+            {ipprefix("::1234/128")},
+
+        });
+    auto expected = makeFlatVector<bool>(
+        {true,
+         false,
+         false,
+         false,
+         false,
+         true,
+         false,
+         false,
+         false,
+         true,
+         false});
+
+    auto result = evaluate(
+        "c0 between c1 and c2",
+        makeRowVector(
+            {inbetween->elements(), left->elements(), right->elements()}));
+    test::assertEqualVectors(result, expected);
+  }
+}
+
 TEST_F(ComparisonsTest, IpAddressType) {
   auto makeIpAdressFromString = [](const std::string& ipAddr) -> int128_t {
     auto ret = ipaddress::tryGetIPv6asInt128FromString(ipAddr);

diff --git a/velox/functions/prestosql/types/IPPrefixType.h b/velox/functions/prestosql/types/IPPrefixType.h
@@ -98,7 +98,7 @@ class IPPrefixType : public RowType {
   IPPrefixType()
       : RowType(
             {ipaddress::kIpRowIndex, ipaddress::kIpPrefixRowIndex},
-            {HUGEINT(), TINYINT()}) {}
+            {IPADDRESS(), TINYINT()}) {}
 
  public:
   static const std::shared_ptr<const IPPrefixType>& get() {