feat(core): range shuffle reads only the bytes a consumer's range covers

ballista/core/proto/ballista.proto

@@ -60,6 +60,7 @@ message BallistaPhysicalPlanNode {
     RangeShuffleReaderExecNode range_shuffle_reader = 12;
     RangeFilterExecNode range_filter = 13;
     PrefixMergeExecNode prefix_merge = 14;
+    RangeShuffleWriterExecNode range_shuffle_writer = 15;
   }
 }
 
@@ -320,6 +321,15 @@ message SortShuffleWriterExecNode {
   optional uint64 memory_limit_per_task_bytes = 9;
 }
 
+// Passthrough shuffle writer that emits the seekable Arrow IPC file format.
+// Carries no output partitioning: like `ShuffleWriterExecNode`'s passthrough
+// case it writes its input's partitioning through, one file per partition.
+message RangeShuffleWriterExecNode {
+  string job_id = 1;
+  uint32 stage_id = 2;
+  datafusion.PhysicalPlanNode input = 3;
+}
+
 message UnresolvedShuffleExecNode {
   uint32 stage_id = 1;
   datafusion_common.Schema schema = 2;
@@ -359,6 +369,11 @@ message RangeShuffleReaderExecNode {
   repeated datafusion.PhysicalSortExprNode merge_ordering = 4;
   // Row limit pushed down by a consuming merge. Absent means read everything.
   optional uint64 fetch = 5;
+  // Half-open value range each output partition covers, halo already applied.
+  // One per output partition, or empty when the scheduler had no cuts. Lets
+  // the reader skip the parts of an indexed source that cannot hold a row the
+  // partition wants; the consuming RangeFilterExec still does the exact trim.
+  repeated RangeBound bounds = 6;
 }
 
 // CoalescePartitionsRule output: groups upstream partitions into coalesced output partitions.
@@ -507,18 +522,65 @@ message ExecutePartition {
   datafusion.PhysicalHashRepartition output_partitioning = 6;
 }
 
+// How a shuffle writer laid its output out on disk. Selects the path shape,
+// and with it how a client parses the index beside the data. It says nothing
+// about the framing of the bytes inside the data file.
+enum ShuffleLayout {
+  // {stage_id}/{partition_id}/data-{file_id}.arrow — one file per output
+  // partition, written by the passthrough and range writers.
+  SHUFFLE_LAYOUT_PASSTHROUGH = 0;
+  // {stage_id}/{file_id}/data.arrow — one file per task holding every
+  // partition, written by the sort-based writer.
+  SHUFFLE_LAYOUT_SORT = 1;
+}
+
+// Which of the files making up a shuffle output is being asked for. Which
+// index sits beside the data, and how to read it, follows from the layout, so
+// it is not spelled out here.
+enum ShuffleFileKind {
+  SHUFFLE_FILE_KIND_DATA = 0;
+  SHUFFLE_FILE_KIND_INDEX = 1;
+}
+
+// A half-open byte range of a file: `[offset, offset + length)`.
+message ByteRange {
+  uint64 offset = 1;
+  uint64 length = 2;
+}
+
 message FetchPartition {
   string job_id = 1;
   uint32 stage_id = 2;
   uint32 partition_id = 3;
   string host = 5;
   uint32 port = 6;
   optional uint64 file_id = 7;
-  bool is_sort_shuffle = 8;
-
+  // How the producing writer laid its output out, which is what turns the
+  // identifiers above into a path.
+  ShuffleLayout layout = 9;
+  // Which file beside that partition to serve.
+  ShuffleFileKind file_kind = 10;
+  // Ranges to return, concatenated in request order, as absolute offsets into
+  // the file.
+  //
+  // Empty asks for whatever the identifiers above address. Under the sort
+  // layout that is the named partition's slice, which the executor resolves
+  // through its own index — one round trip, as it has always been. Under the
+  // passthrough layout it is the whole file.
+  //
+  // With ranges given the executor serves bytes and nothing else: no index
+  // read, no comparison, no notion of what a range means. That is what a
+  // consumer needs against object storage, where there is no executor to ask,
+  // and it is why the range shuffle's index carries byte offsets at all.
+  repeated ByteRange byte_ranges = 11;
+
   // reserved after removing deprecated `path`.
   reserved 4;
   reserved "path";
+  // reserved after `is_sort_shuffle` became `layout`: the bool selected a path
+  // shape, which reads as data-vs-index once a file_kind exists beside it.
+  reserved 8;
+  reserved "is_sort_shuffle";
 }
 
 message PartitionLocation {

ballista/core/src/client.rs

@@ -20,7 +20,9 @@
 use crate::error::{BallistaError, Result as BResult};
 use crate::extension::BallistaConfigGrpcEndpoint;
 use crate::serde::protobuf;
-use crate::serde::scheduler::{Action, PartitionId};
+use crate::serde::scheduler::{
+    Action, ByteRange, PartitionId, ShuffleFileKind, ShuffleLayout,
+};
 use crate::utils::create_grpc_client_endpoint;
 use arrow_flight;
 use arrow_flight::Ticket;
@@ -157,7 +159,7 @@ impl BallistaClient {
         executor_id: &str,
         partition_id: &PartitionId,
         file_id: Option<u64>,
-        is_sort_shuffle: bool,
+        layout: ShuffleLayout,
         flight_transport: bool,
     ) -> BResult<SendableRecordBatchStream> {
         let host = self.host.to_owned();
@@ -166,14 +168,50 @@ impl BallistaClient {
             executor_id,
             partition_id,
             file_id,
-            is_sort_shuffle,
+            layout,
             &host,
             port,
             flight_transport,
         )
         .await
     }
 
+    /// Fetch byte ranges of a shuffle file, as absolute offsets into it.
+    ///
+    /// The executor returns exactly those bytes concatenated in request order,
+    /// having resolved nothing: a caller that knows which bytes it wants — from
+    /// an index it fetched itself — gets them without the executor reading an
+    /// index or comparing a value. Block transport only, since decoding is the
+    /// caller's business.
+    #[allow(clippy::too_many_arguments)]
+    pub async fn fetch_byte_ranges(
+        &mut self,
+        executor_id: &str,
+        partition_id: &PartitionId,
+        file_id: Option<u64>,
+        layout: ShuffleLayout,
+        file_kind: ShuffleFileKind,
+        byte_ranges: Vec<ByteRange>,
+        header: Option<Vec<u8>>,
+    ) -> BResult<SendableRecordBatchStream> {
+        let host = self.host.to_owned();
+        let port = self.port;
+        let action = Action::FetchPartition {
+            job_id: partition_id.job_id.clone(),
+            stage_id: partition_id.stage_id,
+            partition_id: partition_id.partition_id,
+            host,
+            port,
+            file_id,
+            layout,
+            file_kind,
+            byte_ranges,
+        };
+        self.execute_do_action(&action, header)
+            .await
+            .map_err(|error| Self::as_fetch_failed(executor_id, partition_id, error))
+    }
+
     /// Retrieves a partition from an executor.
     ///
     /// Depending on the value of the `flight_transport` parameter, this method will utilize either
@@ -187,48 +225,62 @@ impl BallistaClient {
         executor_id: &str,
         partition_id: &PartitionId,
         file_id: Option<u64>,
-        is_sort_shuffle: bool,
+        layout: ShuffleLayout,
         host: &str,
         port: u16,
         flight_transport: bool,
     ) -> BResult<SendableRecordBatchStream> {
+        // No ranges: asks for whatever these identifiers address, which under
+        // the sort layout is the named partition's slice and under passthrough
+        // is the whole file.
         let action = Action::FetchPartition {
             job_id: partition_id.job_id.clone(),
             stage_id: partition_id.stage_id,
             partition_id: partition_id.partition_id,
             host: host.to_owned(),
             port,
             file_id,
-            is_sort_shuffle,
+            layout,
+            file_kind: ShuffleFileKind::Data,
+            byte_ranges: vec![],
         };
 
         let result = if flight_transport {
             self.execute_do_get(&action).await
         } else {
-            self.execute_do_action(&action).await
+            self.execute_do_action(&action, None).await
         };
 
-        result
-            .map_err(|error| match error {
-                // map grpc connection error to partition fetch error.
-                BallistaError::GrpcActionError(msg) => {
-                    log::warn!(
-                        "grpc client failed to fetch partition: {partition_id:?} , message: {msg:?}"
-                    );
-                    BallistaError::FetchFailed(
-                        executor_id.to_owned(),
-                        partition_id.stage_id,
-                        partition_id.partition_id,
-                        msg,
-                    )
-                }
-                error => {
-                    log::warn!(
-                        "grpc client failed to fetch partition: {partition_id:?} , error: {error:?}"
-                    );
-                    error
-                }
-            })
+        result.map_err(|error| Self::as_fetch_failed(executor_id, partition_id, error))
+    }
+
+    /// Report a transport failure as a partition fetch failure, which is what
+    /// lets the scheduler retry the task rather than fail the query.
+    fn as_fetch_failed(
+        executor_id: &str,
+        partition_id: &PartitionId,
+        error: BallistaError,
+    ) -> BallistaError {
+        match error {
+            // map grpc connection error to partition fetch error.
+            BallistaError::GrpcActionError(msg) => {
+                log::warn!(
+                    "grpc client failed to fetch partition: {partition_id:?} , message: {msg:?}"
+                );
+                BallistaError::FetchFailed(
+                    executor_id.to_owned(),
+                    partition_id.stage_id,
+                    partition_id.partition_id,
+                    msg,
+                )
+            }
+            error => {
+                log::warn!(
+                    "grpc client failed to fetch partition: {partition_id:?} , error: {error:?}"
+                );
+                error
+            }
+        }
     }
 
     #[allow(rustdoc::private_intra_doc_links)]
@@ -321,6 +373,7 @@ impl BallistaClient {
     pub async fn execute_do_action(
         &mut self,
         action: &Action,
+        header: Option<Vec<u8>>,
     ) -> BResult<SendableRecordBatchStream> {
         let serialized_action: protobuf::Action = action.to_owned().try_into()?;
 
@@ -375,6 +428,19 @@ impl BallistaClient {
                 })
             });
 
+            // A caller fetching byte ranges gets batch messages with no schema
+            // ahead of them, because the schema is not in the bytes it asked
+            // for. It supplies the header it already knows.
+            let stream =
+                match header.clone() {
+                    Some(header) => futures::stream::once(async move {
+                        Ok(prost::bytes::Bytes::from(header))
+                    })
+                    .chain(stream)
+                    .boxed(),
+                    None => stream.boxed(),
+                };
+
             return Ok(Box::pin(BlockDataStream::try_new(stream).await?));
         }
         unreachable!("Did not receive schema batch from flight server");

