fix: charge accumulators' construction-time state to the memory pool in ungrouped aggregation

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

@@ -301,6 +301,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

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

@@ -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();