feat: Store shuffles in object store (S3, Azure)

Author: lukekim

Cargo.lock

@@ -40,7 +40,8 @@ datafusion-cli = "51.0.0"
 datafusion-proto = "51.0.0"
 datafusion-proto-common = "51.0.0"
 datafusion-substrait = "51.0.0"
-object_store = "0.12"
+object_store = { version = "0.12", features = ["aws", "azure"] }
+bytes = "1.5"
 prost = "0.14"
 prost-types = "0.14"
 rstest = { version = "0.26" }

Cargo.toml

@@ -45,6 +45,7 @@ arrow-flight = { workspace = true }
 async-trait = { workspace = true }
 aws-config = { version = "1.6.0", optional = true }
 aws-credential-types = { version = "1.2.0", optional = true }
+bytes = { workspace = true }
 chrono = { version = "0.4", default-features = false }
 clap = { workspace = true, optional = true }
 datafusion = { workspace = true }
@@ -54,7 +55,7 @@ futures = { workspace = true }
 itertools = "0.14"
 log = { workspace = true }
 md-5 = { version = "^0.10.0" }
-object_store = { workspace = true, features = ["aws", "http"], optional = true }
+object_store = { workspace = true, features = ["aws", "azure", "http"], optional = true }
 parking_lot = { workspace = true }
 prost = { workspace = true }
 prost-types = { workspace = true }

ballista/core/Cargo.toml

@@ -43,6 +43,10 @@ pub const BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ: &str =
 /// Configuration key to prefer Flight protocol for remote shuffle reads.
 pub const BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT: &str =
     "ballista.shuffle.remote_read_prefer_flight";
+/// Configuration key for shuffle storage type (local, s3, azure).
+pub const BALLISTA_SHUFFLE_STORAGE_TYPE: &str = "ballista.shuffle.storage_type";
+/// Configuration key for shuffle storage base URL/path.
+pub const BALLISTA_SHUFFLE_STORAGE_URL: &str = "ballista.shuffle.storage_url";
 
 /// Configuration key for gRPC client connection timeout in seconds.
 pub const BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS: &str =
@@ -85,6 +89,14 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "Forces the shuffle reader to use flight reader instead of block reader for remote read. Block reader usually has better performance and resource utilization".to_string(),
                          DataType::Boolean,
                          Some((false).to_string())),
+        ConfigEntry::new(BALLISTA_SHUFFLE_STORAGE_TYPE.to_string(),
+                         "Storage type for shuffle data: 'local' (default), 's3', or 'azure'".to_string(),
+                         DataType::Utf8,
+                         Some("local".to_string())),
+        ConfigEntry::new(BALLISTA_SHUFFLE_STORAGE_URL.to_string(),
+                         "Base URL/path for shuffle storage. For local: file path; For S3: s3://bucket/prefix; For Azure: abfs://container@account.dfs.core.windows.net/prefix".to_string(),
+                         DataType::Utf8,
+                         None),
         ConfigEntry::new(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS.to_string(),
                          "Connection timeout for gRPC client in seconds".to_string(),
                          DataType::UInt64,
@@ -264,6 +276,16 @@ impl BallistaConfig {
         self.get_bool_setting(BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT)
     }
 
+    /// Returns the shuffle storage type (local, s3, azure).
+    pub fn shuffle_storage_type(&self) -> String {
+        self.get_string_setting(BALLISTA_SHUFFLE_STORAGE_TYPE)
+    }
+
+    /// Returns the shuffle storage base URL/path if configured.
+    pub fn shuffle_storage_url(&self) -> Option<String> {
+        self.settings.get(BALLISTA_SHUFFLE_STORAGE_URL).cloned()
+    }
+
     fn get_usize_setting(&self, key: &str) -> usize {
         if let Some(v) = self.settings.get(key) {
             // infallible because we validate all configs in the constructor

ballista/core/src/config.rs

@@ -22,12 +22,21 @@ use std::any::Any;
 use std::collections::HashMap;
 use std::fmt::Debug;
 use std::fs::File;
-use std::io::BufReader;
+use std::io::{BufReader, Cursor};
 use std::pin::Pin;
 use std::result;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+#[cfg(feature = "build-binary")]
+use object_store::aws::AmazonS3Builder;
+#[cfg(feature = "build-binary")]
+use object_store::azure::MicrosoftAzureBuilder;
+#[cfg(feature = "build-binary")]
+use object_store::ObjectStore;
+#[cfg(feature = "build-binary")]
+use url::Url;
+
 use crate::client::BallistaClient;
 use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt};
 use crate::serde::scheduler::{PartitionLocation, PartitionStats};
@@ -371,23 +380,34 @@ impl Stream for AbortableReceiverStream {
             .map_err(|e| ArrowError::ExternalError(Box::new(e)))
     }
 }
-/// Splits the provided partition locations into local and remote partitions.
+/// Splits the provided partition locations into local, object store, and remote partitions.
 /// Local partitions are read directly from local Arrow IPC files,
+/// object store partitions are read via the object store client,
 /// 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(
     partition_locations: Vec<PartitionLocation>,
     force_remote_read: bool,
-) -> (Vec<PartitionLocation>, Vec<PartitionLocation>) {
+) -> (Vec<PartitionLocation>, Vec<PartitionLocation>, Vec<PartitionLocation>) {
     if !force_remote_read {
-        partition_locations
+        let (local, non_local): (Vec<_>, Vec<_>) = partition_locations
+            .into_iter()
+            .partition(check_is_local_location);
+        let (object_store, remote): (Vec<_>, Vec<_>) = non_local
             .into_iter()
-            .partition(check_is_local_location)
+            .partition(check_is_object_store_location);
+        (local, object_store, remote)
     } else {
-        (vec![], partition_locations)
+        (vec![], vec![], partition_locations)
     }
 }
 
