Fix collect_set result

facebookincubator · Aug 13, 2024 · c97c479 · c97c479
1 parent de46a38
commit c97c479
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/velox/functions/lib/aggregates/SetBaseAggregate.h b/velox/functions/lib/aggregates/SetBaseAggregate.h
@@ -50,7 +50,13 @@ class SetBaseAggregate : public exec::Aggregate {
     for (auto i = 0; i < numGroups; ++i) {
       auto* group = groups[i];
       if (isNull(group)) {
-        arrayVector->setNull(i, true);
+        if constexpr (ignoreNulls) {
+          // When ignoring nulls, if the group's accumulator is null, the
+          // corresponding result is an empty array.
+          clearNull(rawNulls, i);
+        } else {
+          arrayVector->setNull(i, true);
+        }
       } else {
         clearNull(rawNulls, i);
 

diff --git a/velox/functions/sparksql/aggregates/tests/CollectSetAggregateTest.cpp b/velox/functions/sparksql/aggregates/tests/CollectSetAggregateTest.cpp
@@ -88,6 +88,27 @@ TEST_F(CollectSetAggregateTest, global) {
       {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});
 }
 
+TEST_F(CollectSetAggregateTest, noInputRow) {
+  auto data = makeRowVector({makeFlatVector<int32_t>({})});
+  auto expected = makeRowVector({
+      makeArrayVectorFromJson<int32_t>({"[]"}),
+  });
+  testAggregations({data}, {}, {"collect_set(c0)"}, {}, {expected});
+
+  // No input row passes aggregate filter.
+  const auto size = 1'000;
+  data = makeRowVector({
+      makeFlatVector<int32_t>(size, [](auto row) { return row; }),
+      makeFlatVector<bool>(size, [](auto /*row*/) { return false; }),
+  });
+
+  auto op = exec::test::PlanBuilder()
+                .values({data})
+                .singleAggregation({}, {"collect_set(c0)"}, {"c1"})
+                .planNode();
+  assertQuery(op, expected);
+}
+
 TEST_F(CollectSetAggregateTest, groupBy) {
   auto data = makeRowVector({
       makeFlatVector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),

diff --git a/velox/functions/sparksql/fuzzer/tests/CMakeLists.txt b/velox/functions/sparksql/fuzzer/tests/CMakeLists.txt
@@ -28,6 +28,6 @@ target_link_libraries(
     velox_exec_test_lib
     velox_parse_utils
     velox_vector_test_lib
-    gtest
-    gtest_main
+    GTest::gtest
+    GTest::gtest_main
     Folly::folly)