test: cover accumulator scratch release in clustered single aggregation

The clustered final and single streams also call ClusteredAggregateTable::shrink_to_fit after the memory-pressure emit. Add a Single-mode APPROX_DISTINCT test, and let the clustered test helper take the aggregate function.

Co-authored-by: Isaac <no-reply@databricks.com>

Author: viirya

datafusion/physical-plan/src/aggregates/clustered_single_stream.rs

@@ -762,9 +762,10 @@ mod tests {
     use super::*;
     use crate::aggregates::PhysicalGroupBy;
     use crate::aggregates::clustered_test_util::{
-        assert_completed_groups_emitted_under_limit, assert_output_unchanged_under_limit,
-        assert_resources_exhausted_under_limit, distinct_key_batches, full_batches,
-        held_memory, single_key_batches, sum_over_sorted_input,
+        aggregate_over_sorted_input, assert_completed_groups_emitted_under_limit,
+        assert_output_unchanged_under_limit, assert_resources_exhausted_under_limit,
+        distinct_key_batches, full_batches, held_memory, single_key_batches,
+        sum_over_sorted_input,
     };
     use crate::collect;
     use crate::test::TestMemoryExec;
@@ -776,6 +777,7 @@ mod tests {
         GreedyMemoryPool, MemoryPool, PeakRecordingPool,
     };
     use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+    use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf;
     use datafusion_functions_aggregate::sum::sum_udaf;
     use datafusion_physical_expr::PhysicalSortExpr;
     use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -874,6 +876,22 @@ mod tests {
         assert_output_unchanged_under_limit(&plan, 100 * 1024).await
     }
 
+    /// Like [`completed_groups_release_batch_capacity`], but `APPROX_DISTINCT`
+    /// also keeps the hashes of the last batch (64 KB for 8192 rows), which the
+    /// table must release too. The limit fits the active group, but not those
+    /// hashes.
+    #[tokio::test]
+    async fn completed_groups_release_accumulator_scratch() -> Result<()> {
+        let plan = aggregate_over_sorted_input(
+            AggregateMode::Single,
+            &["k"],
+            distinct_key_batches(),
+            approx_distinct_udaf(),
+        )?;
+        assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full);
+        assert_output_unchanged_under_limit(&plan, 64 * 1024).await
+    }
+
     /// A single key never completes a group before the input ends. The limit
     /// fits only the initial table.
     #[tokio::test]

datafusion/physical-plan/src/aggregates/clustered_test_util.rs

@@ -18,6 +18,7 @@ use datafusion_execution::memory_pool::{
     GreedyMemoryPool, MemoryPool, UnboundedMemoryPool,
 };
 use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+use datafusion_expr::AggregateUDF;
 use datafusion_functions_aggregate::sum::sum_udaf;
 use datafusion_physical_expr::PhysicalSortExpr;
 use datafusion_physical_expr::aggregate::AggregateExprBuilder;
@@ -63,6 +64,17 @@ pub(super) fn sum_over_sorted_input(
     mode: AggregateMode,
     keys: &[&str],
     batches: Vec<Vec<Vec<i64>>>,
+) -> Result<AggregateExec> {
+    aggregate_over_sorted_input(mode, keys, batches, sum_udaf())
+}
+
+/// Like [`sum_over_sorted_input`], but aggregates `v` with `udaf`. For `Final`
+/// mode, the state of `udaf` must be a single `Int64` column.
+pub(super) fn aggregate_over_sorted_input(
+    mode: AggregateMode,
+    keys: &[&str],
+    batches: Vec<Vec<Vec<i64>>>,
+    udaf: Arc<AggregateUDF>,
 ) -> Result<AggregateExec> {
     let fields = keys
         .iter()
@@ -76,9 +88,9 @@ pub(super) fn sum_over_sorted_input(
             .collect::<Result<_>>()?,
     );
     let aggregate_exprs = vec![Arc::new(
-        AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?])
+        AggregateExprBuilder::new(Arc::clone(&udaf), vec![col("v", &schema)?])
             .schema(Arc::clone(&schema))
-            .alias("sum")
+            .alias(udaf.name())
             .build()?,
     )];