Skip to content

Commit 3e7073a

Browse files
committed
add extension system to SessionConfig
This is helpful for application built on top of DataFusion that want to pass certain query-scoped data structures around, e.g. for tracing or caching. This is loosely inspired by: https://github.com/hyperium/http/blob/34a9d6bdab027948d6dea3b36d994f9cbaf96f75/src/extensions.rs#L6-L28 Closes #2939.
1 parent 944ef3d commit 3e7073a

1 file changed

Lines changed: 111 additions & 1 deletion

File tree

‎datafusion/core/src/execution/context.rs‎

Lines changed: 111 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,12 @@ use crate::{
4444
};
4545
pub use datafusion_physical_expr::execution_props::ExecutionProps;
4646
use parking_lot::RwLock;
47-
use std::string::String;
4847
use std::sync::Arc;
48+
use std::{
49+
any::{Any, TypeId},
50+
hash::{BuildHasherDefault, Hasher},
51+
string::String,
52+
};
4953
use 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'
993997
pub 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)]
9971032
pub 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

10241061
impl 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

Comments
 (0)