vortex_array/stats/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Session state for stats APIs.
5
6use std::any::Any;
7use std::sync::Arc;
8
9use parking_lot::RwLock;
10use vortex_session::SessionExt;
11use vortex_session::SessionGuard;
12use vortex_session::SessionVar;
13use vortex_utils::aliases::hash_map::HashMap;
14
15use crate::scalar_fn::ScalarFnId;
16use crate::stats::rewrite::StatsRewriteRule;
17use crate::stats::rewrite::StatsRewriteRuleRef;
18use crate::stats::rewrite::register_builtins;
19
20type StatsRewriteRuleSet = Arc<[StatsRewriteRuleRef]>;
21
22/// Session state for stats APIs.
23#[derive(Clone, Debug)]
24pub struct StatsSession {
25    rewrite_rules: Arc<RwLock<HashMap<ScalarFnId, StatsRewriteRuleSet>>>,
26}
27
28impl Default for StatsSession {
29    fn default() -> Self {
30        let this = Self {
31            rewrite_rules: Arc::new(RwLock::new(HashMap::default())),
32        };
33        register_builtins(&this);
34        this
35    }
36}
37
38impl StatsSession {
39    /// Register a stats rewrite rule.
40    pub fn register_rewrite<R: StatsRewriteRule>(&self, rule: R) {
41        self.register_rewrite_ref(Arc::new(rule));
42    }
43
44    /// Register a shared stats rewrite rule.
45    pub fn register_rewrite_ref(&self, rule: StatsRewriteRuleRef) {
46        let mut rules = self.rewrite_rules.write();
47        let rule_id = rule.scalar_fn_id();
48        let mut updated_rules = rules
49            .get(&rule_id)
50            .map(|rules| rules.iter().cloned().collect::<Vec<_>>())
51            .unwrap_or_default();
52        updated_rules.push(rule);
53        rules.insert(rule_id, updated_rules.into());
54    }
55
56    /// Return the rewrite rules registered for `scalar_fn_id`.
57    pub(crate) fn rewrite_rules_for(
58        &self,
59        scalar_fn_id: ScalarFnId,
60    ) -> Option<StatsRewriteRuleSet> {
61        self.rewrite_rules.read().get(&scalar_fn_id).cloned()
62    }
63}
64
65impl SessionVar for StatsSession {
66    fn as_any(&self) -> &dyn Any {
67        self
68    }
69
70    fn as_any_mut(&mut self) -> &mut dyn Any {
71        self
72    }
73}
74
75/// Extension trait for accessing stats session data.
76pub trait StatsSessionExt: SessionExt {
77    /// Returns the stats session state.
78    fn stats(&self) -> SessionGuard<'_, StatsSession> {
79        self.get::<StatsSession>()
80    }
81}
82impl<S: SessionExt> StatsSessionExt for S {}
vortex_array/stats/session.rs

vortex_array/stats/
session.rs
