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/arrow-datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 8604d0a801 Minor: get mutable ref to `SessionConfig` in `SessionState`
(#10050)
8604d0a801 is described below
commit 8604d0a80155364ff60d3eac5e44f03f6e7e5a8e
Author: LFC <[email protected]>
AuthorDate: Mon Apr 15 18:55:53 2024 +0800
Minor: get mutable ref to `SessionConfig` in `SessionState` (#10050)
* Get mutable `SessionConfig` in `SessionState` so that we can add
extensions to it.
* fix doc tests
* fix: resolve PR comments
---
datafusion/core/src/execution/context/mod.rs | 5 ++++
datafusion/execution/src/config.rs | 43 +++++++++++++++++++++++++++-
2 files changed, 47 insertions(+), 1 deletion(-)
diff --git a/datafusion/core/src/execution/context/mod.rs
b/datafusion/core/src/execution/context/mod.rs
index 5cf8969aa4..fc2cdbb751 100644
--- a/datafusion/core/src/execution/context/mod.rs
+++ b/datafusion/core/src/execution/context/mod.rs
@@ -1952,6 +1952,11 @@ impl SessionState {
&self.config
}
+ /// Return the mutable [`SessionConfig`].
+ pub fn config_mut(&mut self) -> &mut SessionConfig {
+ &mut self.config
+ }
+
/// Return the physical optimizers
pub fn physical_optimizers(&self) -> &[Arc<dyn PhysicalOptimizerRule +
Send + Sync>] {
&self.physical_optimizers.rules
diff --git a/datafusion/execution/src/config.rs
b/datafusion/execution/src/config.rs
index 0a7a87c7d8..e29030e614 100644
--- a/datafusion/execution/src/config.rs
+++ b/datafusion/execution/src/config.rs
@@ -501,13 +501,54 @@ impl SessionConfig {
///
/// [^1]: Compare that to [`ConfigOptions`] which only supports
[`ScalarValue`] payloads.
pub fn with_extension<T>(mut self, ext: Arc<T>) -> Self
+ where
+ T: Send + Sync + 'static,
+ {
+ self.set_extension(ext);
+ self
+ }
+
+ /// Set extension. Pretty much the same as
[`with_extension`](Self::with_extension), but take
+ /// mutable reference instead of owning it. Useful if you want to add
another extension after
+ /// the [`SessionConfig`] is created.
+ ///
+ /// # Example
+ /// ```
+ /// use std::sync::Arc;
+ /// use datafusion_execution::config::SessionConfig;
+ ///
+ /// // application-specific extension types
+ /// struct Ext1(u8);
+ /// struct Ext2(u8);
+ /// struct Ext3(u8);
+ ///
+ /// let ext1a = Arc::new(Ext1(10));
+ /// let ext1b = Arc::new(Ext1(11));
+ /// let ext2 = Arc::new(Ext2(2));
+ ///
+ /// let mut cfg = SessionConfig::default();
+ ///
+ /// // will only remember the last Ext1
+ /// cfg.set_extension(Arc::clone(&ext1a));
+ /// cfg.set_extension(Arc::clone(&ext1b));
+ /// cfg.set_extension(Arc::clone(&ext2));
+ ///
+ /// let ext1_received = cfg.get_extension::<Ext1>().unwrap();
+ /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a));
+ /// assert!(Arc::ptr_eq(&ext1_received, &ext1b));
+ ///
+ /// let ext2_received = cfg.get_extension::<Ext2>().unwrap();
+ /// assert!(Arc::ptr_eq(&ext2_received, &ext2));
+ ///
+ /// assert!(cfg.get_extension::<Ext3>().is_none());
+ /// ```
+ pub fn set_extension<T>(&mut self, ext: Arc<T>)
where
T: Send + Sync + 'static,
{
let ext = ext as Arc<dyn Any + Send + Sync + 'static>;
let id = TypeId::of::<T>();
self.extensions.insert(id, ext);
- self
}
/// Get extension, if any for the specified type `T` exists.
