SortExec (TopK) `output_batches` metric does not match the number of emitted batches

[alamb]

Describe the bug

• Found while reviewing fix: make UnnestExec respect datafusion.execution.batch_size #24384

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.

To Reproduce

The following standalone test (e.g. dropped into datafusion/core/tests/) fails on main:

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

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.

[shinzoxD]

I'll fix SortExec TopK so output_batches matches the batches actually emitted after batch-size splitting.

[shinzoxD]

Opened #24496 to fix SortExec TopK so output_batches matches the batches actually emitted after batch-size splitting.

[shinzoxD]

take