Skip to content

Commit a980792

Browse files
pratikpugaliafacebook-github-bot
authored andcommitted
feat: Add IPADDRESS support to Arbitrary aggregate function
Summary: The arbitrary() aggregate function was throwing NOT_IMPLEMENTED when used with IPADDRESS type. This was caught by the aggregation fuzzer. This change adds a check for isIPAddressType() for HUGEINT type in ArbitraryAggregate, allowing IPADDRESS values to be properly handled. Differential Revision: D92026533
1 parent 6fe8d01 commit a980792

File tree

2 files changed

+31
-1
lines changed

2 files changed

+31
-1
lines changed

velox/functions/prestosql/aggregates/ArbitraryAggregate.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "velox/expression/FunctionSignature.h"
2020
#include "velox/functions/lib/aggregates/SimpleNumericAggregate.h"
2121
#include "velox/functions/lib/aggregates/SingleValueAccumulator.h"
22+
#include "velox/functions/prestosql/types/IPAddressType.h"
2223

2324
using namespace facebook::velox::functions::aggregate;
2425

@@ -446,7 +447,7 @@ void registerArbitraryAggregate(
446447
case TypeKind::BIGINT:
447448
return std::make_unique<ArbitraryAggregate<int64_t>>(inputType);
448449
case TypeKind::HUGEINT:
449-
if (inputType->isLongDecimal()) {
450+
if (inputType->isLongDecimal() || isIPAddressType(inputType)) {
450451
return std::make_unique<ArbitraryAggregate<int128_t>>(inputType);
451452
}
452453
VELOX_NYI();

velox/functions/prestosql/aggregates/tests/ArbitraryTest.cpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "velox/exec/tests/utils/TempDirectoryPath.h"
2020
#include "velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h"
2121
#include "velox/functions/lib/window/tests/WindowTestBase.h"
22+
#include "velox/functions/prestosql/types/IPAddressType.h"
2223

2324
using namespace facebook::velox::exec::test;
2425
using namespace facebook::velox::functions::aggregate::test;
@@ -407,6 +408,34 @@ TEST_F(ArbitraryTest, shortDecimal) {
407408
testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
408409
}
409410

411+
TEST_F(ArbitraryTest, ipAddress) {
412+
auto data = makeRowVector(
413+
{// Grouping key.
414+
makeFlatVector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
415+
// IPADDRESS values
416+
makeNullableFlatVector<int128_t>(
417+
{ipaddress::tryGetIPv6asInt128FromString("192.168.1.1").value(),
418+
ipaddress::tryGetIPv6asInt128FromString("192.168.1.1").value(),
419+
ipaddress::tryGetIPv6asInt128FromString("10.0.0.1").value(),
420+
ipaddress::tryGetIPv6asInt128FromString("10.0.0.1").value(),
421+
std::nullopt,
422+
std::nullopt,
423+
std::nullopt,
424+
ipaddress::tryGetIPv6asInt128FromString("172.16.0.1").value()},
425+
IPADDRESS())});
426+
427+
auto expectedResult = makeRowVector(
428+
{makeFlatVector<int64_t>({1, 2, 3, 4}),
429+
makeNullableFlatVector<int128_t>(
430+
{ipaddress::tryGetIPv6asInt128FromString("192.168.1.1").value(),
431+
ipaddress::tryGetIPv6asInt128FromString("10.0.0.1").value(),
432+
std::nullopt,
433+
ipaddress::tryGetIPv6asInt128FromString("172.16.0.1").value()},
434+
IPADDRESS())});
435+
436+
testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
437+
}
438+
410439
class ArbitraryWindowTest : public WindowTestBase {};
411440

412441
TEST_F(ArbitraryWindowTest, basic) {

0 commit comments

Comments
 (0)