This is an automated email from the ASF dual-hosted git repository.
avantgardnerio 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 06f8f1d9c feat(core): RangeShuffleReaderExec — ordering-preserving
shuffle reader (purely additive) (#2255)
06f8f1d9c is described below
commit 06f8f1d9cfd07d04549404d9ef870d7d42a39b02
Author: Brent Gardner <[email protected]>
AuthorDate: Sat Aug 8 09:46:24 2026 -0600
feat(core): RangeShuffleReaderExec — ordering-preserving shuffle reader
(purely additive) (#2255)
* feat(core): RangeShuffleReaderExec — ordering-preserving shuffle reader
(purely additive)
**Purely additive.** No in-tree callers, no behavior change. The type is
defined, serialized, and round-trip tested; wiring lands in follow-ups.
Extracted from #2223 (a working end-to-end parallel-window branch) as
the first slice of the epic to land parallel windows incrementally,
one reviewable piece at a time.
Sibling to `ShuffleReaderExec`: keeps each upstream source alive as its
own stream and feeds all N into a `StreamingMerge` keyed on the child's
declared output ordering. The regular reader concatenates in arrival
order, which breaks the monotonicity RANGE-frame window operators (and
sort-merge-join build sides) require.
- `range_shuffle_reader.rs` — operator impl + tests
- proto message + oneof slot 12
- serde encode/decode + roundtrip test
- three shuffle_reader helpers promoted to `pub(crate)` so RSR can
reuse them: `local_remote_read_split`, `fetch_partition_local`,
`fetch_partition_remote`
* docs(executor): TODO for RangeShuffleReaderExec late-bind
Marks the follow-up: create_query_stage_exec only downcasts
ShuffleReaderExec, so a task carrying a RangeShuffleReaderExec would
miss work_dir + client_pool the moment a planner rule plants one.
No functional change today — no code path emits the operator yet.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
---------
Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]>
---
ballista/core/proto/ballista.proto | 13 +
ballista/core/src/execution_plans/mod.rs | 2 +
.../src/execution_plans/range_shuffle_reader.rs | 569 +++++++++++++++++++++
.../core/src/execution_plans/shuffle_reader.rs | 6 +-
ballista/core/src/serde/generated/ballista.rs | 22 +-
ballista/core/src/serde/mod.rs | 145 +++++-
ballista/executor/src/execution_engine.rs | 4 +
7 files changed, 755 insertions(+), 6 deletions(-)
diff --git a/ballista/core/proto/ballista.proto
b/ballista/core/proto/ballista.proto
index 285accd52..e8becaba7 100644
--- a/ballista/core/proto/ballista.proto
+++ b/ballista/core/proto/ballista.proto
@@ -57,6 +57,7 @@ message BallistaPhysicalPlanNode {
OrderedRangeRepartitionExecNode ordered_range_repartition = 9;
PerPartitionFilterExecNode per_partition_filter = 10;
PartitionedBoundedWindowAggExecNode partitioned_bounded_window_agg = 11;
+ RangeShuffleReaderExecNode range_shuffle_reader = 12;
}
}
@@ -201,6 +202,18 @@ message ShuffleReaderPartition {
repeated PartitionLocation location = 1;
}
+// Ordering-preserving shuffle reader. Reuses `ShuffleReaderPartition` for the
+// M-shape source layout. Partitioning is derived from `partition.len()`
+// (always `UnknownPartitioning`, range-partitioned by `merge_ordering`).
+message RangeShuffleReaderExecNode {
+ repeated ShuffleReaderPartition partition = 1;
+ datafusion_common.Schema schema = 2;
+ uint32 stage_id = 3;
+ // 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;
+}
+
// CoalescePartitionsRule output: groups upstream partitions into coalesced
output partitions.
// Empty when no coalesce is applied (the optional field on the parent message
is absent).
message CoalescePlan {
diff --git a/ballista/core/src/execution_plans/mod.rs
b/ballista/core/src/execution_plans/mod.rs
index 230b68913..22ce701c0 100644
--- a/ballista/core/src/execution_plans/mod.rs
+++ b/ballista/core/src/execution_plans/mod.rs
@@ -27,6 +27,7 @@ mod partitioned_bounded_window_agg;
mod per_partition_filter;
pub mod plan_algebra;
mod range_repartition_common;
+mod range_shuffle_reader;
mod runtime_stats;
mod shuffle_reader;
mod shuffle_writer;
@@ -46,6 +47,7 @@ pub use
ordered_range_repartition::OrderedRangeRepartitionExec;
pub use partitioned_bounded_window_agg::PartitionedBoundedWindowAggExec;
pub use per_partition_filter::{PerPartitionFilterExec,
range_partition_predicates};
pub use plan_algebra::{preserves_distribution, preserves_partitioning};
+pub use range_shuffle_reader::RangeShuffleReaderExec;
pub use runtime_stats::{
MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats,
collect_reports as collect_runtime_stats_reports, cut_partitions,
diff --git a/ballista/core/src/execution_plans/range_shuffle_reader.rs
b/ballista/core/src/execution_plans/range_shuffle_reader.rs
new file mode 100644
index 000000000..abeb33938
--- /dev/null
+++ b/ballista/core/src/execution_plans/range_shuffle_reader.rs
@@ -0,0 +1,569 @@
+// 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.
+
+//! Ordering-preserving shuffle reader.
+//!
+//! The regular [`ShuffleReaderExec`](super::ShuffleReaderExec) fans every
+//! upstream source for an output partition into one mpsc channel,
`try_flatten`s
+//! them in arrival order, then coalesces batches. Each source file is
+//! internally sorted, but the concatenation is not — the resulting stream
+//! violates the monotonicity that RANGE-frame window operators (and
+//! sort-merge-join build sides) require.
+//!
+//! `RangeShuffleReaderExec` keeps each upstream source alive as its own
+//! stream and feeds all N into a `StreamingMerge` keyed on the child's
+//! declared output ordering. Batches within an output partition are
+//! globally sorted on the merge key.
+//!
+//! Trade-offs vs. the regular reader:
+//!
+//! - No permit-based fetch governor. Backpressure flows from the merge
+//! consumer down through each per-source stream — the merge only polls
+//! the source it needs next, so h2 (for remote) and disk (for local)
+//! throttle naturally.
+//! - No mid-body fetch retry. `fetch_partition_remote` streams directly
+//! without buffering the whole source, so a transport error mid-body
+//! fails the merged stream; the task-level retry re-executes.
+//! - No coalesce or broadcast variants. Ordered outputs come from
+//! `OrderedRangeRepartitionExec`, which is one-to-one and never fanned
+//! into a broadcast.
+//!
+//! # Future work
+//!
+//! Today the reader pulls whole upstream files and lets `StreamingMerge`
+//! do the work. Once value-indexed shuffle files land end-to-end (writer
+//! side in [PR #2204]), the reader will consult per-file ValueIndex
+//! offsets and fetch only the byte ranges covering its target output
+//! partition — dropping the "read the whole file to throw most of it
+//! away" cost that dominates when K is large.
+//!
+//! [PR #2204]: https://github.com/apache/datafusion-ballista/pull/2204
+
+use crate::client_pool::BallistaClientPool;
+use crate::execution_plans::shuffle_reader::{
+ fetch_partition_local, fetch_partition_remote, local_remote_read_split,
+ stats_for_partition,
+};
+use crate::extension::SessionConfigExt;
+use crate::serde::scheduler::PartitionLocation;
+use crate::utils::GrpcClientConfig;
+use datafusion::arrow::datatypes::SchemaRef;
+use datafusion::common::{Result, Statistics};
+use datafusion::error::DataFusionError;
+use datafusion::execution::TaskContext;
+use datafusion::execution::memory_pool::MemoryConsumer;
+use datafusion::physical_expr::{EquivalenceProperties, LexOrdering,
Partitioning};
+use datafusion::physical_plan::metrics::{BaselineMetrics,
ExecutionPlanMetricsSet};
+use datafusion::physical_plan::sorts::streaming_merge::StreamingMergeBuilder;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
+ SendableRecordBatchStream,
+};
+use futures::TryStreamExt;
+use log::debug;
+use std::sync::Arc;
+
+/// Ordering-preserving shuffle reader. See module docs.
+#[derive(Debug)]
+pub struct RangeShuffleReaderExec {
+ /// Upstream stage that produced these files.
+ pub stage_id: usize,
+ pub(crate) schema: SchemaRef,
+ /// M-shape: `partition[k]` lists the upstream partition locations that
+ /// output partition `k` needs to merge.
+ pub partition: Vec<Vec<PartitionLocation>>,
+ /// Sort key the merge preserves. Advertised in
`PlanProperties.eq_properties`
+ /// so downstream operators (BWAG, SMJ build side) see the output ordering.
+ merge_ordering: LexOrdering,
+ metrics: ExecutionPlanMetricsSet,
+ properties: Arc<PlanProperties>,
+ work_dir: Option<String>,
+ client_pool: Option<Arc<dyn BallistaClientPool>>,
+}
+
+impl RangeShuffleReaderExec {
+ /// The output partition count is `partition.len()`. Partitioning is
+ /// range on the merge key — DF has no `Partitioning::Range`, so it's
+ /// reported as `UnknownPartitioning`. Downstream co-partitioning is
+ /// carried by the advertised ordering, not by partitioning kind.
+ pub fn try_new(
+ stage_id: usize,
+ partition: Vec<Vec<PartitionLocation>>,
+ schema: SchemaRef,
+ merge_ordering: LexOrdering,
+ ) -> Result<Self> {
+ let output_partition_count = partition.len();
+ let eq_properties = EquivalenceProperties::new_with_orderings(
+ schema.clone(),
+ vec![merge_ordering.clone()],
+ );
+ let properties = Arc::new(PlanProperties::new(
+ eq_properties,
+ Partitioning::UnknownPartitioning(output_partition_count),
+
datafusion::physical_plan::execution_plan::EmissionType::Incremental,
+ datafusion::physical_plan::execution_plan::Boundedness::Bounded,
+ ));
+ Ok(Self {
+ stage_id,
+ schema,
+ partition,
+ merge_ordering,
+ metrics: ExecutionPlanMetricsSet::new(),
+ properties,
+ work_dir: None,
+ client_pool: None,
+ })
+ }
+
+ /// 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(),
+ }
+ }
+
+ /// 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),
+ }
+ }
+
+ /// Sort key the merge preserves.
+ pub fn merge_ordering(&self) -> &LexOrdering {
+ &self.merge_ordering
+ }
+}
+
+impl DisplayAs for RangeShuffleReaderExec {
+ fn fmt_as(
+ &self,
+ t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ match t {
+ DisplayFormatType::Default | DisplayFormatType::Verbose => {
+ write!(
+ f,
+ "RangeShuffleReaderExec: upstream_stage: {}, partitions:
{}, ordering: {}",
+ self.stage_id,
+ self.partition.len(),
+ self.merge_ordering,
+ )
+ }
+ DisplayFormatType::TreeRender => {
+ writeln!(f, "upstream_stage={}", self.stage_id)?;
+ writeln!(f, "output_partitions={}", self.partition.len())?;
+ writeln!(f, "ordering={}", self.merge_ordering)
+ }
+ }
+ }
+}
+
+impl ExecutionPlan for RangeShuffleReaderExec {
+ fn name(&self) -> &str {
+ "RangeShuffleReaderExec"
+ }
+
+ fn schema(&self) -> SchemaRef {
+ self.schema.clone()
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ if !children.is_empty() {
+ return Err(DataFusionError::Plan(
+ "RangeShuffleReaderExec does not support children
plans".to_owned(),
+ ));
+ }
+ 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(),
+ }))
+ }
+
+ fn execute(
+ &self,
+ output_partition: usize,
+ ctx: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let task_id = ctx
+ .task_id()
+ .unwrap_or_else(|| output_partition.to_string());
+ debug!("RangeShuffleReaderExec::execute({task_id})");
+
+ let config = ctx.session_config();
+ let work_dir = self.work_dir.as_ref().ok_or_else(|| {
+ DataFusionError::Configuration(
+ "RangeShuffleReaderExec work dir should have been set by
executor"
+ .to_owned(),
+ )
+ })?;
+
+ let locations = self.partition[output_partition].clone();
+ let (local_locations, remote_locations) = local_remote_read_split(
+ work_dir,
+ locations,
+ config.ballista_shuffle_reader_force_remote_read(),
+ );
+
+ let mut sub_streams: Vec<SendableRecordBatchStream> =
+ Vec::with_capacity(local_locations.len() + remote_locations.len());
+
+ for loc in local_locations {
+ let stream = fetch_partition_local(work_dir, &loc)
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ sub_streams.push(stream);
+ }
+
+ if !remote_locations.is_empty() {
+ let grpc_config: Arc<GrpcClientConfig> =
+ Arc::new((&config.ballista_config()).into());
+ let customize_endpoint =
+ config.ballista_override_create_grpc_client_endpoint();
+ let prefer_flight =
config.ballista_shuffle_reader_remote_prefer_flight();
+ let client_pool = self.client_pool.clone();
+
+ for loc in remote_locations {
+ let schema = self.schema.clone();
+ let grpc_config = grpc_config.clone();
+ let customize_endpoint = customize_endpoint.clone();
+ let client_pool = client_pool.clone();
+
+ // Lazy connect: the first poll from the merge triggers the
+ // remote fetch; subsequent polls stream batches directly. No
+ // buffering, no retry — task-level retry covers transport
+ // failures.
+ let lazy = futures::stream::once(async move {
+ fetch_partition_remote(
+ &loc,
+ grpc_config,
+ prefer_flight,
+ customize_endpoint,
+ client_pool,
+ )
+ .await
+ .map_err(|e| DataFusionError::External(Box::new(e)))
+ })
+ .try_flatten();
+
+
sub_streams.push(Box::pin(RecordBatchStreamAdapter::new(schema, lazy)));
+ }
+ }
+
+ if sub_streams.is_empty() {
+ return Ok(Box::pin(RecordBatchStreamAdapter::new(
+ self.schema.clone(),
+ futures::stream::empty(),
+ )));
+ }
+
+ let baseline = BaselineMetrics::new(&self.metrics, output_partition);
+ let reservation = MemoryConsumer::new(format!(
+ "RangeShuffleReaderExec[stage={},out={}]",
+ self.stage_id, output_partition,
+ ))
+ .register(ctx.memory_pool());
+
+ let merged = StreamingMergeBuilder::new()
+ .with_streams(sub_streams)
+ .with_schema(self.schema.clone())
+ .with_expressions(&self.merge_ordering)
+ .with_batch_size(config.batch_size())
+ .with_metrics(baseline)
+ .with_reservation(reservation)
+ .build()?;
+
+ Ok(merged)
+ }
+
+ fn metrics(&self) ->
Option<datafusion::physical_plan::metrics::MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+
+ fn partition_statistics(&self, partition: Option<usize>) ->
Result<Arc<Statistics>> {
+ if let Some(idx) = partition {
+ let partition_count = self.partition.len();
+ if idx >= partition_count {
+ return datafusion::common::internal_err!(
+ "Invalid partition index: {}, the partition count is {}",
+ idx,
+ partition_count
+ );
+ }
+ let stats =
+ stats_for_partition(idx, self.schema.fields().len(),
&self.partition)?;
+ return Ok(Arc::new(stats));
+ }
+ Ok(Arc::new(Statistics::new_unknown(&self.schema)))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::serde::scheduler::{
+ ExecutorMetadata, ExecutorOperatingSystemSpecification,
ExecutorSpecification,
+ PartitionId, PartitionStats,
+ };
+ use datafusion::arrow::array::{ArrayRef, Float64Array, RecordBatch};
+ use datafusion::arrow::datatypes::{DataType, Field, Schema};
+ use datafusion::arrow::ipc::writer::StreamWriter;
+ use datafusion::physical_expr::PhysicalSortExpr;
+ use datafusion::physical_expr::expressions::Column;
+ use datafusion::prelude::SessionContext;
+ use std::fs::{File, create_dir_all};
+ use tempfile::tempdir;
+
+ /// Write an Arrow IPC stream containing `batches` at `path`.
+ fn write_ipc_stream(path: &std::path::Path, batches: &[RecordBatch]) {
+ create_dir_all(path.parent().unwrap()).unwrap();
+ let file = File::create(path).unwrap();
+ let mut writer = StreamWriter::try_new(file,
&batches[0].schema()).unwrap();
+ for b in batches {
+ writer.write(b).unwrap();
+ }
+ writer.finish().unwrap();
+ }
+
+ fn sorted_batch(schema: SchemaRef, values: &[f64]) -> RecordBatch {
+ let col: ArrayRef = Arc::new(Float64Array::from(values.to_vec()));
+ RecordBatch::try_new(schema, vec![col]).unwrap()
+ }
+
+ fn make_location(job: &str, stage: usize, partition: usize) ->
PartitionLocation {
+ PartitionLocation {
+ map_partition_id: partition,
+ partition_id: PartitionId {
+ job_id: job.into(),
+ stage_id: stage,
+ partition_id: partition,
+ },
+ executor_meta: ExecutorMetadata {
+ id: "test-executor".to_string(),
+ host: "127.0.0.1".to_string(),
+ port: 0,
+ grpc_port: 0,
+ specification: ExecutorSpecification::default(),
+ os_info: ExecutorOperatingSystemSpecification::default(),
+ },
+ partition_stats: PartitionStats::default(),
+ file_id: None,
+ is_sort_shuffle: false,
+ }
+ }
+
+ /// Two upstream sources, each internally sorted; the range reader must
+ /// interleave them into a single globally-sorted stream. This is the core
+ /// correctness property the regular reader violates.
+ #[tokio::test]
+ async fn merges_two_sorted_sources() {
+ let dir = tempdir().unwrap();
+ let work_dir = dir.path();
+ let job = "job-merges-two";
+ let stage_id = 1usize;
+
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Float64,
false)]));
+
+ // Source 0: [1.0, 3.0, 5.0]
+ // Source 1: [2.0, 4.0, 6.0]
+ // Merged: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
+ let src0 = sorted_batch(schema.clone(), &[1.0, 3.0, 5.0]);
+ let src1 = sorted_batch(schema.clone(), &[2.0, 4.0, 6.0]);
+
+ let loc0 = make_location(job, stage_id, 0);
+ let loc1 = make_location(job, stage_id, 1);
+ let path0 = loc0.path(work_dir.to_str().unwrap()).expect("path0");
+ let path1 = loc1.path(work_dir.to_str().unwrap()).expect("path1");
+ write_ipc_stream(&path0, &[src0]);
+ write_ipc_stream(&path1, &[src1]);
+
+ // Both sources land in output partition 0.
+ let partitions = vec![vec![loc0, loc1]];
+ let merge_ordering =
LexOrdering::new(vec![PhysicalSortExpr::new_default(
+ Arc::new(Column::new("v", 0)),
+ )])
+ .unwrap();
+
+ let reader = RangeShuffleReaderExec::try_new(
+ stage_id,
+ partitions,
+ schema.clone(),
+ merge_ordering,
+ )
+ .unwrap()
+ .with_work_dir(work_dir.to_string_lossy().to_string());
+
+ let ctx = SessionContext::new();
+ let task_ctx = ctx.task_ctx();
+ let stream = reader.execute(0, task_ctx).unwrap();
+ let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+
+ let all: Vec<f64> = batches
+ .iter()
+ .flat_map(|b| {
+ b.column(0)
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect();
+ assert_eq!(all, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
+ }
+
+ /// If the caller only has one source per output partition, the reader
+ /// must still work — StreamingMerge over 1 stream is a valid degenerate
+ /// case.
+ #[tokio::test]
+ async fn passes_through_single_source() {
+ let dir = tempdir().unwrap();
+ let work_dir = dir.path();
+ let job = "job-single-source";
+ let stage_id = 2usize;
+
+ let schema =
+ Arc::new(Schema::new(vec![Field::new("v", DataType::Float64,
false)]));
+ let src = sorted_batch(schema.clone(), &[10.0, 20.0, 30.0]);
+
+ let loc = make_location(job, stage_id, 0);
+ let path = loc.path(work_dir.to_str().unwrap()).expect("path");
+ write_ipc_stream(&path, &[src]);
+
+ let partitions = vec![vec![loc]];
+ let merge_ordering =
LexOrdering::new(vec![PhysicalSortExpr::new_default(
+ Arc::new(Column::new("v", 0)),
+ )])
+ .unwrap();
+
+ let reader = RangeShuffleReaderExec::try_new(
+ stage_id,
+ partitions,
+ schema.clone(),
+ merge_ordering,
+ )
+ .unwrap()
+ .with_work_dir(work_dir.to_string_lossy().to_string());
+
+ let ctx = SessionContext::new();
+ let stream = reader.execute(0, ctx.task_ctx()).unwrap();
+ let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+ let all: Vec<f64> = batches
+ .iter()
+ .flat_map(|b| {
+ b.column(0)
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .unwrap()
+ .values()
+ .to_vec()
+ })
+ .collect();
+ assert_eq!(all, vec![10.0, 20.0, 30.0]);
+ }
+
+ /// Empty partitions.len() > 0 with all-empty inner locations must return
+ /// an empty stream instead of tripping StreamingMerge on zero sub-streams.
+ #[tokio::test]
+ async fn handles_empty_partition() {
+ 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(
+ 7,
+ vec![vec![]],
+ schema.clone(),
+ merge_ordering,
+ )
+ .unwrap()
+ .with_work_dir("/tmp".to_string());
+
+ let ctx = SessionContext::new();
+ let stream = reader.execute(0, ctx.task_ctx()).unwrap();
+ let batches: Vec<RecordBatch> = stream.try_collect().await.unwrap();
+ assert!(batches.is_empty());
+ }
+
+ /// The reader must advertise its merge ordering so downstream operators
+ /// see the sortedness invariant (BWAG's RANGE-frame cursor, SMJ build
side).
+ #[test]
+ fn advertises_merge_ordering() {
+ 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.clone(),
+ )
+ .unwrap();
+
+ let out = reader
+ .properties()
+ .eq_properties
+ .oeq_class()
+ .iter()
+ .next()
+ .cloned();
+ let expected = out.expect("expected an advertised ordering");
+ assert_eq!(expected, merge_ordering);
+
+ assert_eq!(reader.properties().partitioning.partition_count(), 4);
+ }
+}
diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs
b/ballista/core/src/execution_plans/shuffle_reader.rs
index d85f1978f..6f4ec1a6a 100644
--- a/ballista/core/src/execution_plans/shuffle_reader.rs
+++ b/ballista/core/src/execution_plans/shuffle_reader.rs
@@ -692,7 +692,7 @@ impl RecordBatchStream for GovernedStream {
/// Local partitions are read directly from local Arrow IPC files,
/// while remote partitions are fetched using the Arrow Flight client.
/// If `force_remote_read` is true, all partitions are treated as remote.
-fn local_remote_read_split(
+pub(crate) fn local_remote_read_split(
work_dir: &str,
partition_locations: Vec<PartitionLocation>,
force_remote_read: bool,
@@ -1037,7 +1037,7 @@ async fn new_ballista_client(
.await
}
-async fn fetch_partition_remote(
+pub(crate) async fn fetch_partition_remote(
location: &PartitionLocation,
config: Arc<GrpcClientConfig>,
prefer_flight: bool,
@@ -1108,7 +1108,7 @@ async fn fetch_partition_remote(
}
}
-fn fetch_partition_local(
+pub(crate) fn fetch_partition_local(
work_dir: &str,
location: &PartitionLocation,
) -> result::Result<SendableRecordBatchStream, BallistaError> {
diff --git a/ballista/core/src/serde/generated/ballista.rs
b/ballista/core/src/serde/generated/ballista.rs
index c5d3899dd..e01a99aa7 100644
--- a/ballista/core/src/serde/generated/ballista.rs
+++ b/ballista/core/src/serde/generated/ballista.rs
@@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode {
pub struct BallistaPhysicalPlanNode {
#[prost(
oneof = "ballista_physical_plan_node::PhysicalPlanType",
- tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11"
+ tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12"
)]
pub physical_plan_type: ::core::option::Option<
ballista_physical_plan_node::PhysicalPlanType,
@@ -63,6 +63,8 @@ pub mod ballista_physical_plan_node {
PerPartitionFilter(super::PerPartitionFilterExecNode),
#[prost(message, tag = "11")]
PartitionedBoundedWindowAgg(super::PartitionedBoundedWindowAggExecNode),
+ #[prost(message, tag = "12")]
+ RangeShuffleReader(super::RangeShuffleReaderExecNode),
}
}
/// Value-range router over N locally-sorted overlapping input partitions.
@@ -249,6 +251,24 @@ pub struct ShuffleReaderPartition {
#[prost(message, repeated, tag = "1")]
pub location: ::prost::alloc::vec::Vec<PartitionLocation>,
}
+/// Ordering-preserving shuffle reader. Reuses `ShuffleReaderPartition` for the
+/// M-shape source layout. Partitioning is derived from `partition.len()`
+/// (always `UnknownPartitioning`, range-partitioned by `merge_ordering`).
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct RangeShuffleReaderExecNode {
+ #[prost(message, repeated, tag = "1")]
+ pub partition: ::prost::alloc::vec::Vec<ShuffleReaderPartition>,
+ #[prost(message, optional, tag = "2")]
+ pub schema: ::core::option::Option<::datafusion_proto_common::Schema>,
+ #[prost(uint32, tag = "3")]
+ pub stage_id: u32,
+ /// Sort key the reader's k-way merge preserves. Advertised on the reader's
+ /// `PlanProperties.eq_properties` for downstream consumers.
+ #[prost(message, repeated, tag = "4")]
+ pub merge_ordering: ::prost::alloc::vec::Vec<
+ ::datafusion_proto::protobuf::PhysicalSortExprNode,
+ >,
+}
/// CoalescePartitionsRule output: groups upstream partitions into coalesced
output partitions.
/// Empty when no coalesce is applied (the optional field on the parent
message is absent).
#[derive(Clone, PartialEq, ::prost::Message)]
diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs
index eb2f0181c..1a4885fef 100644
--- a/ballista/core/src/serde/mod.rs
+++ b/ballista/core/src/serde/mod.rs
@@ -58,8 +58,8 @@ use crate::execution_plans::sort_shuffle::SortShuffleConfig;
use crate::execution_plans::{
BufferExec, BufferMode, ChaosExec, CoalescePlan,
OrderedRangeRepartitionExec,
PartitionGroup, PartitionedBoundedWindowAggExec, PerPartitionFilterExec,
- RuntimeStatsExec, ShuffleReaderExec, ShuffleWriterExec,
SortShuffleWriterExec,
- UnorderedRangeRepartitionExec, UnresolvedShuffleExec,
+ RangeShuffleReaderExec, RuntimeStatsExec, ShuffleReaderExec,
ShuffleWriterExec,
+ SortShuffleWriterExec, UnorderedRangeRepartitionExec,
UnresolvedShuffleExec,
};
use crate::serde::protobuf::{
ballista_logical_plan_node::LogicalPlanType,
@@ -507,6 +507,46 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
};
Ok(Arc::new(exec))
}
+ PhysicalPlanType::RangeShuffleReader(range_reader) => {
+ let stage_id = range_reader.stage_id as usize;
+ let schema: SchemaRef =
Arc::new(convert_required!(range_reader.schema)?);
+ let partition_location: Vec<Vec<PartitionLocation>> =
range_reader
+ .partition
+ .iter()
+ .map(|p| {
+ p.location
+ .iter()
+ .map(|l| {
+ l.clone().try_into().map_err(|e| {
+ DataFusionError::Internal(format!(
+ "Fail to get partition location due to
{e:?}"
+ ))
+ })
+ })
+ .collect::<Result<Vec<_>, _>>()
+ })
+ .collect::<Result<Vec<_>, DataFusionError>>()?;
+ let merge_ordering_exprs = parse_physical_sort_exprs(
+ &range_reader.merge_ordering,
+ &decode_ctx,
+ schema.as_ref(),
+ &converter,
+ )?;
+ let merge_ordering =
datafusion::physical_expr::LexOrdering::new(
+ merge_ordering_exprs,
+ )
+ .ok_or_else(|| {
+ proto_error(
+ "RangeShuffleReaderExec: merge_ordering must be
non-empty",
+ )
+ })?;
+ Ok(Arc::new(RangeShuffleReaderExec::try_new(
+ stage_id,
+ partition_location,
+ schema,
+ merge_ordering,
+ )?))
+ }
PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => {
let schema: SchemaRef =
Arc::new(convert_required!(unresolved_shuffle.schema)?);
@@ -804,6 +844,46 @@ impl PhysicalExtensionCodec for
BallistaPhysicalExtensionCodec {
))
})?;
+ Ok(())
+ } else if let Some(exec) =
node.downcast_ref::<RangeShuffleReaderExec>() {
+ let stage_id = exec.stage_id as u32;
+ let mut partition = vec![];
+ for location in &exec.partition {
+ partition.push(protobuf::ShuffleReaderPartition {
+ location: location
+ .iter()
+ .map(|l| {
+ l.clone().try_into().map_err(|e| {
+ DataFusionError::Internal(format!(
+ "Fail to get partition location due to
{e:?}"
+ ))
+ })
+ })
+ .collect::<Result<Vec<_>, _>>()?,
+ });
+ }
+ let converter = DefaultPhysicalProtoConverter {};
+ let merge_ordering = serialize_physical_sort_exprs(
+ exec.merge_ordering().iter().cloned(),
+ self.default_codec.as_ref(),
+ &converter,
+ )?;
+ let proto = protobuf::BallistaPhysicalPlanNode {
+ physical_plan_type: Some(PhysicalPlanType::RangeShuffleReader(
+ protobuf::RangeShuffleReaderExecNode {
+ stage_id,
+ partition,
+ schema: Some(exec.schema().as_ref().try_into()?),
+ merge_ordering,
+ },
+ )),
+ };
+ proto.encode(buf).map_err(|e| {
+ DataFusionError::Internal(format!(
+ "failed to encode range shuffle reader execution plan:
{e:?}"
+ ))
+ })?;
+
Ok(())
} else if let Some(exec) =
node.downcast_ref::<UnresolvedShuffleExec>() {
let converter = DefaultPhysicalProtoConverter {};
@@ -1218,6 +1298,67 @@ mod test {
);
}
+ /// `RangeShuffleReaderExec` carries its merge ordering across the wire —
+ /// the reader's k-way merge machinery is inert without it, and the
+ /// scheduler side plants the reader with the child's declared ordering.
+ #[tokio::test]
+ async fn test_range_shuffle_reader_exec_roundtrip() {
+ use datafusion::arrow::compute::SortOptions;
+ use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
+
+ let schema = create_test_schema();
+ let sort_expr = PhysicalSortExpr {
+ expr: col("id", schema.as_ref()).unwrap(),
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ };
+ let merge_ordering =
LexOrdering::new(vec![sort_expr.clone()]).unwrap();
+
+ let original = RangeShuffleReaderExec::try_new(
+ 7,
+ vec![vec![]; 4],
+ schema.clone(),
+ merge_ordering.clone(),
+ )
+ .unwrap();
+
+ let codec = BallistaPhysicalExtensionCodec::default();
+ let mut buf: Vec<u8> = vec![];
+ codec.try_encode(Arc::new(original), &mut buf).unwrap();
+
+ let ctx = SessionContext::new().task_ctx();
+ let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap();
+ let decoded = decoded_plan
+ .downcast_ref::<RangeShuffleReaderExec>()
+ .expect("Expected RangeShuffleReaderExec");
+
+ assert_eq!(decoded.stage_id, 7);
+ assert_eq!(decoded.schema().as_ref(), schema.as_ref());
+ assert_eq!(decoded.partition.len(), 4, "partition shape must survive");
+ assert_eq!(
+ decoded.merge_ordering().len(),
+ 1,
+ "merge_ordering must round-trip"
+ );
+ assert_eq!(
+ decoded.merge_ordering().first().expr.to_string(),
+ sort_expr.expr.to_string(),
+ );
+ // The ordering must land on `PlanProperties.eq_properties` —
downstream
+ // consumers (BWAG, SMJ build side) read it there.
+ let advertised = decoded
+ .properties()
+ .eq_properties
+ .oeq_class()
+ .iter()
+ .next()
+ .cloned()
+ .expect("advertised ordering");
+ assert_eq!(advertised, merge_ordering);
+ }
+
/// The sort shuffle writer's per-task memory budget must survive the
/// protobuf round trip so a session override reaches the executor where
the
/// writer actually runs (issue #2089). A non-default value and the special
diff --git a/ballista/executor/src/execution_engine.rs
b/ballista/executor/src/execution_engine.rs
index 4f8e6663d..fcfd52514 100644
--- a/ballista/executor/src/execution_engine.rs
+++ b/ballista/executor/src/execution_engine.rs
@@ -135,6 +135,10 @@ impl ExecutionEngine for DefaultExecutionEngine {
) -> Result<Arc<dyn QueryStageExecutor>> {
let plan = plan
.transform(|p| {
+ // TODO: RangeShuffleReaderExec needs the same late-bind
+ // (with_work_dir + with_client_pool) once a planner rule
+ // plants it; without it, the first task carrying one will
+ // fail with "work dir should have been set by executor".
if let Some(reader) = p.downcast_ref::<ShuffleReaderExec>() {
match &self.client_pool {
Some(client_pool) => Ok(Transformed::yes(Arc::new(
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
