jayzhan211 commented on code in PR #27129:
URL: https://github.com/apache/datafusion/pull/27129#discussion_r4499931577
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -217,6 +217,20 @@ mod skip_partial;
mod spill;
mod topk;
+/// Bounded batches passing between operators are not charged to the memory
pool.
+/// Partial-state batches larger than this may hold input-sized values and must
+/// be reserved while the final aggregate merges them.
+const MAX_IN_FLIGHT_STATE_BATCH_BYTES: usize = 64 * 1024;
+
+fn incoming_state_batch_reservation_bytes(batch: &RecordBatch) -> usize {
+ let bytes =
datafusion_common::utils::memory::get_record_batch_memory_size(batch);
+ if bytes > MAX_IN_FLIGHT_STATE_BATCH_BYTES {
+ bytes
+ } else {
+ 0
+ }
+}
Review Comment:
`MAX_IN_FLIGHT_STATE_BATCH_BYTES` is below a normal state batch (8192 rows
of two `Int64` columns is 131 KB), so bounded batches get charged and queries
that passed now fail. A `GROUP BY` with 24 `sum`s under a fair pool passes 4/4
on main and fails 4/4 here at `target_partitions=16` with 10-12 MB and at 64
partitions with 48 MB; a `Utf8View` `GROUP BY` at 20 MB goes from 488 spills to
881.
Charging only batches above 1 KiB per row restores those (498 spills) and
still refuses the `median` and `array_agg` states. The fixed-width batches in
`final_aggregate_reserves_incoming_state_batch` and
`final_hash_reserves_incoming_state_batch` become exempt, so they need a batch
with few rows and a large buffer, like the test suggested on
`clustered_final_reserves_incoming_state_batch`.
```suggestion
/// Batches passing between operators are bounded by `batch_size` rows and
are
/// not charged to the memory pool. A partial-state batch is charged only
when
/// its rows carry more than this many bytes on average, which means a row
holds
/// many input values (`median`, `array_agg`), so its size follows the input.
const MAX_IN_FLIGHT_STATE_BYTES_PER_ROW: usize = 1024;
/// Below this size a batch is never charged, whatever its row count.
const MIN_CHARGED_STATE_BATCH_BYTES: usize = 64 * 1024;
fn incoming_state_batch_reservation_bytes(batch: &RecordBatch) -> usize {
let bytes =
datafusion_common::utils::memory::get_record_batch_memory_size(batch);
let in_flight_bytes = batch
.num_rows()
.saturating_mul(MAX_IN_FLIGHT_STATE_BYTES_PER_ROW)
.max(MIN_CHARGED_STATE_BATCH_BYTES);
if bytes > in_flight_bytes { bytes } else { 0 }
}
```
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -701,7 +731,11 @@ impl FinalHashAggregateStream {
));
match resize_result {
- Ok(()) => {}
+ Ok(()) => {
+ if hit_soft_limit {
+ break;
+ }
+ }
Review Comment:
Moving the soft-limit `break` after `try_resize` makes a `DISTINCT ...
LIMIT` fail or read all its input when the pool refuses the batch that reaches
the limit. With a limit of 4000 and two 8192-row batches, main returns 8192
rows at every pool size; this branch errors at 64-128 KB and reads all 16384
rows at 256 KB. Please keep the early `break`: the next `try_resize` in
`produce_output_from_memory` releases the charge, and a test for this shape
would pin it.
```suggestion
if self.hit_soft_group_limit(hash_table) && !spilled {
break;
}
// Replace the input-batch reservation with the table footprint,
// spilling if the table cannot be reserved.
let resize_result =
self.reservation
.try_resize(Self::reservation_size_for_table(
hash_table,
spill_context.as_deref(),
));
match resize_result {
Ok(()) => {}
```
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -768,4 +789,17 @@ mod tests {
let limit = held_memory(&plan).await?;
assert_resources_exhausted_under_limit(&plan, limit).await
}
+
+ #[tokio::test]
+ async fn clustered_final_reserves_incoming_state_batch() -> Result<()> {
+ // Many partial SUM states for one key merge to a small final state.
+ let plan = sum_over_sorted_input(
+ AggregateMode::Final,
+ &["k"],
+ vec![vec![vec![1; 16_384]]],
+ )?;
+ assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full);
+ assert_resources_exhausted_under_limit(&plan, 64 * 1024).await?;
Review Comment:
This test still passes with the new charge turned off, because the table
alone grows past 64 KiB on 16,384 rows, so the clustered reservation has no
test. Please use a batch with few rows that retains a large buffer (the test
below fails with the charge off). The spill retry at lines 334-344 also needs a
`GroupClusteringMode::Partial(vec![0])` test: nothing in physical-plan or
core's `memory_limit` tests reaches it, and the hash retry is reached only
incidentally by two core tests.
<details>
<summary>Clustered reservation test, plus a hash retry test that fails with
the retry or the charge off</summary>
```rs
/// The table stays small: only the incoming batch can exceed the limit.
#[tokio::test]
async fn clustered_final_refuses_large_state_batch() -> Result<()> {
let plan =
sum_over_sorted_input(AggregateMode::Final, &["k"],
vec![vec![vec![1; 64]]])?;
// 64 rows that retain 1 MB per column, like the few rows of a `median`
state.
let state = plan.input().schema();
let values: arrow::array::ArrayRef =
Arc::new(Int64Array::from(vec![1_i64; 131_072]).slice(0, 64));
let batch =
RecordBatch::try_new(Arc::clone(&state), vec![Arc::clone(&values),
values])?;
let ordering =
plan.input().properties().output_ordering().unwrap().clone();
let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&state),
None)?
.try_with_sort_information(vec![ordering])?;
let plan = AggregateExec::try_new(
AggregateMode::Final,
plan.group_expr().clone(),
plan.aggr_expr().to_vec(),
vec![None],
Arc::new(TestMemoryExec::update_cache(&Arc::new(input))),
plan.input_schema(),
)?;
assert_eq!(plan.group_clustering_mode(), &GroupClusteringMode::Full);
assert_resources_exhausted_under_limit(&plan, 1024 * 1024).await?;
assert_output_unchanged_under_limit(&plan, 4 * 1024 * 1024).await
}
```
```rs
/// Spilling the table makes room for a state batch that did not fit next to
it.
#[tokio::test]
async fn final_hash_spills_to_admit_large_state_batch() -> Result<()> {
let schema = Arc::new(Schema::new(vec![
Field::new("k", DataType::Int64, false),
Field::new("v", DataType::Int64, false),
]));
let group_by = PhysicalGroupBy::new_single(vec![(col("k", &schema)?,
"k".into())]);
let aggregate = Arc::new(
AggregateExprBuilder::new(sum_udaf(), vec![col("v", &schema)?])
.schema(Arc::clone(&schema))
.alias("sum")
.build()?,
);
let empty = TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema),
None)?;
let state = AggregateExec::try_new(
AggregateMode::Partial,
group_by.clone(),
vec![Arc::clone(&aggregate)],
vec![None],
empty,
Arc::clone(&schema),
)?
.schema();
let batch = |rows: i64| {
RecordBatch::try_new(
Arc::clone(&state),
vec![
Arc::new(Int64Array::from((0..rows).collect::<Vec<_>>())) as
ArrayRef,
Arc::new(Int64Array::from(vec![1_i64; rows as usize])),
],
)
};
// 64 rows that retain 2 MB, like the few rows of a `median` state.
let large = batch(131_072)?.slice(0, 64);
let input = TestMemoryExec::try_new_exec(
&[vec![batch(4000)?, large]],
Arc::clone(&state),
None,
)?;
let plan = Arc::new(AggregateExec::try_new(
AggregateMode::Final,
group_by.as_final(),
vec![aggregate],
vec![None],
input,
Arc::clone(&state),
)?);
// Fits the large batch, but not next to the table of 4000 groups.
let runtime = RuntimeEnvBuilder::new()
.with_memory_pool(Arc::new(GreedyMemoryPool::new(2200 * 1024)))
.build_arc()?;
let context = Arc::new(
TaskContext::default()
.with_session_config(SessionConfig::new().set_bool(
"datafusion.execution.enable_migration_aggregate",
true,
))
.with_runtime(runtime),
);
let output = collect(plan.execute(0, context)?).await?;
assert!(plan.metrics().unwrap().spill_count().unwrap() > 0);
let rows: usize = output.iter().map(|b| b.num_rows()).sum();
assert_eq!(rows, 4000);
let mut total = 0;
for b in &output {
total += arrow::compute::sum(
b.column(1).as_primitive::<arrow::datatypes::Int64Type>(),
)
.unwrap_or(0);
}
assert_eq!(total, 4064);
Ok(())
}
```
</details>
##########
datafusion/physical-plan/src/aggregates/aggregate_stream.rs:
##########
@@ -372,6 +373,18 @@ impl AggregateStream {
loop {
let result = match this.input.next().await {
Some(Ok(batch)) => {
+ // Partial states can own the entire input. Reserve
the batch
+ // before merging it, then transfer that reservation
to the
+ // accumulator's state below.
+ let previous_size = this.reservation.size();
+ let batch_memory =
incoming_state_batch_reservation_bytes(&batch);
+ if this.mode.input_mode() ==
AggregateInputMode::Partial
+ && batch_memory > 0
+ && let Err(e) =
this.reservation.try_grow(batch_memory)
Review Comment:
Optional cleanup: `Single` and `Partial` modes walk every buffer here and
then ignore the result, so the size can be computed only for partial-state
input.
```suggestion
let batch_memory =
if this.mode.input_mode() ==
AggregateInputMode::Partial {
incoming_state_batch_reservation_bytes(&batch)
} else {
0
};
if batch_memory > 0
&& let Err(e) =
this.reservation.try_grow(batch_memory)
```
##########
datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs:
##########
@@ -446,16 +446,16 @@ fn hide_byte_counts() -> Settings {
/// Runs `case` as Final and as Single and returns the Final run.
///
-/// The Final and Single streams account in the same way, so this asserts that
-/// both make the same pool calls, with the same byte counts, once the stream
-/// name is replaced, and the test snapshots only the Final run.
+/// Final reserves each incoming partial-state batch before merging it, while
+/// Single consumes raw rows. Their pool calls differ, but their results must
+/// still match. The test snapshots the Final run.
async fn run_final_and_single(case: impl Fn(AggregateMode) -> Case) ->
Result<Run> {
let final_run = case(AggregateMode::Final).run().await?;
let single_run = case(AggregateMode::Single).run().await?;
assert_eq!(
Review Comment:
These test edits are no longer needed: with the threshold in place, all 8
tests in this file pass with main's version. Please restore it, so the check
that Final and Single make the same pool calls stays, along with the
34,000-byte limits and the `@r"` snapshots.
```sh
git checkout upstream/main --
datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs
```
##########
datafusion/physical-plan/src/aggregates/hash_stream.rs:
##########
@@ -680,6 +680,37 @@ impl FinalHashAggregateStream {
while let Some(batch) = self.input.next().await.transpose()? {
let _timer = elapsed_compute.timer();
+ // Partial states may retain input-sized buffers. Account for the
+ // incoming batch before merging it into the table.
+ let batch_memory = incoming_state_batch_reservation_bytes(&batch);
+ let admission = if batch_memory > 0 {
+ self.reservation.try_grow(batch_memory)
+ } else {
+ Ok(())
+ };
+ match admission {
+ Ok(()) => {}
+ Err(e @ DataFusionError::ResourcesExhausted(_)) => {
+ // Spill the current table to make room for this batch,
then
+ // retry admission. An individual batch that still cannot
fit
+ // must fail rather than proceed without a reservation.
+ let Some(spill_context) = spill_context.as_mut() else {
+ return Err(e.context(
+ "Final hash aggregate cannot spill because
temporary files are not enabled in the DiskManager",
+ ));
+ };
+ if hash_table.building_group_count() == 0 {
+ return Err(e);
+ }
+ let result = hash_table
Review Comment:
Optional cleanup: this repeats the spill sequence from the resize arm below,
and #26429 is extracting it as `self.spill(...)` in the same loop. A shared
helper here would remove the copy and make the two PRs easier to combine.
--
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]
