From 9f742a43df5ff99b31939fffe845c6783a72281b Mon Sep 17 00:00:00 2001
From: Pradeep Vaka <pvaka@meta.com>
Date: Thu, 6 Feb 2025 12:02:37 -0800
Subject: [PATCH] Register map_union_sum(map<BOOLEAN, . >) function (#12273)

Summary:

- register map_union_sum(map<BOOLEAN, ..>)  function

Reviewed By: Yuhta, singcha

Differential Revision: D69215688
---
 .../aggregates/MapUnionSumAggregate.cpp       |  5 +++-
 .../aggregates/tests/MapUnionSumTest.cpp      | 28 +++++++++++++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp b/velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp
index d86ee3f407d2..431773c455e7 100644
--- a/velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp
+++ b/velox/functions/prestosql/aggregates/MapUnionSumAggregate.cpp
@@ -358,7 +358,8 @@ void registerMapUnionSumAggregate(
       "real",
       "double",
       "varchar",
-      "json"};
+      "json",
+      "boolean"};
   const std::vector<std::string> valueTypes = {
       "tinyint",
       "smallint",
@@ -416,6 +417,8 @@ void registerMapUnionSumAggregate(
           case TypeKind::VARCHAR:
             return createMapUnionSumAggregate<StringView>(
                 valueTypeKind, resultType);
+          case TypeKind::BOOLEAN:
+            return createMapUnionSumAggregate<bool>(valueTypeKind, resultType);
           default:
             VELOX_UNREACHABLE();
         }
diff --git a/velox/functions/prestosql/aggregates/tests/MapUnionSumTest.cpp b/velox/functions/prestosql/aggregates/tests/MapUnionSumTest.cpp
index 77328923d58f..9ac2f9d1f585 100644
--- a/velox/functions/prestosql/aggregates/tests/MapUnionSumTest.cpp
+++ b/velox/functions/prestosql/aggregates/tests/MapUnionSumTest.cpp
@@ -416,6 +416,30 @@ TEST_F(MapUnionSumTest, groupByJsonKey) {
   testAggregations({data}, {"c0"}, {"map_union_sum(c1)"}, {expected});
 }
 
+TEST_F(MapUnionSumTest, groupByBooleanKeys) {
+  auto data = makeRowVector({
+      makeFlatVector<int64_t>({1, 2, 1, 2, 1}),
+      makeNullableMapVector<bool, int64_t>({
+          {}, // empty map
+          std::nullopt, // null map
+          {{{true, 10}, {false, 20}}},
+          {{{true, 11}, {false, 30}, {true, 40}}},
+          {{{false, 28}, {true, 50}, {true, 12}}},
+      }),
+
+  });
+
+  auto expected = makeRowVector({
+      makeFlatVector<int64_t>({1, 2}),
+      makeMapVector<bool, int64_t>({
+          {{{true, 72}, {false, 48}}},
+          {{{true, 51}, {false, 30}}},
+      }),
+  });
+
+  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}),
@@ -442,8 +466,8 @@ TEST_F(MapUnionSumTest, floatingPointKeys) {
 }
 
 TEST_F(MapUnionSumTest, nanKeys) {
-  // Verify that NaNs with different binary representations are considered equal
-  // and deduplicated when used as keys in the output map.
+  // Verify that NaNs with different binary representations are considered
+  // equal and deduplicated when used as keys in the output map.
   constexpr double kNan = std::numeric_limits<double>::quiet_NaN();
   constexpr double kSNaN = std::numeric_limits<double>::signaling_NaN();