viirya commented on code in PR #26782:
URL: https://github.com/apache/datafusion/pull/26782#discussion_r4425941190
##########
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:
Thanks, confirmed: your skewed input (`[0, 1, 1, …, 1]`) failed under 16 KiB
asking for 27.1 KB. In a8b5fb3be, `GroupsAccumulatorAdapter::shrink_to_fit`
also releases the empty `indices` of the remaining states and subtracts them
from `allocation_bytes`. It skips non-empty `indices`, which only a failed
update leaves behind. New test:
`clustered_partial_stream::tests::active_group_releases_adapter_indices`.
##########
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:
Good find, thanks. I missed this in the sweep.
`BooleanBufferBuilder::truncate` keeps the capacity, and `size()` counts it.
Your 131,072-key input failed under 16 KiB asking for 11.1 KB.
In a8b5fb3be, `FirstLastGroupsAccumulator::shrink_to_fit` rebuilds the
scoreboard bitmap at the remaining group count, all false, and clears
`extreme_buf_dirty`. Both users of the bitmap accept that state: the
pre-ordered path expects it all false between batches, and
`get_filtered_extreme_of_each_group` resets it when it starts. The other fields
don't need it: `extreme_of_each_group_buf.0` goes through `take_needed`, and
`is_sets` is rebuilt by `finish()`. New test:
`clustered_partial_stream::tests::completed_groups_release_first_value_bitmap`,
using `FIRST_VALUE(v ORDER BY v)`.
##########
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:
Done. It now uses `collect_under_limit` and asserts `early_emit_count > 0`.
--
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]
