Skip to content

Commit

Permalink
Add ip_subnet_max/min [5/n] (facebookincubator#11515)
Browse files Browse the repository at this point in the history 
Summary:

Add ip_subnet and ip_subnetmax.

Split from facebookincubator#11407

Differential Revision: D65833249
  • Loading branch information
@yuandagits @facebook-github-bot
yuandagits authored and facebook-github-bot committed Dec 18, 2024
1 parent 8991f47 commit 6e6fa66
Show file tree
Hide file tree
Showing 3 changed files with 136 additions and 0 deletions.
16 changes: 16 additions & 0 deletions velox/docs/functions/presto/ipaddress.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}

51 changes: 51 additions & 0 deletions velox/functions/prestosql/IPAddressFunctions.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,13 +75,64 @@ struct IPPrefixFunction {
}
};

template <typename T>
struct IPSubnetMinFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);

FOLLY_ALWAYS_INLINE void call(
out_type<IPAddress>& result,
const arg_type<IPPrefix>& ipPrefix) {
// IPPrefix type stores the smallest(canonical) IP already
result = *ipPrefix.template at<0>();
}
};

template <typename T>
struct IPSubnetMaxFunction {
VELOX_DEFINE_FUNCTION_TYPES(T);

FOLLY_ALWAYS_INLINE void call(
out_type<IPAddress>& result,
const arg_type<IPPrefix>& 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();
registerFunction<IPPrefixFunction, IPPrefix, IPAddress, int64_t>(
{prefix + "ip_prefix"});
registerFunction<IPPrefixFunction, IPPrefix, Varchar, int64_t>(
{prefix + "ip_prefix"});
registerFunction<IPSubnetMinFunction, IPAddress, IPPrefix>(
{prefix + "ip_subnet_min"});
registerFunction<IPSubnetMaxFunction, IPAddress, IPPrefix>(
{prefix + "ip_subnet_max"});
}

} // namespace facebook::velox::functions
69 changes: 69 additions & 0 deletions velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,18 @@ class IPAddressFunctionsTest : public functions::test::FunctionBaseTest {
return evaluateOnce<std::string>(
"cast(ip_prefix(c0, c1) as varchar)", input, mask);
}

std::optional<std::string> ipSubnetMin(
const std::optional<std::string>& input) {
return evaluateOnce<std::string>(
"cast(ip_subnet_min(cast(c0 as ipprefix)) as varchar)", input);
}

std::optional<std::string> ipSubnetMax(
const std::optional<std::string>& input) {
return evaluateOnce<std::string>(
"cast(ip_subnet_max(cast(c0 as ipprefix)) as varchar)", input);
}
};

TEST_F(IPAddressFunctionsTest, ipPrefixFromIpAddress) {
Expand Down Expand Up @@ -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

0 comments on commit 6e6fa66

Please sign in to comment.