sunchao commented on code in PR #25312:
URL: https://github.com/apache/datafusion/pull/25312#discussion_r4027945703
##########
datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs:
##########
@@ -146,199 +172,245 @@ impl OrderedPartialAggregateStream {
.register(context.memory_pool());
Ok(Self {
- schema,
- input,
reservation,
- baseline_metrics,
- reduction_factor,
- table: Some(table),
+ context: OrderedPartialAggregateContext {
+ schema,
+ batch_size,
+ baseline_metrics,
+ reduction_factor,
+ },
+ stage: ExecutionStage::Aggregating(Aggregating { input, table }),
})
}
- pub(crate) fn into_stream(self) -> SendableRecordBatchStream {
- let schema_clone = Arc::clone(&self.schema);
-
- let cloned_metrics = self.baseline_metrics.clone();
- let stream = Box::pin(RecordBatchStreamAdapter::new(
- schema_clone,
- self.create_stream(),
- ));
-
- Box::pin(ObservedStream::new(stream, cloned_metrics, None))
- }
-
- /// Entry point for the ordered partial aggregate state machine.
- ///
- /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas.
+ /// Entry point for the ordered partial aggregate execution stages.
///
- /// State transitions are implemented using the generator pattern; see the
comments in [`async_try_stream`].
+ /// See [`OrderedPartialAggregateStream`] for high-level ideas.
///
- /// Conceptual state-transition graph:
+ /// # Stage transition graph:
///
/// ```text
- /// (start)
- /// -> ReadingInput
- /// The stream starts by polling ordered input and aggregating batches
- /// into the ordered partial aggregate table.
+ /// +----[2]----+ +----[5]----+
+ /// | | | |
+ /// v | v |
+ /// +-------------------+ +-------------------+
+ /// | | | |
+ /// (start)-[1]->| Aggregating |-----[3]---->| Outputting |
+ /// | |<----[6]-----| |
+ /// +-------------------+ +-------------------+
+ /// | [4] | [7]
+ /// | |
+ /// +----------------+----------------+
+ /// |
+ /// v
+ /// +---------+
+ /// | Done |--[8]--> (end)
+ /// +---------+
+ /// ```
+ ///
+ /// ## Stages
+ ///
+ /// - [`Aggregating`]: Aggregates raw input and materializes one batch of
partial
+ /// states.
+ /// - [`Outputting`]: Emits slices of one materialized batch. If the
materialized
+ /// buffers cannot be reserved while slicing, hand off the whole batch,
then
+ /// resume aggregation or finish as described below.
///
- /// ReadingInput
- /// -> ReadingInput
- /// Aggregate one input batch. If the ordering proves some groups are
- /// complete, yield one partial-state batch immediately, then continue
- /// reading input. Otherwise continue directly with the next input
batch.
- /// -> DrainingFinal
- /// Input was exhausted. Mark the table input as done so every
remaining
- /// group is safe to emit.
+ /// ### Incremental output
///
- /// DrainingFinal
- /// -> DrainingFinal
- /// One remaining partial-state batch was yielded; repeat to continue
- /// draining the table.
- /// -> Done
- /// All remaining groups were emitted.
+ /// Consider this query with input ordered only by `k1`:
///
- /// Done
- /// -> (end)
+ /// ```sql
+ /// SELECT k1, k2, AVG(v)
+ /// FROM table_with_order_k1
+ /// GROUP BY k1, k2
/// ```
- fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> {
- async_try_stream(|mut emitter| async move {
- let mut table = self
- .table
- .take()
- .expect("OrderedPartialAggregateStream state should not be
None");
-
- self.handle_reading_input(&mut table, &mut emitter).await?;
-
- // Input has exhausted, move to the final draining stage.
- self.close_input();
- table.input_done();
-
- self.handle_draining_final(table, &mut emitter).await?;
-
+ ///
+ /// Suppose one `k1` value spans 1M rows with distinct, unordered `k2`
+ /// values. Ordering only proves these 1M `(k1, k2)` groups complete when
+ /// `k1` changes, so a single early emission can produce far more than
+ /// `batch_size` rows.
+ ///
+ /// Emitting those groups in small batches through [EmitTo::First] would
+ /// repeatedly remove a prefix from [`GroupValues`]. Because the group
+ /// values are stored contiguously, each removal copies the remaining
values
+ /// and updates their group indexes.
+ ///
+ /// To avoid repeating that work, this stream:
+ ///
+ /// 1. Materializes all completed groups into one large batch.
+ /// 2. Emits `batch_size` slices that share the batch's buffers.
+ ///
+ /// Blocked aggregate state management may simplify this approach:
+ /// <https://github.com/apache/datafusion/issues/24704>
+ ///
+ /// [`GroupValues`]: crate::aggregates::group_values::GroupValues
+ /// [EmitTo::First]: datafusion_expr::EmitTo::First
+ ///
+ ///
+ /// ## Transition Edges
+ ///
+ /// 1. Start.
+ /// 2. Aggregate one input batch. If memory fits and no groups are
complete,
+ /// continue reading input.
+ /// 3. Prepare output:
+ /// - Ordering proves a prefix complete: materialize the entire prefix
once,
+ /// retaining the input and active groups to resume aggregation.
+ /// - On memory pressure with partial ordering, materialize all current
+ /// states instead, including incomplete groups, and reset the table.
+ /// - At EOF, materialize all remaining states and prepare to output.
+ /// 4. Input was exhausted with no remaining groups, directly end.
+ /// 5. Yield one slice without materializing the table again. Keep the
shared
+ /// buffers reserved until handing off the last slice.
+ /// 6. The batch was fully emitted and retained aggregation can resume.
+ /// 7. The output batch was fully emitted.
+ /// 8. End.
+ pub(crate) fn into_stream(self) -> SendableRecordBatchStream {
+ let Self {
+ reservation,
+ context,
+ stage,
+ } = self;
+ let schema = Arc::clone(&context.schema);
+ let metrics = context.baseline_metrics.clone();
+ let stream = async_try_stream(|mut emitter| async move {
+ let mut stage = Some(stage);
+ while let Some(current_stage) = stage {
+ stage = match current_stage {
+ ExecutionStage::Aggregating(aggregating) => {
+ aggregating.handle_stage(&context, &reservation).await?
+ }
+ ExecutionStage::Outputting(outputting) => {
+ outputting
+ .handle_stage(&context, &reservation, &mut emitter)
+ .await?
+ }
+ };
+ }
Ok(())
- })
- }
-
- fn close_input(&mut self) {
- let input_schema = self.input.schema();
- self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
+ });
+ let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream));
+ Box::pin(ObservedStream::new(stream, metrics, None))
}
+}
- /// Consumes one ordered input batch, then immediately emits completed
groups
- /// if the ordering proves any group is ready.
+impl Aggregating {
+ /// Aggregates raw input and materializes one batch of partial states.
///
- /// See comments at [`Self::create_stream`] for details.
- async fn handle_reading_input(
- &mut self,
- table: &mut OrderedAggregateTable<PartialMarker>,
- emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
- ) -> Result<()> {
- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ /// See [`OrderedPartialAggregateStream::into_stream`] for stage
transitions.
+ async fn handle_stage(
+ mut self,
+ context: &OrderedPartialAggregateContext,
+ reservation: &MemoryReservation,
+ ) -> Result<Option<ExecutionStage>> {
+ let elapsed_compute = context.baseline_metrics.elapsed_compute();
while let Some(batch) = self.input.next().await.transpose()? {
- let input_rows = batch.num_rows();
- self.reduction_factor.add_total(input_rows);
-
+ context.reduction_factor.add_total(batch.num_rows());
let timer = elapsed_compute.timer();
-
- table.aggregate_batch(&batch)?;
-
- // Check memory reservation. See function comments for details.
- if let Some(batch) = self.resize_or_take_state_batch(table)? {
- self.reduction_factor.add_part(batch.num_rows());
- drop(timer);
- emitter.emit(batch).await;
- continue;
- }
-
- let Some(batch) = table.next_output_batch()? else {
- // Can't do early emit, continue aggregating.
+ self.table.aggregate_batch(&batch)?;
+
+ let output = match
reservation.try_resize(self.table.memory_size()) {
+ Ok(()) => self.table.take_completed_state_batch()?,
+ Err(oom @ DataFusionError::ResourcesExhausted(_)) => {
+ // Partial ordering may have an unbounded active key range.
+ // The final stage can merge incomplete states emitted
here.
+ if matches!(self.table.group_ordering(),
GroupOrdering::Full(_)) {
+ return Err(oom);
+ }
+ let Some(batch) = self.table.take_state_batch()? else {
+ return Err(oom);
+ };
+ Some(batch)
+ }
+ Err(e) => return Err(e),
+ };
+ let Some(batch) = output else {
continue;
};
- self.reduction_factor.add_part(batch.num_rows());
- self.reservation.try_resize(table.memory_size())?;
-
- drop(timer);
- emitter.emit(batch).await;
- }
-
- Ok(())
- }
-
- /// Update the memory reservation, and:
- /// - If memory reservation succeed, returns `Ok(None)`
- /// - If memory reservation failed,
- /// - If input is partially ordered, materialize all the output, and
- /// directly send them to the final aggregation stage.
- /// Returns `Ok(Some(batch))`
- /// - If input is fully ordered, directly return error. It's not
- /// expected to use more than constant memory.
- /// Returns `Err(..)`
- ///
- /// # Implementation Note
- /// Incrementally output it after the blocked state management is ready,
keep
- /// it simple for now.
- ///
- /// Issue: <https://github.com/apache/datafusion/issues/7065>
- fn resize_or_take_state_batch(
- &mut self,
- table: &mut OrderedAggregateTable<PartialMarker>,
- ) -> Result<Option<RecordBatch>> {
- let oom = match self.reservation.try_resize(table.memory_size()) {
- Ok(()) => return Ok(None),
- Err(e @ DataFusionError::ResourcesExhausted(_)) => e,
- Err(e) => return Err(e),
- };
+ timer.done();
- if matches!(table.group_ordering(), GroupOrdering::Full(_)) {
- return Err(oom);
+ // OOM, do early emit next, and go back to the current state to
continue
+ // aggregating
+ return Ok(Some(ExecutionStage::Outputting(Outputting {
+ batch,
+ resume: Some(self),
+ })));
}
- let Some(batch) = table.take_state_batch()? else {
- return Err(oom);
+ // Release upstream resources before draining the remaining states.
+ drop(self.input);
+ self.table.input_done();
+ let timer = elapsed_compute.timer();
+ let output = self.table.take_completed_state_batch()?;
+ drop(self.table);
+ timer.done();
+
+ let Some(batch) = output else {
+ reservation.try_resize(0)?;
+ return Ok(None);
};
- self.reservation.try_resize(table.memory_size())?;
- Ok(Some(batch))
+ Ok(Some(ExecutionStage::Outputting(Outputting {
+ batch,
+ resume: None,
+ })))
}
+}
- /// Emits one batch after input is exhausted.
- ///
- /// `table.input_done()` has already made every remaining group safe to
emit,
- /// so this state keeps draining until the table is empty.
- ///
- /// See comments at [`Self::create_stream`] for details.
+impl Outputting {
+ /// Emits slices of one materialized batch without touching the hash table.
///
- async fn handle_draining_final(
- &mut self,
- mut table: OrderedAggregateTable<PartialMarker>,
+ /// See [`OrderedPartialAggregateStream::into_stream`] for stage
transitions
+ /// and output memory accounting.
+ async fn handle_stage(
+ self,
+ context: &OrderedPartialAggregateContext,
+ reservation: &MemoryReservation,
emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
- ) -> Result<()> {
- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ ) -> Result<Option<ExecutionStage>> {
+ let Self { mut batch, resume } = self;
+ let elapsed_compute = context.baseline_metrics.elapsed_compute();
let mut timer = elapsed_compute.timer();
-
- while let Some(batch) = table.next_output_batch()? {
- self.reduction_factor.add_part(batch.num_rows());
-
- if table.is_empty() {
- // Clear memory before emitting last batch so we don't have to
wait for next poll to clear
- drop(table);
- let _ = self.reservation.try_resize(0);
- drop(timer);
-
+ let (table_memory, next_stage) = match resume {
+ Some(aggregating) => (
+ aggregating.table.memory_size(),
+ Some(ExecutionStage::Aggregating(aggregating)),
+ ),
+ None => (0, None),
+ };
+ let batch_memory = batch.get_array_memory_size();
+ match reservation.try_resize(table_memory + batch_memory) {
+ Ok(()) => {}
+ Err(DataFusionError::ResourcesExhausted(_)) => {
+ // If we cannot hold the batch while slicing, hand it off
whole.
+ // Only the retained table needs to remain reserved.
+ reservation.try_resize(table_memory)?;
+ context.reduction_factor.add_part(batch.num_rows());
+ timer.done();
emitter.emit(batch).await;
-
- return Ok(());
+ return Ok(next_stage);
}
+ Err(e) => return Err(e),
+ }
- self.reservation.try_resize(table.memory_size())?;
-
+ while batch.num_rows() > context.batch_size {
+ // 1. Emit first `batch_size` rows from `batch`
+ // 2. Update `batch`` with the remaining tail
+ let output = batch.slice(0, context.batch_size);
+ batch =
+ batch.slice(context.batch_size, batch.num_rows() -
context.batch_size);
+ context.reduction_factor.add_part(output.num_rows());
timer.done();
- emitter.emit(batch).await;
+ emitter.emit(output).await;
Review Comment:
Agreed, fixing shared-buffer accounting in the merge is a reasonable way to
address this, and I confirmed that other aggregation paths already slice their
output. The operator in this reproducer is `SortPreservingMergeExec`.
I tested an isolated prototype that changes only `sorts/builder.rs`: use
`RecordBatchMemoryCounter` to count unique backing buffers across the currently
live batches plus the incoming batch, and recompute that total after consumed
batches are pruned. Upstream aggregate and sort-key cursor reservations were
unchanged.
All four previously failing cases pass with that change:
- Integer keys, 1 MiB pool: all 8,192 rows, with disk spilling both disabled
and enabled.
- 128-byte string keys, 3.5 MiB pool: all 4,096 rows, with disk spilling
both disabled and enabled.
Every aggregate count is 2, no spills occur, and all reservations are
released on stream drop. This supports the proposed merge-local fix for these
reproducers; we do not need to wait for the full blocked-storage redesign to
address this finding. The prototype was only a diagnostic: I have not validated
broader behavior or the performance of recomputing the live-buffer set.
The unchanged PR head still fails these cases, so the accounting change and
regression coverage should land before resolving this thread. Accounting should
follow currently live allocations, since a permanent set of previously seen
addresses would be unsafe after buffers are dropped and addresses reused.
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs:
##########
@@ -89,26 +89,22 @@ impl OrderedAggregateTable<PartialMarker> {
)
}
- /// Emits the next batch of partial state rows for groups proven complete
by
- /// the input ordering.
- ///
- /// For example, when the query is `GROUP BY a` and the input is ordered by
- /// `a`, seeing a latest input row with `a = 3` means all groups with `a <
3`
- /// are complete and safe to emit.
- ///
- /// Key steps:
- /// 1. Ask `group_ordering` to decide how many groups can be emitted
eagerly.
- /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and
- /// all `GroupsAccumulator`s.
- ///
- /// This may output small batches. Avoiding tiny batches is left to future
- /// ordered-aggregation optimizations.
- pub(in crate::aggregates) fn next_output_batch(
+ /// Materializes all groups proven complete by the input ordering, leaving
+ /// the active ordered-key range in the table.
+ pub(in crate::aggregates) fn take_completed_state_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- self.next_output_batch_inner(
+ if self.is_empty() {
+ return Ok(None);
+ }
+ let Some(emit_to) = self.group_ordering().emit_to() else {
+ return Ok(None);
+ };
+ self.materialize_groups(
+ emit_to,
HashAggregateAccumulator::state,
AccumulatorPhase::State,
)
Review Comment:
To clarify, this failure happens while constructing the combined output
array, before the stream reaches slicing.
Each input batch has its own valid dictionary. The reproducer has three
batches with 64 distinct strings each (`v000`–`v063`, `v064`–`v127`, and
`v128`–`v191`), and each batch's Int8 keys are only 0–63. There are 192
distinct values across the completed ordered range, but one Int8 dictionary can
address only 128 non-null values.
With `batch_size=32`, base materializes 32 groups at a time and constructs
an independent dictionary for each output. Head calls
`take_completed_state_batch()` for all 192 groups, so dictionary construction
overflows before entering `Outputting`. For the nested grouping key, the path
is `RowsGroupColumn::rows_to_array` -> `encode_array_if_necessary`; its
`expect` turns the overflow into a panic.
I reran both EOF and ordered-boundary cases: they pass on base and fail on
head; widening the dictionary keys to Int16 makes them pass. The independent
`ARRAY_AGG(Dictionary<Int8, Utf8>)` case also fails during state
materialization, with ordinary integer grouping keys and just one value per
group.
A broader fix could normalize the output representation with consistent
schema changes, or materialization could produce multiple independently
representable batches. Simply removing slicing or changing the panic to a
returned error would leave the query failure. The bound needs to account for
dictionary child cardinality as well as outer row count; 32 groups works for
these reproducers, but is not a general bound for nested aggregate states.
--
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]