+/// 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();
+    path.starts_with("s3://") || path.starts_with("abfs://") || path.starts_with("az://")
+}
+
 fn send_fetch_partitions(
     partition_locations: Vec<PartitionLocation>,
     max_request_num: usize,
@@ -401,12 +421,13 @@ fn send_fetch_partitions(
     let semaphore = Arc::new(Semaphore::new(max_request_num));
     let mut spawned_tasks: Vec<SpawnedTask<()>> = vec![];
 
-    let (local_locations, remote_locations): (Vec<_>, Vec<_>) =
+    let (local_locations, object_store_locations, remote_locations): (Vec<_>, Vec<_>, Vec<_>) =
         local_remote_read_split(partition_locations, force_remote_read);
 
     debug!(
-        "local shuffle file counts:{}, remote shuffle file count:{}.",
+        "local shuffle file counts:{}, object store shuffle file count:{}, remote shuffle file count:{}.",
         local_locations.len(),
+        object_store_locations.len(),
         remote_locations.len()
     );
 
@@ -430,6 +451,31 @@ fn send_fetch_partitions(
         }
     }));
 
+    // Handle object store partitions with concurrency control
+    for p in object_store_locations.into_iter() {
+        let semaphore = semaphore.clone();
+        let response_sender = response_sender.clone();
+        spawned_tasks.push(SpawnedTask::spawn(async move {
+            // Block if exceeds max request number.
+            let permit = semaphore.acquire_owned().await.unwrap();
+            let r = PartitionReaderEnum::ObjectStoreRemote
+                .fetch_partition(
+                    &p,
+                    max_message_size,
+                    false, // flight_transport not used for object store
+                    None,  // customize_endpoint not used for object store
+                    false, // use_tls not used for object store
+                )
+                .await;
+            // Block if the channel buffer is full.
+            if let Err(e) = response_sender.send(r).await {
+                error!("Fail to send response event to the channel due to {e}");
+            }
+            // Increase semaphore by dropping existing permits.
+            drop(permit);
+        }));
+    }
+
     for p in remote_locations.into_iter() {
         let semaphore = semaphore.clone();
         let response_sender = response_sender.clone();
@@ -590,14 +636,143 @@ fn fetch_partition_local_inner(
     Ok(reader)
 }
 
+#[cfg(feature = "build-binary")]
+async fn fetch_partition_object_store(
+    location: &PartitionLocation,
+) -> result::Result<SendableRecordBatchStream, BallistaError> {
+    let path = &location.path;
+    let metadata = &location.executor_meta;
+    let partition_id = &location.partition_id;
+
+    debug!("Fetching shuffle partition from object store: {}", path);
+
+    let batches = fetch_partition_object_store_inner(path).await.map_err(|e| {
+        // return BallistaError::FetchFailed may let scheduler retry this task.
+        BallistaError::FetchFailed(
+            metadata.id.clone(),
+            partition_id.stage_id,
+            partition_id.partition_id,
+            e.to_string(),
+        )
+    })?;
+
+    if batches.is_empty() {
+        return Err(BallistaError::General(format!(
+            "No batches found in shuffle partition at {}",
+            path
+        )));
+    }
+
+    let schema = batches[0].schema();
+    let stream = futures::stream::iter(batches.into_iter().map(Ok));
+    Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
+}
+
+#[cfg(not(feature = "build-binary"))]
 async fn fetch_partition_object_store(
     _location: &PartitionLocation,
 ) -> result::Result<SendableRecordBatchStream, BallistaError> {
     Err(BallistaError::NotImplemented(
-        "Should not use ObjectStorePartitionReader".to_string(),
+        "Object store support requires 'build-binary' feature".to_string(),
     ))
 }
 
+#[cfg(feature = "build-binary")]
+async fn fetch_partition_object_store_inner(
+    path: &str,
+) -> result::Result<Vec<RecordBatch>, BallistaError> {
+    use object_store::path::Path as ObjectPath;
+
+    let url = Url::parse(path).map_err(|e| {
+        BallistaError::General(format!("Failed to parse object store URL '{}': {:?}", path, e))
+    })?;
+
+    let scheme = url.scheme();
+    let store: Arc<dyn ObjectStore> = match scheme {
+        "s3" => {
+            let bucket = url.host_str().ok_or_else(|| {
+                BallistaError::General(format!("No bucket in S3 URL: {}", path))
+            })?;
+            let builder = AmazonS3Builder::from_env().with_bucket_name(bucket);
+            Arc::new(builder.build().map_err(|e| {
+                BallistaError::General(format!("Failed to create S3 client: {:?}", e))
+            })?)
+        }
+        "abfs" | "az" => {
+            // Parse Azure URL: abfs://container@account.dfs.core.windows.net/path
+            let host = url.host_str().ok_or_else(|| {
+                BallistaError::General(format!("No host in Azure URL: {}", path))
+            })?;
+            
+            // Extract container from username portion
+            let container = url.username();
+            if container.is_empty() {
+                return Err(BallistaError::General(format!(
+                    "No container in Azure URL. Expected format: abfs://container@account.dfs.core.windows.net/path. Got: {}",
+                    path
+                )));
+            }
+
+            // Extract account from host (account.dfs.core.windows.net)
+            let account = host.split('.').next().ok_or_else(|| {
+                BallistaError::General(format!("No account in Azure URL: {}", path))
+            })?;
+
+            let builder = MicrosoftAzureBuilder::from_env()
+                .with_account(account)
+                .with_container_name(container);
+            Arc::new(builder.build().map_err(|e| {
+                BallistaError::General(format!("Failed to create Azure client: {:?}", e))
+            })?)
+        }
+        _ => {
+            return Err(BallistaError::General(format!(
+                "Unsupported object store scheme: {}. Supported: s3, abfs, az",
+                scheme
+            )));
+        }
+    };
+
+    // Extract the object path from the URL
+    let object_path = ObjectPath::from(url.path().trim_start_matches('/'));
+
+    debug!("Reading object from path: {:?}", object_path);
+
+    let get_result = store.get(&object_path).await.map_err(|e| {
+        BallistaError::General(format!(
+            "Failed to read object from {}: {:?}",
+            path, e
+        ))
+    })?;
+
+    let bytes = get_result.bytes().await.map_err(|e| {
+        BallistaError::General(format!(
+            "Failed to read bytes from {}: {:?}",
+            path, e
+        ))
+    })?;
+
+    let cursor = Cursor::new(bytes.to_vec());
+    let stream_reader = StreamReader::try_new(cursor, None).map_err(|e| {
+        BallistaError::General(format!(
+            "Failed to create Arrow stream reader for {}: {:?}",
+            path, e
+        ))
+    })?;
+
+    let mut batches = Vec::new();
+    for batch_result in stream_reader {
+        batches.push(batch_result.map_err(|e| {
+            BallistaError::General(format!(
+                "Failed to read batch from {}: {:?}",
+                path, e
+            ))
+        })?);
+    }
+
+    Ok(batches)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -955,14 +1130,16 @@ mod tests {
         let partition_locations =
             get_test_partition_locations(1, file_path.to_str().unwrap().to_string());
 
-        let (local, remote) = local_remote_read_split(partition_locations.clone(), false);
+        let (local, object_store, remote) = local_remote_read_split(partition_locations.clone(), false);
 
         assert!(!local.is_empty());
+        assert!(object_store.is_empty());
         assert!(remote.is_empty());
 
-        let (local, remote) = local_remote_read_split(partition_locations, true);
+        let (local, object_store, remote) = local_remote_read_split(partition_locations, true);
 
         assert!(local.is_empty());
+        assert!(object_store.is_empty());
         assert!(!remote.is_empty());
     }