Add ip_subnet_max/min [5/n] (#11515)

Summary: Add ip_subnet and ip_subnetmax. Split from #11407 Differential Revision: D65833249
facebookincubator · Jan 9, 2025 · 4995804 · 4995804
1 parent ff943fd
commit 4995804
Show file tree

Hide file tree

Showing 4 changed files with 137 additions and 0 deletions.
diff --git 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/expression/fuzzer/ExpressionFuzzer.cpp b/velox/expression/fuzzer/ExpressionFuzzer.cpp
@@ -676,6 +676,7 @@ bool ExpressionFuzzer::isSupportedSignature(
   if (usesTypeName(signature, "opaque") ||
       usesTypeName(signature, "timestamp with time zone") ||
       usesTypeName(signature, "interval day to second") ||
+      usesTypeName(signature, "ipprefix") ||
       (!options_.enableDecimalType && usesTypeName(signature, "decimal")) ||
       (!options_.enableComplexTypes && useComplexType) ||
       (options_.enableComplexTypes && usesTypeName(signature, "unknown"))) {

diff --git a/velox/functions/prestosql/IPAddressFunctions.h b/velox/functions/prestosql/IPAddressFunctions.h
@@ -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
diff --git 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<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) {
@@ -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