From d2d01f5ad53360b6f208c5346fd07dfa3435471f Mon Sep 17 00:00:00 2001 From: Yenda Li Date: Thu, 5 Dec 2024 10:34:48 -0800 Subject: [PATCH] Support comparison for ipprefix [3/n] (#11546) Summary: We need to support comparison for IPPrefix, which has underlying type of RowVector. 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. 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 --- .../ComparisonFunctionsRegistration.cpp | 5 + .../prestosql/tests/ComparisonsTest.cpp | 283 ++++++++++++++++++ .../functions/prestosql/types/IPPrefixType.h | 2 +- 3 files changed, 289 insertions(+), 1 deletion(-) diff --git a/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp index 2870a3ff787f0..153d7d140a841 100644 --- 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& aliases) { registerFunction( aliases); registerFunction(aliases); + registerFunction(aliases); } } // namespace @@ -40,6 +42,7 @@ void registerComparisonFunctions(const std::string& prefix) { // independent of DateTimeFunctions registerTimestampWithTimeZoneType(); registerIPAddressType(); + registerIPPrefixType(); registerNonSimdizableScalar({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( {prefix + "between"}); + registerFunction( + {prefix + "between"}); } } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/tests/ComparisonsTest.cpp b/velox/functions/prestosql/tests/ComparisonsTest.cpp index 707a61d108fa1..f10cd2cff7854 100644 --- 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,288 @@ 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( + {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( + {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( + {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( + {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")}}); + 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( + {true, false, false, false, false, true, false, 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( + {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( + {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( + {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 index 5ece386a4a3ab..d345e1c17af4f 100644 --- 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& get() {