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::Ref;
11use vortex_session::SessionExt;
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(Debug)]
24pub struct StatsSession {
25    rewrite_rules: RwLock<HashMap<ScalarFnId, StatsRewriteRuleSet>>,
26}
27
28impl Default for StatsSession {
29    fn default() -> Self {
30        let this = Self {
31            rewrite_rules: RwLock::new(HashMap::default()),
32        };
33        register_builtins(&this);
34        this
35    }
36}
37
38impl StatsSession {
39    /// Register a stats rewrite rule.
40    #[allow(dead_code)]
41    pub(crate) fn register_rewrite<R: StatsRewriteRule>(&self, rule: R) {
42        self.register_rewrite_ref(Arc::new(rule));
43    }
44
45    /// Register a shared stats rewrite rule.
46    #[allow(dead_code)]
47    pub(crate) fn register_rewrite_ref(&self, rule: StatsRewriteRuleRef) {
48        let mut rules = self.rewrite_rules.write();
49        let rule_id = rule.scalar_fn_id();
50        let mut updated_rules = rules
51            .get(&rule_id)
52            .map(|rules| rules.iter().cloned().collect::<Vec<_>>())
53            .unwrap_or_default();
54        updated_rules.push(rule);
55        rules.insert(rule_id, updated_rules.into());
56    }
57
58    /// Return the rewrite rules registered for `scalar_fn_id`.
59    pub(crate) fn rewrite_rules_for(
60        &self,
61        scalar_fn_id: ScalarFnId,
62    ) -> Option<StatsRewriteRuleSet> {
63        self.rewrite_rules.read().get(&scalar_fn_id).cloned()
64    }
65}
66
67impl SessionVar for StatsSession {
68    fn as_any(&self) -> &dyn Any {
69        self
70    }
71
72    fn as_any_mut(&mut self) -> &mut dyn Any {
73        self
74    }
75}
76
77/// Extension trait for accessing stats session data.
78pub(crate) trait StatsSessionExt: SessionExt {
79    /// Returns the stats session state.
80    fn stats(&self) -> Ref<'_, StatsSession> {
81        self.get::<StatsSession>()
82    }
83}
84impl<S: SessionExt> StatsSessionExt for S {}
vortex_array/stats/session.rs

vortex_array/stats/
session.rs
