Skip to content

WindowAggExec and grouped TopK can exceed batch_size #26772

Open

Description

@yashrb24

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:

  1. 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)
    }
  2. // 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:

  1. Operators should respect batch_size instead of Emit::All #19481 - attempts to cover all blocking operators emitting oversized final batches
  2. perf: emit WindowAggExec output in batch_size chunks #23206 - originally identifies the gap for WindowAggExec but 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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions