comphead commented on code in PR #26637:
URL: https://github.com/apache/datafusion/pull/26637#discussion_r4361499205
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -335,7 +336,27 @@ impl ClusteredSingleAggregateStream {
let Some(spill_context) = spill_context else {
Review Comment:
The new `None` branch repeats the tail of the success path below:
`take_completed_result_batch`, `reservation_size_for_table`,
`start_outputting(.., ReadingInput { table, spill_context })`. It also adds two
match arms that both call `break_with_err`, with the second `e` shadowing the
OOM error.
With `spill_context == None` the `has_spills` check below is always false,
so the success path already does the right thing. Keeping the refused error in
an `Option` and falling through would remove the duplicate block. The only
extra logic is in the `Ok(None)` arm, which returns the saved error instead of
continuing. The same shape applies to
`ClusteredFinalAggregateStream::Aggregating::handle_stage`, where it saves a
third copy of the `Outputting { batch, resume: Some(self) }` construction.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -732,4 +744,87 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ /// Runs a final `SUM(k) GROUP BY k` over two partial-state batches sorted
on
+ /// `k` that share a boundary key. Returns the output and the most memory
the
+ /// stream held while handing off output.
+ async fn run_fully_sorted(
+ pool: Arc<dyn MemoryPool>,
+ ) -> Result<(Vec<RecordBatch>, usize)> {
+ let schema = Arc::new(Schema::new(vec![Field::new("k",
DataType::Int64, false)]));
+ let groups = PhysicalGroupBy::new_single(vec![(col("k", &schema)?,
"k".into())]);
+ let expressions = vec![Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("k", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("sum")
+ .build()?,
+ )];
+ let empty = TestMemoryExec::try_new_exec(&[vec![]],
Arc::clone(&schema), None)?;
+ let partial = AggregateExec::try_new(
+ AggregateMode::Partial,
+ groups.clone(),
+ expressions.clone(),
+ vec![None],
+ empty,
+ Arc::clone(&schema),
+ )?;
+ let partial_schema = partial.schema();
+ let batches = [0..=256, 256..=512]
+ .into_iter()
+ .map(|keys| {
+ let k = Int64Array::from_iter_values(keys);
+ RecordBatch::try_new(
+ Arc::clone(&partial_schema),
+ vec![Arc::new(k.clone()), Arc::new(k)],
+ )
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+ let ordering =
+ LexOrdering::new([PhysicalSortExpr::new_default(col("k",
&partial_schema)?)])
+ .unwrap();
+ let input =
+ TestMemoryExec::try_new(&[batches], Arc::clone(&partial_schema),
None)?
+ .try_with_sort_information(vec![ordering])?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::Final,
+ groups.as_final(),
+ expressions,
+ vec![None],
+ Arc::new(TestMemoryExec::update_cache(&Arc::new(input))),
+ schema,
+ )?;
+ assert!(matches!(
+ aggregate.group_clustering_mode(),
+ GroupClusteringMode::Full
+ ));
+ let context = Arc::new(
+ TaskContext::default().with_runtime(
+ RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&pool))
+ .build_arc()?,
+ ),
+ );
+ let mut stream =
+ ClusteredFinalAggregateStream::new(&aggregate, &context,
0)?.into_stream();
+ // Each output batch is a single slice, so the stream releases it
before
+ // handing it off and holds only the table.
+ let mut held = pool.reserved();
+ let mut output = vec![];
+ while let Some(batch) = stream.next().await.transpose()? {
+ held = held.max(pool.reserved());
+ output.push(batch);
+ }
+ Ok((output, held))
+ }
+
+ /// The limit fits the table with only the active group, but not the table
+ /// before the completed groups are emitted. The result must not change.
+ #[tokio::test]
+ async fn completed_groups_under_memory_pressure() -> Result<()> {
Review Comment:
The test only checks that a tight limit now succeeds. Nothing checks the
other half of the contract in the docs: the error is still returned when no
group can be released. A rerun with a limit just below `held` (or a single-key
input) that expects `ResourcesExhausted` would show the fix does not swallow
genuine exhaustion and exercises the `ok_or(oom)` and `Ok(None)` arms.
##########
datafusion/physical-plan/src/aggregates/clustered_single_stream.rs:
##########
@@ -851,4 +872,70 @@ mod tests {
);
Ok(())
}
+
+ /// Runs a single `SUM(k) GROUP BY k` over two batches sorted on `k` that
+ /// share a boundary key. Returns the output and the most memory the stream
+ /// held while handing off output.
+ async fn run_fully_sorted(
Review Comment:
This driver (context with `RuntimeEnvBuilder`, the `held` sampling loop,
unlimited run then `GreedyMemoryPool::new(held)` run, `batches_to_string`
comparison) is now the third copy, after `clustered_partial_stream.rs` on
`main` and `run_fully_sorted` in `clustered_final_stream.rs`. Only the plan
construction differs, and the final one needs partial-state input.
`AggregateExec::execute` already dispatches to the clustered stream by mode,
since `enable_migration_aggregate` defaults to `true`. One shared helper that
takes `Arc<dyn ExecutionPlan>`, runs it on both pools and returns the output
and `held` would leave each module with only its plan builder. The helper could
live next to the existing aggregate test utilities in `aggregates/mod.rs`.
The unit-test form makes sense here since a SQL memory limit cannot target
one stream's table size. If a SQL form is feasible,
`aggregate_memory_spill.slt` already sets a 1M limit and notes that sorted
`generate_series` takes the streaming path. I have not tried it, so it may be
too fragile.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -732,4 +744,87 @@ mod tests {
assert_eq!(pool.reserved(), 0);
Ok(())
}
+
+ /// Runs a final `SUM(k) GROUP BY k` over two partial-state batches sorted
on
+ /// `k` that share a boundary key. Returns the output and the most memory
the
+ /// stream held while handing off output.
+ async fn run_fully_sorted(
+ pool: Arc<dyn MemoryPool>,
+ ) -> Result<(Vec<RecordBatch>, usize)> {
+ let schema = Arc::new(Schema::new(vec![Field::new("k",
DataType::Int64, false)]));
+ let groups = PhysicalGroupBy::new_single(vec![(col("k", &schema)?,
"k".into())]);
+ let expressions = vec![Arc::new(
+ AggregateExprBuilder::new(sum_udaf(), vec![col("k", &schema)?])
+ .schema(Arc::clone(&schema))
+ .alias("sum")
+ .build()?,
+ )];
+ let empty = TestMemoryExec::try_new_exec(&[vec![]],
Arc::clone(&schema), None)?;
+ let partial = AggregateExec::try_new(
+ AggregateMode::Partial,
+ groups.clone(),
+ expressions.clone(),
+ vec![None],
+ empty,
+ Arc::clone(&schema),
+ )?;
+ let partial_schema = partial.schema();
+ let batches = [0..=256, 256..=512]
+ .into_iter()
+ .map(|keys| {
+ let k = Int64Array::from_iter_values(keys);
+ RecordBatch::try_new(
+ Arc::clone(&partial_schema),
+ vec![Arc::new(k.clone()), Arc::new(k)],
+ )
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+ let ordering =
+ LexOrdering::new([PhysicalSortExpr::new_default(col("k",
&partial_schema)?)])
+ .unwrap();
+ let input =
+ TestMemoryExec::try_new(&[batches], Arc::clone(&partial_schema),
None)?
+ .try_with_sort_information(vec![ordering])?;
+ let aggregate = AggregateExec::try_new(
+ AggregateMode::Final,
+ groups.as_final(),
+ expressions,
+ vec![None],
+ Arc::new(TestMemoryExec::update_cache(&Arc::new(input))),
+ schema,
+ )?;
+ assert!(matches!(
+ aggregate.group_clustering_mode(),
+ GroupClusteringMode::Full
Review Comment:
Both new tests assert `GroupClusteringMode::Full`. The changed branch is
`spill_context == None`, which is also true for `Partial` clustering with
temporary files disabled (`DiskManager` without tmp files). That case reaches
the new code with an active key range defined by the sorted prefix, and it is
not covered. A case with a two-column key sorted on the first column and tmp
files disabled would cover it, for the final and the single stream.
##########
datafusion/physical-plan/src/aggregates/clustered_final_stream.rs:
##########
@@ -340,7 +341,15 @@ impl Aggregating {
Ok(()) => {}
Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
let Some(spill_context) = self.spill_context.as_mut() else
{
- return Err(oom);
+ // Without spilling, emitting completed groups is the
only
Review Comment:
Design question for both streams. The fix uses a refused `try_resize` as the
trigger and then emits. The issue points at the legacy
`GroupedHashAggregateStream`, which emits completed groups before it updates
the reservation. For the non-spilling case, taking the completed groups first
and then sizing the reservation once would not need an error path, and a
failure would then mean the active group itself does not fit. Is there a reason
to prefer the error-driven fallback, for example keeping the success path
untouched? If so, a short comment saying why would help. Note that neither
variant lowers the real peak, since `aggregate_batch` has already allocated the
new groups. It only makes the accounting consistent.
--
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]
