diff --git a/velox/docs/functions/presto/ipaddress.rst b/velox/docs/functions/presto/ipaddress.rst index 6b9140dec883..b497c1043dac 100644 --- a/velox/docs/functions/presto/ipaddress.rst +++ b/velox/docs/functions/presto/ipaddress.rst @@ -10,3 +10,19 @@ IP Functions 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} + diff --git a/velox/functions/prestosql/IPAddressFunctions.h b/velox/functions/prestosql/IPAddressFunctions.h index 6ce0acbc0ef7..0d14eba00145 100644 --- a/velox/functions/prestosql/IPAddressFunctions.h +++ b/velox/functions/prestosql/IPAddressFunctions.h @@ -75,6 +75,53 @@ struct IPPrefixFunction { } }; +template +struct IPSubnetMinFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix) { + // IPPrefix type stores the smallest(canonical) IP already + result = *ipPrefix.template at<0>(); + } +}; + +template +struct IPSubnetMaxFunction { + VELOX_DEFINE_FUNCTION_TYPES(T); + + FOLLY_ALWAYS_INLINE void call( + out_type& result, + const arg_type& ipPrefix) { + result = + getIPSubnetMax(*ipPrefix.template at<0>(), *ipPrefix.template at<1>()); + } + + private: + static int128_t getIPSubnetMax(int128_t ip, uint8_t prefix) { + auto tryIsIpv4 = ipaddress::tryIpPrefixLengthFromIPAddressType(ip); + // This check should never fail because we're taking a pre-existing + // IPPrefix. + VELOX_DCHECK(tryIsIpv4.hasValue()); + + const bool isIpV4 = tryIsIpv4.value() == ipaddress::kIPV4Bits; + uint128_t mask = 1; + if (isIpV4) { + ip |= (mask << (ipaddress::kIPV4Bits - prefix)) - 1; + return ip; + } + + // Special case: Overflow to all 0 subtracting 1 does not work. + if (prefix == 0) { + return -1; + } + + ip |= (mask << (ipaddress::kIPV6Bits - prefix)) - 1; + return ip; + } +}; + void registerIPAddressFunctions(const std::string& prefix) { registerIPAddressType(); registerIPPrefixType(); @@ -82,6 +129,10 @@ void registerIPAddressFunctions(const std::string& prefix) { {prefix + "ip_prefix"}); registerFunction( {prefix + "ip_prefix"}); + registerFunction( + {prefix + "ip_subnet_min"}); + registerFunction( + {prefix + "ip_subnet_max"}); } } // namespace facebook::velox::functions diff --git a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp index 532da9e02409..95f1dd2f8010 100644 --- a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp +++ b/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp @@ -33,6 +33,18 @@ class IPAddressFunctionsTest : public functions::test::FunctionBaseTest { return evaluateOnce( "cast(ip_prefix(c0, c1) as varchar)", input, mask); } + + std::optional ipSubnetMin( + const std::optional& input) { + return evaluateOnce( + "cast(ip_subnet_min(cast(c0 as ipprefix)) as varchar)", input); + } + + std::optional ipSubnetMax( + const std::optional& input) { + return evaluateOnce( + "cast(ip_subnet_max(cast(c0 as ipprefix)) as varchar)", input); + } }; TEST_F(IPAddressFunctionsTest, ipPrefixFromIpAddress) { @@ -111,4 +123,61 @@ TEST_F(IPAddressFunctionsTest, ipPrefixFromVarChar) { "Cannot cast value to IPADDRESS: 123.456.789.012"); } +TEST_F(IPAddressFunctionsTest, ipSubnetMin) { + ASSERT_EQ(ipSubnetMin("1.2.3.4/24"), "1.2.3.0"); + ASSERT_EQ(ipSubnetMin("1.2.3.4/32"), "1.2.3.4"); + ASSERT_EQ(ipSubnetMin("64:ff9b::17/64"), "64:ff9b::"); + ASSERT_EQ(ipSubnetMin("64:ff9b::17/127"), "64:ff9b::16"); + ASSERT_EQ(ipSubnetMin("64:ff9b::17/128"), "64:ff9b::17"); + ASSERT_EQ(ipSubnetMin("64:ff9b::17/0"), "::"); + ASSERT_EQ(ipSubnetMin("192.64.1.1/9"), "192.0.0.0"); + ASSERT_EQ(ipSubnetMin("192.64.1.1/0"), "0.0.0.0"); + ASSERT_EQ(ipSubnetMin("192.64.1.1/1"), "128.0.0.0"); + ASSERT_EQ(ipSubnetMin("192.64.1.1/31"), "192.64.1.0"); + ASSERT_EQ(ipSubnetMin("192.64.1.1/32"), "192.64.1.1"); + ASSERT_EQ( + ipSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48"), + "2001:db8:85a3::"); + ASSERT_EQ(ipSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0"), "::"); + ASSERT_EQ(ipSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1"), "::"); + ASSERT_EQ( + ipSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127"), + "2001:db8:85a3:1:1:8a2e:370:7334"); + ASSERT_EQ( + ipSubnetMin("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128"), + "2001:db8:85a3:1:1:8a2e:370:7334"); +} + +TEST_F(IPAddressFunctionsTest, ipSubnetMax) { + ASSERT_EQ(ipSubnetMax("1.2.3.128/26"), "1.2.3.191"); + ASSERT_EQ(ipSubnetMax("192.168.128.4/32"), "192.168.128.4"); + ASSERT_EQ(ipSubnetMax("10.1.16.3/9"), "10.127.255.255"); + ASSERT_EQ(ipSubnetMax("2001:db8::16/127"), "2001:db8::17"); + ASSERT_EQ(ipSubnetMax("2001:db8::16/128"), "2001:db8::16"); + ASSERT_EQ(ipSubnetMax("64:ff9b::17/64"), "64:ff9b::ffff:ffff:ffff:ffff"); + ASSERT_EQ(ipSubnetMax("64:ff9b::17/72"), "64:ff9b::ff:ffff:ffff:ffff"); + ASSERT_EQ( + ipSubnetMax("64:ff9b::17/0"), "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff"); + ASSERT_EQ(ipSubnetMax("192.64.1.1/9"), "192.127.255.255"); + ASSERT_EQ(ipSubnetMax("192.64.1.1/0"), "255.255.255.255"); + ASSERT_EQ(ipSubnetMax("192.64.1.1/1"), "255.255.255.255"); + ASSERT_EQ(ipSubnetMax("192.64.1.1/31"), "192.64.1.1"); + ASSERT_EQ(ipSubnetMax("192.64.1.1/32"), "192.64.1.1"); + ASSERT_EQ( + ipSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/48"), + "2001:db8:85a3:ffff:ffff:ffff:ffff:ffff"); + ASSERT_EQ( + ipSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/0"), + "ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff"); + ASSERT_EQ( + ipSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/1"), + "7fff:ffff:ffff:ffff:ffff:ffff:ffff:ffff"); + ASSERT_EQ( + ipSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/127"), + "2001:db8:85a3:1:1:8a2e:370:7335"); + ASSERT_EQ( + ipSubnetMax("2001:0db8:85a3:0001:0001:8a2e:0370:7334/128"), + "2001:db8:85a3:1:1:8a2e:370:7334"); +} + } // namespace facebook::velox::functions::prestosql