[arrow-ballista] branch main updated: Add support for GCS data sources (#805)

[dh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dheres pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-ballista.git


The following commit(s) were added to refs/heads/main by this push:
     new 8dbf9d22 Add support for GCS data sources (#805)
8dbf9d22 is described below

commit 8dbf9d22425a53bc0b0651fdae8a5f29f61e4da7
Author: Xin Hao <ha...@gmail.com>
AuthorDate: Fri Jun 9 18:42:57 2023 +0800

    Add support for GCS data sources (#805)
    
    * Add support for GCS data sources
    
    closes #804
    
    * Add prefix gcs://
---
 ballista/core/src/utils.rs | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs
index f3859386..2d19889c 100644
--- a/ballista/core/src/utils.rs
+++ b/ballista/core/src/utils.rs
@@ -56,6 +56,8 @@ use log::error;
 use object_store::aws::AmazonS3Builder;
 #[cfg(feature = "azure")]
 use object_store::azure::MicrosoftAzureBuilder;
+#[cfg(feature = "gcs")]
+use object_store::gcp::GoogleCloudStorageBuilder;
 use object_store::ObjectStore;
 use std::io::{BufWriter, Write};
 use std::marker::PhantomData;
@@ -148,6 +150,22 @@ impl BallistaObjectStoreRegistry {
             }
         }
 
+        #[cfg(feature = "gcs")]
+        {
+            if url.to_string().starts_with("gs://")
+                || url.to_string().starts_with("gcs://")
+            {
+                if let Some(bucket_name) = url.host_str() {
+                    let store = Arc::new(
+                        GoogleCloudStorageBuilder::from_env()
+                            .with_bucket_name(bucket_name)
+                            .build()?,
+                    );
+                    return Ok(store);
+                }
+            }
+        }
+
         Err(DataFusionError::Execution(format!(
             "No object store available for: {url}"
         )))