[arrow-ballista] branch master updated: Add support for Azure (#599)

[ag...@apache.org]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 7f285339 Add support for Azure (#599)
7f285339 is described below

commit 7f285339ac7ef4840fc78bdf6da0b14bacac2984
Author: Aidan Kovacic <95...@users.noreply.github.com>
AuthorDate: Sat Jan 14 10:02:31 2023 -0500

    Add support for Azure (#599)
---
 ballista/client/Cargo.toml |  1 +
 ballista/core/Cargo.toml   |  1 +
 ballista/core/src/utils.rs | 14 ++++++++++++++
 3 files changed, 16 insertions(+)

diff --git a/ballista/client/Cargo.toml b/ballista/client/Cargo.toml
index cc870c87..eaf7365d 100644
--- a/ballista/client/Cargo.toml
+++ b/ballista/client/Cargo.toml
@@ -45,4 +45,5 @@ default = []
 hdfs = ["ballista-core/hdfs"]
 hdfs3 = ["ballista-core/hdfs3"]
 s3 = ["ballista-core/s3"]
+azure = ["ballista-core/azure"]
 standalone = ["ballista-executor", "ballista-scheduler"]
diff --git a/ballista/core/Cargo.toml b/ballista/core/Cargo.toml
index 0d8621f0..d39675e0 100644
--- a/ballista/core/Cargo.toml
+++ b/ballista/core/Cargo.toml
@@ -40,6 +40,7 @@ force_hash_collisions = ["datafusion/force_hash_collisions"]
 hdfs = ["datafusion-objectstore-hdfs/hdfs"]
 hdfs3 = ["datafusion-objectstore-hdfs/hdfs3"]
 s3 = ["object_store/aws"]
+azure = ["object_store/azure"]
 simd = ["datafusion/simd"]
 
 [dependencies]
diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs
index 9fb62cd2..c83e94aa 100644
--- a/ballista/core/src/utils.rs
+++ b/ballista/core/src/utils.rs
@@ -52,6 +52,8 @@ use futures::StreamExt;
 use log::error;
 #[cfg(feature = "s3")]
 use object_store::aws::AmazonS3Builder;
+#[cfg(feature = "azure")]
+use object_store::azure::MicrosoftAzureBuilder;
 use object_store::ObjectStore;
 use std::io::{BufWriter, Write};
 use std::marker::PhantomData;
@@ -120,6 +122,18 @@ impl ObjectStoreProvider for FeatureBasedObjectStoreProvider {
             }
         }
 
+        #[cfg(feature = "azure")]
+        {
+            if url.to_string().starts_with("azure://") {
+                if let Some(bucket_name) = url.host_str() {
+                    let store = MicrosoftAzureBuilder::from_env()
+                        .with_container_name(bucket_name)
+                        .build()?;
+                    return Ok(Arc::new(store));
+                }
+            }
+        }
+
         Err(DataFusionError::Execution(format!(
             "No object store available for {}",
             url