diff --git a/velox/docs/develop/types.rst b/velox/docs/develop/types.rst index 39166de9373e..2eb03ce5bc3f 100644 --- a/velox/docs/develop/types.rst +++ b/velox/docs/develop/types.rst @@ -136,6 +136,8 @@ HYPERLOGLOG VARBINARY JSON VARCHAR TIMESTAMP WITH TIME ZONE BIGINT UUID HUGEINT +IPADDRESS HUGEINT +IPPREFIX (HUGEINT, BIGINT) ======================== ===================== TIMESTAMP WITH TIME ZONE represents a time point in milliseconds precision diff --git a/velox/docs/functions/presto/ipaddress.rst b/velox/docs/functions/presto/ipaddress.rst new file mode 100644 index 000000000000..dcf5d4923b7b --- /dev/null +++ b/velox/docs/functions/presto/ipaddress.rst @@ -0,0 +1,51 @@ +=================== +IP Functions +=================== + +.. function:: ip_prefix(ip_address, prefix_bits) -> ipprefix + + Returns the IP prefix of a given ``ip_address`` with subnet size of ``prefix_bits``. + ``ip_address`` can be either of type ``VARCHAR`` or type ``IPADDRESS``. :: + + SELECT ip_prefix(CAST('192.168.255.255' AS IPADDRESS), 9); -- {192.128.0.0/9} + SELECT ip_prefix('2001:0db8:85a3:0001:0001:8a2e:0370:7334', 48); -- {2001:db8:85a3::/48} + +.. function:: ip_subnet_min(ip_prefix) -> ip_address + + Returns the smallest IP address of type ``IPADDRESS`` in the subnet + specified by ``ip_prefix``. :: + + SELECT ip_subnet_min(IPPREFIX '192.168.255.255/9'); -- {192.128.0.0} + SELECT ip_subnet_min(IPPREFIX '2001:0db8:85a3:0001:0001:8a2e:0370:7334/48'); -- {2001:db8:85a3::} + +.. function:: ip_subnet_max(ip_prefix) -> ip_address + + Returns the largest IP address of type ``IPADDRESS`` in the subnet + specified by ``ip_prefix``. :: + + SELECT ip_subnet_max(IPPREFIX '192.64.0.0/9'); -- {192.127.255.255} + SELECT ip_subnet_max(IPPREFIX '2001:0db8:85a3:0001:0001:8a2e:0370:7334/48'); -- {2001:db8:85a3:ffff:ffff:ffff:ffff:ffff} + +.. function:: ip_subnet_range(ip_prefix) -> array(ip_address) + + Return an array of 2 IP addresses. + The array contains the smallest and the largest IP address + in the subnet specified by ``ip_prefix``. :: + + SELECT ip_subnet_range(IPPREFIX '1.2.3.160/24'); -- [{1.2.3.0}, {1.2.3.255}] + SELECT ip_subnet_range(IPPREFIX '64:ff9b::52f4/120'); -- [{64:ff9b::5200}, {64:ff9b::52ff}] + +.. function:: is_subnet_of(ip_prefix, ip_address) -> boolean + + Returns ``true`` if the ``ip_address`` is in the subnet of ``ip_prefix``. :: + + SELECT is_subnet_of(IPPREFIX '1.2.3.128/26', IPADDRESS '1.2.3.129'); -- true + SELECT is_subnet_of(IPPREFIX '64:fa9b::17/64', IPADDRESS '64:ffff::17'); -- false + +.. function:: is_subnet_of(ip_prefix1, ip_prefix2) -> boolean + + Returns ``true`` if ``ip_prefix2`` is a subnet of ``ip_prefix1``. :: + + SELECT is_subnet_of(IPPREFIX '192.168.3.131/26', IPPREFIX '192.168.3.144/30'); -- true + SELECT is_subnet_of(IPPREFIX '64:ff9b::17/64', IPPREFIX '64:ffff::17/64'); -- false + SELECT is_subnet_of(IPPREFIX '192.168.3.131/26', IPPREFIX '192.168.3.131/26'); -- true diff --git a/velox/expression/tests/CustomTypeTest.cpp b/velox/expression/tests/CustomTypeTest.cpp index 01db687746d6..26540cd2fdd6 100644 --- a/velox/expression/tests/CustomTypeTest.cpp +++ b/velox/expression/tests/CustomTypeTest.cpp @@ -216,7 +216,8 @@ TEST_F(CustomTypeTest, getCustomTypeNames) { "HYPERLOGLOG", "TIMESTAMP WITH TIME ZONE", "UUID", - }), + "IPADDRESS", + "IPPREFIX"}), names); ASSERT_TRUE(registerCustomType( @@ -229,6 +230,8 @@ TEST_F(CustomTypeTest, getCustomTypeNames) { "HYPERLOGLOG", "TIMESTAMP WITH TIME ZONE", "UUID", + "IPADDRESS", + "IPPREFIX", "FANCY_INT", }), names); diff --git a/velox/functions/prestosql/IPAddressFunctions.h b/velox/functions/prestosql/IPAddressFunctions.h new file mode 100644 index 000000000000..2a8397e595d9 --- /dev/null +++ b/velox/functions/prestosql/IPAddressFunctions.h @@ -0,0 +1,185 @@ +/* + * 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. + */ +#pragma once + +#include "velox/functions/Macros.h" +#include "velox/functions/Registerer.h" +#include "velox/functions/lib/string/StringImpl.h" +#include "velox/functions/prestosql/types/IPPrefixType.h" + +namespace facebook::velox::functions { + +inline bool isIPV4(int128_t ip) { + int128_t ipV4 = 0x0000FFFF00000000; + uint128_t mask = 0xFFFFFFFFFFFFFFFF; + mask = (mask << 64) | 0xFFFFFFFF00000000; + return (ip & mask) == ipV4; +} + +template +struct IPPrefixFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ip, + const arg_type prefixBits) { + // Presto stores prefixBits in one signed byte. Cast to unsigned + uint8_t prefix = (uint8_t)prefixBits; + folly::ByteArray16 addrBytes; + memcpy(&addrBytes, &ip, 16); + bigEndianByteArray(addrBytes); + + // All IPs are stored as V6 + folly::IPAddressV6 v6Addr(addrBytes); + + // For return + folly::ByteArray16 canonicalBytes; + int128_t canonicalAddrInt; + + if (v6Addr.isIPv4Mapped()) { + canonicalBytes = + v6Addr.createIPv4().mask(prefix).createIPv6().toByteArray(); + } else { + canonicalBytes = v6Addr.mask(prefix).toByteArray(); + } + bigEndianByteArray(canonicalBytes); + memcpy(&canonicalAddrInt, &canonicalBytes, 16); + + result = std::make_shared(canonicalAddrInt, prefix); + } + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ip, + const arg_type prefixBits) { + int128_t intAddr; + folly::IPAddress addr(ip); + auto addrBytes = folly::IPAddress::createIPv6(addr).toByteArray(); + + bigEndianByteArray(addrBytes); + memcpy(&intAddr, &addrBytes, 16); + + call(result, intAddr, prefixBits); + } +}; + +template +struct IPSubnetMinFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix) { + // IPPrefix type should store the smallest(canonical) IP already + memcpy(&result, &ipPrefix->ip, 16); + } +}; + +inline int128_t getIPSubnetMax(int128_t ip, uint8_t prefix) { + uint128_t mask = 1; + int128_t result; + memcpy(&result, &ip, 16); + + if (isIPV4(ip)) { + result |= (mask << (32 - prefix)) - 1; + } else { + // Special case: Overflow to all 0 subtracting 1 does not work. + if (prefix == 0) { + result = -1; + } else { + result |= (mask << (128 - prefix)) - 1; + } + } + return result; +} + +template +struct IPSubnetMaxFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix) { + result = getIPSubnetMax(ipPrefix->ip, (uint8_t)ipPrefix->prefix); + } +}; + +template +struct IPSubnetRangeFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type>& result, + const arg_type& ipPrefix) { + result.push_back(ipPrefix->ip); + result.push_back(getIPSubnetMax(ipPrefix->ip, (uint8_t)ipPrefix->prefix)); + } +}; + +template +struct IPSubnetOfFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix, + const arg_type& ip) { + uint128_t mask = 1; + uint8_t prefix = (uint8_t)ipPrefix->prefix; + int128_t checkIP = ip; + + if (isIPV4(ipPrefix->ip)) { + checkIP &= ((mask << (32 - prefix)) - 1) ^ -1; + } else { + // Special case: Overflow to all 0 subtracting 1 does not work. + if (prefix == 0) { + checkIP = 0; + } else { + checkIP &= ((mask << (128 - prefix)) - 1) ^ -1; + } + } + result = (ipPrefix->ip == checkIP); + } + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix, + const arg_type& ipPrefix2) { + call(result, ipPrefix, ipPrefix2->ip); + result = result && (ipPrefix2->prefix >= ipPrefix->prefix); + } +}; + +void registerIPAddressFunctions(const std::string& prefix) { + registerIPAddressType(); + registerIPPrefixType(); + registerFunction( + {prefix + "ip_prefix"}); + registerFunction( + {prefix + "ip_prefix"}); + registerFunction( + {prefix + "ip_subnet_min"}); + registerFunction( + {prefix + "ip_subnet_max"}); + registerFunction, TheIPPrefix>( + {prefix + "ip_subnet_range"}); + registerFunction( + {prefix + "is_subnet_of"}); + registerFunction( + {prefix + "is_subnet_of"}); +} + +} // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/TypeOf.cpp b/velox/functions/prestosql/TypeOf.cpp index 048c8d031c6a..48667bb2249e 100644 --- a/velox/functions/prestosql/TypeOf.cpp +++ b/velox/functions/prestosql/TypeOf.cpp @@ -16,6 +16,8 @@ #include "velox/expression/VectorFunction.h" #include "velox/functions/prestosql/types/HyperLogLogType.h" #include "velox/functions/prestosql/types/JsonType.h" +#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/IPPrefixType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/functions/prestosql/types/UuidType.h" @@ -54,6 +56,8 @@ std::string typeName(const TypePtr& type) { case TypeKind::HUGEINT: { if (isUuidType(type)) { return "uuid"; + } else if (isIPAddressType(type)) { + return "ipaddress"; } VELOX_USER_CHECK( type->isDecimal(), @@ -104,6 +108,11 @@ std::string typeName(const TypePtr& type) { } case TypeKind::UNKNOWN: return "unknown"; + case TypeKind::OPAQUE: + if (isIPPrefixType(type)) { + return "ipprefix"; + } + return "opaque"; default: VELOX_UNSUPPORTED("Unsupported type: {}", type->toString()) } diff --git a/velox/functions/prestosql/registration/RegistrationFunctions.cpp b/velox/functions/prestosql/registration/RegistrationFunctions.cpp index 3e1b8bb647e1..8ba86c72345e 100644 --- a/velox/functions/prestosql/registration/RegistrationFunctions.cpp +++ b/velox/functions/prestosql/registration/RegistrationFunctions.cpp @@ -15,6 +15,7 @@ */ #include #include "velox/functions/prestosql/UuidFunctions.h" +#include "velox/functions/prestosql/IPAddressFunctions.h" namespace facebook::velox::functions { @@ -104,6 +105,7 @@ void registerAllScalarFunctions(const std::string& prefix) { registerGeneralFunctions(prefix); registerDateTimeFunctions(prefix); registerURLFunctions(prefix); + registerIPAddressFunctions(prefix); registerStringFunctions(prefix); registerBinaryFunctions(prefix); registerBitwiseFunctions(prefix); diff --git a/velox/functions/prestosql/tests/CMakeLists.txt b/velox/functions/prestosql/tests/CMakeLists.txt index 8fff437472c5..22a716778c78 100644 --- a/velox/functions/prestosql/tests/CMakeLists.txt +++ b/velox/functions/prestosql/tests/CMakeLists.txt @@ -63,6 +63,7 @@ add_executable( GreatestLeastTest.cpp HyperLogLogCastTest.cpp HyperLogLogFunctionsTest.cpp + IPAddressFunctionsTest.cpp InPredicateTest.cpp JsonCastTest.cpp JsonExtractScalarTest.cpp diff --git a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp new file mode 100644 index 000000000000..b2600e7936e6 --- /dev/null +++ b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp @@ -0,0 +1,463 @@ +/* + * 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/common/base/tests/GTestUtils.h" +#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" + +namespace facebook::velox::functions::prestosql { + +namespace { + +class IPAddressTest : public functions::test::FunctionBaseTest { + protected: + std::optional castIPAddress( + const std::optional input) { + auto result = evaluateOnce( + "cast(cast(c0 as ipaddress) as varchar)", input); + return result; + } + + std::optional castIPPrefix( + const std::optional input) { + auto result = evaluateOnce( + "cast(cast(c0 as ipprefix) as varchar)", input); + return result; + } + + std::optional getIPPrefix( + const std::optional input, + std::optional mask) { + auto result = evaluateOnce( + "cast(ip_prefix(cast(c0 as ipaddress), c1) as varchar)", input, mask); + return result; + } + + std::optional getIPPrefixUsingVarchar( + const std::optional input, + std::optional mask) { + auto result = evaluateOnce( + "cast(ip_prefix(c0, c1) as varchar)", input, mask); + return result; + } + + std::optional getIPSubnetMin( + const std::optional input) { + auto result = evaluateOnce( + "cast(ip_subnet_min(cast(c0 as ipprefix)) as varchar)", input); + return result; + } + + std::optional getIPSubnetMax( + const std::optional input) { + auto result = evaluateOnce( + "cast(ip_subnet_max(cast(c0 as ipprefix)) as varchar)", input); + return result; + } + + std::optional getIPSubnetRangeMin( + const std::optional input) { + auto result = evaluateOnce( + "cast(ip_subnet_range(cast(c0 as ipprefix))[1] as varchar)", input); + return result; + } + + std::optional getIPSubnetRangeMax( + const std::optional input) { + auto result = evaluateOnce( + "cast(ip_subnet_range(cast(c0 as ipprefix))[2] as varchar)", input); + return result; + } + + std::optional getIsSubnetOfIP( + const std::optional prefix, + const std::optional ip) { + auto result = evaluateOnce( + "is_subnet_of(cast(c0 as ipprefix), cast(c1 as ipaddress))", + prefix, + ip); + return result; + } + + std::optional getIsSubnetOfIPPrefix( + const std::optional prefix, + const std::optional prefix2) { + auto result = evaluateOnce( + "is_subnet_of(cast(c0 as ipprefix), cast(c1 as ipprefix))", + prefix, + prefix2); + return result; + } +}; + +TEST_F(IPAddressTest, castCircle) { + EXPECT_EQ(castIPAddress("::ffff:ffff:ffff"), "255.255.255.255"); +} + +TEST_F(IPAddressTest, castFail) { + EXPECT_THROW(castIPAddress("12.483.09.1"), VeloxUserError); + EXPECT_THROW(castIPAddress("10.135.23.12.12"), VeloxUserError); + EXPECT_THROW(castIPAddress("10.135.23"), VeloxUserError); + EXPECT_THROW( + castIPAddress("q001:0db8:85a3:0001:0001:8a2e:0370:7334"), VeloxUserError); + EXPECT_THROW( + castIPAddress("2001:0db8:85a3:542e:0001:0001:8a2e:0370:7334"), + VeloxUserError); + EXPECT_THROW( + castIPAddress("2001:0db8:85a3:0001:0001:8a2e:0370"), VeloxUserError); + + EXPECT_THROW(castIPPrefix("12.135.23.12/-1"), VeloxUserError); + EXPECT_THROW(castIPPrefix("10.135.23.12/33"), VeloxUserError); + EXPECT_THROW(castIPPrefix("::ffff:ffff:ffff/-1"), VeloxUserError); + EXPECT_THROW(castIPPrefix("::ffff:ffff:ffff/33"), VeloxUserError); + EXPECT_THROW(castIPPrefix("64:ff9b::10/-1"), VeloxUserError); + EXPECT_THROW(castIPPrefix("64:ff9b::10/129"), VeloxUserError); + EXPECT_THROW(castIPPrefix("::ffff:1.2.3.4/-1"), VeloxUserError); + EXPECT_THROW(castIPPrefix("::ffff:1.2.3.4/33"), VeloxUserError); + EXPECT_THROW(castIPPrefix("::ffff:909:909/-1"), VeloxUserError); + EXPECT_THROW(castIPPrefix("::ffff:909:909/33"), VeloxUserError); + EXPECT_THROW(castIPPrefix("64:ff9b::10/-1"), VeloxUserError); + EXPECT_THROW(castIPPrefix("64:ff9b::10/129"), VeloxUserError); + EXPECT_THROW(castIPPrefix("localhost/24"), VeloxUserError); + EXPECT_THROW(castIPPrefix("64::ff9b::10/24"), VeloxUserError); + EXPECT_THROW(castIPPrefix("64:face:book::10/24"), VeloxUserError); + EXPECT_THROW(castIPPrefix("123.456.789.012/24"), VeloxUserError); +} + +TEST_F(IPAddressTest, castRoundTrip) { + auto strings = makeFlatVector( + {"87a0:ce14:8989:44c9:826e:b4d8:73f9:1542", + "7cd6:bcec:1216:5c20:4b67:b1bd:173:ced", + "192.128.0.0"}); + + auto ipaddresses = + evaluate("cast(c0 as ipaddress)", makeRowVector({strings})); + auto stringsCopy = + evaluate("cast(c0 as varchar)", makeRowVector({ipaddresses})); + auto ipaddressesCopy = + evaluate("cast(c0 as ipaddress)", makeRowVector({stringsCopy})); + + velox::test::assertEqualVectors(strings, stringsCopy); + velox::test::assertEqualVectors(ipaddresses, ipaddressesCopy); +} + +TEST_F(IPAddressTest, IPPrefixv4) { + // EXPECT_EQ("10.0.0.0/8", getIPPrefixUsingVarchar("10.135.23.12", 8)); + + EXPECT_EQ("10.0.0.0/8", getIPPrefix("10.135.23.12", 8)); + EXPECT_EQ("192.128.0.0/9", getIPPrefix("192.168.255.255", 9)); + EXPECT_EQ("192.168.255.255/32", getIPPrefix("192.168.255.255", 32)); + EXPECT_EQ("0.0.0.0/0", getIPPrefix("192.168.255.255", 0)); + + EXPECT_THROW(getIPPrefix("12.483.09.1", 8), VeloxUserError); + EXPECT_THROW(getIPPrefix("10.135.23.12.12", 8), VeloxUserError); + EXPECT_THROW(getIPPrefix("10.135.23", 8), VeloxUserError); + EXPECT_THROW(getIPPrefix("12.135.23.12", -1), VeloxUserError); + EXPECT_THROW(getIPPrefix("10.135.23.12", 33), VeloxUserError); +} + +TEST_F(IPAddressTest, IPPrefixv4UsingVarchar) { + EXPECT_EQ("10.0.0.0/8", getIPPrefixUsingVarchar("10.135.23.12", 8)); + EXPECT_EQ("192.128.0.0/9", getIPPrefixUsingVarchar("192.168.255.255", 9)); + EXPECT_EQ( + "192.168.255.255/32", getIPPrefixUsingVarchar("192.168.255.255", 32)); + EXPECT_EQ("0.0.0.0/0", getIPPrefixUsingVarchar("192.168.255.255", 0)); + + EXPECT_THROW(getIPPrefixUsingVarchar("12.483.09.1", 8), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("10.135.23.12.12", 8), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("10.135.23", 8), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("12.135.23.12", -1), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("10.135.23.12", 33), VeloxUserError); +} + +TEST_F(IPAddressTest, IPPrefixv6) { + EXPECT_EQ( + "2001:db8:85a3::/48", + getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 48)); + EXPECT_EQ( + "2001:db8:85a3::/52", + getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 52)); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334/128", + getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 128)); + EXPECT_EQ("::/0", getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 0)); + + EXPECT_THROW( + getIPPrefix("q001:0db8:85a3:0001:0001:8a2e:0370:7334", 8), + VeloxUserError); + EXPECT_THROW( + getIPPrefix("2001:0db8:85a3:542e:0001:0001:8a2e:0370:7334", 8), + VeloxUserError); + EXPECT_THROW( + getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370", 8), VeloxUserError); + EXPECT_THROW( + getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", -1), + VeloxUserError); + EXPECT_THROW( + getIPPrefix("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 140), + VeloxUserError); +} + +TEST_F(IPAddressTest, IPPrefixv6UsingVarchar) { + EXPECT_EQ( + "2001:db8:85a3::/48", + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 48)); + EXPECT_EQ( + "2001:db8:85a3::/52", + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 52)); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334/128", + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 128)); + EXPECT_EQ( + "::/0", + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 0)); + + EXPECT_THROW( + getIPPrefixUsingVarchar("q001:0db8:85a3:0001:0001:8a2e:0370:7334", 8), + VeloxUserError); + EXPECT_THROW( + getIPPrefixUsingVarchar( + "2001:0db8:85a3:542e:0001:0001:8a2e:0370:7334", 8), + VeloxUserError); + EXPECT_THROW( + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370", 8), + VeloxUserError); + EXPECT_THROW( + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", -1), + VeloxUserError); + EXPECT_THROW( + getIPPrefixUsingVarchar("2001:0db8:85a3:0001:0001:8a2e:0370:7334", 140), + VeloxUserError); +} + +TEST_F(IPAddressTest, IPPrefixPrestoTests) { + EXPECT_EQ(getIPPrefix("1.2.3.4", 24), "1.2.3.0/24"); + EXPECT_EQ(getIPPrefix("1.2.3.4", 32), "1.2.3.4/32"); + EXPECT_EQ(getIPPrefix("1.2.3.4", 0), "0.0.0.0/0"); + EXPECT_EQ(getIPPrefix("::ffff:1.2.3.4", 24), "1.2.3.0/24"); + EXPECT_EQ(getIPPrefix("64:ff9b::17", 64), "64:ff9b::/64"); + EXPECT_EQ(getIPPrefix("64:ff9b::17", 127), "64:ff9b::16/127"); + EXPECT_EQ(getIPPrefix("64:ff9b::17", 128), "64:ff9b::17/128"); + EXPECT_EQ(getIPPrefix("64:ff9b::17", 0), "::/0"); + EXPECT_THROW(getIPPrefix("::ffff:1.2.3.4", -1), VeloxUserError); + EXPECT_THROW(getIPPrefix("::ffff:1.2.3.4", 33), VeloxUserError); + EXPECT_THROW(getIPPrefix("64:ff9b::10", -1), VeloxUserError); + EXPECT_THROW(getIPPrefix("64:ff9b::10", 129), VeloxUserError); +} + +TEST_F(IPAddressTest, IPPrefixVarcharPrestoTests) { + EXPECT_EQ(getIPPrefixUsingVarchar("1.2.3.4", 24), "1.2.3.0/24"); + EXPECT_EQ(getIPPrefixUsingVarchar("1.2.3.4", 32), "1.2.3.4/32"); + EXPECT_EQ(getIPPrefixUsingVarchar("1.2.3.4", 0), "0.0.0.0/0"); + EXPECT_EQ(getIPPrefixUsingVarchar("::ffff:1.2.3.4", 24), "1.2.3.0/24"); + EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 64), "64:ff9b::/64"); + EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 127), "64:ff9b::16/127"); + EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 128), "64:ff9b::17/128"); + EXPECT_EQ(getIPPrefixUsingVarchar("64:ff9b::17", 0), "::/0"); + EXPECT_THROW(getIPPrefixUsingVarchar("::ffff:1.2.3.4", -1), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("::ffff:1.2.3.4", 33), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("64:ff9b::10", -1), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("64:ff9b::10", 129), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("localhost", 24), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("64::ff9b::10", 24), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("64:face:book::10", 24), VeloxUserError); + EXPECT_THROW(getIPPrefixUsingVarchar("123.456.789.012", 24), VeloxUserError); +} + +TEST_F(IPAddressTest, castRoundTripPrefix) { + auto strings = makeFlatVector( + {"87a0:ce14:8989::/48", "7800::/5", "192.0.0.0/5"}); + + auto ipprefixes = evaluate("cast(c0 as ipprefix)", makeRowVector({strings})); + auto stringsCopy = + evaluate("cast(c0 as varchar)", makeRowVector({ipprefixes})); + auto ipprefixesCopy = + evaluate("cast(c0 as ipprefix)", makeRowVector({stringsCopy})); + + velox::test::assertEqualVectors(strings, stringsCopy); +} + +TEST_F(IPAddressTest, IPSubnetMin) { + EXPECT_EQ("192.0.0.0", getIPSubnetMin("192.64.1.1/9")); + EXPECT_EQ("0.0.0.0", getIPSubnetMin("192.64.1.1/0")); + EXPECT_EQ("128.0.0.0", getIPSubnetMin("192.64.1.1/1")); + EXPECT_EQ("192.64.1.0", getIPSubnetMin("192.64.1.1/31")); + EXPECT_EQ("192.64.1.1", getIPSubnetMin("192.64.1.1/32")); + + EXPECT_EQ( + "2001:db8:85a3::", + getIPSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48")); + EXPECT_EQ("::", getIPSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0")); + EXPECT_EQ("::", getIPSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + getIPSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + getIPSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128")); +} + +TEST_F(IPAddressTest, IPSubnetMinPrestoTests) { + EXPECT_EQ(getIPSubnetMin("1.2.3.4/24"), "1.2.3.0"); + EXPECT_EQ(getIPSubnetMin("1.2.3.4/32"), "1.2.3.4"); + EXPECT_EQ(getIPSubnetMin("64:ff9b::17/64"), "64:ff9b::"); + EXPECT_EQ(getIPSubnetMin("64:ff9b::17/127"), "64:ff9b::16"); + EXPECT_EQ(getIPSubnetMin("64:ff9b::17/128"), "64:ff9b::17"); + EXPECT_EQ(getIPSubnetMin("64:ff9b::17/0"), "::"); +} + +TEST_F(IPAddressTest, IPSubnetMax) { + EXPECT_EQ("192.127.255.255", getIPSubnetMax("192.64.1.1/9")); + EXPECT_EQ("255.255.255.255", getIPSubnetMax("192.64.1.1/0")); + EXPECT_EQ("255.255.255.255", getIPSubnetMax("192.64.1.1/1")); + EXPECT_EQ("192.64.1.1", getIPSubnetMax("192.64.1.1/31")); + EXPECT_EQ("192.64.1.1", getIPSubnetMax("192.64.1.1/32")); + + EXPECT_EQ( + "2001:db8:85a3:ffff:ffff:ffff:ffff:ffff", + getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48")); + EXPECT_EQ( + "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff", + getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0")); + EXPECT_EQ( + "7fff:ffff:ffff:ffff:ffff:ffff:ffff:ffff", + getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7335", + getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + getIPSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128")); +} + +TEST_F(IPAddressTest, IPSubnetMaxPrestoTests) { + EXPECT_EQ(getIPSubnetMax("1.2.3.128/26"), "1.2.3.191"); + EXPECT_EQ(getIPSubnetMax("192.168.128.4/32"), "192.168.128.4"); + EXPECT_EQ(getIPSubnetMax("10.1.16.3/9"), "10.127.255.255"); + EXPECT_EQ(getIPSubnetMax("2001:db8::16/127"), "2001:db8::17"); + EXPECT_EQ(getIPSubnetMax("2001:db8::16/128"), "2001:db8::16"); + EXPECT_EQ(getIPSubnetMax("64:ff9b::17/64"), "64:ff9b::ffff:ffff:ffff:ffff"); + EXPECT_EQ(getIPSubnetMax("64:ff9b::17/72"), "64:ff9b::ff:ffff:ffff:ffff"); + EXPECT_EQ( + getIPSubnetMax("64:ff9b::17/0"), + "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff"); +} + +TEST_F(IPAddressTest, IPSubnetRange) { + EXPECT_EQ("192.0.0.0", getIPSubnetRangeMin("192.64.1.1/9")); + EXPECT_EQ("192.127.255.255", getIPSubnetRangeMax("192.64.1.1/9")); + + EXPECT_EQ("0.0.0.0", getIPSubnetRangeMin("192.64.1.1/0")); + EXPECT_EQ("255.255.255.255", getIPSubnetRangeMax("192.64.1.1/0")); + + EXPECT_EQ("128.0.0.0", getIPSubnetRangeMin("192.64.1.1/1")); + EXPECT_EQ("255.255.255.255", getIPSubnetRangeMax("192.64.1.1/1")); + + EXPECT_EQ("192.64.1.0", getIPSubnetRangeMin("192.64.1.1/31")); + EXPECT_EQ("192.64.1.1", getIPSubnetRangeMax("192.64.1.1/31")); + + EXPECT_EQ("192.64.1.1", getIPSubnetRangeMin("192.64.1.1/32")); + EXPECT_EQ("192.64.1.1", getIPSubnetRangeMax("192.64.1.1/32")); + + EXPECT_EQ( + "2001:db8:85a3::", + getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48")); + EXPECT_EQ( + "2001:db8:85a3:ffff:ffff:ffff:ffff:ffff", + getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48")); + + EXPECT_EQ( + "::", getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0")); + EXPECT_EQ( + "::", getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + getIPSubnetRangeMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128")); + + EXPECT_EQ( + "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff", + getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0")); + EXPECT_EQ( + "7fff:ffff:ffff:ffff:ffff:ffff:ffff:ffff", + getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7335", + getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127")); + EXPECT_EQ( + "2001:db8:85a3:1:1:8a2e:370:7334", + getIPSubnetRangeMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128")); +} + +TEST_F(IPAddressTest, IPSubnetRangePrestoTests) { + EXPECT_EQ(getIPSubnetRangeMin("1.2.3.160/24"), "1.2.3.0"); + EXPECT_EQ(getIPSubnetRangeMin("1.2.3.128/31"), "1.2.3.128"); + EXPECT_EQ(getIPSubnetRangeMin("10.1.6.46/32"), "10.1.6.46"); + EXPECT_EQ(getIPSubnetRangeMin("10.1.6.46/0"), "0.0.0.0"); + EXPECT_EQ(getIPSubnetRangeMin("64:ff9b::17/64"), "64:ff9b::"); + EXPECT_EQ(getIPSubnetRangeMin("64:ff9b::52f4/120"), "64:ff9b::5200"); + EXPECT_EQ(getIPSubnetRangeMin("64:ff9b::17/128"), "64:ff9b::17"); + + EXPECT_EQ(getIPSubnetRangeMax("1.2.3.160/24"), "1.2.3.255"); + EXPECT_EQ(getIPSubnetRangeMax("1.2.3.128/31"), "1.2.3.129"); + EXPECT_EQ(getIPSubnetRangeMax("10.1.6.46/32"), "10.1.6.46"); + EXPECT_EQ(getIPSubnetRangeMax("10.1.6.46/0"), "255.255.255.255"); + EXPECT_EQ( + getIPSubnetRangeMax("64:ff9b::17/64"), "64:ff9b::ffff:ffff:ffff:ffff"); + EXPECT_EQ(getIPSubnetRangeMax("64:ff9b::52f4/120"), "64:ff9b::52ff"); + EXPECT_EQ(getIPSubnetRangeMax("64:ff9b::17/128"), "64:ff9b::17"); +} + +TEST_F(IPAddressTest, IPSubnetOfIPAddress) { + EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/26", "1.2.3.129"), true); + EXPECT_EQ(getIsSubnetOfIP("64:fa9b::17/64", "64:ffff::17"), false); +} + +TEST_F(IPAddressTest, IPSubnetOfIPPrefix) { + EXPECT_EQ( + getIsSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.144/30"), true); + EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/64", "64:ffff::17/64"), false); + EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/32", "64:ffff::17/24"), false); + EXPECT_EQ(getIsSubnetOfIPPrefix("64:ffff::17/24", "64:ff9b::17/32"), true); + EXPECT_EQ( + getIsSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.131/26"), true); +} + +TEST_F(IPAddressTest, IPSubnetOfPrestoTests) { + EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/26", "1.2.3.129"), true); + EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/26", "1.2.5.1"), false); + EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/32", "1.2.3.128"), true); + EXPECT_EQ(getIsSubnetOfIP("1.2.3.128/0", "192.168.5.1"), true); + EXPECT_EQ(getIsSubnetOfIP("64:ff9b::17/64", "64:ff9b::ffff:ff"), true); + EXPECT_EQ(getIsSubnetOfIP("64:ff9b::17/64", "64:ffff::17"), false); + + EXPECT_EQ( + getIsSubnetOfIPPrefix("192.168.3.131/26", "192.168.3.144/30"), true); + EXPECT_EQ(getIsSubnetOfIPPrefix("1.2.3.128/26", "1.2.5.1/30"), false); + EXPECT_EQ(getIsSubnetOfIPPrefix("1.2.3.128/26", "1.2.3.128/26"), true); + EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/64", "64:ff9b::ff:25/80"), true); + EXPECT_EQ(getIsSubnetOfIPPrefix("64:ff9b::17/64", "64:ffff::17/64"), false); + EXPECT_EQ( + getIsSubnetOfIPPrefix("2804:431:b000::/37", "2804:431:b000::/38"), true); + EXPECT_EQ( + getIsSubnetOfIPPrefix("2804:431:b000::/38", "2804:431:b000::/37"), false); + EXPECT_EQ(getIsSubnetOfIPPrefix("170.0.52.0/22", "170.0.52.0/24"), true); + EXPECT_EQ(getIsSubnetOfIPPrefix("170.0.52.0/24", "170.0.52.0/22"), false); +} +} // namespace + +} // namespace facebook::velox::functions::prestosql diff --git a/velox/functions/prestosql/types/CMakeLists.txt b/velox/functions/prestosql/types/CMakeLists.txt index b3481493a3c0..8b46bdf2fa5c 100644 --- a/velox/functions/prestosql/types/CMakeLists.txt +++ b/velox/functions/prestosql/types/CMakeLists.txt @@ -11,8 +11,10 @@ # 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. -add_library(velox_presto_types HyperLogLogType.cpp JsonType.cpp - TimestampWithTimeZoneType.cpp UuidType.cpp) +add_library( + velox_presto_types + HyperLogLogType.cpp JsonType.cpp TimestampWithTimeZoneType.cpp UuidType.cpp + IPAddressType.cpp IPPrefixType.cpp) target_link_libraries( velox_presto_types velox_memory velox_expression velox_functions_util diff --git a/velox/functions/prestosql/types/IPAddressType.cpp b/velox/functions/prestosql/types/IPAddressType.cpp new file mode 100644 index 000000000000..35c2ef758b75 --- /dev/null +++ b/velox/functions/prestosql/types/IPAddressType.cpp @@ -0,0 +1,141 @@ +/* + * 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/prestosql/types/IPAddressType.h" +#include + +namespace facebook::velox { + +namespace { + +class IPAddressCastOperator : public exec::CastOperator { + public: + bool isSupportedFromType(const TypePtr& other) const override { + return VARCHAR()->equivalent(*other); + } + + bool isSupportedToType(const TypePtr& other) const override { + return VARCHAR()->equivalent(*other); + } + + void castTo( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + const TypePtr& resultType, + VectorPtr& result) const override { + context.ensureWritable(rows, resultType, result); + + if (input.typeKind() == TypeKind::VARCHAR) { + castFromString(input, context, rows, *result); + } else { + VELOX_UNSUPPORTED( + "Cast from {} to IPAddress not yet supported", + resultType->toString()); + } + } + + void castFrom( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + const TypePtr& resultType, + VectorPtr& result) const override { + context.ensureWritable(rows, resultType, result); + + if (resultType->kind() == TypeKind::VARCHAR) { + castToString(input, context, rows, *result); + } else { + VELOX_UNSUPPORTED( + "Cast from IPAddress to {} not yet supported", + resultType->toString()); + } + } + + private: + static void castToString( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + BaseVector& result) { + auto* flatResult = result.as>(); + const auto* ipaddresses = input.as>(); + + context.applyToSelectedNoThrow(rows, [&](auto row) { + const auto intAddr = ipaddresses->valueAt(row); + folly::ByteArray16 addrBytes; + std::string s; + memcpy(&addrBytes, &intAddr, 16); + + bigEndianByteArray(addrBytes); + folly::IPAddressV6 v6Addr(addrBytes); + + if (v6Addr.isIPv4Mapped()) { + s = v6Addr.createIPv4().str(); + } else { + s = v6Addr.str(); + } + + exec::StringWriter result(flatResult, row); + result.append(s); + result.finalize(); + }); + } + + static void castFromString( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + BaseVector& result) { + auto* flatResult = result.as>(); + const auto* ipAddressStrings = input.as>(); + + context.applyToSelectedNoThrow(rows, [&](auto row) { + const auto ipAddressString = ipAddressStrings->valueAt(row); + int128_t intAddr; + folly::IPAddress addr(ipAddressString); + + auto addrBytes = folly::IPAddress::createIPv6(addr).toByteArray(); + + bigEndianByteArray(addrBytes); + memcpy(&intAddr, &addrBytes, 16); + + flatResult->set(row, intAddr); + }); + } +}; + +class IPAddressTypeFactories : public CustomTypeFactories { + public: + IPAddressTypeFactories() = default; + + TypePtr getType() const override { + return IPADDRESS(); + } + + exec::CastOperatorPtr getCastOperator() const override { + return std::make_shared(); + } +}; + +} // namespace + +void registerIPAddressType() { + registerCustomType( + "ipaddress", std::make_unique()); +} + +} // namespace facebook::velox diff --git a/velox/functions/prestosql/types/IPAddressType.h b/velox/functions/prestosql/types/IPAddressType.h new file mode 100644 index 000000000000..935529552f2f --- /dev/null +++ b/velox/functions/prestosql/types/IPAddressType.h @@ -0,0 +1,86 @@ +/* + * 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. + */ +#pragma once + +#include +#include +#include "velox/expression/CastExpr.h" +#include "velox/type/SimpleFunctionApi.h" +#include "velox/type/Type.h" + +namespace facebook::velox { + +// Converts BigEndian <-> native byte array +// NOOP if system is Big Endian already +inline void bigEndianByteArray(folly::ByteArray16& addrBytes) { + if (folly::kIsLittleEndian) { + std::reverse(addrBytes.begin(), addrBytes.end()); + } +} + +class IPAddressType : public HugeintType { + IPAddressType() = default; + + public: + static const std::shared_ptr& get() { + static const std::shared_ptr instance{ + new IPAddressType()}; + + return instance; + } + + bool equivalent(const Type& other) const override { + // Pointer comparison works since this type is a singleton. + return this == &other; + } + + const char* name() const override { + return "IPADDRESS"; + } + + std::string toString() const override { + return name(); + } + + folly::dynamic serialize() const override { + folly::dynamic obj = folly::dynamic::object; + obj["name"] = "Type"; + obj["type"] = name(); + return obj; + } +}; + +FOLLY_ALWAYS_INLINE bool isIPAddressType(const TypePtr& type) { + // Pointer comparison works since this type is a singleton. + return IPAddressType::get() == type; +} + +FOLLY_ALWAYS_INLINE std::shared_ptr IPADDRESS() { + return IPAddressType::get(); +} + +// Type used for function registration. +struct IPAddressT { + using type = int128_t; + static constexpr const char* typeName = "ipaddress"; +}; + +using IPAddress = CustomType; + +void registerIPAddressType(); + +} // namespace facebook::velox + diff --git a/velox/functions/prestosql/types/IPPrefixType.cpp b/velox/functions/prestosql/types/IPPrefixType.cpp new file mode 100644 index 000000000000..98ec6ac237c6 --- /dev/null +++ b/velox/functions/prestosql/types/IPPrefixType.cpp @@ -0,0 +1,151 @@ +/* + * 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/prestosql/types/IPPrefixType.h" +#include + +namespace facebook::velox { + +namespace { + +class IPPrefixCastOperator : public exec::CastOperator { + public: + bool isSupportedFromType(const TypePtr& other) const override { + return VARCHAR()->equivalent(*other); + } + + bool isSupportedToType(const TypePtr& other) const override { + return VARCHAR()->equivalent(*other); + } + + void castTo( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + const TypePtr& resultType, + VectorPtr& result) const override { + context.ensureWritable(rows, resultType, result); + + if (input.typeKind() == TypeKind::VARCHAR) { + castFromString(input, context, rows, *result); + } else { + VELOX_UNSUPPORTED( + "Cast from {} to IPPrefix not yet supported", resultType->toString()); + } + } + + void castFrom( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + const TypePtr& resultType, + VectorPtr& result) const override { + context.ensureWritable(rows, resultType, result); + + if (resultType->kind() == TypeKind::VARCHAR) { + castToString(input, context, rows, *result); + } else { + VELOX_UNSUPPORTED( + "Cast from IPPrefix to {} not yet supported", resultType->toString()); + } + } + + private: + static void castToString( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + BaseVector& result) { + auto* flatResult = result.as>(); + const auto* ipaddresses = input.as>>(); + + context.applyToSelectedNoThrow(rows, [&](auto row) { + const auto intAddr = + std::static_pointer_cast(ipaddresses->valueAt(row)); + folly::ByteArray16 addrBytes; + std::string s; + + memcpy(&addrBytes, &intAddr->ip, 16); + bigEndianByteArray(addrBytes); + folly::IPAddressV6 v6Addr(addrBytes); + + if (v6Addr.isIPv4Mapped()) { + s = v6Addr.createIPv4().str(); + } else { + s = v6Addr.str(); + } + s += "/" + std::to_string((uint8_t)intAddr->prefix); + exec::StringWriter result(flatResult, row); + result.append(s); + result.finalize(); + }); + } + + static void castFromString( + const BaseVector& input, + exec::EvalCtx& context, + const SelectivityVector& rows, + BaseVector& result) { + auto* flatResult = result.as>>(); + const auto* ipAddressStrings = input.as>(); + + context.applyToSelectedNoThrow(rows, [&](auto row) { + auto ipAddressString = ipAddressStrings->valueAt(row); + folly::CIDRNetwork net = + folly::IPAddress::createNetwork(ipAddressString, -1, false); + IPPrefix res(0, 0); + folly::ByteArray16 addrBytes; + + if (net.first.isIPv4Mapped() || net.first.isV4()) { + addrBytes = folly::IPAddress::createIPv4(net.first) + .mask(net.second) + .createIPv6() + .toByteArray(); + } else { + addrBytes = folly::IPAddress::createIPv6(net.first) + .mask(net.second) + .toByteArray(); + } + res.prefix = (uint8_t)net.second; + + bigEndianByteArray(addrBytes); + memcpy(&res.ip, &addrBytes, 16); + + flatResult->set( + row, std::make_shared(res.ip, (uint8_t)res.prefix)); + }); + } +}; + +class IPPrefixTypeFactories : public CustomTypeFactories { + public: + TypePtr getType() const override { + return IPPrefixType::get(); + } + + exec::CastOperatorPtr getCastOperator() const override { + return std::make_shared(); + } +}; + +} // namespace + +void registerIPPrefixType() { + registerCustomType( + "ipprefix", std::make_unique()); +} + +} // namespace facebook::velox diff --git a/velox/functions/prestosql/types/IPPrefixType.h b/velox/functions/prestosql/types/IPPrefixType.h new file mode 100644 index 000000000000..a55dd62029a3 --- /dev/null +++ b/velox/functions/prestosql/types/IPPrefixType.h @@ -0,0 +1,87 @@ +/* + * 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. + */ +#pragma once + +#include +#include +#include +#include +#include "velox/expression/CastExpr.h" +#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/type/OpaqueCustomTypes.h" +#include "velox/type/SimpleFunctionApi.h" +#include "velox/type/Type.h" + +#include + +namespace facebook::velox { + +struct IPPrefix { + int128_t ip; + int8_t prefix; + explicit IPPrefix(int128_t _ip, int8_t _prefix) : ip{_ip}, prefix{_prefix} {} +}; + +class IPPrefixType : public OpaqueType { + IPPrefixType() : OpaqueType(std::type_index(typeid(IPPrefix))) {} + + public: + static const std::shared_ptr& get() { + static const std::shared_ptr instance{ + new IPPrefixType()}; + + return instance; + } + + bool equivalent(const Type& other) const override { + // Pointer comparison works since this type is a singleton. + return this == &other; + } + + const char* name() const override { + return "IPPREFIX"; + } + + std::string toString() const override { + return name(); + } + + folly::dynamic serialize() const override { + folly::dynamic obj = folly::dynamic::object; + obj["name"] = "Type"; + obj["type"] = name(); + return obj; + } +}; + +FOLLY_ALWAYS_INLINE bool isIPPrefixType(const TypePtr& type) { + // Pointer comparison works since this type is a singleton. + return IPPrefixType::get() == type; +} + +FOLLY_ALWAYS_INLINE std::shared_ptr IPPREFIX() { + return IPPrefixType::get(); +} + +struct IPPrefixT { + using type = std::shared_ptr; + static constexpr const char* typeName = "ipprefix"; +}; +using TheIPPrefix = CustomType; + +void registerIPPrefixType(); + +} // namespace facebook::velox diff --git a/velox/functions/prestosql/types/tests/CMakeLists.txt b/velox/functions/prestosql/types/tests/CMakeLists.txt index b7cbaa4f4336..100a049e709f 100644 --- a/velox/functions/prestosql/types/tests/CMakeLists.txt +++ b/velox/functions/prestosql/types/tests/CMakeLists.txt @@ -14,8 +14,13 @@ add_executable( velox_presto_types_test - HyperLogLogTypeTest.cpp JsonTypeTest.cpp TimestampWithTimeZoneTypeTest.cpp - TypeTestBase.cpp UuidTypeTest.cpp) + HyperLogLogTypeTest.cpp + JsonTypeTest.cpp + TimestampWithTimeZoneTypeTest.cpp + TypeTestBase.cpp + UuidTypeTest.cpp + IPAddressTypeTest.cpp + IPPrefixTypeTest.cpp) add_test(velox_presto_types_test velox_presto_types_test) diff --git a/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp b/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp new file mode 100644 index 000000000000..a24ce1fec3a6 --- /dev/null +++ b/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp @@ -0,0 +1,41 @@ +/* + * 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/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/tests/TypeTestBase.h" + +namespace facebook::velox::test { + +class IPAddressTypeTest : public testing::Test, public TypeTestBase { + public: + IPAddressTypeTest() { + registerIPAddressType(); + } +}; + +TEST_F(IPAddressTypeTest, basic) { + ASSERT_EQ(IPADDRESS()->name(), "IPADDRESS"); + ASSERT_EQ(IPADDRESS()->kindName(), "HUGEINT"); + ASSERT_TRUE(IPADDRESS()->parameters().empty()); + ASSERT_EQ(IPADDRESS()->toString(), "IPADDRESS"); + + ASSERT_TRUE(hasType("IPADDRESS")); + ASSERT_EQ(*getType("IPADDRESS", {}), *IPADDRESS()); +} + +TEST_F(IPAddressTypeTest, serde) { + testTypeSerde(IPADDRESS()); +} +} // namespace facebook::velox::test diff --git a/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp b/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp new file mode 100644 index 000000000000..a8f92d99dc5a --- /dev/null +++ b/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp @@ -0,0 +1,40 @@ +/* + * 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/prestosql/types/IPPrefixType.h" +#include "velox/functions/prestosql/types/tests/TypeTestBase.h" + +namespace facebook::velox::test { + +class IPPrefixTypeTest : public testing::Test, public TypeTestBase { + public: + IPPrefixTypeTest() { + registerIPPrefixType(); + } +}; + +TEST_F(IPPrefixTypeTest, basic) { + ASSERT_EQ(IPPREFIX()->name(), "IPPREFIX"); + ASSERT_TRUE(IPPREFIX()->parameters().empty()); + ASSERT_EQ(IPPREFIX()->toString(), "IPPREFIX"); + + ASSERT_TRUE(hasType("IPPREFIX")); + ASSERT_EQ(*getType("IPPREFIX", {}), *IPPREFIX()); +} + +TEST_F(IPPrefixTypeTest, serde) { + testTypeSerde(IPPREFIX()); +} +} // namespace facebook::velox::test