diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs index 4c8b8f9c01122..992c954f09b25 100644 --- a/datafusion/core/tests/sql/explain_analyze.rs +++ b/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::>(); + 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 diff --git a/datafusion/physical-expr-common/src/metrics/mod.rs b/datafusion/physical-expr-common/src/metrics/mod.rs index 146c039c75f6a..fcdd33dbe7f06 100644 --- a/datafusion/physical-expr-common/src/metrics/mod.rs +++ b/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> { self.metrics.iter() diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 189cdcf225e98..b7dacd0cd08d0 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -2128,7 +2128,7 @@ impl ExecutionPlan for AggregateExec { } fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + Some(self.metrics.clone_inner().with_output_rows_skew()) } fn child_stats_requests(&self, partition: Option) -> Vec { diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 84ed96bb52170..b05114a4fd573 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -1580,7 +1580,7 @@ impl ExecutionPlan for HashJoinExec { } fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + Some(self.metrics.clone_inner().with_output_rows_skew()) } fn child_stats_requests(&self, partition: Option) -> Vec { diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 063954a72a094..ecb3926d2ef72 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1768,7 +1768,7 @@ impl ExecutionPlan for RepartitionExec { } fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + Some(self.metrics.clone_inner().with_output_rows_skew()) } fn child_stats_requests(&self, _partition: Option) -> Vec { diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index d4c98009ba70d..cfb6d0110f725 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -528,7 +528,7 @@ impl ExecutionPlan for BoundedWindowAggExec { } fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + Some(self.metrics.clone_inner().with_output_rows_skew()) } fn child_stats_requests(&self, partition: Option) -> Vec { diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index d794e7df9d0a9..c6865a01bf3f4 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -320,7 +320,7 @@ impl ExecutionPlan for WindowAggExec { } fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) + Some(self.metrics.clone_inner().with_output_rows_skew()) } fn child_stats_requests(&self, partition: Option) -> Vec { diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 7a706836f44d6..54123002d0aa5 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/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=, output_bytes=, output_batches=1, build_mem_used=, array_map_created_count=0, build_input_batches=0, build_input_rows=0, input_batches=1, input_rows=2, build_time=, join_time=, 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=, output_bytes=, output_batches=1, build_mem_used=, array_map_created_count=0, build_input_batches=0, build_input_rows=0, input_batches=1, input_rows=2, build_time=, join_time=, 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=, output_bytes=, output_batches=0, expr_0_eval_time=] 03)----FilterExec: column1@0 != 1, metrics=[output_rows=0, elapsed_compute=, output_bytes=, output_batches=0, selectivity=0% (0/1)] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]