This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 7c5a81bfaa fix: count UnnestExec output batches after splitting
(#24482)
7c5a81bfaa is described below
commit 7c5a81bfaae4e1a431cc46013fabdb80114212fb
Author: Goutam Adwant <[email protected]>
AuthorDate: Thu Aug 20 18:03:15 2026 +0000
fix: count UnnestExec output batches after splitting (#24482)
## Which issue does this PR close?
- Closes #24469.
## Rationale for this change
`UnnestExec` currently records its baseline output metrics before
`BatchSplitStream` divides an oversized result. For example, unnesting
25 rows with a batch size of 10 emits batches containing `[10, 10, 5]`
rows, but reports `output_batches=1` instead of `3`.
This makes the metric disagree with the batches observed by downstream
operators and by users inspecting execution metrics.
## What changes are included in this PR?
- Move baseline output observation to the final stream boundary by
wrapping `BatchSplitStream` with the existing `ObservedStream`.
- Keep `elapsed_compute` measurement inside `UnnestStream`, where the
unnest work is performed.
- Add a regression test that verifies the emitted batch sizes and the
corresponding `output_batches` and `output_rows` metrics.
As a result, baseline output metrics, including `output_bytes`, are
recorded from the batches emitted by the final stream. As documented for
`BaselineMetrics`, `output_bytes` may overestimate memory when sliced
batches share underlying buffers.
## Are these changes tested?
Yes.
- The regression test fails before the production change with
`output_batches=1` and passes afterward with `output_batches=3`.
- All 13 `unnest` tests pass.
- All 1,760 `datafusion-physical-plan` unit tests and its documentation
tests pass.
- The extended workspace test command from `AGENTS.md` passes, including
all 503 SQL logic test files.
- All 107 `datafusion-cli` unit and integration tests pass.
- `cargo fmt --all -- --check` passes.
- `cargo clippy --all-targets --all-features -- -D warnings` passes.
- Rust documentation builds with warnings denied.
## Are there any user-facing changes?
Yes. `UnnestExec` now reports baseline output metrics at the
consumer-visible output boundary, so `output_batches` matches the number
of emitted batches. Query results, row ordering, batch boundaries,
plans, and public APIs are unchanged.
Co-authored-by: Andrew Lamb <[email protected]>
---
datafusion/physical-plan/src/unnest.rs | 65 +++++++++++++++++++++++++++++++---
1 file changed, 60 insertions(+), 5 deletions(-)
diff --git a/datafusion/physical-plan/src/unnest.rs
b/datafusion/physical-plan/src/unnest.rs
index 6dfc2a0e53..964225d71c 100644
--- a/datafusion/physical-plan/src/unnest.rs
+++ b/datafusion/physical-plan/src/unnest.rs
@@ -23,10 +23,10 @@ use std::task::{Poll, ready};
use super::metrics::{
self, BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder,
MetricCategory,
- MetricsSet, RecordOutput, SplitMetrics,
+ MetricsSet, SplitMetrics,
};
use super::{DisplayAs, ExecutionPlanProperties, PlanProperties};
-use crate::stream::{BatchSplitStream, EmptyRecordBatchStream};
+use crate::stream::{BatchSplitStream, EmptyRecordBatchStream, ObservedStream};
use crate::{
ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan,
RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream,
@@ -296,6 +296,7 @@ impl ExecutionPlan for UnnestExec {
let batch_size = context.session_config().batch_size();
let input = self.input.execute(partition, context)?;
let metrics = UnnestMetrics::new(partition, &self.metrics);
+ let baseline_metrics = metrics.baseline_metrics.clone();
let stream = Box::pin(UnnestStream {
input,
@@ -311,10 +312,15 @@ impl ExecutionPlan for UnnestExec {
// Chunking the input bounds each build to roughly `batch_size` rows,
but two cases
// can still produce an oversized batch (see `predict_output_lens`),
so the output
// goes through the shared splitter to make the bound unconditional.
- Ok(Box::pin(BatchSplitStream::new(
+ let stream = Box::pin(BatchSplitStream::new(
stream,
batch_size,
SplitMetrics::new(&self.metrics, partition),
+ ));
+ Ok(Box::pin(ObservedStream::new(
+ stream,
+ baseline_metrics,
+ None,
)))
}
@@ -651,7 +657,6 @@ impl UnnestStream {
// that with `None` rather than an empty batch.
if let Some(batch) = result? {
debug_assert!(batch.num_rows() > 0);
- (&batch).record_output(&self.metrics.baseline_metrics);
return Poll::Ready(Some(Ok(batch)));
}
continue;
@@ -1435,6 +1440,7 @@ mod tests {
use arrow::datatypes::{Field, Int32Type};
use datafusion_common::NullHandling;
use datafusion_common::test_util::batches_to_string;
+ use datafusion_physical_expr_common::metrics::MetricValue;
use insta::assert_snapshot;
// Create a GenericListArray with the following list values:
@@ -2213,6 +2219,19 @@ mod tests {
options: UnnestOptions,
depth: usize,
) -> Result<Vec<RecordBatch>> {
+ Ok(
+ unnest_at_depth_with_metrics(input, batch_size, options, depth)
+ .await?
+ .0,
+ )
+ }
+
+ async fn unnest_at_depth_with_metrics(
+ input: Vec<RecordBatch>,
+ batch_size: usize,
+ options: UnnestOptions,
+ depth: usize,
+ ) -> Result<(Vec<RecordBatch>, MetricsSet)> {
let input_schema = input[0].schema();
let output_schema =
Arc::new(Schema::new(vec![Field::new("l", DataType::Int32,
true)]));
@@ -2234,7 +2253,9 @@ mod tests {
.with_batch_size(batch_size),
),
);
- crate::common::collect(unnest.execute(0, task_ctx)?).await
+ let batches = crate::common::collect(unnest.execute(0,
task_ctx)?).await?;
+ let metrics = unnest.metrics().expect("UnnestExec exposes metrics");
+ Ok((batches, metrics))
}
/// The values an unnest produces, flattened across all output batches.
@@ -2329,6 +2350,40 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_unnest_stream_output_metrics_after_split() -> Result<()> {
+ let (batches, metrics) = unnest_at_depth_with_metrics(
+ vec![list_batch(&[Some(25)])],
+ 10,
+ UnnestOptions::default(),
+ 1,
+ )
+ .await?;
+
+ assert_eq!(
+ batches
+ .iter()
+ .map(RecordBatch::num_rows)
+ .collect::<Vec<_>>(),
+ vec![10, 10, 5]
+ );
+ let output_batches = metrics
+ .sum(|metric| matches!(metric.value(),
MetricValue::OutputBatches(_)))
+ .expect("output_batches metric exists")
+ .as_usize();
+ assert_eq!(output_batches, batches.len());
+ let output_rows = metrics
+ .sum(|metric| matches!(metric.value(), MetricValue::OutputRows(_)))
+ .expect("output_rows metric exists")
+ .as_usize();
+ assert_eq!(
+ output_rows,
+ batches.iter().map(RecordBatch::num_rows).sum::<usize>()
+ );
+
+ Ok(())
+ }
+
#[tokio::test]
async fn test_unnest_stream_chunking_preserves_null_handling() ->
Result<()> {
// NULL and empty lists each contribute one NULL output row under
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
