Fix partial aggregation (#10)

Summary:
HashAggregation operator could get into an invalid state where it is not
blocked, doesn't need input, but getOutput() returns null. This happened when
operator processed a "distinct" aggregation and reached a limit on the amount
of memory to use for partial aggregation. This state confused the Driver loop
and caused it to complete prematurely. This would result in either wrong
results or a crash.

The fix is to flush partial aggregation state when reached the limit.

Pull Request resolved: #10

Reviewed By: pedroerp

Differential Revision: D30203061

Pulled By: mbasmanova

fbshipit-source-id: 12a381717462e77d2254490ca01c98920a9b2ad7

velox/core/QueryCtx.h

@@ -98,7 +98,8 @@ class QueryCtx : public Context {
   }
 
   uint64_t maxPartialAggregationMemoryUsage() const {
-    return 1L << 24; // 16MB
+    return get<uint64_t>(
+        kMaxPartialAggregationMemory, kMaxPartialAggregationMemoryDefault);
   }
 
   uint64_t maxPartitionedOutputBufferSize() const {
@@ -195,6 +196,9 @@ class QueryCtx : public Context {
   static constexpr const char* kMaxLocalExchangeBufferSize =
       "max_local_exchange_buffer_size";
 
+  static constexpr const char* kMaxPartialAggregationMemory =
+      "max_partial_aggregation_memory";
+
   // Overrides the previous configuration. Note that this function is NOT
   // thread-safe and should probably only be used in tests.
   void setConfigOverridesUnsafe(
@@ -222,6 +226,9 @@ class QueryCtx : public Context {
 
   static constexpr uint64_t kMaxLocalExchangeBufferSizeDefault = 32UL << 20;
 
+  // 16MB
+  static constexpr uint64_t kMaxPartialAggregationMemoryDefault = 1L << 24;
+
   CancelPoolPtr cancelPool_;
   std::unique_ptr<memory::MemoryPool> pool_;
   memory::MappedMemory* mappedMemory_;

velox/exec/HashAggregation.cpp

@@ -408,6 +408,11 @@ RowVectorPtr HashAggregation::getOutput() {
     // Drop reference to input_ to make it singly-referenced at the producer and
     // allow for memory reuse.
     input_ = nullptr;
+
+    if (partialFull_) {
+      groupingSet_->resetPartial();
+      partialFull_ = false;
+    }
     return output;
   }
 

velox/exec/tests/AggregationTest.cpp

@@ -391,5 +391,45 @@ TEST_F(AggregationTest, allKeyTypes) {
       " GROUP BY c0, C1, C2, c3, C4, C5");
 }
 
+TEST_F(AggregationTest, partialAggregationMemoryLimit) {
+  auto vectors = {
+      makeRowVector({makeFlatVector<int32_t>(
+          100, [](auto row) { return row; }, nullEvery(5))}),
+      makeRowVector({makeFlatVector<int32_t>(
+          110, [](auto row) { return row + 29; }, nullEvery(7))}),
+      makeRowVector({makeFlatVector<int32_t>(
+          90, [](auto row) { return row - 71; }, nullEvery(7))}),
+  };
+
+  createDuckDbTable(vectors);
+
+  // Set an artificially low limit on the amount of data to accumulate in
+  // the partial aggregation.
+  CursorParameters params;
+  params.queryCtx = core::QueryCtx::create();
+
+  params.queryCtx->setConfigOverridesUnsafe({
+      {core::QueryCtx::kMaxPartialAggregationMemory, "100"},
+  });
+
+  // Distinct aggregation.
+  params.planNode = PlanBuilder()
+                        .values(vectors)
+                        .partialAggregation({0}, {})
+                        .finalAggregation({0}, {})
+                        .planNode();
+
+  assertQuery(params, "SELECT distinct c0 FROM tmp");
+
+  // Count aggregation.
+  params.planNode = PlanBuilder()
+                        .values(vectors)
+                        .partialAggregation({0}, {"count(1)"})
+                        .finalAggregation({0}, {"sum(a0)"})
+                        .planNode();
+
+  assertQuery(params, "SELECT c0, count(1) FROM tmp GROUP BY 1");
+}
+
 } // namespace
 } // namespace facebook::velox::exec::test

velox/exec/tests/OperatorTestBase.h

@@ -48,6 +48,13 @@ class OperatorTestBase : public testing::Test {
     return assertQuery(plan, splits, duckDbSql, sortingKeys);
   }
 
+  std::shared_ptr<Task> assertQuery(
+      const CursorParameters& params,
+      const std::string& duckDbSql) {
+    return test::assertQuery(
+        params, [&](exec::Task* /*task*/) {}, duckDbSql, duckDbQueryRunner_);
+  }
+
   std::shared_ptr<Task> assertQuery(
       const std::shared_ptr<const core::PlanNode>& plan,
       const std::string& duckDbSql) {

velox/exec/tests/QueryAssertions.cpp

@@ -524,7 +524,7 @@ std::shared_ptr<Task> assertQuery(
 }
 
 std::shared_ptr<Task> assertQuery(
-    CursorParameters& params,
+    const CursorParameters& params,
     std::function<void(exec::Task*)> addSplits,
     const std::string& duckDbSql,
     DuckDbQueryRunner& duckDbQueryRunner,

velox/exec/tests/QueryAssertions.h

@@ -84,7 +84,7 @@ std::shared_ptr<Task> assertQuery(
     std::optional<std::vector<uint32_t>> sortingKeys = std::nullopt);
 
 std::shared_ptr<Task> assertQuery(
-    CursorParameters& params,
+    const CursorParameters& params,
     std::function<void(exec::Task*)> addSplits,
     const std::string& duckDbSql,
     DuckDbQueryRunner& duckDbQueryRunner,