@@ -44,8 +44,12 @@ use crate::{
4444} ;
4545pub use datafusion_physical_expr:: execution_props:: ExecutionProps ;
4646use parking_lot:: RwLock ;
47- use std:: string:: String ;
4847use std:: sync:: Arc ;
48+ use std:: {
49+ any:: { Any , TypeId } ,
50+ hash:: { BuildHasherDefault , Hasher } ,
51+ string:: String ,
52+ } ;
4953use std:: {
5054 collections:: { HashMap , HashSet } ,
5155 fmt:: Debug ,
@@ -992,6 +996,37 @@ pub const REPARTITION_WINDOWS: &str = "repartition_windows";
992996/// Session Configuration entry name for 'PARQUET_PRUNING'
993997pub const PARQUET_PRUNING : & str = "parquet_pruning" ;
994998
999+ /// Map that holds opaque objects indexed by their type.
1000+ ///
1001+ /// Data is wrapped into an [`Arc`] to enable [`Clone`] while still being [object safe].
1002+ ///
1003+ /// [object safe]: https://doc.rust-lang.org/reference/items/traits.html#object-safety
1004+ type AnyMap =
1005+ HashMap < TypeId , Arc < dyn Any + Send + Sync + ' static > , BuildHasherDefault < IdHasher > > ;
1006+
1007+ /// Hasher for [`AnyMap`].
1008+ ///
1009+ /// With [`TypeId`}s as keys, there's no need to hash them. They are already hashes themselves, coming from the compiler.
1010+ /// The [`IdHasher`} just holds the [`u64`} of the [`TypeId`}, and then returns it, instead of doing any bit fiddling.
1011+ #[ derive( Default ) ]
1012+ struct IdHasher ( u64 ) ;
1013+
1014+ impl Hasher for IdHasher {
1015+ fn write ( & mut self , _: & [ u8 ] ) {
1016+ unreachable ! ( "TypeId calls write_u64" ) ;
1017+ }
1018+
1019+ #[ inline]
1020+ fn write_u64 ( & mut self , id : u64 ) {
1021+ self . 0 = id;
1022+ }
1023+
1024+ #[ inline]
1025+ fn finish ( & self ) -> u64 {
1026+ self . 0
1027+ }
1028+ }
1029+
9951030/// Configuration options for session context
9961031#[ derive( Clone ) ]
9971032pub struct SessionConfig {
@@ -1019,6 +1054,8 @@ pub struct SessionConfig {
10191054 pub parquet_pruning : bool ,
10201055 /// Configuration options
10211056 pub config_options : ConfigOptions ,
1057+ /// Opaque extensions.
1058+ extensions : AnyMap ,
10221059}
10231060
10241061impl Default for SessionConfig {
@@ -1034,6 +1071,11 @@ impl Default for SessionConfig {
10341071 repartition_windows : true ,
10351072 parquet_pruning : true ,
10361073 config_options : ConfigOptions :: new ( ) ,
1074+ // Assume no extensions by default.
1075+ extensions : HashMap :: with_capacity_and_hasher (
1076+ 0 ,
1077+ BuildHasherDefault :: default ( ) ,
1078+ ) ,
10371079 }
10381080 }
10391081}
@@ -1173,6 +1215,74 @@ impl SessionConfig {
11731215 ) ;
11741216 map
11751217 }
1218+
1219+ /// Add extensions.
1220+ ///
1221+ /// Extensions can be used to attach extra data to the session config -- e.g. tracing information or caches.
1222+ /// Extensions are opaque and the types are unknown to DataFusion itself, which makes them extremely flexible. [^1]
1223+ ///
1224+ /// Extensions are stored within an [`Arc`] so they do NOT require [`Clone`]. The are immutable. If you need to
1225+ /// modify their state over their lifetime -- e.g. for caches -- you need to establish some for of interior mutability.
1226+ ///
1227+ /// Extensions are indexed by their type `T`. If multiple values of the same type are provided, only the last one
1228+ /// will be kept.
1229+ ///
1230+ /// You may use [`get_extension`](Self::get_extension) to retrieve extensions.
1231+ ///
1232+ /// # Example
1233+ /// ```
1234+ /// use std::sync::Arc;
1235+ /// use datafusion::execution::context::SessionConfig;
1236+ ///
1237+ /// // application-specific extension types
1238+ /// struct Ext1(u8);
1239+ /// struct Ext2(u8);
1240+ /// struct Ext3(u8);
1241+ ///
1242+ /// let ext1a = Arc::new(Ext1(10));
1243+ /// let ext1b = Arc::new(Ext1(11));
1244+ /// let ext2 = Arc::new(Ext2(2));
1245+ ///
1246+ /// let cfg = SessionConfig::default()
1247+ /// // will only remember the last Ext1
1248+ /// .with_extension(Arc::clone(&ext1a))
1249+ /// .with_extension(Arc::clone(&ext1b))
1250+ /// .with_extension(Arc::clone(&ext2));
1251+ ///
1252+ /// let ext1_received = cfg.get_extension::<Ext1>().unwrap();
1253+ /// assert!(!Arc::ptr_eq(&ext1_received, &ext1a));
1254+ /// assert!(Arc::ptr_eq(&ext1_received, &ext1b));
1255+ ///
1256+ /// let ext2_received = cfg.get_extension::<Ext2>().unwrap();
1257+ /// assert!(Arc::ptr_eq(&ext2_received, &ext2));
1258+ ///
1259+ /// assert!(cfg.get_extension::<Ext3>().is_none());
1260+ /// ```
1261+ ///
1262+ /// [^1]: Compare that to [`ConfigOptions`] which only supports [`ScalarValue`] payloads.
1263+ pub fn with_extension < T > ( mut self , ext : Arc < T > ) -> Self
1264+ where
1265+ T : Send + Sync + ' static ,
1266+ {
1267+ let ext = ext as Arc < dyn Any + Send + Sync + ' static > ;
1268+ let id = TypeId :: of :: < T > ( ) ;
1269+ self . extensions . insert ( id, ext) ;
1270+ self
1271+ }
1272+
1273+ /// Get extension, if any for the specified type `T` exists.
1274+ ///
1275+ /// See [`with_extension`](Self::with_extension) on how to add attach extensions.
1276+ pub fn get_extension < T > ( & self ) -> Option < Arc < T > >
1277+ where
1278+ T : Send + Sync + ' static ,
1279+ {
1280+ let id = TypeId :: of :: < T > ( ) ;
1281+ self . extensions
1282+ . get ( & id)
1283+ . cloned ( )
1284+ . map ( |ext| Arc :: downcast ( ext) . expect ( "TypeId unique" ) )
1285+ }
11761286}
11771287
11781288/// Execution context for registering data sources and executing queries
0 commit comments