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).

[shinzoxD]

I'll add ExecutionPlan metrics to PartitionedTopKExec so EXPLAIN ANALYZE reports output_rows/output_batches, matching other sort/window operators.

[shinzoxD]

Opened #24495 to expose PartitionedTopKExec metrics (output_rows / output_batches) so EXPLAIN ANALYZE can report them.

[shinzoxD]

take