Add skewness metric to nodes that execute in partitioned mode (#25977)

datafusion/core/tests/sql/explain_analyze.rs

@@ -19,6 +19,7 @@ use super::*;
 use insta::assert_snapshot;
 use rstest::rstest;
 
+use datafusion::catalog::MemTable;
 use datafusion::config::ConfigOptions;
 use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::physical_plan::metrics::Timestamp;
@@ -907,6 +908,83 @@ async fn parquet_explain_analyze() {
     );
 }
 
+/// Operators that run once per partition report `output_rows_skew`. A single
+/// join/group/window key hashes every row to one partition, so each
+/// hash-partitioned operator must report the maximum skew.
+#[tokio::test]
+#[cfg_attr(coverage, ignore)]
+async fn explain_analyze_output_rows_skew_on_hash_partitioned_operators() {
+    let mut config = SessionConfig::new().with_target_partitions(4);
+    // Force a partitioned (hash repartitioned) join instead of `CollectLeft`
+    let options = config.options_mut();
+    options.optimizer.hash_join_single_partition_threshold = 0;
+    options.optimizer.hash_join_single_partition_threshold_rows = 0;
+    let ctx = SessionContext::new_with_config(config);
+    // Spread the rows over 4 input partitions so they get hash repartitioned
+    // (a single input partition already satisfies the hash distribution)
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("k", DataType::Int64, false),
+        Field::new("v", DataType::Int64, false),
+    ]));
+    let partitions = (0..4)
+        .map(|p| {
+            let v = (p * 25..(p + 1) * 25).collect::<Vec<i64>>();
+            let batch = RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![
+                    Arc::new(Int64Array::from(vec![1; v.len()])),
+                    Arc::new(Int64Array::from(v)),
+                ],
+            )
+            .unwrap();
+            vec![batch]
+        })
+        .collect();
+    let table = MemTable::try_new(schema, partitions).unwrap();
+    ctx.register_table("t", Arc::new(table)).unwrap();
+
+    let sql = "EXPLAIN ANALYZE \
+               SELECT a.k, count(*) FROM t a JOIN t b ON a.k = b.k GROUP BY a.k";
+    let formatted = explain_analyze_output(&ctx, sql).await;
+    for operator in [
+        "RepartitionExec: partitioning=Hash",
+        "HashJoinExec: mode=Partitioned",
+        "AggregateExec: mode=SinglePartitioned",
+    ] {
+        assert_operator_metric(&formatted, operator, "output_rows_skew=100%");
+    }
+
+    let sql = "EXPLAIN ANALYZE \
+               SELECT k, \
+                 row_number() OVER (PARTITION BY k ORDER BY v), \
+                 sum(v) OVER (PARTITION BY k) \
+               FROM t";
+    let formatted = explain_analyze_output(&ctx, sql).await;
+    for operator in ["BoundedWindowAggExec", "WindowAggExec"] {
+        assert_operator_metric(&formatted, operator, "output_rows_skew=100%");
+    }
+}
+
+async fn explain_analyze_output(ctx: &SessionContext, sql: &str) -> String {
+    let actual = execute_to_batches(ctx, sql).await;
+    arrow::util::pretty::pretty_format_batches(&actual)
+        .unwrap()
+        .to_string()
+}
+
+fn assert_operator_metric(formatted: &str, operator: &str, metric: &str) {
+    // Leading space so `WindowAggExec` does not match `BoundedWindowAggExec`
+    let needle = format!(" {operator}");
+    let line = formatted
+        .lines()
+        .find(|line| line.contains(&needle))
+        .unwrap_or_else(|| panic!("{operator} not found in plan:\n{formatted}"));
+    assert!(
+        line.contains(metric),
+        "expected {metric} on {operator}, got:\n{line}"
+    );
+}
+
 // This test reproduces the behavior described in
 // https://github.com/apache/datafusion/issues/16684 where projection
 // pushdown with recursive CTEs could fail to remove unused columns

datafusion/physical-expr-common/src/metrics/mod.rs

