andygrove commented on PR #2196:
URL:
https://github.com/apache/datafusion-ballista/pull/2196#issuecomment-5133343868
Really like the shape of this, and the "a rule only has to splice two
operators" goal is the right one. I dug into one thing and it looks like a real
gap, so I built a repro rather than just asserting it.
## The issue
`update_stage_progress` does the duplication first and attaches the filter
second:
```rust
let remapped = cut_partitions(partitions, reports, &routing.cuts)?; //
straddlers duplicated
self.planner.set_repartition_routing(stage_id, routing)?; // may
silently no-op
```
`set_repartition_routing` returns `Ok(())` when the stage's boundary root is
an `AdaptiveDatafusionExec` rather than an `ExchangeExec`. That is documented
as intentional, but by then `cut_partitions` has already pushed each straddling
file into every overlapping downstream partition. With no routing parked,
`BallistaAdapter` never injects `PerPartitionFilterExec`, so nothing trims them
back. In the final-stage case it goes straight to the user via:
```rust
self.output_locations = partitions.into_iter().flatten().collect();
```
## The reachable shape is the canonical one
I initially assumed you would need something exotic to hit this. You do not.
It is the exact `RuntimeStatsExec -> URRE` splice the PR tells rule authors to
emit, sitting at the root of the plan.
`DistributedExchangeRule::optimize` runs `transform_up` first and only wraps
the result in `AdaptiveDatafusionExec` afterwards. So at transform time the
chain top has no parent for the new branch to visit, and the branch inspects
`execution_plan.children()`. The only node visited is the `RuntimeStatsExec`
itself, whose child is the URRE, which is not a chain top.
## Repro
Three tests against `a92f0617`, all passing.
**Leg 1**, dropped into the `distributed_exchange.rs` test module next to
your existing ones:
```rust
#[test]
fn repro_root_level_chain_top_gets_no_exchange() {
let rule = DistributedExchangeRule::default();
let root = stats_over_urre_over_leaf(); // your existing helper
let result = rule.optimize(root, &config()).unwrap();
let adaptive = result.downcast_ref::<AdaptiveDatafusionExec>().unwrap();
assert!(adaptive.input().is::<RuntimeStatsExec>());
assert_eq!(count_exchanges(result.as_ref()), 0);
}
```
```
AdaptiveDatafusionExec: is_final=false, plan_id=0, stage_id=pending,
stage_resolved=false
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc → 4 partitions
StatisticsExec: col_count=1, row_count=Absent
exchanges inserted: 0
```
**Leg 2**, composing `cut_partitions` and `set_repartition_routing` in the
order `update_stage_progress` calls them, through a real `AdaptivePlanner`. One
producer file, sketch over `[5, 15, 25]`, single cut at 15:
```
after cut_partitions: 2 downstream partitions, file counts = [1, 1]
set_repartition_routing returned: Ok(())
--- plan after set_repartition_routing ---
AdaptiveDatafusionExec: is_final=true, plan_id=0, stage_id=0,
stage_resolved=false
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc → 2 partitions
DataSourceExec: partitions=1, partition_sizes=[0]
output_locations file_ids = [7, 7]
```
No `range_repartition_cuts` on the plan, so no `PerPartitionFilterExec`
downstream. File 7 appears twice. Three rows would be read as six.
**Leg 3**, positive control so this is not just me misreading the mechanism.
Same splice, same calls, only a `FilterExec` parent added:
```
AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending,
stage_resolved=false
FilterExec: v@0 > 0
ExchangeExec: partitioning=None, plan_id=0, stage_id=0,
stage_resolved=false, range_repartition_cuts=1
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc → 2 partitions
```
Exchange inserted, cuts parked, everything works. Position in the plan is
the only variable between leg 2 and leg 3.
<details>
<summary>Full leg 2 and leg 3 test file</summary>
Add `datafusion-functions-aggregate-common = { workspace = true }` to
`ballista/scheduler` dev-dependencies, drop this in as
`ballista/scheduler/src/state/aqe/test/repro_range_repartition.rs`, and
register `mod repro_range_repartition;` in that directory's `mod.rs`.
```rust
use crate::state::aqe::execution_plan::RangeRepartitionRouting;
use crate::state::aqe::planner::AdaptivePlanner;
use ballista_core::execution_plans::{
RuntimeStatsExec, UnorderedRangeRepartitionExec, cut_partitions,
};
use ballista_core::extension::SessionConfigExt;
use ballista_core::serde::protobuf::{RuntimeStatsPartitionEntry,
RuntimeStatsReport};
use ballista_core::serde::scheduler::{
ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
PartitionId, PartitionLocation, PartitionStats,
};
use datafusion::arrow::compute::SortOptions;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::source::DataSourceExec;
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::expressions::col;
use datafusion::prelude::SessionConfig;
use std::sync::Arc;
fn v_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)]))
}
/// The canonical splice a range-repartition rule is told to emit,
/// here at the root of the plan.
fn stats_over_urre_root() -> Arc<dyn ExecutionPlan> {
let schema = v_schema();
let source: Arc<dyn ExecutionPlan> =
Arc::new(DataSourceExec::new(Arc::new(
MemorySourceConfig::try_new(&[vec![]], schema.clone(),
None).unwrap(),
)));
let sort_expr = PhysicalSortExpr {
expr: col("v", schema.as_ref()).unwrap(),
options: SortOptions { descending: false, nulls_first: false },
};
let urre: Arc<dyn ExecutionPlan> = Arc::new(
UnorderedRangeRepartitionExec::try_new(source,
vec![sort_expr.clone()], 2).unwrap(),
);
Arc::new(RuntimeStatsExec::try_new(urre, Some(vec![sort_expr])).unwrap())
}
/// Same splice with an ordinary parent above it. A FilterExec survives the
/// optimizer pipeline. RoundRobin RepartitionExec and CoalescePartitionsExec
/// do NOT, both get stripped.
fn stats_over_urre_with_parent() -> Arc<dyn ExecutionPlan> {
let pred: Arc<dyn datafusion::physical_expr::PhysicalExpr> =
Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new(
col("v", v_schema().as_ref()).unwrap(),
datafusion::logical_expr::Operator::Gt,
Arc::new(datafusion::physical_expr::expressions::Literal::new(
datafusion::scalar::ScalarValue::Float64(Some(0.0)),
)),
));
Arc::new(
datafusion::physical_plan::filter::FilterExec::try_new(pred,
stats_over_urre_root())
.unwrap(),
)
}
fn location(sub_part_id: usize, producer_task_id: usize, rows: u64) ->
PartitionLocation {
PartitionLocation {
map_partition_id: 0,
partition_id: PartitionId {
job_id: "repro-job".into(),
stage_id: 0,
partition_id: sub_part_id,
},
executor_meta: ExecutorMetadata {
id: format!("exec-{producer_task_id}"),
host: "".to_string(),
port: 0,
grpc_port: 0,
specification: ExecutorSpecification::default().with_vcores(0),
os_info: ExecutorOperatingSystemSpecification::default(),
},
partition_stats: PartitionStats::new(Some(rows), None, None),
file_id: Some(producer_task_id as u64),
is_sort_shuffle: false,
}
}
#[tokio::test]
async fn repro_routing_dropped_after_partitions_already_duplicated()
-> datafusion::error::Result<()> {
let config = SessionConfig::new_with_ballista();
let mut planner =
AdaptivePlanner::try_from_plan(&config, stats_over_urre_root(),
"repro-job".into())?;
let stages = planner.runnable_stages()?;
let stage_id = stages
.as_ref()
.and_then(|s| s.first())
.map(|e| e.plan.stage_id())
.expect("a runnable stage must exist");
// Step 1: the remap, exactly as update_stage_progress calls it.
// Producer task 7, sub-part 0, sketched [5, 25], straddles the cut at
15.
let reports = vec![ballista_core::execution_plans::TaskRuntimeStats {
producer_task_id: 7,
report: RuntimeStatsReport {
order_by: vec![],
partitions: vec![RuntimeStatsPartitionEntry {
partition_id: 0,
row_count: 3,
sketch: Some(ballista_core::execution_plans::sketch_to_proto(
&datafusion_functions_aggregate_common::tdigest::TDigest::new(100)
.merge_unsorted_f64(vec![5.0, 15.0, 25.0]),
)?),
}],
},
}];
let cuts = vec![15.0];
let remapped = cut_partitions(vec![vec![location(0, 7, 3)]], &reports,
&cuts)?;
assert_eq!(remapped[0].len(), 1, "straddler duplicated into partition
0");
assert_eq!(remapped[1].len(), 1, "straddler duplicated into partition
1");
// Step 2: park the routing so a read-side filter can trim it.
let routing = RangeRepartitionRouting {
cuts: cuts.clone(),
routing_expr: col("v", v_schema().as_ref()).unwrap(),
};
let parked = planner.set_repartition_routing(stage_id as usize, routing);
assert!(parked.is_ok(), "silently succeeds even though nothing was
parked");
// Step 3: confirm nothing was actually parked.
let plan_str = format!(
"{}",
datafusion::physical_plan::displayable(planner.current_plan()).indent(true)
);
assert!(
!plan_str.contains("range_repartition_cuts"),
"no ExchangeExec carries the cuts, so no PerPartitionFilterExec
downstream"
);
// What update_stage_progress assigns to output_locations.
let file_ids: Vec<u64> = remapped
.into_iter()
.flatten()
.map(|l| l.file_id.unwrap())
.collect();
assert_eq!(file_ids, vec![7, 7], "same file twice, 3 rows read as 6");
Ok(())
}
#[tokio::test]
async fn control_routing_is_parked_when_chain_top_has_a_parent()
-> datafusion::error::Result<()> {
let config = SessionConfig::new_with_ballista();
let mut planner = AdaptivePlanner::try_from_plan(
&config,
stats_over_urre_with_parent(),
"control-job".into(),
)?;
let stages = planner.runnable_stages()?;
let stage_id = stages
.as_ref()
.and_then(|s| s.first())
.map(|e| e.plan.stage_id())
.expect("a runnable stage must exist");
let routing = RangeRepartitionRouting {
cuts: vec![15.0],
routing_expr: col("v", v_schema().as_ref()).unwrap(),
};
planner.set_repartition_routing(stage_id as usize, routing)?;
let plan_str = format!(
"{}",
datafusion::physical_plan::displayable(planner.current_plan()).indent(true)
);
assert!(
plan_str.contains("range_repartition_cuts=1"),
"control: the cuts must be parked on the boundary ExchangeExec"
);
Ok(())
}
```
</details>
## A side finding that might matter more
I first wrote the control with a `RoundRobinBatch` `RepartitionExec` parent,
then with a `CoalescePartitionsExec` parent. Both were stripped by the
optimizer pipeline, collapsing back to the root-level shape and making my
control fail before I switched to `FilterExec`.
That is worth flagging on its own.
`range_repartition_chain_top_gets_exchange_inserted` builds a `RoundRobinBatch`
parent and calls `rule.optimize()` in isolation, where the parent survives and
the exchange gets inserted. Through the real `AdaptivePlanner` pipeline that
same parent is gone by the time DER runs. So the test passes for a plan shape
the pipeline does not actually produce, and a rule author could write something
that looks protected and lose the protection in production. Might be worth
driving those DER tests through the planner rather than the rule alone.
## What I did not prove
I composed `cut_partitions` and `set_repartition_routing` directly in the
order `update_stage_progress` calls them, rather than driving a full query to a
wrong answer. Since no rule inserts a range repartition yet that is not
constructible on this branch. The link I reasoned about but did not execute is
the graph-side gate, `stages.get(&stage_id)` being `Running` with non-empty
`runtime_stats_reports`. Reading `revive()` it promotes every Resolved stage to
Running including the final one, and the `RuntimeStatsExec` in that stage's
plan is what `collect_reports` walks, so I think it holds. Happy to be told I
have that wrong.
## Suggested fix
Make `set_repartition_routing` error when it cannot find an `ExchangeExec`
to park on, or move it ahead of `cut_partitions` so a failure to attach the
filter aborts before the duplication is committed. Everything else in this code
errors hard on invariant breaks like `missing file_id`, `no usable sketch`, and
`merge_reports returned N groups`, so the silent `Ok(())` is the outlier.
Separately `is_range_repartition_chain_top` probably wants to fire when the
chain top is itself the plan root.
Nothing else in the review turned up anything. No public API breaks, the
`GlobalPartitionMap` and `AdaptivePlanner` renames are both behind private
modules, and the machinery really is inert until a rule lands.
Investigation and repro here were LLM assisted, with the test output and
plan dumps above coming from actual runs against the PR head.
--
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]
