This is an automated email from the ASF dual-hosted git repository.
Dandandan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git
The following commit(s) were added to refs/heads/main by this push:
new fcd8a5a04 perf(core): push a consumer's row limit into
RangeShuffleReaderExec (#2334)
fcd8a5a04 is described below
commit fcd8a5a0486a4e337c698654091e19576d312258
Author: Daniël Heres <[email protected]>
AuthorDate: Mon Aug 17 11:19:52 2026 +0200
perf(core): push a consumer's row limit into RangeShuffleReaderExec (#2334)
* perf(core): push a consumer's row limit into RangeShuffleReaderExec
`RangeShuffleReaderExec` merges all of its input streams to completion even
when the `SortPreservingMergeExec` above it only wants the first few rows.
Its
`StreamingMerge` was built without a fetch, so a `fetch=20` top-N over a
wide
fan-in merged everything and then discarded almost all of it.
Both operators merge on the same ordering, so the consumer's first `n` rows
can
only come from the reader's first `n`. The reader now carries an optional
`fetch` and passes it to `StreamingMergeBuilder`, and the adapter pushes the
limit down when it plants the reader under such a merge.
DataFusion's own limit pushdown can't do this: the reader is planted at
adapt
time, after the optimizer chain has run.
This reduces merge work. It does not lower the merge's peak memory, which is
set by fan-in times batch size and is paid before the first row is emitted.
Co-Authored-By: Claude Opus 5 <[email protected]>
* refactor(scheduler): apply the row limit where the reader is built
The limit was applied by a second walk of the plan that re-visited a node
the
adapter had just created. The adapter already walks parent-first, so it can
handle the merge and its exchange together and build the reader with the
limit
already set.
Extracts the exchange-to-reader conversion into `build_reader` so both paths
share it, and drops the extra pass.
Co-Authored-By: Claude Opus 5 <[email protected]>
* fix(core): carry the pushed-down fetch through the wire format
`RangeShuffleReaderExec` gained a `fetch` field, but the proto message had
no matching field, so the scheduler's limit was dropped on encode and the
executor merged every row anyway. Add `fetch` to the proto, write it in the
encoder, restore it on decode, and preserve it when the task builder
restricts the reader to a partition subset.
Also simplify the operator (derive `Clone`, update via struct update
syntax), show `fetch` in EXPLAIN, and reuse `replace_children_if_necessary`
in the adapter.
Co-Authored-By: Claude Opus 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01DZS73f3mPrwQ8EPbswFJVa
* test(scheduler): cover the limit reaching the planted range reader
The adapter rebuilds the merge around a new reader, so a limit that fails
to reach the reader is silent — nothing in the plan shape changes.
Co-Authored-By: Claude Opus 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01DZS73f3mPrwQ8EPbswFJVa
---------
Co-authored-by: Claude Opus 5 <[email protected]>
---
ballista/core/proto/ballista.proto | 2 +
.../src/execution_plans/range_shuffle_reader.rs | 91 ++++++---
ballista/core/src/serde/generated/ballista.rs | 3 +
ballista/core/src/serde/mod.rs | 18 +-
ballista/scheduler/src/state/aqe/adapter.rs | 222 +++++++++++++--------
ballista/scheduler/src/state/task_builder.rs | 2 +-
6 files changed, 230 insertions(+), 108 deletions(-)
diff --git a/ballista/core/proto/ballista.proto
b/ballista/core/proto/ballista.proto
index bd20ed65d..ae8098bce 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -235,6 +235,8 @@ message RangeShuffleReaderExecNode {
// Sort key the reader's k-way merge preserves. Advertised on the reader's
// `PlanProperties.eq_properties` for downstream consumers.
repeated datafusion.PhysicalSortExprNode merge_ordering = 4;
+ // Row limit pushed down by a consuming merge. Absent means read everything.
+ optional uint64 fetch = 5;
}
// CoalescePartitionsRule output: groups upstream partitions into coalesced
output partitions.
diff --git a/ballista/core/src/execution_plans/range_shuffle_reader.rs
b/ballista/core/src/execution_plans/range_shuffle_reader.rs
index be8fe8f71..ed43ed271 100644
--- a/ballista/core/src/execution_plans/range_shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/range_shuffle_reader.rs
@@ -82,7 +82,7 @@ use log::debug;
use std::sync::Arc;
/// Ordering-preserving shuffle reader. See module docs.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
pub struct RangeShuffleReaderExec {
/// Upstream stage that produced these files.
pub stage_id: usize,
@@ -93,6 +93,8 @@ pub struct RangeShuffleReaderExec {
/// Sort key the merge preserves. Advertised in
`PlanProperties.eq_properties`
/// so downstream operators (BWAG, SMJ build side) see the output ordering.
merge_ordering: LexOrdering,
+ /// Row limit pushed down by a consumer. `None` means read everything.
+ fetch: Option<usize>,
metrics: ExecutionPlanMetricsSet,
properties: Arc<PlanProperties>,
work_dir: Option<String>,
@@ -126,6 +128,7 @@ impl RangeShuffleReaderExec {
schema,
partition,
merge_ordering,
+ fetch: None,
metrics: ExecutionPlanMetricsSet::new(),
properties,
work_dir: None,
@@ -133,31 +136,25 @@ impl RangeShuffleReaderExec {
})
}
+ /// Set the row limit at construction, without cloning the reader.
+ pub fn with_fetch_limit(mut self, fetch: Option<usize>) -> Self {
+ self.fetch = fetch;
+ self
+ }
+
/// Late-bound by the executor.
pub fn with_work_dir(&self, work_dir: String) -> Self {
Self {
- stage_id: self.stage_id,
- schema: self.schema.clone(),
- partition: self.partition.clone(),
- merge_ordering: self.merge_ordering.clone(),
- metrics: self.metrics.clone(),
- properties: self.properties.clone(),
work_dir: Some(work_dir),
- client_pool: self.client_pool.clone(),
+ ..self.clone()
}
}
/// Late-bound by the executor.
pub fn with_client_pool(&self, client_pool: Arc<dyn BallistaClientPool>)
-> Self {
Self {
- stage_id: self.stage_id,
- schema: self.schema.clone(),
- partition: self.partition.clone(),
- merge_ordering: self.merge_ordering.clone(),
- metrics: self.metrics.clone(),
- properties: self.properties.clone(),
- work_dir: self.work_dir.clone(),
client_pool: Some(client_pool),
+ ..self.clone()
}
}
@@ -181,12 +178,20 @@ impl DisplayAs for RangeShuffleReaderExec {
self.stage_id,
self.partition.len(),
self.merge_ordering,
- )
+ )?;
+ if let Some(fetch) = self.fetch {
+ write!(f, ", fetch: {fetch}")?;
+ }
+ Ok(())
}
DisplayFormatType::TreeRender => {
writeln!(f, "upstream_stage={}", self.stage_id)?;
writeln!(f, "output_partitions={}", self.partition.len())?;
- writeln!(f, "ordering={}", self.merge_ordering)
+ writeln!(f, "ordering={}", self.merge_ordering)?;
+ if let Some(fetch) = self.fetch {
+ writeln!(f, "fetch={fetch}")?;
+ }
+ Ok(())
}
}
}
@@ -231,14 +236,8 @@ impl ExecutionPlan for RangeShuffleReaderExec {
));
}
Ok(Arc::new(Self {
- stage_id: self.stage_id,
- schema: self.schema.clone(),
- partition: self.partition.clone(),
- merge_ordering: self.merge_ordering.clone(),
metrics: ExecutionPlanMetricsSet::new(),
- properties: self.properties.clone(),
- work_dir: self.work_dir.clone(),
- client_pool: self.client_pool.clone(),
+ ..self.as_ref().clone()
}))
}
@@ -330,6 +329,7 @@ impl ExecutionPlan for RangeShuffleReaderExec {
.with_schema(self.schema.clone())
.with_expressions(&self.merge_ordering)
.with_batch_size(config.batch_size())
+ .with_fetch(self.fetch)
.with_metrics(baseline)
.with_reservation(reservation)
.build()?;
@@ -337,6 +337,19 @@ impl ExecutionPlan for RangeShuffleReaderExec {
Ok(merged)
}
+ fn fetch(&self) -> Option<usize> {
+ self.fetch
+ }
+
+ /// Output is sorted on `merge_ordering`, so the first `n` rows can only
+ /// come from the first `n` of each input.
+ fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn
ExecutionPlan>> {
+ Some(Arc::new(Self {
+ fetch: limit,
+ ..self.clone()
+ }))
+ }
+
fn metrics(&self) ->
Option<datafusion::physical_plan::metrics::MetricsSet> {
Some(self.metrics.clone_inner())
}
@@ -371,6 +384,7 @@ mod tests {
use datafusion::arrow::ipc::writer::StreamWriter;
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_expr::expressions::Column;
+ use datafusion::physical_plan::{ChildrenPropertiesMode,
ReplaceChildrenOptions};
use datafusion::prelude::SessionContext;
use std::fs::{File, create_dir_all};
use tempfile::tempdir;
@@ -550,6 +564,35 @@ mod tests {
assert!(batches.is_empty());
}
+ /// A limit must survive a rebuild, or the merge quietly goes back to
+ /// reading everything.
+ #[test]
+ fn fetch_roundtrips_through_with_fetch() {
+ use datafusion::physical_plan::ExecutionPlan;
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Float64,
false)]));
+ let merge_ordering =
LexOrdering::new(vec![PhysicalSortExpr::new_default(
+ Arc::new(Column::new("v", 0)),
+ )])
+ .unwrap();
+ let reader =
+ RangeShuffleReaderExec::try_new(3, vec![vec![]; 4], schema,
merge_ordering)
+ .unwrap();
+ assert_eq!(reader.fetch(), None, "reader starts unlimited");
+
+ let limited = ExecutionPlan::with_fetch(&reader, Some(20))
+ .expect("RangeShuffleReaderExec must accept a pushed-down limit");
+ assert_eq!(limited.fetch(), Some(20));
+
+ let rebuilt = Arc::clone(&limited)
+ .replace_children(
+ vec![],
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ .expect("rebuild");
+ assert_eq!(rebuilt.fetch(), Some(20), "fetch must survive a rebuild");
+ }
+
/// The reader must advertise its merge ordering so downstream operators
/// see the sortedness invariant (BWAG's RANGE-frame cursor, SMJ build
side).
#[test]
diff --git a/ballista/core/src/serde/generated/ballista.rs
b/ballista/core/src/serde/generated/ballista.rs
index a868ff2b3..9e05f2ed2 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -300,6 +300,9 @@ pub struct RangeShuffleReaderExecNode {
pub merge_ordering: ::prost::alloc::vec::Vec<
::datafusion_proto::protobuf::PhysicalSortExprNode,
>,
+ /// Row limit pushed down by a consuming merge. Absent means read
everything.
+ #[prost(uint64, optional, tag = "5")]
+ pub fetch: ::core::option::Option<u64>,
}
/// CoalescePartitionsRule output: groups upstream partitions into coalesced
output partitions.
/// Empty when no coalesce is applied (the optional field on the parent
message is absent).
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index 642feec2c..56d2b8eb8 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -542,12 +542,15 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
"RangeShuffleReaderExec: merge_ordering must be
non-empty",
)
})?;
- Ok(Arc::new(RangeShuffleReaderExec::try_new(
+ let reader = RangeShuffleReaderExec::try_new(
stage_id,
partition_location,
schema,
merge_ordering,
- )?))
+ )?;
+ Ok(Arc::new(
+ reader.with_fetch_limit(range_reader.fetch.map(|f| f as
usize)),
+ ))
}
PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => {
let schema: SchemaRef =
@@ -934,6 +937,7 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
partition,
schema: Some(exec.schema().as_ref().try_into()?),
merge_ordering,
+ fetch: exec.fetch().map(|f| f as u64),
},
)),
};
@@ -1430,7 +1434,8 @@ mod test {
schema.clone(),
merge_ordering.clone(),
)
- .unwrap();
+ .unwrap()
+ .with_fetch_limit(Some(20));
let codec = BallistaPhysicalExtensionCodec::default();
let mut buf: Vec<u8> = vec![];
@@ -1462,6 +1467,13 @@ mod test {
decoded.merge_ordering().first().expr.to_string(),
sort_expr.expr.to_string(),
);
+ // Dropped on the wire, the limit would silently stop applying: the
+ // reader only ever executes after being decoded on an executor.
+ assert_eq!(
+ ExecutionPlan::fetch(decoded),
+ Some(20),
+ "fetch must round-trip"
+ );
// The ordering must land on `PlanProperties.eq_properties` —
downstream
// consumers (BWAG, SMJ build side) read it there.
let advertised = decoded
diff --git a/ballista/scheduler/src/state/aqe/adapter.rs
b/ballista/scheduler/src/state/aqe/adapter.rs
index 29d137615..9e39c5ffb 100644
--- a/ballista/scheduler/src/state/aqe/adapter.rs
+++ b/ballista/scheduler/src/state/aqe/adapter.rs
@@ -28,7 +28,10 @@ use ballista_core::execution_plans::{
use datafusion::common::exec_err;
use datafusion::config::ConfigOptions;
use datafusion::error::DataFusionError;
-use datafusion::physical_plan::{ExecutionPlanProperties, Partitioning};
+use
datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion::physical_plan::{
+ ExecutionPlanProperties, Partitioning, replace_children_if_necessary,
+};
use datafusion::scalar::ScalarValue;
use datafusion::{
common::tree_node::{Transformed, TreeNode, TreeNodeRecursion},
@@ -48,94 +51,120 @@ pub(crate) struct BallistaAdapter {
/// ShuffleWriterExec/SortShuffleWriterExec and [ShuffleReaderExec]
///
impl BallistaAdapter {
- fn transform_children(
+ /// Build the reader that replaces `exchange`, recording its upstream
+ /// stage as an input of this stage. `fetch` is a row limit from the
+ /// consumer; only the ordered reader can honor it.
+ fn build_reader(
&mut self,
- plan: Arc<dyn ExecutionPlan>,
- ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
- if let Some(exchange) = plan.downcast_ref::<ExchangeExec>() {
- let schema = exchange.schema().clone();
- let partitions = exchange.shuffle_partitions().ok_or_else(|| {
- DataFusionError::Execution(
- "partitions have to be resolved at this point".to_string(),
- )
- })?;
+ exchange: &ExchangeExec,
+ fetch: Option<usize>,
+ ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
+ let schema = exchange.schema().clone();
+ let partitions = exchange.shuffle_partitions().ok_or_else(|| {
+ DataFusionError::Execution(
+ "partitions have to be resolved at this point".to_string(),
+ )
+ })?;
- let stage_id = exchange.stage_id().ok_or_else(|| {
- DataFusionError::Execution(
- "stage ID has to be generated at this point".to_string(),
- )
- })?;
- let mut stage_output = StageOutput::new();
- for partition in partitions.iter().flatten().cloned() {
- stage_output.add_partition(partition);
- }
- stage_output.complete = true;
- self.inputs.insert(stage_id, stage_output);
- let partitioning = exchange.properties().partitioning.clone();
+ let stage_id = exchange.stage_id().ok_or_else(|| {
+ DataFusionError::Execution(
+ "stage ID has to be generated at this point".to_string(),
+ )
+ })?;
+ let mut stage_output = StageOutput::new();
+ for partition in partitions.iter().flatten().cloned() {
+ stage_output.add_partition(partition);
+ }
+ stage_output.complete = true;
+ self.inputs.insert(stage_id, stage_output);
+ let partitioning = exchange.properties().partitioning.clone();
- let reader: Arc<dyn ExecutionPlan> =
- match (exchange.coalesce(), exchange.broadcast) {
- (Some(cp), false) => {
- // Concatenate M-shape locations into K-shape per
CoalescePlan.groups.
- let k_shape: Vec<Vec<_>> = cp
- .groups
- .iter()
- .map(|pg| {
- let mut concat = Vec::new();
- for &idx in &pg.upstream_indices {
- if let Some(inner) = partitions.get(idx as
usize) {
- concat.extend_from_slice(inner);
- }
- }
- concat
- })
- .collect();
- let new_partitioning = match &partitioning {
- Partitioning::Hash(keys, _m) => {
- Partitioning::Hash(keys.clone(),
cp.groups.len())
+ Ok(match (exchange.coalesce(), exchange.broadcast) {
+ (Some(cp), false) => {
+ // Concatenate M-shape locations into K-shape per
CoalescePlan.groups.
+ let k_shape: Vec<Vec<_>> = cp
+ .groups
+ .iter()
+ .map(|pg| {
+ let mut concat = Vec::new();
+ for &idx in &pg.upstream_indices {
+ if let Some(inner) = partitions.get(idx as usize) {
+ concat.extend_from_slice(inner);
}
- _ =>
Partitioning::UnknownPartitioning(cp.groups.len()),
- };
- Arc::new(ShuffleReaderExec::try_new_coalesced(
- stage_id,
- k_shape,
- (*cp).clone(),
- schema,
- new_partitioning,
- )?)
- }
- (None, false) => {
- // Ordered-writer path: when the child declared an
output
- // ordering, preserve it across the shuffle boundary
with a
- // k-way merge instead of the arrival-order concat
that the
- // regular reader does.
- if let Some(ordering) =
exchange.input().output_ordering() {
- Arc::new(RangeShuffleReaderExec::try_new(
- stage_id,
- partitions,
- schema,
- ordering.clone(),
- )?)
- } else {
- Arc::new(ShuffleReaderExec::try_new(
- stage_id,
- partitions,
- schema,
- partitioning,
- )?)
}
+ concat
+ })
+ .collect();
+ let new_partitioning = match &partitioning {
+ Partitioning::Hash(keys, _m) => {
+ Partitioning::Hash(keys.clone(), cp.groups.len())
}
- (_, true) => Arc::new(ShuffleReaderExec::try_new_broadcast(
+ _ => Partitioning::UnknownPartitioning(cp.groups.len()),
+ };
+ Arc::new(ShuffleReaderExec::try_new_coalesced(
+ stage_id,
+ k_shape,
+ (*cp).clone(),
+ schema,
+ new_partitioning,
+ )?)
+ }
+ (None, false) => {
+ // Ordered-writer path: when the child declared an output
+ // ordering, preserve it across the shuffle boundary with a
+ // k-way merge instead of the arrival-order concat that the
+ // regular reader does.
+ if let Some(ordering) = exchange.input().output_ordering() {
+ Arc::new(
+ RangeShuffleReaderExec::try_new(
+ stage_id,
+ partitions,
+ schema,
+ ordering.clone(),
+ )?
+ .with_fetch_limit(fetch),
+ )
+ } else {
+ Arc::new(ShuffleReaderExec::try_new(
stage_id,
- exchange.shuffle_partitions_flattened(),
+ partitions,
schema,
-
exchange.input().output_partitioning().partition_count(),
- )?),
- };
- Ok(Transformed::yes(reader))
- } else {
- Ok(Transformed::no(plan))
+ partitioning,
+ )?)
+ }
+ }
+ (_, true) => Arc::new(ShuffleReaderExec::try_new_broadcast(
+ stage_id,
+ exchange.shuffle_partitions_flattened(),
+ schema,
+ exchange.input().output_partitioning().partition_count(),
+ )?),
+ })
+ }
+
+ fn transform_children(
+ &mut self,
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> datafusion::error::Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ // A merge with a row limit on top of an exchange: build the reader
+ // with the limit already set. Both merge on the same ordering, so the
+ // consumer's first `n` rows come from the reader's first `n`.
+ if let Some(spm) = plan.downcast_ref::<SortPreservingMergeExec>()
+ && let Some(fetch) = spm.fetch()
+ && let Some(exchange) = spm.input().downcast_ref::<ExchangeExec>()
+ {
+ let reader = self.build_reader(exchange, Some(fetch))?;
+ return Ok(Transformed::yes(replace_children_if_necessary(
+ plan,
+ vec![reader],
+ )?));
}
+
+ if let Some(exchange) = plan.downcast_ref::<ExchangeExec>() {
+ return Ok(Transformed::yes(self.build_reader(exchange, None)?));
+ }
+
+ Ok(Transformed::no(plan))
}
/// Converts Adaptive plan to plan which ballista expects
@@ -384,4 +413,37 @@ mod tests {
);
assert!(out.downcast_ref::<RangeShuffleReaderExec>().is_none());
}
+
+ /// The consuming merge's row limit must reach the reader. The merge is
+ /// rebuilt around the new reader, so a dropped limit is silent.
+ #[test]
+ fn pushes_consumer_limit_into_range_reader() {
+ let schema = f64_schema();
+ let empty: Vec<Vec<RecordBatch>> = vec![vec![]];
+ let source =
+ MemorySourceConfig::try_new_exec(&empty, schema.clone(),
None).unwrap();
+ let sort_lex = LexOrdering::new(vec![asc(&schema, "v")]).unwrap();
+ let sorted = Arc::new(
+ SortExec::new(sort_lex.clone(),
source).with_preserve_partitioning(true),
+ ) as Arc<dyn ExecutionPlan>;
+
+ let exchange = ExchangeExec::new(sorted, None, 0);
+ exchange.set_stage_id(1);
+ exchange.resolve_shuffle_partitions(vec![vec![]]);
+ let merge = Arc::new(
+ SortPreservingMergeExec::new(sort_lex, Arc::new(exchange))
+ .with_fetch(Some(7)),
+ ) as Arc<dyn ExecutionPlan>;
+
+ let mut adapter = BallistaAdapter::default();
+ let out = adapter.transform_children(merge).unwrap().data;
+
+ let children = out.children();
+ let reader = children[0]
+ .downcast_ref::<RangeShuffleReaderExec>()
+ .unwrap_or_else(|| {
+ panic!("expected a range reader, got {}", children[0].name())
+ });
+ assert_eq!(reader.fetch(), Some(7));
+ }
}
diff --git a/ballista/scheduler/src/state/task_builder.rs
b/ballista/scheduler/src/state/task_builder.rs
index f1ed152a2..a9ab5e502 100644
--- a/ballista/scheduler/src/state/task_builder.rs
+++ b/ballista/scheduler/src/state/task_builder.rs
@@ -336,7 +336,7 @@ fn select_output_partitions(
) else {
return Ok(None);
};
- return Ok(Some(Arc::new(restricted)));
+ return Ok(Some(Arc::new(restricted.with_fetch_limit(reader.fetch()))));
}
// DataSourceExec: file-backed or in-memory scans.
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
