Skip to content

Commit 49c0aa3

Browse files
authored
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.
1 parent 7570366 commit 49c0aa3

1 file changed

Lines changed: 42 additions & 1 deletion

File tree

  • datafusion/physical-plan/src/topk

‎datafusion/physical-plan/src/topk/mod.rs‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1876,7 +1876,6 @@ impl PartitionedTopKRank {
18761876
coalescer.push_batch(batch)?;
18771877
}
18781878
for tie in ties {
1879-
(&tie.batch).record_output(&metrics.baseline);
18801879
coalescer.push_batch(tie.batch)?;
18811880
}
18821881
}
@@ -4138,6 +4137,48 @@ mod tests {
41384137
Ok(())
41394138
}
41404139

4140+
/// Tie rows are emitted through the same coalescer as heap rows, so they
4141+
/// must be counted in `output_rows` once, not once as a tie batch and
4142+
/// again as part of the coalesced output batch.
4143+
#[tokio::test]
4144+
async fn test_partitioned_topk_rank_output_rows_counts_ties_once() -> Result<()> {
4145+
let schema = pk_val_schema(false);
4146+
let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?;
4147+
let pk_sort_expr = PhysicalSortExpr {
4148+
expr: Arc::clone(&pk_expr),
4149+
options: SortOptions::default(),
4150+
};
4151+
let val_sort_expr = PhysicalSortExpr {
4152+
expr: col("val", schema.as_ref())?,
4153+
options: SortOptions::default(),
4154+
};
4155+
let metrics = ExecutionPlanMetricsSet::new();
4156+
let mut state = PartitionedTopKRank::try_new(
4157+
0,
4158+
Arc::clone(&schema),
4159+
vec![pk_expr],
4160+
build_sort_fields(&[pk_sort_expr], &schema)?,
4161+
LexOrdering::from([val_sort_expr]),
4162+
2,
4163+
8, // batch_size
4164+
&Arc::new(RuntimeEnv::default()),
4165+
&metrics,
4166+
)?;
4167+
4168+
// Two 5s fill the heap, the third 5 is retained as a tie.
4169+
let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 5, 10, 5])?;
4170+
state.insert_batch(&batch)?;
4171+
4172+
let results: Vec<RecordBatch> = state.emit()?.try_collect().await?;
4173+
let emitted_rows: usize = results.iter().map(|b| b.num_rows()).sum();
4174+
assert_eq!(emitted_rows, 3);
4175+
assert_eq!(
4176+
output_batches_and_rows(&metrics),
4177+
(results.len(), emitted_rows)
4178+
);
4179+
Ok(())
4180+
}
4181+
41414182
/// RANK-specific: heap fills with K rows tied at value V, equal_indices
41424183
/// accumulate at V, then a strictly-better row arrives whose admission
41434184
/// shifts the boundary strictly below V. The boundary-changed branch

0 commit comments

Comments
 (0)