andygrove commented on code in PR #2196: URL: https://github.com/apache/datafusion-ballista/pull/2196#discussion_r3705634907
########## ballista/core/src/execution_plans/plan_algebra.rs: ########## @@ -0,0 +1,70 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Algebraic properties of physical plan nodes — do they preserve +//! partitioning, do they preserve the distribution of row values, etc. +//! +//! DataFusion doesn't expose these as trait methods on `ExecutionPlan` +//! (nice-to-haves like `ExecutionPlan::affects_partitioning()` / +//! `ExecutionPlan::affects_distribution()` that would hopefully land one +//! day), so we downcast against a hand-maintained whitelist. Being +//! conservative is the safety net: unrecognized node → property assumed +//! false → caller falls back to the safer path. + +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; + +use crate::execution_plans::{ + BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, +}; + +/// Whitelisted ops preserve the routing key's row set, values, and +/// partitioning — an upstream sketch remains valid after the operator. +pub fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool { + // Buffered batches replayed verbatim. + plan.downcast_ref::<BufferExec>().is_some() + // Per-partition sort: rows reorder within a partition, row set + // and counts unchanged. `preserve_partitioning=false` collapses + // N→1 (like SortPreservingMergeExec), so gate on the flag. + || plan + .downcast_ref::<SortExec>() + .is_some_and(|sort| sort.preserve_partitioning()) + // Stage-boundary writers: batches to disk unchanged. + || plan.downcast_ref::<ShuffleWriterExec>().is_some() + || plan.downcast_ref::<SortShuffleWriterExec>().is_some() + // Pure row-annotation: one input row → one output row with an + // added column (window fn result); values, partitioning, count preserved. + || plan.downcast_ref::<BoundedWindowAggExec>().is_some() + || plan.downcast_ref::<WindowAggExec>().is_some() +} + +/// Looser sibling of [`preserves_distribution`]: partitioning survives, +/// but rows and values within a partition are fair game. +pub fn preserves_partitioning(plan: &dyn ExecutionPlan) -> bool { + // Distribution-preserving is strictly stronger; compose to keep the + // whitelist deduplicated. + preserves_distribution(plan) + // Drops rows, but per-partition — no rows migrate. + || plan.downcast_ref::<FilterExec>().is_some() + // Rewrites columns; partition boundaries untouched. + || plan.downcast_ref::<ProjectionExec>().is_some() Review Comment: A `ProjectionExec` here means the read-side filter can silently route on the wrong column. `is_range_repartitioned` gates on `preserves_partitioning`, so a `ProjectionExec` between the URRE and the boundary is accepted and DER splices the `ExchangeExec` above it. But `repartition_routing_expr` recovers `order_by[0].expr` from the URRE, which is a `Column` indexed into the pre-projection schema, while `PerPartitionFilterExec` evaluates it against the post-projection schema. The old `preserves_distribution` list excluded `ProjectionExec` for exactly this reason ("might compute a new column that shadows or replaces the routing key"), and that exclusion is still what the routing expression needs. I ran a probe with a projection that swaps two `Float64` columns above the URRE: ``` is_range_repartitioned(proj) = true AdaptiveDatafusionExec: is_final=false, plan_id=1 ExchangeExec: partitioning=None, plan_id=0 ProjectionExec: expr=[tag@1 as tag, k@0 as k] UnorderedRangeRepartitionExec: routing=k@0 asc -> 4 partitions StatisticsExec: col_count=2, row_count=Absent recovered routing_expr = k@0 read-side schema = [tag: Float64, k: Float64] predicate[1] = k@0 >= 10 AND k@0 < 20 PerPartitionFilterExec::try_new -> Ok ``` `try_new`'s Boolean check passes because both columns are `Float64`, so index 0 now resolves to `tag` and this lands as wrong answers rather than an error. The probe, if it is useful, dropped into the test module in `distributed_exchange.rs`: ```rust #[test] fn probe_projection_reindexes_routing_column() { use ballista_core::execution_plans::{ PerPartitionFilterExec, range_partition_predicates, repartition_routing_expr, }; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::projection::ProjectionExec; let schema = Schema::new(vec![ Field::new("k", DataType::Float64, false), Field::new("tag", DataType::Float64, false), ]); let stats = Statistics { num_rows: Default::default(), total_byte_size: Default::default(), column_statistics: vec![ ColumnStatistics::new_unknown(), ColumnStatistics::new_unknown(), ], }; let leaf: Arc<dyn ExecutionPlan> = Arc::new(StatisticsExec::new(stats, schema)); let sort_expr = PhysicalSortExpr { expr: Arc::new(Column::new("k", 0)), options: SortOptions { descending: false, nulls_first: false }, }; let urre: Arc<dyn ExecutionPlan> = Arc::new( UnorderedRangeRepartitionExec::try_new(leaf, vec![sort_expr], 4).unwrap(), ); // Projection swaps the columns: output is (tag, k). let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![ (Arc::new(Column::new("tag", 1)), "tag".to_string()), (Arc::new(Column::new("k", 0)), "k".to_string()), ]; let proj: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(proj_exprs, urre).unwrap()); println!("is_range_repartitioned = {}", is_range_repartitioned(&proj)); let optimized = DistributedExchangeRule::default() .optimize(proj.clone(), &config()) .unwrap(); println!("{}", display_plan(&optimized)); let routing = repartition_routing_expr(proj.as_ref()).unwrap().unwrap(); println!("routing_expr = {routing}"); println!("read-side schema = {:?}", proj.schema()); let preds = range_partition_predicates(routing, &[10.0, 20.0, 30.0]); println!("predicate[1] = {}", preds[1]); println!("{:?}", PerPartitionFilterExec::try_new(proj, preds).is_ok()); } ``` To be clear, `preserves_distribution` is fine as-is, and `collect_reachable_stats` using it is correct. It is only this looser sibling being used as the gate in `is_range_repartitioned` that opens the hole. Options I can see are gating `is_range_repartitioned` on the strict predicate instead, or rebasing the routing expression onto the boundary schema when the cuts get parked. ########## ballista/scheduler/src/state/aqe/adapter.rs: ########## @@ -105,7 +108,24 @@ impl BallistaAdapter { )?, }; - Ok(Transformed::yes(Arc::new(reader))) + let reader: Arc<dyn ExecutionPlan> = Arc::new(reader); + // Without a per-partition filter, straddling sub-parts from a + // range-repartitioned upstream would feed multiple downstream + // partitions and `FinalPartitioned` would split their partial sums. + if let Some(routing) = exchange.range_repartition_routing() { Review Comment: The coalesce slot and the routing slot can both be set on the same `ExchangeExec`, and this branch does not account for it. `CoalescePartitionsRule` collects every leaf `ExchangeExec` in the subtree. It bails for `broadcast`, but there is no equivalent bail for a range-repartitioned exchange, so both slots land on the same node: ``` coalesce groups = 2, cuts = 3 ``` When that happens, the `(Some(cp), false)` arm above builds a reader with `cp.groups.len()` partitions, and then this block hands `cuts.len() + 1` predicates to `PerPartitionFilterExec::try_new`, which rejects the count mismatch. So it is a hard error at plan time rather than wrong data, which makes it the lesser of the two issues, but it does mean the AQE coalesce path and any range-repartition rule are mutually exclusive with nothing saying so. Worth noting the coalesce rule groups neighbouring upstream partitions only, so contiguous-group coalescing is actually compatible with range partitioning. Merging the cuts alongside the groups looks like a real fix rather than just a bail, if you want to go that way. A bail matching the broadcast one would be fine for now too, as long as it is explicit. Repro: ```rust #[test] fn probe_coalesce_and_routing_on_same_exchange() { use ballista_core::execution_plans::{CoalescePlan, PartitionGroup}; use crate::state::aqe::execution_plan::RangeRepartitionRouting; let ex = ExchangeExec::new(stats_over_urre_over_leaf(), None, 0); ex.set_coalesce(Arc::new(CoalescePlan { groups: vec![ PartitionGroup { upstream_indices: vec![0, 1] }, PartitionGroup { upstream_indices: vec![2, 3] }, ], upstream_partition_count: 4, })); ex.resolve_range_repartition_routing(RangeRepartitionRouting { cuts: vec![10.0, 20.0, 30.0], routing_expr: Arc::new(Column::new("v", 0)), }); println!( "coalesce groups = {}, cuts = {}", ex.coalesce().unwrap().groups.len(), ex.range_repartition_routing().unwrap().cuts.len() ); } ``` ########## ballista/core/src/execution_plans/runtime_stats.rs: ########## @@ -1357,7 +1462,7 @@ mod merge_tests { /// getting silently dropped. #[test] fn merge_reports_propagates_sketch_decode_errors() { - use crate::serde::protobuf::QuantileSketchState; + use QuantileSketchState; Review Comment: Leftover from the import rewrite. This compiles under uniform paths, but `use super::*` at the top of the module already brings `QuantileSketchState` in, so the line can go. ########## ballista/scheduler/src/state/aqe/execution_plan/exchange.rs: ########## @@ -32,6 +33,26 @@ use parking_lot::Mutex; use std::ops::Deref; use std::sync::{Arc, atomic::AtomicI64}; +/// Range-partition boundaries recovered from an +/// `UnorderedRangeRepartitionExec` / `OrderedRangeRepartitionExec` upstream +/// of this exchange. Written after the range-repartition-producing stage +/// completes and its runtime-stats sketches are merged; read at +/// task-specialization time to build per-downstream-partition range filters +/// (see `PerPartitionFilterExec`). +/// +/// `cuts` are `K - 1` monotone `f64` boundaries expressed in the value space +/// of `routing_expr`; downstream partition `k` owns `[cuts[k-1], cuts[k])` +/// with virtual `-∞`/`+∞` sentinels on the ends (matching the range +/// repartition's write-side convention). `routing_expr` is the same +/// expression the range repartition routes on — a `CAST(order_by[0] AS Review Comment: Small doc nit. `repartition_routing_expr` returns `order_by[0].expr` verbatim, not a cast of it. Both range repartition `try_new`s already require the expression to be `Float64`, so the code is right and this comment is just stale. -- 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]
