feat: Add HUGEINT key support to MapUnionSumAggregate

velox/functions/prestosql/aggregates/ArbitraryAggregate.cpp

@@ -19,6 +19,7 @@
 #include "velox/expression/FunctionSignature.h"
 #include "velox/functions/lib/aggregates/SimpleNumericAggregate.h"
 #include "velox/functions/lib/aggregates/SingleValueAccumulator.h"
+#include "velox/functions/prestosql/types/IPAddressType.h"
 
 using namespace facebook::velox::functions::aggregate;
 
@@ -446,7 +447,7 @@ void registerArbitraryAggregate(
           case TypeKind::BIGINT:
             return std::make_unique<ArbitraryAggregate<int64_t>>(inputType);
           case TypeKind::HUGEINT:
-            if (inputType->isLongDecimal()) {
+            if (inputType->isLongDecimal() || isIPAddressType(inputType)) {
               return std::make_unique<ArbitraryAggregate<int128_t>>(inputType);
             }
             VELOX_NYI();

velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp

@@ -534,6 +534,9 @@ void registerMapUnionSumAggregate(
           case TypeKind::BIGINT:
             return createMapUnionSumAggregate<int64_t>(
                 valueTypeKind, resultType);
+          case TypeKind::HUGEINT:
+            return createMapUnionSumAggregate<int128_t>(
+                valueTypeKind, resultType);
           case TypeKind::REAL:
             return createMapUnionSumAggregate<float>(valueTypeKind, resultType);
           case TypeKind::DOUBLE:

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

@@ -19,6 +19,7 @@
 #include "velox/exec/tests/utils/TempDirectoryPath.h"
 #include "velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h"
 #include "velox/functions/lib/window/tests/WindowTestBase.h"
+#include "velox/functions/prestosql/types/IPAddressType.h"
 
 using namespace facebook::velox::exec::test;
 using namespace facebook::velox::functions::aggregate::test;
@@ -407,6 +408,34 @@ TEST_F(ArbitraryTest, shortDecimal) {
   testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
 }
 
+TEST_F(ArbitraryTest, ipAddress) {
+  auto data = makeRowVector(
+      {// Grouping key.
+       makeFlatVector<int64_t>({1, 1, 2, 2, 3, 3, 4, 4}),
+       // IPADDRESS values
+       makeNullableFlatVector<int128_t>(
+           {ipaddress::tryGetIPv6asInt128FromString("192.168.1.1").value(),
+            ipaddress::tryGetIPv6asInt128FromString("192.168.1.1").value(),
+            ipaddress::tryGetIPv6asInt128FromString("10.0.0.1").value(),
+            ipaddress::tryGetIPv6asInt128FromString("10.0.0.1").value(),
+            std::nullopt,
+            std::nullopt,
+            std::nullopt,
+            ipaddress::tryGetIPv6asInt128FromString("172.16.0.1").value()},
+           IPADDRESS())});
+
+  auto expectedResult = makeRowVector(
+      {makeFlatVector<int64_t>({1, 2, 3, 4}),
+       makeNullableFlatVector<int128_t>(
+           {ipaddress::tryGetIPv6asInt128FromString("192.168.1.1").value(),
+            ipaddress::tryGetIPv6asInt128FromString("10.0.0.1").value(),
+            std::nullopt,
+            ipaddress::tryGetIPv6asInt128FromString("172.16.0.1").value()},
+           IPADDRESS())});
+
+  testAggregations({data}, {"c0"}, {"arbitrary(c1)"}, {expectedResult});
+}
+
 class ArbitraryWindowTest : public WindowTestBase {};
 
 TEST_F(ArbitraryWindowTest, basic) {

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

@@ -17,6 +17,7 @@
 #include "velox/common/testutil/OptionalEmpty.h"
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/functions/lib/aggregates/tests/utils/AggregationTestBase.h"
+#include "velox/functions/prestosql/types/IPAddressType.h"
 
 using namespace facebook::velox::exec;
 using namespace facebook::velox::exec::test;
@@ -571,5 +572,51 @@ TEST_F(MapUnionSumTest, decimalKey) {
       {makeRowVector({data})}, {}, {"map_union_sum(c0)"}, {expected});
 }
 
+TEST_F(MapUnionSumTest, ipAddressKey) {
+  // Helper to convert IP string to int128_t.
+  auto ipToInt128 = [](std::string_view ipAddr) -> int128_t {
+    return ipaddress::tryGetIPv6asInt128FromString(ipAddr).value();
+  };
+
+  // Test on nulls.
+  auto nullData =
+      makeRowVector({makeAllNullMapVector(3, IPADDRESS(), BIGINT())});
+  testAggregations({nullData}, {}, {"map_union_sum(c0)"}, {"VALUES (NULL)"});
+
+  auto ip1 = ipToInt128("192.168.1.1");
+  auto ip2 = ipToInt128("10.0.0.1");
+  auto ip3 = ipToInt128("172.16.0.1");
+
+  // Test global aggregation with IPADDRESS keys.
+  auto data = makeMapVector<int128_t, int64_t>(
+      {{{ip1, 10}, {ip2, 20}}, {{ip1, 5}, {ip3, 15}}},
+      MAP(IPADDRESS(), BIGINT()));
+  auto expected = makeRowVector({makeMapVector<int128_t, int64_t>(
+      {{{ip1, 15}, {ip2, 20}, {ip3, 15}}}, MAP(IPADDRESS(), BIGINT()))});
+  testAggregations(
+      {makeRowVector({data})}, {}, {"map_union_sum(c0)"}, {expected});
+
+  // Test group-by aggregation with IPADDRESS keys.
+  auto groupByData = makeRowVector({
+      makeFlatVector<int64_t>({1, 2, 1, 2}),
+      makeMapVector<int128_t, int64_t>(
+          {{{ip1, 10}, {ip2, 20}},
+           {{ip1, 5}, {ip3, 15}},
+           {{ip1, 3}, {ip2, 7}},
+           {{ip3, 25}}},
+          MAP(IPADDRESS(), BIGINT())),
+  });
+
+  auto groupByExpected = makeRowVector({
+      makeFlatVector<int64_t>({1, 2}),
+      makeMapVector<int128_t, int64_t>(
+          {{{ip1, 13}, {ip2, 27}}, {{ip1, 5}, {ip3, 40}}},
+          MAP(IPADDRESS(), BIGINT())),
+  });
+
+  testAggregations(
+      {groupByData}, {"c0"}, {"map_union_sum(c1)"}, {groupByExpected});
+}
+
 } // namespace
 } // namespace facebook::velox::aggregate::test