Skip to content

Commit

Permalink
Add ipprefix cast operators for ipaddr [2/n]
Browse files Browse the repository at this point in the history 
Summary:
Support casting ipprefix -> ipaddress, and ipaddress -> ipprefix.

We just need one new helper function, tryIpPrefixLengthFromIPAddressType, which takes in int128_t, gets the ippaddress, and then we check if the ippaddress is ipv4 or ipv6 mapped. We return 32 for ipv4 and 128 for ipv6 mapped.

Differential Revision: D65641994
  • Loading branch information
@yuandagits @facebook-github-bot
yuandagits authored and facebook-github-bot committed Nov 8, 2024
1 parent 31aa90c commit 9e09303
Show file tree
Hide file tree
Showing 5 changed files with 197 additions and 1 deletion.
25 changes: 25 additions & 0 deletions velox/functions/prestosql/tests/IPAddressCastTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,13 @@ class IPAddressCastTest : public functions::test::FunctionBaseTest {
input);
return result;
}

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

int128_t stringToInt128(const std::string& value) {
Expand All @@ -53,6 +60,24 @@ int128_t stringToInt128(const std::string& value) {
return res;
}

TEST_F(IPAddressCastTest, castToIPPrefix) {
EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("1.2.3.4"), "1.2.3.4");
EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("::ffff:1.2.3.4"), "1.2.3.4");
EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("::ffff:102:304"), "1.2.3.4");
EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("192.168.0.0"), "192.168.0.0");
EXPECT_EQ(
castToIPPrefixAndBackToIpVarchar(
"2001:0db8:0000:0000:0000:ff00:0042:8329"),
"2001:db8::ff00:42:8329");
EXPECT_EQ(
castToIPPrefixAndBackToIpVarchar("2001:db8:0:0:1:0:0:1"),
"2001:db8::1:0:0:1");
EXPECT_EQ(castToIPPrefixAndBackToIpVarchar("::1"), "::1");
EXPECT_EQ(
castToIPPrefixAndBackToIpVarchar("2001:db8::ff00:42:8329"),
"2001:db8::ff00:42:8329");
}

TEST_F(IPAddressCastTest, castToVarchar) {
EXPECT_EQ(castToVarchar("::ffff:1.2.3.4"), "1.2.3.4");
EXPECT_EQ(castToVarchar("0:0:0:0:0:0:13.1.68.3"), "::13.1.68.3");
Expand Down
39 changes: 39 additions & 0 deletions velox/functions/prestosql/tests/IPPrefixCastTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,47 @@ class IPPrefixTypeTest : public functions::test::FunctionBaseTest {
"cast(cast(c0 as ipprefix) as varchar)", input);
return result;
}

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

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

TEST_F(IPPrefixTypeTest, castFromIpAddress) {
EXPECT_EQ(castFromIPAddress(std::nullopt), std::nullopt);
EXPECT_EQ(castFromIPAddress("1.2.3.4"), "1.2.3.4/32");
EXPECT_EQ(castFromIPAddress("::ffff:1.2.3.4"), "1.2.3.4/32");
EXPECT_EQ(castFromIPAddress("::ffff:102:304"), "1.2.3.4/32");
EXPECT_EQ(castFromIPAddress("192.168.0.0"), "192.168.0.0/32");
EXPECT_EQ(
castFromIPAddress("2001:0db8:0000:0000:0000:ff00:0042:8329"),
"2001:db8::ff00:42:8329/128");
EXPECT_EQ(castFromIPAddress("2001:db8:0:0:1:0:0:1"), "2001:db8::1:0:0:1/128");
EXPECT_EQ(castFromIPAddress("::1"), "::1/128");
EXPECT_EQ(
castFromIPAddress("2001:db8::ff00:42:8329"),
"2001:db8::ff00:42:8329/128");
EXPECT_EQ(castFromIPAddress("2001:db8::"), "2001:db8::/128");
}

