Address PR comments

Author: lukekim

ballista/client/src/builder.rs

@@ -39,7 +39,7 @@
 //!     .build()?;
 //!
 //! let ctx = BallistaBuilder::new()
-//!     .with_object_store("s3://my-bucket", Arc::new(s3_store))
+//!     .add_object_store("s3://my-bucket", Arc::new(s3_store))
 //!     .standalone()
 //!     .await?;
 //! # Ok(())
@@ -79,7 +79,7 @@ use url::Url;
 ///     .build()?;
 ///
 /// let ctx = BallistaBuilder::new()
-///     .with_object_store("s3://my-bucket", Arc::new(s3_store))
+///     .add_object_store("s3://my-bucket", Arc::new(s3_store))
 ///     .remote("df://localhost:50050")
 ///     .await?;
 /// # Ok(())
@@ -159,32 +159,28 @@ impl BallistaBuilder {
     ///     .build()?;
     ///
     /// let ctx = BallistaBuilder::new()
-    ///     .with_object_store("s3://my-bucket", Arc::new(s3_store))
+    ///     .add_object_store("s3://my-bucket", Arc::new(s3_store))
     ///     .standalone()
     ///     .await?;
     /// # Ok(())
     /// # }
     /// ```
-    pub fn with_object_store(mut self, url: &str, store: Arc<dyn ObjectStore>) -> Self {
+    pub fn add_object_store(mut self, url: &str, store: Arc<dyn ObjectStore>) -> Self {
         // Parse the URL, or store it for later error handling during build
         if let Ok(parsed_url) = Url::parse(url) {
             self.object_stores.push((parsed_url, store));
         } else {
             // We'll handle invalid URLs during build
-            log::warn!("Invalid object store URL: {}", url);
+            log::warn!("Invalid object store URL: {url}");
         }
         self
     }
 
     /// Registers a pre-created object store for a given URL.
     ///
-    /// This is the same as [`with_object_store`](Self::with_object_store) but takes
+    /// This is the same as [`add_object_store`](Self::add_object_store) but takes
     /// a pre-parsed [`Url`] instead of a string.
-    pub fn with_object_store_url(
-        mut self,
-        url: Url,
-        store: Arc<dyn ObjectStore>,
-    ) -> Self {
+    pub fn add_object_store_url(mut self, url: Url, store: Arc<dyn ObjectStore>) -> Self {
         self.object_stores.push((url, store));
         self
     }

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -876,24 +876,19 @@ async fn fetch_partition_object_store_with_runtime(
     let metadata = &location.executor_meta;
     let partition_id = &location.partition_id;
 
-    debug!(
-        "Fetching shuffle partition from object store using runtime_env: {}",
-        path
-    );
+    debug!("Fetching shuffle partition from object store using runtime_env: {path}");
 
     let url = Url::parse(path).map_err(|e| {
         BallistaError::General(format!(
-            "Failed to parse object store URL '{}': {:?}",
-            path, e
+            "Failed to parse object store URL '{path}': {e:?}"
         ))
     })?;
 
     // Get the object store from the RuntimeEnv's registry
     // This uses the credentials configured in the runtime (e.g., SpiceObjectStoreRegistry)
     let object_store_url = ObjectStoreUrl::parse(&url).map_err(|e| {
         BallistaError::General(format!(
-            "Failed to parse object store URL '{}': {:?}",
-            path, e
+            "Failed to parse object store URL '{path}': {e:?}"
         ))
     })?;
 
@@ -902,21 +897,21 @@ async fn fetch_partition_object_store_with_runtime(
             metadata.id.clone(),
             partition_id.stage_id,
             partition_id.partition_id,
-            format!("Failed to get object store for URL '{}': {:?}", path, e),
+            format!("Failed to get object store for URL '{path}': {e:?}"),
         )
     })?;
 
     // Extract the object path from the URL
     let object_path = ObjectPath::from(url.path().trim_start_matches('/'));
 
-    debug!("Reading object from path: {:?}", object_path);
+    debug!("Reading object from path: {object_path:?}");
 
     let get_result = store.get(&object_path).await.map_err(|e| {
         BallistaError::FetchFailed(
             metadata.id.clone(),
             partition_id.stage_id,
             partition_id.partition_id,
-            format!("Failed to read object from {}: {:?}", path, e),
+            format!("Failed to read object from {path}: {e:?}"),
         )
     })?;
 
@@ -925,7 +920,7 @@ async fn fetch_partition_object_store_with_runtime(
             metadata.id.clone(),
             partition_id.stage_id,
             partition_id.partition_id,
-            format!("Failed to read bytes from {}: {:?}", path, e),
+            format!("Failed to read bytes from {path}: {e:?}"),
         )
     })?;
 
