alamb opened a new issue, #24468:
URL: https://github.com/apache/datafusion/issues/24468
### Describe the bug
When `SortExec` runs in TopK mode (`ORDER BY ... LIMIT k` on unsorted input)
and the k result rows span more than one output batch, its `output_batches`
metric reports `1` regardless of how many batches the operator actually emits.
For example, a top-25 query at `datafusion.execution.batch_size = 10` emits
three batches (`[10, 10, 5]` rows) but reports `output_batches=1`. The
`output_rows` metric is correct.
### To Reproduce
The following standalone test (e.g. dropped into `datafusion/core/tests/`)
fails on `main`:
```rust
use std::sync::Arc;
use datafusion::common::Result;
use datafusion::physical_plan::metrics::MetricValue;
use datafusion::physical_plan::sorts::sort::SortExec;
use datafusion::physical_plan::{ExecutionPlan, collect};
use datafusion::prelude::*;
/// The `output_batches` metric of a TopK sort should equal the number of
/// batches the operator emits to its consumer.
#[tokio::test]
async fn topk_output_batches_metric_matches_emitted_batches() -> Result<()> {
// A top-25 over 100 unsorted rows, at batch_size 10, so the TopK result
// must be emitted as multiple batches.
let config = SessionConfig::new()
.with_batch_size(10)
.with_target_partitions(1);
let ctx = SessionContext::new_with_config(config);
ctx.sql("CREATE TABLE t AS SELECT value FROM range(0, 100)")
.await?
.collect()
.await?;
let df = ctx
.sql("SELECT value FROM t ORDER BY value DESC LIMIT 25")
.await?;
let plan = df.create_physical_plan().await?;
let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await?;
let emitted_sizes: Vec<usize> = batches.iter().map(|b|
b.num_rows()).collect();
let sort = find_sort(&plan).expect("plan should contain SortExec");
let metrics = sort.metrics().expect("SortExec should have metrics");
let output_batches = metrics
.sum(|m| matches!(m.value(), MetricValue::OutputBatches(_)))
.expect("output_batches metric should be present")
.as_usize();
// The 25 result rows arrive as three batches of at most batch_size rows
assert_eq!(emitted_sizes, vec![10, 10, 5]);
// ... so the metric should report three output batches
assert_eq!(
output_batches,
emitted_sizes.len(),
"output_batches metric disagrees with the number of emitted batches"
);
Ok(())
}
fn find_sort(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn
ExecutionPlan>> {
if plan.downcast_ref::<SortExec>().is_some() {
return Some(Arc::clone(plan));
}
plan.children().into_iter().find_map(find_sort)
}
```
Output:
```
assertion `left == right` failed: output_batches metric disagrees with the
number of emitted batches
left: 1
right: 3
```
The first assertion passes, i.e. the operator really does emit three
batches; only the metric disagrees.
### Expected behavior
`output_batches` matches the number of batches the operator emits to its
consumer (`3` in the reproducer above).
### Additional context
Only the TopK path is affected. `SortExec` without a `fetch`, and `SortExec`
with a `fetch` over already-sorted input (the `LimitStream` path), report
`output_batches` correctly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