TEST_F(IPPrefixTypeTest, castToIpAddress) {
EXPECT_EQ(castToIpAddress(std::nullopt), std::nullopt);
EXPECT_EQ(castToIpAddress("1.2.3.4/32"), "1.2.3.4");
EXPECT_EQ(castToIpAddress("1.2.3.4/24"), "1.2.3.0");
EXPECT_EQ(castToIpAddress("::1/128"), "::1");
EXPECT_EQ(
castToIpAddress("2001:db8::ff00:42:8329/128"), "2001:db8::ff00:42:8329");
EXPECT_EQ(castToIpAddress("2001:db8::ff00:42:8329/64"), "2001:db8::");
}

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");
Expand Down
59 changes: 59 additions & 0 deletions velox/functions/prestosql/types/IPAddressType.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

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

namespace facebook::velox {

Expand All @@ -28,6 +30,11 @@ class IPAddressCastOperator : public exec::CastOperator {
case TypeKind::VARBINARY:
case TypeKind::VARCHAR:
return true;
case TypeKind::ROW:
if (isIPPrefixType(other)) {
return true;
}
[[fallthrough]];
default:
return false;
}
Expand All @@ -38,6 +45,11 @@ class IPAddressCastOperator : public exec::CastOperator {
case TypeKind::VARBINARY:
case TypeKind::VARCHAR:
return true;
case TypeKind::ROW:
if (isIPPrefixType(other)) {
return true;
}
[[fallthrough]];
default:
return false;
}
Expand All @@ -55,6 +67,9 @@ class IPAddressCastOperator : public exec::CastOperator {
castFromString(input, context, rows, *result);
} else if (input.typeKind() == TypeKind::VARBINARY) {
castFromVarbinary(input, context, rows, *result);
} else if (
input.typeKind() == TypeKind::ROW && isIPPrefixType(input.type())) {
castFromIPPrefix(input, context, rows, *result);
} else {
VELOX_UNSUPPORTED(
"Cast from {} to IPAddress not supported", resultType->toString());
Expand All @@ -73,6 +88,9 @@ class IPAddressCastOperator : public exec::CastOperator {
castToString(input, context, rows, *result);
} else if (resultType->kind() == TypeKind::VARBINARY) {
castToVarbinary(input, context, rows, *result);
} else if (
resultType->kind() == TypeKind::ROW && isIPPrefixType(resultType)) {
castToIPPrefix(input, context, rows, *result);
} else {
VELOX_UNSUPPORTED(
"Cast from IPAddress to {} not supported", resultType->toString());
Expand Down Expand Up @@ -155,6 +173,47 @@ class IPAddressCastOperator : public exec::CastOperator {
});
}

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

context.applyToSelectedNoThrow(rows, [&](auto row) {
const auto ipAddrVal = ipaddresses->valueAt(row);
const auto tryPrefixLength =
ipaddress::tryIpPrefixLengthFromIPAddressType(ipAddrVal);
if (tryPrefixLength.hasError()) {
context.setStatus(row, std::move(tryPrefixLength).error());
return;
}

auto writer = exec::VectorWriter<Row<int128_t, int8_t>>();
writer.init(*rowVectorResult);
writer.setOffset(row);
auto& rowWriter = writer.current();
rowWriter.get_writer_at<0>() = ipAddrVal;
rowWriter.get_writer_at<1>() = tryPrefixLength.value();
writer.commit();
});
}

static void castFromIPPrefix(
const BaseVector& input,
exec::EvalCtx& context,
const SelectivityVector& rows,
BaseVector& result) {
auto* flatResult = result.as<FlatVector<int128_t>>();
const auto* ipprefix = input.as<RowVector>();
const auto* ipaddr =
ipprefix->childAt(ipaddress::kIpRowIndex)->as<SimpleVector<int128_t>>();

context.applyToSelectedNoThrow(
rows, [&](auto row) { flatResult->set(row, ipaddr->valueAt(row)); });
}

