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

vortex_array/stats/
session.rs
