Add ip_prefix function [4/n] (#11514)

Summary: 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<int128_t, int8_t> which corresponds to the ipaddress and prefix for IPAddressPrefix type. Split from #11407 Reviewed By: xiaoxmeng, spershin Differential Revision: D65802211
facebookincubator · Dec 19, 2024 · 861f388 · 861f388
1 parent 0ee82f3
commit 861f388
Show file tree

Hide file tree

Showing 3 changed files with 185 additions and 0 deletions.
diff --git a/velox/docs/functions/presto/ipaddress.rst 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
@@ -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 <typename T>
+struct IPPrefixFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<IPPrefix>& result,
+      const arg_type<IPAddress>& ip,
+      const arg_type<int64_t>& 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<IPPrefix>& result,
+      const arg_type<Varchar>& ipString,
+      const arg_type<int64_t>& 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<int128_t, int8_t> makeIPPrefix(
+      const folly::IPAddressV6& v6Addr,
+      int64_t prefixBits) {
+    if (v6Addr.isIPv4Mapped()) {
+      VELOX_USER_CHECK(
+          0 <= prefixBits && prefixBits <= ipaddress::kIPV4Bits,
+          "IPv4 subnet size must be in range [0, 32]");
+    } else {
+      VELOX_USER_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<int8_t>(prefixBits));
+  }
+};
+
 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"});
 }
 
 } // namespace facebook::velox::functions
diff --git a/velox/functions/prestosql/tests/IPAddressFunctionsTest.cpp 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<std::string> ipPrefixFunctionFromIpAddress(
+      const std::optional<std::string>& input,
+      const std::optional<int64_t>& mask) {
+    return evaluateOnce<std::string>(
+        "cast(ip_prefix(cast(c0 as ipaddress), c1) as varchar)", input, mask);
+  }
+
+  std::optional<std::string> ipPrefixFromVarChar(
+      const std::optional<std::string>& input,
+      const std::optional<int64_t>& mask) {
+    return evaluateOnce<std::string>(
+        "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