Skip to content

Commit

Permalink
Add ipprefix cast operators for varchar [1/n]
Browse files Browse the repository at this point in the history
Summary:
Support cast for ipprefix for varchar. Based off of facebookincubator#11122

Will do IPAddress in follow up diff

Differential Revision: D65449935
  • Loading branch information
yuandagits authored and facebook-github-bot committed Nov 6, 2024
1 parent c3023b6 commit 050bcfb
Show file tree
Hide file tree
Showing 3 changed files with 348 additions and 8 deletions.
110 changes: 110 additions & 0 deletions velox/functions/prestosql/tests/IPPrefixCastTest.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
/*
* 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/functions/prestosql/tests/utils/FunctionBaseTest.h"

namespace facebook::velox::functions::prestosql {

class IPPrefixTypeTest : public functions::test::FunctionBaseTest {
protected:
std::optional<std::string> castToVarchar(
const std::optional<std::string>& input) {
auto result = evaluateOnce<std::string>(
"cast(cast(c0 as ipprefix) as varchar)", input);
return result;
}
};

TEST_F(IPPrefixTypeTest, castToVarchar) {
EXPECT_EQ(castToVarchar("::ffff:1.2.3.4/24"), "1.2.3.0/24");
EXPECT_EQ(castToVarchar("192.168.0.0/24"), "192.168.0.0/24");
EXPECT_EQ(castToVarchar("255.2.3.4/0"), "0.0.0.0/0");
EXPECT_EQ(castToVarchar("255.2.3.4/1"), "128.0.0.0/1");
EXPECT_EQ(castToVarchar("255.2.3.4/2"), "192.0.0.0/2");
EXPECT_EQ(castToVarchar("255.2.3.4/4"), "240.0.0.0/4");
EXPECT_EQ(castToVarchar("1.2.3.4/8"), "1.0.0.0/8");
EXPECT_EQ(castToVarchar("1.2.3.4/16"), "1.2.0.0/16");
EXPECT_EQ(castToVarchar("1.2.3.4/24"), "1.2.3.0/24");
EXPECT_EQ(castToVarchar("1.2.3.255/25"), "1.2.3.128/25");
EXPECT_EQ(castToVarchar("1.2.3.255/26"), "1.2.3.192/26");
EXPECT_EQ(castToVarchar("1.2.3.255/28"), "1.2.3.240/28");
EXPECT_EQ(castToVarchar("1.2.3.255/30"), "1.2.3.252/30");
EXPECT_EQ(castToVarchar("1.2.3.255/32"), "1.2.3.255/32");
EXPECT_EQ(
castToVarchar("2001:0db8:0000:0000:0000:ff00:0042:8329/128"),
"2001:db8::ff00:42:8329/128");
EXPECT_EQ(
castToVarchar("2001:db8::ff00:42:8329/128"),
"2001:db8::ff00:42:8329/128");
EXPECT_EQ(castToVarchar("2001:db8:0:0:1:0:0:1/128"), "2001:db8::1:0:0:1/128");
EXPECT_EQ(castToVarchar("2001:db8:0:0:1::1/128"), "2001:db8::1:0:0:1/128");
EXPECT_EQ(castToVarchar("2001:db8::1:0:0:1/128"), "2001:db8::1:0:0:1/128");
EXPECT_EQ(
castToVarchar("2001:DB8::FF00:ABCD:12EF/128"),
"2001:db8::ff00:abcd:12ef/128");
EXPECT_EQ(castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/0"), "::/0");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/1"), "8000::/1");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/2"), "c000::/2");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/4"), "f000::/4");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/8"), "ff00::/8");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/16"), "ffff::/16");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/32"),
"ffff:ffff::/32");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/48"),
"ffff:ffff:ffff::/48");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/64"),
"ffff:ffff:ffff:ffff::/64");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/80"),
"ffff:ffff:ffff:ffff:ffff::/80");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/96"),
"ffff:ffff:ffff:ffff:ffff:ffff::/96");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/112"),
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:0/112");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/120"),
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:ff00/120");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/124"),
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:fff0/124");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/126"),
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:fffc/126");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/127"),
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:fffe/127");
EXPECT_EQ(
castToVarchar("ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/128"),
"ffff:ffff:ffff:ffff:ffff:ffff:ffff:ffff/128");
EXPECT_THROW(castToVarchar("facebook.com/32"), VeloxUserError);
EXPECT_THROW(castToVarchar("localhost/32"), VeloxUserError);
EXPECT_THROW(castToVarchar("2001:db8::1::1/128"), VeloxUserError);
EXPECT_THROW(castToVarchar("2001:zxy::1::1/128"), VeloxUserError);
EXPECT_THROW(castToVarchar("789.1.1.1/32"), VeloxUserError);
EXPECT_THROW(castToVarchar("192.1.1.1"), VeloxUserError);
EXPECT_THROW(castToVarchar("192.1.1.1/128"), VeloxUserError);
}
} // namespace facebook::velox::functions::prestosql
112 changes: 105 additions & 7 deletions velox/functions/prestosql/types/IPPrefixType.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/

#include <folly/small_vector.h>

#include "velox/expression/CastExpr.h"
#include "velox/expression/VectorWriters.h"
#include "velox/functions/prestosql/types/IPPrefixType.h"

namespace facebook::velox {
Expand All @@ -26,11 +26,21 @@ namespace {
class IPPrefixCastOperator : public exec::CastOperator {
public:
bool isSupportedFromType(const TypePtr& other) const override {
return false;
switch (other->kind()) {
case TypeKind::VARCHAR:
return true;
default:
return false;
}
}

bool isSupportedToType(const TypePtr& other) const override {
return false;
switch (other->kind()) {
case TypeKind::VARCHAR:
return true;
default:
return false;
}
}

void castTo(
Expand All @@ -40,8 +50,14 @@ class IPPrefixCastOperator : public exec::CastOperator {
const TypePtr& resultType,
VectorPtr& result) const override {
context.ensureWritable(rows, resultType, result);
VELOX_NYI(
"Cast from {} to IPPrefix not yet supported", input.type()->toString());
switch (input.typeKind()) {
case TypeKind::VARCHAR:
return castFromString(input, context, rows, *result);
default:
VELOX_NYI(
"Cast from {} to IPPrefix not yet supported",
input.type()->toString());
}
}

void castFrom(
Expand All @@ -51,8 +67,90 @@ class IPPrefixCastOperator : public exec::CastOperator {
const TypePtr& resultType,
VectorPtr& result) const override {
context.ensureWritable(rows, resultType, result);
VELOX_NYI(
"Cast from IPPrefix to {} not yet supported", resultType->toString());
switch (resultType->kind()) {
case TypeKind::VARCHAR:
return castToString(input, context, rows, *result);
default:
VELOX_NYI(
"Cast from IPPrefix to {} not yet supported",
resultType->toString());
}
}

private:
static void castToString(
const BaseVector& input,
exec::EvalCtx& context,
const SelectivityVector& rows,
BaseVector& result) {
auto* flatResult = result.as<FlatVector<StringView>>();
auto rowVector = input.as<RowVector>();
auto rowType = rowVector->type();
context.applyToSelectedNoThrow(rows, [&](auto row) {
const auto* ipaddr = rowVector->childAt(ipaddress::kIpRowIndex)
->as<SimpleVector<int128_t>>();
const auto* prefix = rowVector->childAt(ipaddress::kIpPrefixRowIndex)
->as<SimpleVector<int8_t>>();

const auto ipAddrVal = ipaddr->valueAt(row);
// The string representation of the last byte needs
// to be unsigned
const uint8_t prefixVal = prefix->valueAt(row);

// Copy the first 16 bytes into a ByteArray16.
folly::ByteArray16 addrBytes;
memcpy(&addrBytes, &ipAddrVal, ipaddress::kIPAddressBytes);
// Reverse the bytes to get the correct order. Similar to
// IPAddressType. We assume we're ALWAYS on a little endian machine.
// Note: for big endian, we should not reverse the bytes.
std::reverse(addrBytes.begin(), addrBytes.end());
// // Construct a V6 address from the ByteArray16.
folly::IPAddressV6 v6Addr(addrBytes);

// Inline func to get string for ipv4 or ipv6 string
auto ipString = [&]() {
if (v6Addr.isIPv4Mapped()) {
auto v4Addr = v6Addr.createIPv4();
return v4Addr.str();
}
return v6Addr.str();
};

// Format of string is {ipString}/{mask}
auto stringRet = fmt::format("{}/{}", ipString(), prefixVal);

// Write the string to the result vector
exec::StringWriter<false> result(flatResult, row);
result.append(stringRet);
result.finalize();
});
}

static void castFromString(
const BaseVector& input,
exec::EvalCtx& context,
const SelectivityVector& rows,
BaseVector& result) {
auto* rowVectorResult = result.as<RowVector>();
const auto* ipPrefixStrings = input.as<SimpleVector<StringView>>();

context.applyToSelectedNoThrow(rows, [&](auto row) {
auto ipAddressStringView = ipPrefixStrings->valueAt(row);
auto ipAddressString = ipAddressStringView.str();
auto tryIpPrefix = ipaddress::tryParseIpPrefixString(ipAddressString);
if (tryIpPrefix.hasError()) {
context.setStatus(row, std::move(tryIpPrefix.error()));
}

auto ipPrefix = std::move(tryIpPrefix.value());
auto writer = exec::VectorWriter<Row<int128_t, int8_t>>();
writer.init(*rowVectorResult);
writer.setOffset(row);
auto& rowWriter = writer.current();
rowWriter.get_writer_at<0>() = ipPrefix.first;
rowWriter.get_writer_at<1>() = ipPrefix.second;
writer.commit();
});
}
};

Expand Down
Loading

0 comments on commit 050bcfb

Please sign in to comment.