PartitionedTopKExec does not expose any metrics (`metrics()` returns `None`)

[alamb]

Describe the bug

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

PartitionedTopKExec (the operator that enable_window_topn substitutes for per-partition top-K window queries) does not expose any metrics: ExecutionPlan::metrics() returns None, and EXPLAIN ANALYZE shows no metrics for it. This also means its output_batches / output_rows cannot be checked against what it actually emits.

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::partitioned_topk::PartitionedTopKExec;
use datafusion::physical_plan::{ExecutionPlan, collect};
use datafusion::prelude::*;

/// The `output_batches` metric of `PartitionedTopKExec` should equal the
/// number of batches the operator emits to its consumer.
#[tokio::test]
async fn partitioned_topk_output_batches_metric_matches_emitted_batches() -> Result<()> {
    // Top-1 per partition over 50 partition keys, at batch_size 10, so the 50
    // result rows are emitted as several coalesced batches.
    let mut config = SessionConfig::new()
        .with_batch_size(10)
        .with_target_partitions(1);
    config.options_mut().optimizer.enable_window_topn = true;
    let ctx = SessionContext::new_with_config(config);
    ctx.sql("CREATE TABLE t AS SELECT value % 50 AS pk, value AS val FROM range(0, 150)")
        .await?
        .collect()
        .await?;
    let df = ctx
        .sql(
            "SELECT * FROM ( \
                 SELECT pk, val, row_number() OVER (PARTITION BY pk ORDER BY val) AS rn \
                 FROM t \
             ) WHERE rn <= 1",
        )
        .await?;
    let plan = df.create_physical_plan().await?;

    // `PartitionedTopKExec` sits below the window operator, so execute it
    // directly to observe the batches it emits.
    let topk = find_partitioned_topk(&plan).expect("plan should contain PartitionedTopKExec");
    let batches = collect(Arc::clone(&topk), ctx.task_ctx()).await?;
    let emitted_sizes: Vec<usize> = batches.iter().map(|b| b.num_rows()).collect();

    // The 50 result rows arrive as five coalesced batches of batch_size rows
    assert_eq!(emitted_sizes, vec![10, 10, 10, 10, 10]);

    // The operator should expose its metrics (e.g. for EXPLAIN ANALYZE) ...
    let metrics = topk
        .metrics()
        .expect("PartitionedTopKExec should expose metrics");
    // ... and its output_batches metric should match the emitted batches
    let output_batches = metrics
        .sum(|m| matches!(m.value(), MetricValue::OutputBatches(_)))
        .expect("output_batches metric should be present")
        .as_usize();
    assert_eq!(
        output_batches,
        emitted_sizes.len(),
        "output_batches metric disagrees with the number of emitted batches"
    );
    Ok(())
}

fn find_partitioned_topk(plan: &Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> {
    if plan.downcast_ref::<PartitionedTopKExec>().is_some() {
        return Some(Arc::clone(plan));
    }
    plan.children().into_iter().find_map(find_partitioned_topk)
}

Output:

thread ... panicked:
PartitionedTopKExec should expose metrics

Expected behavior

PartitionedTopKExec::metrics() returns its metrics, and its output_batches metric matches the number of batches the operator emits to its consumer (5 in the reproducer above).

Additional context

• The operator is behind datafusion.optimizer.enable_window_topn, which is false by default.
• Related to SortExec (TopK) output_batches metric does not match the number of emitted batches #24468 (output_batches mismatch in the TopK path of SortExec).