Repository navigation
WindowAggExec and grouped TopK can exceed batch_size #26772
Description
WindowAggExec and grouped TopK return their entire result in one RecordBatch. They ignore the datafusion.execution.batch_size setting and can cause a single poll to do much more work than expected. Another side-effect being the next operator can receive a large batch.
References:
datafusion/datafusion/physical-plan/src/windows/window_agg_exec.rs
Lines 607 to 627 in 350ca76
loop { return Poll::Ready(Some(match ready!(self.input.poll_next_unpin(cx)) { Some(Ok(batch)) => { self.batches.push(batch); continue; } Some(Err(e)) => Err(e), None => { // Release the input pipeline's resources before computing // the final aggregates. let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); let Some(result) = self.compute_aggregates()? else { return Poll::Ready(None); }; self.finished = true; // Empty record batches should not be emitted. // They need to be treated as [`Option<RecordBatch>`]es and handled separately debug_assert!(result.num_rows() > 0); Ok(result) } datafusion/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
Lines 268 to 293 in 350ca76
// inner is done, emit all rows and switch to producing output None => { // Release the input pipeline's resources before emitting. let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); if self.priority_map.is_empty() && !self.null_group_seen { trace!("partition {} emit None", self.partition); self.done = true; return Poll::Ready(None); } let group_by_metrics = self.group_by_metrics.clone(); let batch = group_by_metrics.time_emitting(|| { let cols = self.emit_columns()?; RecordBatch::try_new(Arc::clone(&self.schema), cols) })?; let batch = batch.record_output(&self.baseline_metrics); trace!( "partition {} emit batch with {} rows", self.partition, batch.num_rows() ); if log::log_enabled!(Level::Trace) { print_batches(std::slice::from_ref(&batch))?; } self.done = true; return Poll::Ready(Some(Ok(batch)));
Related:
- Operators should respect batch_size instead of Emit::All #19481 - attempts to cover all blocking operators emitting oversized final batches
- perf: emit WindowAggExec output in batch_size chunks #23206 - originally identifies the gap for
WindowAggExecbut was closed without merging.
I believe there's a simpler fix here for the above 2 cases, both already materialize their final result so we can just use RecordBatch::slice method and it should be sufficient for these paths. Something along the lines of #23206 implemented
Activity