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-25535-6574a8c3b2c93ef9689a677c1cc1c0a4a844617f in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit b1467ea5c2b16cd68d78d2cf4c994e9cc55dd3b5 Author: Jay Zhan <[email protected]> AuthorDate: Mon Sep 21 12:16:00 2026 +0000 fix: PartitionedTopKRank counts tie rows twice in output_rows (#25535) ## Which issue does this PR close? - N/A — no issue filed; found while reading the code. ## Rationale for this change For a `RANK()` window Top-N (`PartitionedTopKExec: fn=rank`), the `output_rows` and `output_batches` metrics shown by `EXPLAIN ANALYZE` are too high whenever rows tie at the K-th value. `PartitionedTopKRank::emit` calls `record_output` on every tie batch before pushing it into the `BatchCoalescer`, and then calls `record_output` again on each completed batch coming out of that coalescer — which already contains those tie rows. Heap rows are only counted at the second site, so each tie row is counted twice and each tie batch adds a phantom output batch. Example: K = 2, one partition with values `5, 5, 10, 5`. Three rows are emitted in one batch, but the metrics report `output_rows=4`, `output_batches=2`. `PartitionedTopK` (`ROW_NUMBER`) and `PartitionedTopKDenseRank` only record at the coalescer output and are not affected. ## What changes are included in this PR? Remove the extra `record_output` call on tie batches so every emitted row is counted once, at the coalescer output, as in the other two operators. Query results are unchanged; only the metrics are corrected. ## What is the testing strategy for this PR? New unit test `test_partitioned_topk_rank_output_rows_counts_ties_once`, which emits heap rows plus a boundary tie and asserts `output_rows` / `output_batches` equal what the stream actually produced. It fails on `main` with `(2, 4)` vs `(1, 3)`. ## Are there any user-facing changes? `output_rows` / `output_batches` reported for `PartitionedTopKExec` with `fn=rank` are now accurate when ties are present. No API changes. --- datafusion/physical-plan/src/topk/mod.rs | 43 +++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 0ca700cb37..cec3681de9 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1876,7 +1876,6 @@ impl PartitionedTopKRank { coalescer.push_batch(batch)?; } for tie in ties { - (&tie.batch).record_output(&metrics.baseline); coalescer.push_batch(tie.batch)?; } } @@ -4138,6 +4137,48 @@ mod tests { Ok(()) } + /// Tie rows are emitted through the same coalescer as heap rows, so they + /// must be counted in `output_rows` once, not once as a tie batch and + /// again as part of the coalesced output batch. + #[tokio::test] + async fn test_partitioned_topk_rank_output_rows_counts_ties_once() -> Result<()> { + let schema = pk_val_schema(false); + let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?; + let pk_sort_expr = PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }; + let val_sort_expr = PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: SortOptions::default(), + }; + let metrics = ExecutionPlanMetricsSet::new(); + let mut state = PartitionedTopKRank::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + build_sort_fields(&[pk_sort_expr], &schema)?, + LexOrdering::from([val_sort_expr]), + 2, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &metrics, + )?; + + // Two 5s fill the heap, the third 5 is retained as a tie. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 5, 10, 5])?; + state.insert_batch(&batch)?; + + let results: Vec<RecordBatch> = state.emit()?.try_collect().await?; + let emitted_rows: usize = results.iter().map(|b| b.num_rows()).sum(); + assert_eq!(emitted_rows, 3); + assert_eq!( + output_batches_and_rows(&metrics), + (results.len(), emitted_rows) + ); + Ok(()) + } + /// RANK-specific: heap fills with K rows tied at value V, equal_indices /// accumulate at V, then a strictly-better row arrives whose admission /// shifts the boundary strictly below V. The boundary-changed branch --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
