This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new a35d007574 Add method SessionStateBuilder::with_object_store (#12578)
a35d007574 is described below
commit a35d0075744a058f81bd9ebed747e2e597434019
Author: OussamaSaoudi <[email protected]>
AuthorDate: Mon Sep 23 03:24:09 2024 -0700
Add method SessionStateBuilder::with_object_store (#12578)
This PR adds the the method `SessionStateBuilder::with_object_store`,
which registers an object store with a specified URL to the RuntimeEnv.
---
datafusion/core/src/execution/session_state.rs | 37 ++++++++++++++++++++++++++
1 file changed, 37 insertions(+)
diff --git a/datafusion/core/src/execution/session_state.rs
b/datafusion/core/src/execution/session_state.rs
index a46e5c3615..e447220a1f 100644
--- a/datafusion/core/src/execution/session_state.rs
+++ b/datafusion/core/src/execution/session_state.rs
@@ -68,6 +68,7 @@ use datafusion_sql::parser::{DFParser, Statement};
use datafusion_sql::planner::{ContextProvider, ParserOptions, PlannerContext,
SqlToRel};
use itertools::Itertools;
use log::{debug, info};
+use object_store::ObjectStore;
use sqlparser::ast::Expr as SQLExpr;
use sqlparser::dialect::dialect_from_str;
use std::any::Any;
@@ -75,6 +76,7 @@ use std::collections::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;
use std::sync::Arc;
+use url::Url;
use uuid::Uuid;
/// `SessionState` contains all the necessary state to plan and execute
queries,
@@ -1229,6 +1231,41 @@ impl SessionStateBuilder {
self
}
+ /// Register an `ObjectStore` to the [`RuntimeEnv`]. See
[`RuntimeEnv::register_object_store`]
+ /// for more details.
+ ///
+ /// Note that this creates a default [`RuntimeEnv`] if there isn't one
passed in already.
+ ///
+ /// ```
+ /// # use datafusion::prelude::*;
+ /// # use datafusion::execution::session_state::SessionStateBuilder;
+ /// # use datafusion_execution::runtime_env::RuntimeEnv;
+ /// # use url::Url;
+ /// # use std::sync::Arc;
+ /// # let http_store = object_store::local::LocalFileSystem::new();
+ /// let url = Url::try_from("file://").unwrap();
+ /// let object_store = object_store::local::LocalFileSystem::new();
+ /// let state = SessionStateBuilder::new()
+ /// .with_config(SessionConfig::new())
+ /// .with_object_store(&url, Arc::new(object_store))
+ /// .with_default_features()
+ /// .build();
+ /// ```
+ pub fn with_object_store(
+ mut self,
+ url: &Url,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Self {
+ if self.runtime_env.is_none() {
+ self.runtime_env = Some(Arc::new(RuntimeEnv::default()));
+ }
+ self.runtime_env
+ .as_ref()
+ .unwrap()
+ .register_object_store(url, object_store);
+ self
+ }
+
/// Builds a [`SessionState`] with the current configuration.
///
/// Note that there is an explicit option for enabling catalog and schema
defaults
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