static void castFromVarbinary(
const BaseVector& input,
exec::EvalCtx& context,
Expand Down
57 changes: 56 additions & 1 deletion velox/functions/prestosql/types/IPPrefixType.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ class IPPrefixCastOperator : public exec::CastOperator {
switch (other->kind()) {
case TypeKind::VARCHAR:
return true;
case TypeKind::HUGEINT:
return isIPAddressType(other);
default:
return false;
}
Expand All @@ -38,6 +40,8 @@ class IPPrefixCastOperator : public exec::CastOperator {
switch (other->kind()) {
case TypeKind::VARCHAR:
return true;
case TypeKind::HUGEINT:
return isIPAddressType(other);
default:
return false;
}
Expand All @@ -53,6 +57,12 @@ class IPPrefixCastOperator : public exec::CastOperator {
switch (input.typeKind()) {
case TypeKind::VARCHAR:
return castFromString(input, context, rows, *result);
case TypeKind::HUGEINT: {
if (isIPAddressType(input.type())) {
return castFromIpAddress(input, context, rows, *result);
}
[[fallthrough]];
}
default:
VELOX_NYI(
"Cast from {} to IPPrefix not yet supported",
Expand All @@ -70,6 +80,12 @@ class IPPrefixCastOperator : public exec::CastOperator {
switch (resultType->kind()) {
case TypeKind::VARCHAR:
return castToString(input, context, rows, *result);
case TypeKind::HUGEINT: {
if (isIPAddressType(input.type())) {
return castToIpAddress(input, context, rows, *result);
}
[[fallthrough]];
}
default:
VELOX_NYI(
"Cast from IPPrefix to {} not yet supported",
Expand All @@ -78,14 +94,26 @@ class IPPrefixCastOperator : public exec::CastOperator {
}

private:
static void castToIpAddress(
const BaseVector& input,
exec::EvalCtx& context,
const SelectivityVector& rows,
BaseVector& result) {
auto* flatResult = result.as<FlatVector<int128_t>>();
auto rowVector = input.as<RowVector>();
const auto* ipaddr = rowVector->childAt(ipaddress::kIpRowIndex)
->as<SimpleVector<int128_t>>();
context.applyToSelectedNoThrow(
rows, [&](auto row) { flatResult->set(row, ipaddr->valueAt(row)); });
}

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();
const auto* ipaddr = rowVector->childAt(ipaddress::kIpRowIndex)
->as<SimpleVector<int128_t>>();
const auto* prefix = rowVector->childAt(ipaddress::kIpPrefixRowIndex)
Expand Down Expand Up @@ -120,6 +148,33 @@ class IPPrefixCastOperator : public exec::CastOperator {
});
}

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

context.applyToSelectedNoThrow(rows, [&](auto row) {
auto intIpAddr = ipAddrVector->valueAt(row);
const auto tryPrefixLength =
ipaddress::tryIpPrefixLengthFromIPAddressType(intIpAddr);
if (tryPrefixLength.hasError()) {
context.setStatus(row, std::move(tryPrefixLength).error());
return;
}

auto writer = exec::VectorWriter<Row<int128_t, int8_t>>();
writer.init(*rowVectorResult);
writer.setOffset(row);
auto& rowWriter = writer.current();
rowWriter.get_writer_at<0>() = intIpAddr;
rowWriter.get_writer_at<1>() = tryPrefixLength.value();
writer.commit();
});
}

static void castFromString(
const BaseVector& input,
exec::EvalCtx& context,
Expand Down
18 changes: 18 additions & 0 deletions velox/functions/prestosql/types/IPPrefixType.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,24 @@ Status handleFailedToCreateNetworkError(
}
} // namespace

inline folly::Expected<int8_t, Status> 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 {}", tryV6Addr.error()));
}

return tryV6Addr.value().isIPv4Mapped() ? ipaddress::kIPV4Bits
: ipaddress::kIPV6Bits;
}

inline folly::Expected<std::pair<int128_t, int8_t>, Status>
tryParseIpPrefixString(folly::StringPiece ipprefixString) {
// Ensure '/' is present
Expand Down

0 comments on commit 9e09303

Please sign in to comment.