From 297a8d6f8e8858c6b6d347d1d6cb0d0e5d41b81b Mon Sep 17 00:00:00 2001 From: luofucong Date: Fri, 12 Apr 2024 13:49:26 +0800 Subject: [PATCH 1/3] Get mutable `SessionConfig` in `SessionState` so that we can add extensions to it. --- datafusion/core/src/execution/context/mod.rs | 5 +++++ datafusion/execution/src/config.rs | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 5cf8969aa46d..fc2cdbb7518d 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] { &self.physical_optimizers.rules diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs index 0a7a87c7d81a..c9a1fa686a63 100644 --- a/datafusion/execution/src/config.rs +++ b/datafusion/execution/src/config.rs @@ -500,7 +500,7 @@ impl SessionConfig { /// ``` /// /// [^1]: Compare that to [`ConfigOptions`] which only supports [`ScalarValue`] payloads. - pub fn with_extension(mut self, ext: Arc) -> Self + pub fn with_extension(&mut self, ext: Arc) -> &mut Self where T: Send + Sync + 'static, { From 1b96dc0ebedfebd9150842c1dc5f9c59cc4c46e0 Mon Sep 17 00:00:00 2001 From: luofucong Date: Fri, 12 Apr 2024 14:32:29 +0800 Subject: [PATCH 2/3] fix doc tests --- datafusion/execution/src/config.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs index c9a1fa686a63..7db043ab83e5 100644 --- a/datafusion/execution/src/config.rs +++ b/datafusion/execution/src/config.rs @@ -483,11 +483,11 @@ impl SessionConfig { /// let ext1b = Arc::new(Ext1(11)); /// let ext2 = Arc::new(Ext2(2)); /// - /// let cfg = SessionConfig::default() - /// // will only remember the last Ext1 - /// .with_extension(Arc::clone(&ext1a)) - /// .with_extension(Arc::clone(&ext1b)) - /// .with_extension(Arc::clone(&ext2)); + /// let mut cfg = SessionConfig::default(); + /// // will only remember the last Ext1 + /// cfg.with_extension(Arc::clone(&ext1a)) + /// .with_extension(Arc::clone(&ext1b)) + /// .with_extension(Arc::clone(&ext2)); /// /// let ext1_received = cfg.get_extension::().unwrap(); /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a)); From 916606b5d94e4aa6dd739f9044ea2df437ed9d2c Mon Sep 17 00:00:00 2001 From: luofucong Date: Sun, 14 Apr 2024 22:39:30 +0800 Subject: [PATCH 3/3] fix: resolve PR comments --- datafusion/execution/src/config.rs | 55 ++++++++++++++++++++++++++---- 1 file changed, 48 insertions(+), 7 deletions(-) diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs index 7db043ab83e5..e29030e61457 100644 --- a/datafusion/execution/src/config.rs +++ b/datafusion/execution/src/config.rs @@ -483,11 +483,11 @@ impl SessionConfig { /// let ext1b = Arc::new(Ext1(11)); /// let ext2 = Arc::new(Ext2(2)); /// - /// let mut cfg = SessionConfig::default(); - /// // will only remember the last Ext1 - /// cfg.with_extension(Arc::clone(&ext1a)) - /// .with_extension(Arc::clone(&ext1b)) - /// .with_extension(Arc::clone(&ext2)); + /// let cfg = SessionConfig::default() + /// // will only remember the last Ext1 + /// .with_extension(Arc::clone(&ext1a)) + /// .with_extension(Arc::clone(&ext1b)) + /// .with_extension(Arc::clone(&ext2)); /// /// let ext1_received = cfg.get_extension::().unwrap(); /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a)); @@ -500,14 +500,55 @@ impl SessionConfig { /// ``` /// /// [^1]: Compare that to [`ConfigOptions`] which only supports [`ScalarValue`] payloads. - pub fn with_extension(&mut self, ext: Arc) -> &mut Self + pub fn with_extension(mut self, ext: Arc) -> 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::().unwrap(); + /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a)); + /// assert!(Arc::ptr_eq(&ext1_received, &ext1b)); + /// + /// let ext2_received = cfg.get_extension::().unwrap(); + /// assert!(Arc::ptr_eq(&ext2_received, &ext2)); + /// + /// assert!(cfg.get_extension::().is_none()); + /// ``` + pub fn set_extension(&mut self, ext: Arc) where T: Send + Sync + 'static, { let ext = ext as Arc; let id = TypeId::of::(); self.extensions.insert(id, ext); - self } /// Get extension, if any for the specified type `T` exists.