fix: register object store with datafusion (#107)

--------- Co-authored-by: Shiyan Xu <2701446+xushiyan@users.noreply.github.com>
apache · Aug 12, 2024 · 3359e10 · 3359e10
1 parent b3ffa32
commit 3359e10
Show file tree

Hide file tree

Showing 7 changed files with 40 additions and 3 deletions.
diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml
@@ -60,5 +60,19 @@ dashmap = { workspace = true }
 futures = { workspace = true }
 tokio = { workspace = true }
 
+# datafusion
+datafusion = { workspace = true, optional = true }
+datafusion-expr = { workspace = true, optional = true }
+datafusion-common = { workspace = true, optional = true }
+datafusion-physical-expr = { workspace = true, optional = true }
+
 [dev-dependencies]
 hudi-tests = { path = "../tests" }
+
+[features]
+datafusion = [
+    "dep:datafusion",
+    "datafusion-expr",
+    "datafusion-common",
+    "datafusion-physical-expr",
+]
diff --git a/crates/core/src/storage/mod.rs b/crates/core/src/storage/mod.rs
@@ -58,6 +58,14 @@ impl Storage {
         }
     }
 
+    #[cfg(feature = "datafusion")]
+    pub fn register_object_store(
+        &self,
+        runtime_env: Arc<datafusion::execution::runtime_env::RuntimeEnv>,
+    ) {
+        runtime_env.register_object_store(self.base_url.as_ref(), self.object_store.clone());
+    }
+
     #[cfg(test)]
     async fn get_file_info(&self, relative_path: &str) -> Result<FileInfo> {
         let obj_url = join_url_segments(&self.base_url, &[relative_path])?;

diff --git a/crates/core/src/table/fs_view.rs b/crates/core/src/table/fs_view.rs
@@ -34,7 +34,7 @@ use crate::storage::{get_leaf_dirs, Storage};
 #[allow(dead_code)]
 pub struct FileSystemView {
     configs: Arc<HudiConfigs>,
-    storage: Arc<Storage>,
+    pub(crate) storage: Arc<Storage>,
     partition_to_file_groups: Arc<DashMap<String, Vec<FileGroup>>>,
 }
 

diff --git a/crates/core/src/table/mod.rs b/crates/core/src/table/mod.rs
@@ -93,6 +93,19 @@ impl Table {
         })
     }
 
+    #[cfg(feature = "datafusion")]
+    pub fn register_storage(
+        &self,
+        runtime_env: Arc<datafusion::execution::runtime_env::RuntimeEnv>,
+    ) {
+        self.timeline
+            .storage
+            .register_object_store(runtime_env.clone());
+        self.file_system_view
+            .storage
+            .register_object_store(runtime_env.clone());
+    }
+
     async fn load_configs<I, K, V>(
         base_url: Arc<Url>,
         all_options: I,

diff --git a/crates/core/src/table/timeline.rs b/crates/core/src/table/timeline.rs
@@ -92,7 +92,7 @@ impl Instant {
 #[allow(dead_code)]
 pub struct Timeline {
     configs: Arc<HudiConfigs>,
-    storage: Arc<Storage>,
+    pub(crate) storage: Arc<Storage>,
     pub instants: Vec<Instant>,
 }
 

diff --git a/crates/datafusion/Cargo.toml b/crates/datafusion/Cargo.toml
@@ -28,7 +28,7 @@ homepage.workspace = true
 repository.workspace = true
 
 [dependencies]
-hudi-core = { version = "0.2.0", path = "../core" }
+hudi-core = { version = "0.2.0", path = "../core", features = ["datafusion"] }
 # arrow
 arrow-schema = { workspace = true }
 

diff --git a/crates/datafusion/src/lib.rs b/crates/datafusion/src/lib.rs
@@ -97,6 +97,8 @@ impl TableProvider for HudiDataSource {
         filters: &[Expr],
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        self.table.register_storage(state.runtime_env().clone());
+
         let file_slices = self
             .table
             .split_file_slices(self.get_input_partitions())