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