ballista/core/src/execution_plans/distributed_query.rs

@@ -27,6 +27,7 @@ use crate::serde::protobuf::{
     scheduler_grpc_client::SchedulerGrpcClient,
 };
 use crate::serde::protobuf::{ExecutorMetadata, SuccessfulJob};
+use crate::serde::scheduler::ShuffleLayout;
 use crate::utils::{GrpcClientConfig, create_grpc_client_endpoint};
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::arrow::record_batch::RecordBatch;
@@ -877,7 +878,11 @@ async fn fetch_partition(
             &metadata.id,
             &partition_id.into(),
             location.file_id,
-            location.is_sort_shuffle,
+            if location.is_sort_shuffle {
+                ShuffleLayout::Sort
+            } else {
+                ShuffleLayout::Passthrough
+            },
             host,
             port,
             flight_transport,

ballista/core/src/execution_plans/mod.rs

@@ -29,6 +29,7 @@ pub mod plan_algebra;
 mod prefix_merge;
 mod range_filter;
 mod range_repartition_common;
+pub mod range_shuffle;
 mod range_shuffle_reader;
 mod runtime_stats;
 mod shuffle_reader;
@@ -52,6 +53,7 @@ pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicate
 pub use plan_algebra::{preserves_distribution, preserves_partitioning};
 pub use prefix_merge::{FinalizedPartitionState, PrefixMergeExec, ScalarOp, WindowApply};
 pub use range_filter::{InputOrder, RangeBound, RangeFilterExec, WidenedBound};
+pub use range_shuffle::RangeShuffleWriterExec;
 pub use range_shuffle_reader::RangeShuffleReaderExec;
 pub use runtime_stats::{
     MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats,

ballista/core/src/execution_plans/plan_algebra.rs

@@ -35,8 +35,8 @@ use datafusion::physical_plan::sorts::sort::SortExec;
 use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
 
 use crate::execution_plans::{
-    BufferExec, PrefixMergeExec, RangeFilterExec, RuntimeStatsExec, ShuffleWriterExec,
-    SortShuffleWriterExec,
+    BufferExec, PrefixMergeExec, RangeFilterExec, RangeShuffleWriterExec,
+    RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec,
 };
 
 /// Whitelisted ops preserve the routing key's row set, values, and
@@ -52,6 +52,7 @@ pub fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool {
             .is_some_and(|sort| sort.preserve_partitioning())
         // Stage-boundary writers: batches to disk unchanged.
         || plan.downcast_ref::<ShuffleWriterExec>().is_some()
+        || plan.downcast_ref::<RangeShuffleWriterExec>().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.