You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/30 13:13:46 UTC
(datafusion) 01/01: Fix metrics for repartition when `preserve_order=true` (#20924)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-20924-4dadbbd58dcf0f2f8269ba77971fa32707e4922e
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit d1c875ae8aa515b9dd5741fa7dfe3ec91adc9f56
Author: Xander <za...@googlemail.com>
AuthorDate: Tue Jun 30 18:43:23 2026 +0530
Fix metrics for repartition when `preserve_order=true` (#20924)
## Which issue does this PR close?
Found this when working on
https://github.com/apache/datafusion/pull/20875/
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes #.
## Rationale for this change
Metric reporting was previously incorrect for `RepartitionExec` if
preserve order was set to true.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
Create new metrics before creating `PerPartitionStream`
## Are these changes tested?
Yes and confirmed that this fails on main:
```
thread 'repartition::test::test_preserve_order_output_rows_not_double_counted' (12487869)
panicked at datafusion/physical-plan/src/repartition/mod.rs:3007:9:
assertion `left == right` failed: metrics output_rows (8) should match actual rows
collected (4), not double-count
left: 8
right: 4
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
```
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example,
are they covered by existing tests)?
-->
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---
datafusion/physical-plan/src/repartition/mod.rs | 72 ++++++++++++++++++++++---
1 file changed, 64 insertions(+), 8 deletions(-)
diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
index 1059f0225f..1617af3a68 100644
--- a/datafusion/physical-plan/src/repartition/mod.rs
+++ b/datafusion/physical-plan/src/repartition/mod.rs
@@ -1299,6 +1299,12 @@ impl ExecutionPlan for RepartitionExec {
if preserve_order {
// Store streams from all the input partitions:
// Each input partition gets its own spill reader to maintain proper FIFO ordering
+ //
+ // Pass None for metrics here — these intermediate streams feed into
+ // StreamingMerge which is the actual output. Only the merge's
+ // BaselineMetrics should contribute to the operator's reported
+ // output_rows. Without this, every row would be counted twice
+ // (once by PerPartitionStream, once by StreamingMerge).
let input_streams = rx
.into_iter()
.zip(spill_readers)
@@ -1311,7 +1317,7 @@ impl ExecutionPlan for RepartitionExec {
Arc::clone(&reservation),
spill_stream,
1, // Each receiver handles one input partition
- BaselineMetrics::new(&metrics, partition),
+ None,
)) as SendableRecordBatchStream
})
.collect::<Vec<_>>();
@@ -1349,7 +1355,7 @@ impl ExecutionPlan for RepartitionExec {
reservation,
spill_stream,
num_input_partitions,
- BaselineMetrics::new(&metrics, partition),
+ Some(BaselineMetrics::new(&metrics, partition)),
)) as SendableRecordBatchStream)
}
})
@@ -1862,8 +1868,8 @@ struct PerPartitionStream {
/// each sending None when complete. We must wait for all of them.
remaining_partitions: usize,
- /// Execution metrics
- baseline_metrics: BaselineMetrics,
+ /// Execution metrics (None in preserve-order mode where StreamingMerge owns the metrics)
+ baseline_metrics: Option<BaselineMetrics>,
}
impl PerPartitionStream {
@@ -1874,7 +1880,7 @@ impl PerPartitionStream {
reservation: SharedMemoryReservation,
spill_stream: SendableRecordBatchStream,
num_input_partitions: usize,
- baseline_metrics: BaselineMetrics,
+ baseline_metrics: Option<BaselineMetrics>,
) -> Self {
Self {
schema,
@@ -1893,8 +1899,11 @@ impl PerPartitionStream {
cx: &mut Context<'_>,
) -> Poll<Option<Result<RecordBatch>>> {
use futures::StreamExt;
- let cloned_time = self.baseline_metrics.elapsed_compute().clone();
- let _timer = cloned_time.timer();
+ let elapsed = self
+ .baseline_metrics
+ .as_ref()
+ .map(|m| m.elapsed_compute().clone());
+ let _timer = elapsed.as_ref().map(|t| t.timer());
loop {
match self.state {
@@ -1980,7 +1989,11 @@ impl Stream for PerPartitionStream {
cx: &mut Context<'_>,
) -> Poll<Option<Self::Item>> {
let poll = self.poll_next_inner(cx);
- self.baseline_metrics.record_poll(poll)
+ if let Some(metrics) = &self.baseline_metrics {
+ metrics.record_poll(poll)
+ } else {
+ poll
+ }
}
}
@@ -3294,4 +3307,47 @@ mod test {
let exec = Arc::new(exec);
Arc::new(TestMemoryExec::update_cache(&exec))
}
+
+ /// preserve_order repartition should not double-count
+ /// output rows.
+ #[tokio::test]
+ async fn test_preserve_order_output_rows_not_double_counted() -> Result<()> {
+ use datafusion_execution::TaskContext;
+
+ // Two sorted input partitions, 2 rows each (4 total)
+ let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap();
+ let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap();
+ let schema = batch1.schema();
+ let sort_exprs = sort_exprs(&schema);
+
+ let input_partitions = vec![vec![batch1], vec![batch2]];
+ let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?
+ .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?;
+ let exec = Arc::new(exec);
+ let exec = Arc::new(TestMemoryExec::update_cache(&exec));
+
+ let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))?
+ .with_preserve_order();
+
+ let task_ctx = Arc::new(TaskContext::default());
+ let mut total_rows = 0;
+ for i in 0..exec.partitioning().partition_count() {
+ let mut stream = exec.execute(i, Arc::clone(&task_ctx))?;
+ while let Some(result) = stream.next().await {
+ total_rows += result?.num_rows();
+ }
+ }
+
+ assert_eq!(total_rows, 4, "actual rows collected should be 4");
+
+ let metrics = exec.metrics().unwrap();
+ let reported_output_rows = metrics.output_rows().unwrap();
+ assert_eq!(
+ reported_output_rows, total_rows,
+ "metrics output_rows ({reported_output_rows}) should match \
+ actual rows collected ({total_rows}), not double-count"
+ );
+
+ Ok(())
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org