Fix TypedDistinctAggregations::extractValues to only call addSingleGr…

…oupRawInput when data exists.
kgpai · Feb 22, 2024 · 35bf30f · 35bf30f
1 parent 51007ea
commit 35bf30f
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 5 deletions.
diff --git a/velox/exec/DistinctAggregations.cpp b/velox/exec/DistinctAggregations.cpp
@@ -122,11 +122,13 @@ class TypedDistinctAggregations : public DistinctAggregations {
           accumulator->extractValues(*(data->template as<FlatVector<T>>()), 0);
         }
 
-        rows.resize(data->size());
-        std::vector<VectorPtr> inputForAggregation_ =
-            makeInputForAggregation(data);
-        aggregate.function->addSingleGroupRawInput(
-            group, rows, inputForAggregation_, false);
+        if (data->size() > 0) {
+          rows.resize(data->size());
+          std::vector<VectorPtr> inputForAggregation =
+              makeInputForAggregation(data);
+          aggregate.function->addSingleGroupRawInput(
+              group, rows, inputForAggregation, false);
+        }
       }
 
       aggregate.function->extractValues(

diff --git a/velox/functions/prestosql/aggregates/tests/AverageAggregationTest.cpp b/velox/functions/prestosql/aggregates/tests/AverageAggregationTest.cpp
@@ -707,5 +707,28 @@ TEST_F(AverageAggregationTest, companionFunctionsWithNonFlatAndLazyInputs) {
   testFunction("simple_avg");
 }
 
+/// We can get 0 as the count of a group when
+/// try and do a single aggregation over distinct values.
+///  In this case presto returns null as avg and not 'NaN'.
+TEST_F(AverageAggregationTest, zeroCounts) {
+  auto data = makeRowVector(
+      {makeNullableFlatVector<int64_t>({std::nullopt, 1}),
+       makeNullableFlatVector<int64_t>({2, 1}),
+       makeFlatVector<bool>({true, false})});
+
+  auto expected = makeRowVector({
+      makeNullableFlatVector<int64_t>({std::nullopt, 1}),
+      makeNullableFlatVector<double>({2.0, std::nullopt}),
+  });
+
+  auto op = PlanBuilder()
+                .values({data})
+                .project({"c0", "c1", "c2"})
+                .singleAggregation({"c0"}, {"avg(distinct c1)"}, {{"c2"}})
+                .planNode();
+
+  assertQuery(op, expected);
+}
+
 } // namespace
 } // namespace facebook::velox::aggregate::test