@@ -229,6 +229,22 @@ impl MetricsSet {
         self.metrics.push(metric)
     }
 
+    /// Returns this set with the derived `output_rows_skew` metric appended,
+    /// or unchanged if no partition reported `output_rows` yet.
+    ///
+    /// This is typically used in `ExecutionPlan::metrics` of operators that
+    /// execute in multiple partitions (e.g. repartitions, partitioned joins and
+    /// aggregations), where it shows how evenly their output rows are spread
+    /// across partitions. With a single partition the skew is always `0%`.
+    ///
+    /// See [`BaselineMetrics::output_rows_skew_metric`] for how skew is computed.
+    pub fn with_output_rows_skew(mut self) -> Self {
+        if let Some(output_rows_skew) = BaselineMetrics::output_rows_skew_metric(&self) {
+            self.push(output_rows_skew);
+        }
+        self
+    }
+
     /// Returns an iterator across all metrics
     pub fn iter(&self) -> impl Iterator<Item = &Arc<Metric>> {
         self.metrics.iter()

datafusion/physical-plan/src/aggregates/mod.rs

@@ -2128,7 +2128,7 @@ impl ExecutionPlan for AggregateExec {
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
-        Some(self.metrics.clone_inner())
+        Some(self.metrics.clone_inner().with_output_rows_skew())
     }
 
     fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -1580,7 +1580,7 @@ impl ExecutionPlan for HashJoinExec {
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
-        Some(self.metrics.clone_inner())
+        Some(self.metrics.clone_inner().with_output_rows_skew())
     }
 
     fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {

datafusion/physical-plan/src/repartition/mod.rs

@@ -1768,7 +1768,7 @@ impl ExecutionPlan for RepartitionExec {
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
-        Some(self.metrics.clone_inner())
+        Some(self.metrics.clone_inner().with_output_rows_skew())
     }
 
     fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> {

datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs

@@ -528,7 +528,7 @@ impl ExecutionPlan for BoundedWindowAggExec {
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
-        Some(self.metrics.clone_inner())
+        Some(self.metrics.clone_inner().with_output_rows_skew())
     }
 
     fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {

datafusion/physical-plan/src/windows/window_agg_exec.rs

@@ -320,7 +320,7 @@ impl ExecutionPlan for WindowAggExec {
     }
 
     fn metrics(&self) -> Option<MetricsSet> {
-        Some(self.metrics.clone_inner())
+        Some(self.metrics.clone_inner().with_output_rows_skew())
     }
 
     fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> {

datafusion/sqllogictest/test_files/joins.slt

@@ -5185,7 +5185,7 @@ LEFT ANTI JOIN (
 ) t2 ON t1.k = t2.k;
 ----
 Plan with Metrics
-01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(k@0, k@0)], metrics=[output_rows=2, elapsed_compute=<slt:ignore>, output_bytes=<slt:ignore>, output_batches=1, build_mem_used=<slt:ignore>, array_map_created_count=0, build_input_batches=0, build_input_rows=0, input_batches=1, input_rows=2, build_time=<slt:ignore>, join_time=<slt:ignore>, avg_fanout=N/A (0/0), probe_hit_rate=0% (0/2)]
+01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(k@0, k@0)], metrics=[output_rows=2, elapsed_compute=<slt:ignore>, output_bytes=<slt:ignore>, output_batches=1, build_mem_used=<slt:ignore>, array_map_created_count=0, build_input_batches=0, build_input_rows=0, input_batches=1, input_rows=2, build_time=<slt:ignore>, join_time=<slt:ignore>, avg_fanout=N/A (0/0), output_rows_skew=0%, probe_hit_rate=0% (0/2)]
 02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_rows=0, elapsed_compute=<slt:ignore>, output_bytes=<slt:ignore>, output_batches=0, expr_0_eval_time=<slt:ignore>]
 03)----FilterExec: column1@0 != 1, metrics=[output_rows=0, elapsed_compute=<slt:ignore>, output_bytes=<slt:ignore>, output_batches=0, selectivity=0% (0/1)]
 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]