sunchao commented on code in PR #26782:
URL: https://github.com/apache/datafusion/pull/26782#discussion_r4419021806
##########
datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs:
##########
@@ -653,6 +658,10 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
self.allocation_bytes
}
+ fn shrink_to_fit(&mut self) {
+ self.free_scratch();
Review Comment:
**[P2] Release the active group’s empty row-index buffer as well**
`free_scratch()` releases the adapter-wide `Scratch`, but each surviving
`AccumulatorState` also has an `indices: Vec<u32>`. `invoke_per_accumulator`
charges its capacity as it grows, then calls `indices.clear()` after consuming
the rows. The new hook leaves that empty allocation charged.
I reproduced this with a fully sorted partial `ANY_VALUE(v) GROUP BY k`, one
8,192-row batch, two keys, and a 16 KiB greedy pool. With keys `[0, 1, 1, …,
1]`, group 0 completes and the active group retains 8,192 index slots. The
stream then fails reserving another 27.1 KiB. Reversing the frequencies to `[0,
0, …, 0, 1]` returns both groups under the same limit. Independently, an
adapter-only probe measured 33,032 bytes after shrinking, including 32,768
bytes of empty indices. There is a completed group to emit here, which
distinguishes this from #26911. In a control build that released those empty
indices and updated the accounting, the original input returned both groups at
16 KiB.
Could `shrink_to_fit()` release unused index capacity in the remaining
states and subtract the released bytes from `allocation_bytes`? A skewed
active-group case would complement the current distinct-key fixture.
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -1648,4 +1649,63 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ /// Every row is a distinct group, so under memory pressure the stream
+ /// emits every state. After the emit, `APPROX_DISTINCT` must release the
+ /// hashes of the last batch (64 KB for 8192 rows), or reserving the empty
+ /// table fails the limit.
+ #[tokio::test]
+ async fn emitted_states_release_accumulator_scratch() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("k", DataType::Int64, false),
+ Field::new("v", DataType::Int64, false),
+ ]));
+ let batches = [0..8192, 8192..16384, 16384..24576]
+ .into_iter()
+ .map(|range| {
+ let values: ArrayRef =
Arc::new(Int64Array::from_iter_values(range));
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::clone(&values), values],
+ )
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+ let input = TestMemoryExec::try_new_exec(&[batches],
Arc::clone(&schema), None)?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::Partial,
+ PhysicalGroupBy::new_single(vec![(col("k", &schema)?,
"k".into())]),
+ vec![Arc::new(
+ AggregateExprBuilder::new(
+ approx_distinct_udaf(),
+ vec![col("v", &schema)?],
+ )
+ .schema(Arc::clone(&schema))
+ .alias("approx_distinct")
+ .build()?,
+ )],
+ vec![None],
+ input,
+ schema,
+ )?;
+ assert_eq!(
+ aggregate.group_clustering_mode(),
+ &GroupClusteringMode::None
+ );
+ let context =
+ Arc::new(
+ TaskContext::default()
+ .with_session_config(SessionConfig::new().set_bool(
+ "datafusion.execution.enable_migration_aggregate",
+ true,
+ ))
+ .with_runtime(
+ RuntimeEnvBuilder::new()
+
.with_memory_pool(Arc::new(GreedyMemoryPool::new(64 * 1024)))
+ .build_arc()?,
+ ),
+ );
+ let output = collect(aggregate.execute(0, context)?).await?;
Review Comment:
**Optional test cleanup.** The new `collect_under_limit(&aggregate, 64 *
1024)` already constructs a migration-enabled context and collects the stream.
Could this reuse it and assert `early_emit_count > 0`, as the existing
hash-stream test does? The row count alone could also pass if a future
skip-aggregation threshold sends the input straight through.
##########
datafusion/expr-common/src/groups_accumulator.rs:
##########
@@ -368,6 +368,14 @@ pub trait GroupsAccumulator: Send + std::any::Any {
///
/// May be expensive; check the implementation before calling on hot paths.
fn size(&self) -> usize;
+
+ /// Keeps the stored groups but releases capacity they do not need, such as
+ /// per-batch scratch buffers, so that [`Self::size`] reflects the stored
+ /// groups. For example, after [`EmitTo::First`] removes most groups, an
+ /// accumulator may keep a buffer it sized for the last input batch.
+ ///
+ /// The default does nothing.
+ fn shrink_to_fit(&mut self) {}
Review Comment:
**[P2] Ordered `FIRST_VALUE` / `LAST_VALUE` also retain counted batch
scratch**
`FirstLastGroupsAccumulator` uses the default hook. Its `take_state`
truncates `extreme_of_each_group_buf.1`, a `BooleanBufferBuilder` used to track
each batch's winners. Truncation preserves capacity, and `size()` still charges
that capacity. The bitmap reached 16 KiB in my test even though only one group
remained.
I reproduced this with a fully clustered partial `FIRST_VALUE(v ORDER BY v)
GROUP BY k` over one batch of 131,072 distinct sorted `Int64` keys (`v = k`).
The 16 KiB greedy pool refused another 11.1 KiB after the completed prefix was
emitted. With 32 KiB it returned all 131,072 rows. The 8,192-key control also
succeeded at 16 KiB. Rebuilding just the winner bitmap at the surviving group
count let the original 131,072-key input return all rows at 16 KiB.
Could the native first/last accumulator release its batch scratch through
this hook and cover the ordered case? The `ORDER BY` matters because it selects
the native implementation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
