From 60f34ea7103e1c09d94b697badba4ac12c02301c Mon Sep 17 00:00:00 2001 From: Yenda Li Date: Wed, 11 Dec 2024 08:16:45 -0800 Subject: [PATCH] Add ip_prefix function [4/n] (#11514) Summary: Pull Request resolved: https://github.com/facebookincubator/velox/pull/11514 Add ip_prefix function which takes in varchar or ipaddress along with a prefixBits. Use folly lib to convert the varchar or ipaddress type to folly::IPAddressV6, and then apply the mask. We return a tuple which corresponds to the ipaddress and prefix for IPAddressPrefix type. Split from https://github.com/facebookincubator/velox/pull/11407 Differential Revision: D65802211 --- velox/docs/functions/presto/ipaddress.rst | 12 ++ .../functions/prestosql/IPAddressFunctions.h | 59 +++++++++ .../tests/IPAddressFunctionsTest.cpp | 114 ++++++++++++++++++ 3 files changed, 185 insertions(+) create mode 100644 velox/docs/functions/presto/ipaddress.rst create mode 100644 velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp diff --git a/velox/docs/functions/presto/ipaddress.rst b/velox/docs/functions/presto/ipaddress.rst new file mode 100644 index 000000000000..6b9140dec883 --- /dev/null +++ b/velox/docs/functions/presto/ipaddress.rst @@ -0,0 +1,12 @@ +=================== +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} + diff --git a/velox/functions/prestosql/IPAddressFunctions.h b/velox/functions/prestosql/IPAddressFunctions.h index a08ef92de534..256ee61bd0af 100644 --- a/velox/functions/prestosql/IPAddressFunctions.h +++ b/velox/functions/prestosql/IPAddressFunctions.h @@ -15,14 +15,73 @@ */ #pragma once +#include "velox/functions/Macros.h" +#include "velox/functions/Registerer.h" #include "velox/functions/prestosql/types/IPAddressType.h" #include "velox/functions/prestosql/types/IPPrefixType.h" namespace facebook::velox::functions { +template +struct IPPrefixFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ip, + const arg_type& prefixBits) { + folly::ByteArray16 addrBytes; + memcpy(&addrBytes, &ip, ipaddress::kIPAddressBytes); + std::reverse(addrBytes.begin(), addrBytes.end()); + + result = makeIPPrefix(folly::IPAddressV6(addrBytes), prefixBits); + } + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipString, + const arg_type& prefixBits) { + auto tryIp = folly::IPAddress::tryFromString(ipString); + if (tryIp.hasError()) { + VELOX_USER_FAIL("Cannot cast value to IPADDRESS: {}", ipString); + } + + result = makeIPPrefix( + folly::IPAddress::createIPv6(folly::IPAddress(tryIp.value())), + prefixBits); + } + + private: + static std::tuple makeIPPrefix( + folly::IPAddressV6 v6Addr, + int64_t prefixBits) { + if (v6Addr.isIPv4Mapped()) { + VELOX_CHECK( + 0 <= prefixBits && prefixBits <= ipaddress::kIPV4Bits, + "IPv4 subnet size must be in range [0, 32]"); + } else { + VELOX_CHECK( + 0 <= prefixBits && prefixBits <= ipaddress::kIPV6Bits, + "IPv6 subnet size must be in range [0, 128]"); + } + auto canonicalBytes = v6Addr.isIPv4Mapped() + ? v6Addr.createIPv4().mask(prefixBits).createIPv6().toByteArray() + : v6Addr.mask(prefixBits).toByteArray(); + + int128_t intAddr; + std::reverse(canonicalBytes.begin(), canonicalBytes.end()); + memcpy(&intAddr, &canonicalBytes, ipaddress::kIPAddressBytes); + return std::make_tuple(intAddr, static_cast(prefixBits)); + } +}; + void registerIPAddressFunctions(const std::string& prefix) { registerIPAddressType(); registerIPPrefixType(); + registerFunction( + {prefix + "ip_prefix"}); + registerFunction( + {prefix + "ip_prefix"}); } } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp new file mode 100644 index 000000000000..532da9e02409 --- /dev/null +++ b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp @@ -0,0 +1,114 @@ +/* + * 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 { +class IPAddressFunctionsTest : public functions::test::FunctionBaseTest { + protected: + std::optional ipPrefixFunctionFromIpAddress( + const std::optional& input, + const std::optional& mask) { + return evaluateOnce( + "cast(ip_prefix(cast(c0 as ipaddress), c1) as varchar)", input, mask); + } + + std::optional ipPrefixFromVarChar( + const std::optional& input, + const std::optional& mask) { + return evaluateOnce( + "cast(ip_prefix(c0, c1) as varchar)", input, mask); + } +}; + +TEST_F(IPAddressFunctionsTest, ipPrefixFromIpAddress) { + ASSERT_EQ(ipPrefixFunctionFromIpAddress("1.2.3.4", 24), "1.2.3.0/24"); + ASSERT_EQ(ipPrefixFunctionFromIpAddress("1.2.3.4", 32), "1.2.3.4/32"); + ASSERT_EQ(ipPrefixFunctionFromIpAddress("1.2.3.4", 0), "0.0.0.0/0"); + ASSERT_EQ(ipPrefixFunctionFromIpAddress("::ffff:1.2.3.4", 24), "1.2.3.0/24"); + ASSERT_EQ(ipPrefixFunctionFromIpAddress("64:ff9b::17", 64), "64:ff9b::/64"); + ASSERT_EQ( + ipPrefixFunctionFromIpAddress("64:ff9b::17", 127), "64:ff9b::16/127"); + ASSERT_EQ( + ipPrefixFunctionFromIpAddress("64:ff9b::17", 128), "64:ff9b::17/128"); + ASSERT_EQ(ipPrefixFunctionFromIpAddress("64:ff9b::17", 0), "::/0"); + ASSERT_EQ( + ipPrefixFunctionFromIpAddress( + "2001:0db8:85a3:0001:0001:8a2e:0370:7334", 48), + "2001:db8:85a3::/48"); + ASSERT_EQ( + ipPrefixFunctionFromIpAddress( + "2001:0db8:85a3:0001:0001:8a2e:0370:7334", 52), + "2001:db8:85a3::/52"); + ASSERT_EQ( + ipPrefixFunctionFromIpAddress( + "2001:0db8:85a3:0001:0001:8a2e:0370:7334", 128), + "2001:db8:85a3:1:1:8a2e:370:7334/128"); + ASSERT_EQ( + ipPrefixFunctionFromIpAddress( + "2001:0db8:85a3:0001:0001:8a2e:0370:7334", 0), + "::/0"); + VELOX_ASSERT_THROW( + ipPrefixFunctionFromIpAddress("::ffff:1.2.3.4", -1), + "IPv4 subnet size must be in range [0, 32]"); + VELOX_ASSERT_THROW( + ipPrefixFunctionFromIpAddress("::ffff:1.2.3.4", 33), + "IPv4 subnet size must be in range [0, 32]"); + VELOX_ASSERT_THROW( + ipPrefixFunctionFromIpAddress("64:ff9b::10", -1), + "IPv6 subnet size must be in range [0, 128]"); + VELOX_ASSERT_THROW( + ipPrefixFunctionFromIpAddress("64:ff9b::10", 129), + "IPv6 subnet size must be in range [0, 128]"); +} + +TEST_F(IPAddressFunctionsTest, ipPrefixFromVarChar) { + ASSERT_EQ(ipPrefixFromVarChar("1.2.3.4", 24), "1.2.3.0/24"); + ASSERT_EQ(ipPrefixFromVarChar("1.2.3.4", 32), "1.2.3.4/32"); + ASSERT_EQ(ipPrefixFromVarChar("1.2.3.4", 0), "0.0.0.0/0"); + ASSERT_EQ(ipPrefixFromVarChar("::ffff:1.2.3.4", 24), "1.2.3.0/24"); + ASSERT_EQ(ipPrefixFromVarChar("64:ff9b::17", 64), "64:ff9b::/64"); + ASSERT_EQ(ipPrefixFromVarChar("64:ff9b::17", 127), "64:ff9b::16/127"); + ASSERT_EQ(ipPrefixFromVarChar("64:ff9b::17", 128), "64:ff9b::17/128"); + ASSERT_EQ(ipPrefixFromVarChar("64:ff9b::17", 0), "::/0"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("::ffff:1.2.3.4", -1), + "IPv4 subnet size must be in range [0, 32]"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("::ffff:1.2.3.4", 33), + "IPv4 subnet size must be in range [0, 32]"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("64:ff9b::10", -1), + "IPv6 subnet size must be in range [0, 128]"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("64:ff9b::10", 129), + "IPv6 subnet size must be in range [0, 128]"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("localhost", 24), + "Cannot cast value to IPADDRESS: localhost"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("64::ff9b::10", 24), + "Cannot cast value to IPADDRESS: 64::ff9b::10"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("64:face:book::10", 24), + "Cannot cast value to IPADDRESS: 64:face:book::10"); + VELOX_ASSERT_THROW( + ipPrefixFromVarChar("123.456.789.012", 24), + "Cannot cast value to IPADDRESS: 123.456.789.012"); +} + +} // namespace facebook::velox::functions::prestosql