vortex_array/aggregate_fn/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use parking_lot::RwLock;
7use vortex_session::Ref;
8use vortex_session::SessionExt;
9use vortex_session::registry::Registry;
10use vortex_utils::aliases::hash_map::HashMap;
11
12use crate::aggregate_fn::AggregateFnId;
13use crate::aggregate_fn::AggregateFnPluginRef;
14use crate::aggregate_fn::AggregateFnVTable;
15use crate::aggregate_fn::kernels::DynAggregateKernel;
16use crate::aggregate_fn::kernels::DynGroupedAggregateKernel;
17use crate::arrays::ChunkedVTable;
18use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate;
19use crate::vtable::ArrayId;
20
21/// Registry of aggregate function vtables.
22pub type AggregateFnRegistry = Registry<AggregateFnPluginRef>;
23
24/// Session state for aggregate function vtables.
25#[derive(Debug)]
26pub struct AggregateFnSession {
27    registry: AggregateFnRegistry,
28
29    pub(super) kernels: RwLock<HashMap<KernelKey, &'static dyn DynAggregateKernel>>,
30    pub(super) grouped_kernels: RwLock<HashMap<KernelKey, &'static dyn DynGroupedAggregateKernel>>,
31}
32
33type KernelKey = (ArrayId, Option<AggregateFnId>);
34
35impl Default for AggregateFnSession {
36    fn default() -> Self {
37        let this = Self {
38            registry: AggregateFnRegistry::default(),
39            kernels: RwLock::new(HashMap::default()),
40            grouped_kernels: RwLock::new(HashMap::default()),
41        };
42
43        // Register the built-in aggregate kernels.
44        this.register_aggregate_kernel(ChunkedVTable::ID, None, &ChunkedArrayAggregate);
45
46        this
47    }
48}
49
50impl AggregateFnSession {
51    /// Returns the aggregate function registry.
52    pub fn registry(&self) -> &AggregateFnRegistry {
53        &self.registry
54    }
55
56    /// Register an aggregate function vtable in the session, replacing any existing vtable with
57    /// the same ID.
58    pub fn register<V: AggregateFnVTable>(&self, vtable: V) {
59        self.registry
60            .register(vtable.id(), Arc::new(vtable) as AggregateFnPluginRef);
61    }
62
63    /// Register an aggregate function kernel for a specific aggregate function and array type.
64    pub fn register_aggregate_kernel(
65        &self,
66        array_id: ArrayId,
67        agg_fn_id: Option<AggregateFnId>,
68        kernel: &'static dyn DynAggregateKernel,
69    ) {
70        self.kernels.write().insert((array_id, agg_fn_id), kernel);
71    }
72}
73
74/// Extension trait for accessing aggregate function session data.
75pub trait AggregateFnSessionExt: SessionExt {
76    /// Returns the aggregate function vtable registry.
77    fn aggregate_fns(&self) -> Ref<'_, AggregateFnSession> {
78        self.get::<AggregateFnSession>()
79    }
80}
81impl<S: SessionExt> AggregateFnSessionExt for S {}
vortex_array/aggregate_fn/session.rs

vortex_array/aggregate_fn/
session.rs
