From abf9f72ccde130369f83cd55aad318f4b77b57a8 Mon Sep 17 00:00:00 2001 From: mohsaka <135669458+mohsaka@users.noreply.github.com> Date: Tue, 10 Dec 2024 08:07:50 -0800 Subject: [PATCH] Refactor --- velox/expression/FunctionSignature.cpp | 2 +- .../functions/prestosql/IPAddressFunctions.h | 4 +- velox/functions/prestosql/TypeOf.cpp | 4 +- .../prestosql/aggregates/PrestoHasher.cpp | 2 +- .../aggregates/tests/PrestoHasherTest.cpp | 2 +- .../ComparisonFunctionsRegistration.cpp | 2 +- .../GeneralFunctionsRegistration.cpp | 2 +- .../prestosql/tests/ComparisonsTest.cpp | 2 +- .../prestosql/types/IPAddressType.cpp | 4 +- .../prestosql/types/IPPrefixType.cpp | 2 +- .../types/{ => headers}/IPAddressType.h | 0 .../prestosql/types/headers/IPPrefixType.h | 159 ++++++++++++++++++ .../types/tests/IPAddressTypeTest.cpp | 2 +- .../types/tests/IPPrefixTypeTest.cpp | 2 +- .../functions/tests/FunctionRegistryTest.cpp | 2 +- velox/functions/tests/RegistryTestUtil.h | 2 +- 16 files changed, 176 insertions(+), 17 deletions(-) rename velox/functions/prestosql/types/{ => headers}/IPAddressType.h (100%) create mode 100644 velox/functions/prestosql/types/headers/IPPrefixType.h diff --git a/velox/expression/FunctionSignature.cpp b/velox/expression/FunctionSignature.cpp index c996e6cc2e5d9..b85b9b76ce924 100644 --- a/velox/expression/FunctionSignature.cpp +++ b/velox/expression/FunctionSignature.cpp @@ -17,7 +17,7 @@ #include #include #include "velox/common/base/Exceptions.h" -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" #include "velox/type/Type.h" namespace facebook::velox::exec { diff --git a/velox/functions/prestosql/IPAddressFunctions.h b/velox/functions/prestosql/IPAddressFunctions.h index 295ca68a6fa11..1588f8a8b0739 100644 --- a/velox/functions/prestosql/IPAddressFunctions.h +++ b/velox/functions/prestosql/IPAddressFunctions.h @@ -15,8 +15,8 @@ */ #pragma once -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" namespace facebook::velox::functions { diff --git a/velox/functions/prestosql/TypeOf.cpp b/velox/functions/prestosql/TypeOf.cpp index c1465b8c5f236..81b98894bb321 100644 --- a/velox/functions/prestosql/TypeOf.cpp +++ b/velox/functions/prestosql/TypeOf.cpp @@ -14,9 +14,9 @@ * limitations under the License. */ #include "velox/expression/VectorFunction.h" -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" #include "velox/functions/prestosql/types/HyperLogLogType.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/JsonType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/functions/prestosql/types/UuidType.h" diff --git a/velox/functions/prestosql/aggregates/PrestoHasher.cpp b/velox/functions/prestosql/aggregates/PrestoHasher.cpp index 0df9c7756ebde..97c99a14afc31 100644 --- a/velox/functions/prestosql/aggregates/PrestoHasher.cpp +++ b/velox/functions/prestosql/aggregates/PrestoHasher.cpp @@ -21,7 +21,7 @@ #include #include "velox/functions/lib/RowsTranslationUtil.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" namespace facebook::velox::aggregate { diff --git a/velox/functions/prestosql/aggregates/tests/PrestoHasherTest.cpp b/velox/functions/prestosql/aggregates/tests/PrestoHasherTest.cpp index 6b1c004b9283f..dbeec84d7cf16 100644 --- a/velox/functions/prestosql/aggregates/tests/PrestoHasherTest.cpp +++ b/velox/functions/prestosql/aggregates/tests/PrestoHasherTest.cpp @@ -17,7 +17,7 @@ #include #include "velox/functions/prestosql/aggregates/PrestoHasher.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/type/tz/TimeZoneMap.h" #include "velox/vector/tests/utils/VectorTestBase.h" diff --git a/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp b/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp index 2870a3ff787f0..fbaac212bee98 100644 --- a/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/ComparisonFunctionsRegistration.cpp @@ -16,7 +16,7 @@ #include "velox/functions/Registerer.h" #include "velox/functions/lib/RegistrationHelpers.h" #include "velox/functions/prestosql/Comparisons.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/type/Type.h" diff --git a/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp b/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp index a33d0a62a6bf7..36f860049c67e 100644 --- a/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp +++ b/velox/functions/prestosql/registration/GeneralFunctionsRegistration.cpp @@ -21,7 +21,7 @@ #include "velox/functions/prestosql/GreatestLeast.h" #include "velox/functions/prestosql/InPredicate.h" #include "velox/functions/prestosql/Reduce.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" namespace facebook::velox::functions { diff --git a/velox/functions/prestosql/tests/ComparisonsTest.cpp b/velox/functions/prestosql/tests/ComparisonsTest.cpp index 707a61d108fa1..a6fce73c52803 100644 --- a/velox/functions/prestosql/tests/ComparisonsTest.cpp +++ b/velox/functions/prestosql/tests/ComparisonsTest.cpp @@ -19,7 +19,7 @@ #include "velox/functions/Udf.h" #include "velox/functions/lib/RegistrationHelpers.h" #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/TimestampWithTimeZoneType.h" #include "velox/type/tests/utils/CustomTypesForTesting.h" #include "velox/type/tz/TimeZoneMap.h" diff --git a/velox/functions/prestosql/types/IPAddressType.cpp b/velox/functions/prestosql/types/IPAddressType.cpp index e402bacb99063..44349b7b4085c 100644 --- a/velox/functions/prestosql/types/IPAddressType.cpp +++ b/velox/functions/prestosql/types/IPAddressType.cpp @@ -14,10 +14,10 @@ * limitations under the License. */ -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/expression/CastExpr.h" #include "velox/expression/VectorWriters.h" -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" namespace facebook::velox { diff --git a/velox/functions/prestosql/types/IPPrefixType.cpp b/velox/functions/prestosql/types/IPPrefixType.cpp index 5d25df58af55d..30fbe07f58a66 100644 --- a/velox/functions/prestosql/types/IPPrefixType.cpp +++ b/velox/functions/prestosql/types/IPPrefixType.cpp @@ -16,7 +16,7 @@ #include #include "velox/expression/CastExpr.h" -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" namespace facebook::velox { diff --git a/velox/functions/prestosql/types/IPAddressType.h b/velox/functions/prestosql/types/headers/IPAddressType.h similarity index 100% rename from velox/functions/prestosql/types/IPAddressType.h rename to velox/functions/prestosql/types/headers/IPAddressType.h diff --git a/velox/functions/prestosql/types/headers/IPPrefixType.h b/velox/functions/prestosql/types/headers/IPPrefixType.h new file mode 100644 index 0000000000000..bfd7bf278ba2d --- /dev/null +++ b/velox/functions/prestosql/types/headers/IPPrefixType.h @@ -0,0 +1,159 @@ +/* + * 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. + */ +#pragma once + +#include + +#include "velox/common/base/Status.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" +#include "velox/type/SimpleFunctionApi.h" +#include "velox/type/Type.h" + +namespace facebook::velox { + +namespace ipaddress { +constexpr uint8_t kIPV4Bits = 32; +constexpr uint8_t kIPV6Bits = 128; +constexpr int kIPPrefixLengthIndex = 16; +constexpr int kIPPrefixBytes = 17; +constexpr auto kIpRowIndex = "ip"; +constexpr auto kIpPrefixRowIndex = "prefix"; + +inline folly::Expected tryIpPrefixLengthFromIPAddressType( + const int128_t& intIpAddr) { + folly::ByteArray16 addrBytes = {0}; + memcpy(&addrBytes, &intIpAddr, sizeof(intIpAddr)); + std::reverse(addrBytes.begin(), addrBytes.end()); + auto tryV6Addr = folly::IPAddressV6::tryFromBinary(addrBytes); + if (tryV6Addr.hasError()) { + return folly::makeUnexpected( + threadSkipErrorDetails() + ? Status::UserError() + : Status::UserError("Received invalid ip address")); + } + + return tryV6Addr.value().isIPv4Mapped() ? ipaddress::kIPV4Bits + : ipaddress::kIPV6Bits; +} + +inline folly::Expected, Status> +tryParseIpPrefixString(folly::StringPiece ipprefixString) { + // Ensure '/' is present + if (ipprefixString.find('/') == std::string::npos) { + return folly::makeUnexpected( + threadSkipErrorDetails() + ? Status::UserError() + : Status::UserError( + "Cannot cast value to IPPREFIX: {}", ipprefixString)); + } + + auto tryCdirNetwork = folly::IPAddress::tryCreateNetwork( + ipprefixString, /*defaultCidr*/ -1, /*applyMask*/ false); + + if (tryCdirNetwork.hasError()) { + return folly::makeUnexpected( + threadSkipErrorDetails() + ? Status::UserError() + : Status::UserError( + "Cannot cast value to IPPREFIX: {}", ipprefixString)); + } + + auto [ip, prefix] = tryCdirNetwork.value(); + if (prefix > ((ip.isIPv4Mapped() || ip.isV4()) ? ipaddress::kIPV4Bits + : ipaddress::kIPV6Bits)) { + return folly::makeUnexpected( + threadSkipErrorDetails() + ? Status::UserError() + : Status::UserError( + "Cannot cast value to IPPREFIX: {}", ipprefixString)); + } + + auto addrBytes = (ip.isIPv4Mapped() || ip.isV4()) + ? folly::IPAddress::createIPv4(ip).mask(prefix).createIPv6().toByteArray() + : folly::IPAddress::createIPv6(ip).mask(prefix).toByteArray(); + + std::reverse(addrBytes.begin(), addrBytes.end()); + + int128_t intAddr = 0; + static_assert(sizeof(intAddr) == ipaddress::kIPAddressBytes); + memcpy(&intAddr, &addrBytes, ipaddress::kIPAddressBytes); + return std::make_pair(intAddr, prefix); +} +}; // namespace ipaddress + +class IPPrefixType : public RowType { + IPPrefixType() + : RowType( + {ipaddress::kIpRowIndex, ipaddress::kIpPrefixRowIndex}, + {HUGEINT(), TINYINT()}) {} + + public: + static const std::shared_ptr& get() { + static const std::shared_ptr instance{ + new IPPrefixType()}; + + return instance; + } + + bool equivalent(const Type& other) const override { + // Pointer comparison works since this type is a singleton. + return this == &other; + } + + const char* name() const override { + return "IPPREFIX"; + } + + std::string toString() const override { + return name(); + } + + folly::dynamic serialize() const override { + folly::dynamic obj = folly::dynamic::object; + obj["name"] = "Type"; + obj["type"] = name(); + return obj; + } + + const std::vector& parameters() const override { + static const std::vector kEmpty = {}; + return kEmpty; + } +}; + +FOLLY_ALWAYS_INLINE bool isIPPrefixType(const TypePtr& type) { + // Pointer comparison works since this type is a singleton. + return IPPrefixType::get() == type; +} + +FOLLY_ALWAYS_INLINE bool isIPPrefixName(const std::string& name) { + return (name == IPPrefixType::get()->name()); +} + +FOLLY_ALWAYS_INLINE std::shared_ptr IPPREFIX() { + return IPPrefixType::get(); +} + +struct IPPrefixT { + using type = Row; + static constexpr const char* typeName = "ipprefix"; +}; + +using IPPrefix = CustomType; + +void registerIPPrefixType(); + +} // namespace facebook::velox diff --git a/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp b/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp index f65691e1c05f2..b0de1bd9567d9 100644 --- a/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp +++ b/velox/functions/prestosql/types/tests/IPAddressTypeTest.cpp @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "velox/functions/prestosql/types/IPAddressType.h" +#include "velox/functions/prestosql/types/headers/IPAddressType.h" #include "velox/functions/prestosql/types/tests/TypeTestBase.h" namespace facebook::velox::test { diff --git a/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp b/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp index e1bb851e993f7..4b3b5058dc6e3 100644 --- a/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp +++ b/velox/functions/prestosql/types/tests/IPPrefixTypeTest.cpp @@ -13,7 +13,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" #include "velox/functions/prestosql/types/tests/TypeTestBase.h" namespace facebook::velox::test { diff --git a/velox/functions/tests/FunctionRegistryTest.cpp b/velox/functions/tests/FunctionRegistryTest.cpp index 24d1b617853f6..c85a2db0ef37f 100644 --- a/velox/functions/tests/FunctionRegistryTest.cpp +++ b/velox/functions/tests/FunctionRegistryTest.cpp @@ -26,7 +26,7 @@ #include "velox/functions/Registerer.h" #include "velox/functions/prestosql/registration/RegistrationFunctions.h" #include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h" -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" #include "velox/functions/tests/RegistryTestUtil.h" #include "velox/type/Type.h" diff --git a/velox/functions/tests/RegistryTestUtil.h b/velox/functions/tests/RegistryTestUtil.h index 80c3fcb205c45..b0753e8bf7c23 100644 --- a/velox/functions/tests/RegistryTestUtil.h +++ b/velox/functions/tests/RegistryTestUtil.h @@ -18,7 +18,7 @@ #include "velox/expression/FunctionSignature.h" #include "velox/expression/VectorFunction.h" #include "velox/functions/Macros.h" -#include "velox/functions/prestosql/type_headers/IPPrefixType.h" +#include "velox/functions/prestosql/types/headers/IPPrefixType.h" namespace facebook::velox {