@@ -935,7 +930,7 @@ async fn fetch_partition_object_store_with_runtime(
             metadata.id.clone(),
             partition_id.stage_id,
             partition_id.partition_id,
-            format!("Failed to create Arrow stream reader for {}: {:?}", path, e),
+            format!("Failed to create Arrow stream reader for {path}: {e:?}"),
         )
     })?;
 
@@ -946,15 +941,14 @@ async fn fetch_partition_object_store_with_runtime(
                 metadata.id.clone(),
                 partition_id.stage_id,
                 partition_id.partition_id,
-                format!("Failed to read batch from {}: {:?}", path, e),
+                format!("Failed to read batch from {path}: {e:?}"),
             )
         })?);
     }
 
     if batches.is_empty() {
         return Err(BallistaError::General(format!(
-            "No batches found in shuffle partition at {}",
-            path
+            "No batches found in shuffle partition at {path}"
         )));
     }
 
@@ -981,7 +975,7 @@ async fn fetch_partition_object_store(
     let metadata = &location.executor_meta;
     let partition_id = &location.partition_id;
 
-    debug!("Fetching shuffle partition from object store: {}", path);
+    debug!("Fetching shuffle partition from object store: {path}");
 
     let batches = fetch_partition_object_store_inner(path)
         .await
@@ -997,8 +991,7 @@ async fn fetch_partition_object_store(
 
     if batches.is_empty() {
         return Err(BallistaError::General(format!(
-            "No batches found in shuffle partition at {}",
-            path
+            "No batches found in shuffle partition at {path}"
         )));
     }
 
@@ -1014,82 +1007,78 @@ async fn fetch_partition_object_store_inner(
 
     let url = Url::parse(path).map_err(|e| {
         BallistaError::General(format!(
-            "Failed to parse object store URL '{}': {:?}",
-            path, e
+            "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))
+                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))
+                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))
+                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
+                    "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))
+                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))
+                BallistaError::General(format!("Failed to create Azure client: {e:?}"))
             })?)
         }
         _ => {
             return Err(BallistaError::General(format!(
-                "Unsupported object store scheme: {}. Supported: s3, abfs, az",
-                scheme
+                "Unsupported object store scheme: {scheme}. Supported: s3, abfs, az"
             )));
         }
     };
 
     // Extract the object path from the URL
     let object_path = ObjectPath::from(url.path().trim_start_matches('/'));
 
-    debug!("Reading object from path: {:?}", object_path);
+    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))
+        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))
+        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
+            "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))
+            BallistaError::General(format!("Failed to read batch from {path}: {e:?}"))
         })?);
     }
 

examples/examples/object-store-builder.rs

@@ -96,7 +96,7 @@ async fn main() -> Result<()> {
     // Use BallistaBuilder to create a context with the pre-created object store
     let ctx = BallistaBuilder::new()
         .with_job_name("Object Store Builder Example")
-        .with_object_store(&format!("s3://{S3_BUCKET}"), Arc::new(s3_store))
+        .add_object_store(&format!("s3://{S3_BUCKET}"), Arc::new(s3_store))
         .standalone()
         .await?;