Add JSON to Map Union Sum (#12208)

Summary: Add support for JSON for map_union_sum Differential Revision: D68867704
facebookincubator · Feb 3, 2025 · 1e30d79 · 1e30d79
1 parent e484679
commit 1e30d79
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 1 deletion.
diff --git a/velox/functions/prestosql/aggregates/CMakeLists.txt b/velox/functions/prestosql/aggregates/CMakeLists.txt
@@ -57,6 +57,7 @@ velox_link_libraries(
   velox_functions_aggregates
   velox_functions_lib
   velox_functions_util
+  velox_presto_types
   Folly::folly)
 
 if(${VELOX_BUILD_TESTING})

diff --git a/velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp b/velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp
@@ -358,7 +358,7 @@ void registerMapUnionSumAggregate(
       "real",
       "double",
       "varchar",
-  };
+      "json"};
   const std::vector<std::string> valueTypes = {
       "tinyint",
       "smallint",

diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
@@ -15,6 +15,7 @@
  */
 #include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h"
 #include "velox/exec/Aggregate.h"
+#include "velox/functions/prestosql/types/JsonType.h"
 
 namespace facebook::velox::aggregate::prestosql {
 
@@ -155,6 +156,7 @@ void registerAllAggregateFunctions(
     bool withCompanionFunctions,
     bool onlyPrestoSignatures,
     bool overwrite) {
+  registerJsonType();
   registerApproxDistinctAggregates(prefix, withCompanionFunctions, overwrite);
   registerApproxMostFrequentAggregate(
       prefix, withCompanionFunctions, overwrite);

diff --git a/velox/functions/prestosql/aggregates/tests/MapUnionSumTest.cpp b/velox/functions/prestosql/aggregates/tests/MapUnionSumTest.cpp
@@ -378,6 +378,44 @@ TEST_F(MapUnionSumTest, groupByVarcharKey) {
   testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
 }
 
+TEST_F(MapUnionSumTest, groupByJsonKey) {
+  std::vector<std::string> jsonKeyStrings = {
+      "\"key1\"",
+      "\"key2\"",
+      "\"key3\"",
+      "\"key4\"",
+      "\"key5\"",
+  };
+  std::vector<StringView> jsonKeys;
+  jsonKeys.reserve(jsonKeyStrings.size());
+  for (const auto& key : jsonKeyStrings) {
+    jsonKeys.emplace_back(key);
+  }
+  auto data = makeRowVector({
+      makeFlatVector<int64_t>({1, 2, 1, 2, 1}),
+      makeNullableMapVector<StringView, int64_t>({
+          {},
+          std::nullopt,
+          {{{jsonKeys[0], 10}, {jsonKeys[1], 20}}},
+          {{{jsonKeys[0], 11}, {jsonKeys[2], 30}, {jsonKeys[3], 40}}},
+          {{{jsonKeys[2], 30}, {jsonKeys[4], 50}, {jsonKeys[0], 12}}},
+      }),
+  });
+
+  auto expected = makeRowVector({
+      makeFlatVector<int64_t>({1, 2}),
+      makeMapVector<StringView, int64_t>({
+          {{jsonKeys[0], 22},
+           {jsonKeys[1], 20},
+           {jsonKeys[2], 30},
+           {jsonKeys[4], 50}},
+          {{jsonKeys[0], 11}, {jsonKeys[2], 30}, {jsonKeys[3], 40}},
+      }),
+  });
+
+  testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
+}
+
 TEST_F(MapUnionSumTest, floatingPointKeys) {
   auto data = makeRowVector({
       makeFlatVector<int32_t>({1, 2, 1, 2, 1, 1, 2, 2}),