You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/29 09:21:10 UTC
(datafusion) branch gh-readonly-queue/main/pr-24692-0f55902640ac36ed09f2c7acccda968726891e6e created (now e1d7fe2498)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24692-0f55902640ac36ed09f2c7acccda968726891e6e
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at e1d7fe2498 fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation (#24692)
This branch includes the following new commits:
new e1d7fe2498 fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation (#24692)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation (#24692)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24692-0f55902640ac36ed09f2c7acccda968726891e6e
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e1d7fe249824b89176114aeebfa8aa7a21d09716
Author: ranflarion <ra...@flarion.io>
AuthorDate: Sat Aug 29 09:20:47 2026 +0000
fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation (#24692)
## Which issue does this PR close?
- Closes #24691.
## Rationale for this change
An ungrouped aggregation can hold arbitrarily large accumulator state
that the memory pool never sees. `AggregateStream` grows its reservation
only by per-batch `size()` deltas (`size_post.saturating_sub(size_pre)`
in `aggregate_batch`), so an accumulator that allocates its retained
state in its constructor and never resizes it is charged nothing for the
stream's lifetime. The grouped path already handles this:
`GroupsAccumulatorAdapter` charges `state.size()` when it creates each
accumulator. Embedding DataFusion 54.1.0 in a Spark accelerator we
measured Spark's runtime-filter `bloom_filter_agg` (an 8 MiB bit array
zero-filled at construction, one per concurrent task) running entirely
unaccounted, so a fair-spill pool could neither fail admission nor
pressure other consumers to spill.
## What changes are included in this PR?
One charge in `AggregateStream::new`: after registering the reservation,
`try_grow` the sum of the accumulators' initial `size()`. Per-batch
deltas compose on top of it, so there is no double counting.
`test_oom`'s ungrouped arm is updated for the earlier failure point:
with its 1-byte limit the median accumulator's initial 48 bytes now fail
admission at `execute_typed`, so that arm asserts the construction-time
`ResourcesExhausted` instead of collecting first; the grouped arms are
unchanged.
## Are these changes tested?
The updated `test_oom` is the regression pin: on `main` without the fix
the ungrouped stream constructs successfully and the arm fails, with the
fix it passes. The full `datafusion-physical-plan` lib suite passes
(1784 tests).
## Are there any user-facing changes?
Yes, behavioral: an ungrouped aggregation whose accumulators'
construction-time state exceeds the memory limit now fails with
`ResourcesExhausted` at `execute` instead of silently running past the
configured limit, and memory-limited setups that relied on that
overshoot may need limits raised to what the query actually uses. No API
changes.
---
.../src/aggregates/aggregate_stream.rs | 5 +
datafusion/physical-plan/src/aggregates/mod.rs | 148 ++++++++++++++++++++-
2 files changed, 148 insertions(+), 5 deletions(-)
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs
index 7666fc6da7..a3fc731cb8 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs
@@ -324,6 +324,11 @@ impl AggregateStream {
let reservation = MemoryConsumer::new(format!("AggregateStream[{partition}]"))
.register(context.memory_pool());
+ // Accumulators that allocate their state at construction are invisible to the
+ // per-batch size deltas in `aggregate_batch`, so charge their initial sizes up
+ // front (mirroring `GroupsAccumulatorAdapter`, which charges `state.size()` when
+ // it creates each accumulator).
+ reservation.try_grow(accumulators.iter().map(|accum| accum.size()).sum())?;
// Enable dynamic filter if:
// 1. AggregateExec did the check and ensure it supports the dynamic filter
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 69a24ce17e..375e392e6a 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -3235,7 +3235,7 @@ mod tests {
use crate::test::TestMemoryExec;
use crate::test::assert_is_pending;
use crate::test::exec::{
- BlockingExec, StatisticsExec, assert_strong_count_converges_to_zero,
+ BlockingExec, PanicExec, StatisticsExec, assert_strong_count_converges_to_zero,
};
use arrow::array::{
@@ -4011,21 +4011,32 @@ mod tests {
Arc::clone(&input_schema),
)?);
- let stream = partial_aggregate.execute_typed(0, &task_ctx)?;
+ let stream = partial_aggregate.execute_typed(0, &task_ctx);
// ensure that we really got the version we wanted
- match version {
+ let stream = match version {
0 => {
- assert!(matches!(stream, StreamType::AggregateStream(_)));
+ // the ungrouped stream charges its accumulators' construction-time
+ // state up front, so admission fails before any input is read
+ let err = stream.err().unwrap();
+ assert!(
+ matches!(err.find_root(), DataFusionError::ResourcesExhausted(_)),
+ "Wrong error type: {err}",
+ );
+ continue;
}
1 => {
+ let stream = stream?;
assert!(matches!(stream, StreamType::GroupedHash(_)));
+ stream
}
2 => {
+ let stream = stream?;
assert!(matches!(stream, StreamType::SingleHash(_)));
+ stream
}
_ => panic!("Unknown version: {version}"),
- }
+ };
let stream: SendableRecordBatchStream = stream.into();
let err = collect(stream).await.unwrap_err();
@@ -4041,6 +4052,133 @@ mod tests {
Ok(())
}
+ /// An accumulator that allocates its retained state in its constructor and keeps that
+ /// size across `update_batch` (Spark's `bloom_filter_agg` shape) is charged at
+ /// `execute`, so it is rejected before any input is polled and, when it fits, is
+ /// visible to the pool for the stream's lifetime.
+ #[tokio::test]
+ async fn test_ungrouped_agg_charges_constructor_allocated_state() -> Result<()> {
+ const STATE_BYTES: usize = 8 * 1024 * 1024;
+
+ let schema: SchemaRef =
+ Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)]));
+ // panics when polled, so nothing here can pass by reading the input first
+ let input: Arc<dyn ExecutionPlan> =
+ Arc::new(PanicExec::new(Arc::clone(&schema), 1));
+ let udaf = Arc::new(AggregateUDF::from(ConstructorAllocatingUdaf::new(
+ STATE_BYTES,
+ )));
+ let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new(
+ AggregateExprBuilder::new(udaf, vec![col("b", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("constructor_allocating(b)")
+ .build()?,
+ )];
+ let aggregate = Arc::new(AggregateExec::try_new(
+ AggregateMode::Single,
+ PhysicalGroupBy::default(),
+ aggregates,
+ vec![None],
+ input,
+ Arc::clone(&schema),
+ )?);
+
+ let too_small = RuntimeEnvBuilder::new()
+ .with_memory_limit(STATE_BYTES / 2, 1.0)
+ .build_arc()?;
+ let task_ctx = Arc::new(TaskContext::default().with_runtime(too_small));
+ let err = aggregate
+ .execute_typed(0, &task_ctx)
+ .err()
+ .expect("initial accumulator state should not fit the pool");
+ assert!(
+ matches!(err.find_root(), DataFusionError::ResourcesExhausted(_)),
+ "Wrong error type: {err}",
+ );
+
+ let fits = RuntimeEnvBuilder::new()
+ .with_memory_limit(4 * STATE_BYTES, 1.0)
+ .build_arc()?;
+ let pool = Arc::clone(&fits.memory_pool);
+ let task_ctx = Arc::new(TaskContext::default().with_runtime(fits));
+ let stream = aggregate.execute_typed(0, &task_ctx)?;
+ assert!(matches!(stream, StreamType::AggregateStream(_)));
+ assert!(
+ pool.reserved() >= STATE_BYTES,
+ "initial accumulator state not charged: {} bytes reserved",
+ pool.reserved(),
+ );
+
+ Ok(())
+ }
+
+ /// UDAF whose accumulator allocates all of its state at construction, the shape the
+ /// per-batch `size()` deltas in `aggregate_batch` cannot observe.
+ #[derive(Debug, PartialEq, Eq, Hash)]
+ struct ConstructorAllocatingUdaf {
+ signature: Signature,
+ state_bytes: usize,
+ }
+
+ impl ConstructorAllocatingUdaf {
+ fn new(state_bytes: usize) -> Self {
+ Self {
+ signature: Signature::any(1, Volatility::Immutable),
+ state_bytes,
+ }
+ }
+ }
+
+ impl AggregateUDFImpl for ConstructorAllocatingUdaf {
+ fn name(&self) -> &str {
+ "constructor_allocating"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ Ok(DataType::Int64)
+ }
+
+ fn accumulator(
+ &self,
+ _acc_args: AccumulatorArgs,
+ ) -> Result<Box<dyn Accumulator>> {
+ Ok(Box::new(ConstructorAllocatingAccumulator {
+ state: vec![0; self.state_bytes],
+ }))
+ }
+ }
+
+ #[derive(Debug)]
+ struct ConstructorAllocatingAccumulator {
+ state: Vec<u8>,
+ }
+
+ impl Accumulator for ConstructorAllocatingAccumulator {
+ fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> {
+ Ok(())
+ }
+
+ fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> {
+ Ok(())
+ }
+
+ fn evaluate(&mut self) -> Result<ScalarValue> {
+ Ok(ScalarValue::Int64(Some(self.state.len() as i64)))
+ }
+
+ fn state(&mut self) -> Result<Vec<ScalarValue>> {
+ self.evaluate().map(|value| vec![value])
+ }
+
+ fn size(&self) -> usize {
+ size_of_val(self) + self.state.capacity()
+ }
+ }
+
#[tokio::test]
async fn partial_grouped_aggregate_uses_raw_partial_stream() -> Result<()> {
let (schema, batches) = some_data();
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org