andygrove opened a new pull request, #2664: URL: https://github.com/apache/datafusion-ballista/pull/2664
# Which issue does this PR close? Closes #1960. # Rationale for this change #1916 collects per-column null counts on the executor and aggregates them onto the completed stage (`SuccessfulStage.output_column_stats`), but nothing consumed them: a resolved exchange still reported `ColumnStatistics::new_unknown()` for every column, so the planner never saw them. Only the sort shuffle writer collected counts in the first place. This is also the plumbing that distinct-count sketches need. Without distinct counts, DataFusion estimates a join downstream of a shuffle as its smaller input's row count, which is what led to #2661 (see #2662). # What changes are included in this PR? As the comment on #1960 notes, the planner has been refactored since the issue was written. AQE now reads statistics from `ExchangeExec`, and the `ShuffleReaderExec` built from it at task time is executed, not re-planned, so the statistics go on the exchange and `ShuffleReaderExec` and its proto are unchanged. **Executor** - `ShuffleWriterExec` (passthrough) and `RangeShuffleWriterExec` count per-column logical nulls over everything they write, through a new `ColumnNullCounts` stream wrapper shared by a writer's clones, and expose them as `column_null_counts()`, like `SortShuffleWriterExec`. - `DefaultQueryStageExec::collect_column_stats` reports them for all three writer variants. **Scheduler** - `RunningStage::output_column_stats` is factored out of `to_successful`, so the aggregate is available when the stage completes, before it becomes a `SuccessfulStage`. - `AdaptivePlanner::resolve_stage_partitions` takes the stage's column statistics and stores them on its `ExchangeExec` before resolving the partitions, so the replan that follows sees them. - `ExchangeExec` holds them in a slot shared across `with_new_children` and `to_broadcast`, like `shuffle_partitions`. `partition_statistics(None)` reports them for the exchange as a whole. Per-partition statistics stay unknown, because the counts are stage totals and say nothing about how nulls fall across partitions. Statistics whose length doesn't match the schema are ignored. Out of scope: distinct counts. The protocol already has `ColumnSketch` / `SketchType::HLL`; once the writers fill a sketch, `output_column_stats` can set `distinct_count` and it will flow through the same path. **Tests** - `shuffle_writer::tests::column_null_counts_fold_across_partitions` and `range_shuffle::writer::tests::reports_a_null_count_per_written_column`. - `exchange::column_statistics_tests`: whole-exchange statistics, unknown per-partition statistics, mismatched length ignored, unknown before the shuffle resolves, and kept across `with_new_children` and `to_broadcast`. - `aqe::test::column_statistics::completed_stage_column_statistics_reach_its_exchange`: through `AdaptivePlanner`, a completed stage's statistics are visible on its resolved exchange after the replan. - The existing `to_successful` aggregation tests cover `output_column_stats`. # Are there any user-facing changes? No configuration changes. Public API: `AdaptivePlanner::resolve_stage_partitions` takes a third argument, `column_statistics: Vec<ColumnStatistics>`. New public methods `ExchangeExec::resolve_column_statistics`, `RunningStage::output_column_stats`, `ShuffleWriterExec::column_null_counts` and `RangeShuffleWriterExec::column_null_counts`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
