fix(shuffle): route s3 partition fetches through object store and stream per-partition IPC as a single Arrow stream

Two independent object-store shuffle correctness bugs that were latent until
PR #42 fixed the writer-side prefix bug and let intermediate shuffles actually
land in S3:

1. BallistaClient::fetch_partition was always going via the gRPC
   Action::FetchPartition, whose executor-side handler does
   `tokio::fs::File::open(path)` and only understands local paths and
   `memory://`. With object-store shuffle enabled, writers reported
   `location.path = s3://...` and the gRPC handler then failed with
   `Failed to open file: No such file or directory`. Dispatch on the path
   scheme (s3 / abfs / az / gs) and route those to the existing object-store
   reader, sharing it via two new pub(crate) helpers
   (path_is_object_store, fetch_object_store_partition_stream).

2. Both the no-repart and hash-repart Arrow IPC writer paths serialised
   each batch via serialize_batch_to_ipc_bytes — every call writes a
   complete IPC stream WITH a StreamWriter::finish() EOS marker — and
   concatenated those streams into one S3 object. StreamReader stops at the
   first EOS, so any partition holding more than one batch came back as
   `Arrow error: Ipc error: Unexpected EOS`. Replace with a long-lived
   StreamingMultipartIpcUploader that writes the header once on
   construction, appends each batch's bytes into the multipart upload, and
   emits the EOS marker exactly once on finish().

Verified end-to-end against a 1 scheduler + 2 executor local Spice cluster
with a multi-segment-prefix `shuffle_location: s3://bucket/<prefix>`:
- `SELECT COUNT(*) FROM t` (single-stage; exercises Bug A): succeeds.
- `SELECT col, COUNT(*) FROM t GROUP BY col ORDER BY ... LIMIT 5`
  (hash-repartition shuffle; exercises Bug B): succeeds.

Author: phillipleblanc

ballista/core/src/client.rs

@@ -164,6 +164,20 @@ impl BallistaClient {
         port: u16,
         flight_transport: bool,
     ) -> BResult<SendableRecordBatchStream> {
+        // When the writer-side stored this partition in object store (s3 / abfs / az / gs),
+        // skip the gRPC FetchPartition path entirely — the executor's handler only knows
+        // local paths and `memory://` and would `tokio::fs::File::open("s3://...")`, failing
+        // with `No such file or directory`. Read straight from object store instead.
+        if crate::execution_plans::shuffle_reader::path_is_object_store(path) {
+            return crate::execution_plans::shuffle_reader::fetch_object_store_partition_stream(
+                path,
+                executor_id,
+                partition_id.stage_id,
+                partition_id.partition_id,
+            )
+            .await;
+        }
+
         let action = Action::FetchPartition {
             job_id: partition_id.job_id.clone(),
             stage_id: partition_id.stage_id,

ballista/core/src/execution_plans/distributed_query.rs

@@ -769,6 +769,9 @@ async fn fetch_partition(
     .await
     .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
 
+    // `BallistaClient::fetch_partition` dispatches object-store URLs (s3 / abfs / az / gs)
+    // to the object-store reader, bypassing the gRPC FetchPartition path that only
+    // understands local files and `memory://`.
     let stream = ballista_client
         .fetch_partition(
             &metadata.id,

ballista/core/src/execution_plans/mod.rs

@@ -21,7 +21,7 @@
 mod distributed_explain_analyze;
 mod distributed_query;
 mod shuffle_manager;
-mod shuffle_reader;
+pub(crate) mod shuffle_reader;
 mod shuffle_writer;
 mod shuffle_writer_trait;
 pub mod sort_shuffle;

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -932,11 +932,26 @@ impl RecordBatchStream for InMemoryShuffleStream {
 async fn fetch_partition_object_store_streaming(
     location: &PartitionLocation,
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {
-    use object_store::path::Path as ObjectPath;
+    fetch_object_store_partition_stream(
+        &location.path,
+        &location.executor_meta.id,
+        location.partition_id.stage_id,
+        location.partition_id.partition_id,
+    )
+    .await
+}
 
-    let path = &location.path;
-    let metadata = &location.executor_meta;
-    let partition_id = &location.partition_id;
+/// Streams a shuffle / result partition directly from object store. Used by both the
+/// intermediate shuffle reader and the driver-side final-stage fetch in
+/// `distributed_query`; the two call sites have differently-shaped `PartitionLocation`
+/// types but only need the path and identifiers, so this helper takes plain args.
+pub(crate) async fn fetch_object_store_partition_stream(
+    path: &str,
+    executor_id: &str,
+    stage_id: usize,
+    partition_id: usize,
+) -> result::Result<SendableRecordBatchStream, BallistaError> {
+    use object_store::path::Path as ObjectPath;
 
     debug!("Fetching shuffle partition from object store: {path}");
 
@@ -948,9 +963,9 @@ async fn fetch_partition_object_store_streaming(
 
     let store = build_shuffle_object_store(&url).map_err(|e| {
         BallistaError::FetchFailed(
-            metadata.id.clone(),
-            partition_id.stage_id,
-            partition_id.partition_id,
+            executor_id.to_owned(),
+            stage_id,
+            partition_id,
             format!("Failed to build object store client for '{path}': {e:?}"),
         )
     })?;
@@ -961,15 +976,15 @@ async fn fetch_partition_object_store_streaming(
 
     let get_result = store.get(&object_path).await.map_err(|e| {
         BallistaError::FetchFailed(
-            metadata.id.clone(),
-            partition_id.stage_id,
-            partition_id.partition_id,
+            executor_id.to_owned(),
+            stage_id,
+            partition_id,
             format!("Failed to read object from {path}: {e:?}"),
         )
     })?;
 
     let byte_stream = get_result.into_stream();
-    let stream = ObjectStoreShuffleStream::try_new(byte_stream, path.clone()).await?;
+    let stream = ObjectStoreShuffleStream::try_new(byte_stream, path.to_owned()).await?;
 
     Ok(Box::pin(stream))
 }
@@ -1178,15 +1193,22 @@ impl RecordBatchStream for ObjectStoreShuffleStream {
     }
 }
 
-/// Check if the location is an object store path (S3 or Azure).
-fn check_is_object_store_location(location: &PartitionLocation) -> bool {
-    let path = location.path.as_str();
+/// Returns true when the given path is an object-store URL (s3 / abfs / az / gs).
+/// Exposed so non-shuffle-reader callers (e.g. the driver-side final-stage fetch in
+/// `distributed_query`) can route around the gRPC `FetchPartition` path, which only
+/// understands local files and `memory://`.
+pub(crate) fn path_is_object_store(path: &str) -> bool {
     path.starts_with("s3://")
         || path.starts_with("abfs://")
         || path.starts_with("az://")
         || path.starts_with("gs://")
 }
 
+/// Check if the location is an object store path (S3 or Azure).
+fn check_is_object_store_location(location: &PartitionLocation) -> bool {
+    path_is_object_store(location.path.as_str())
+}
+
 async fn fetch_partition_object_store(
     location: &PartitionLocation,
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {