vortex_array/aggregate_fn/
session.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::sync::Arc;
6
7use vortex_session::ArcSwapMap;
8use vortex_session::SessionExt;
9use vortex_session::SessionGuard;
10use vortex_session::SessionVar;
11
12use crate::aggregate_fn::AggregateFnId;
13use crate::aggregate_fn::AggregateFnPluginRef;
14use crate::aggregate_fn::AggregateFnRef;
15use crate::aggregate_fn::AggregateFnVTable;
16use crate::aggregate_fn::fns::all_nan::AllNan;
17use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct;
18use crate::aggregate_fn::fns::all_non_nan::AllNonNan;
19use crate::aggregate_fn::fns::all_non_null::AllNonNull;
20use crate::aggregate_fn::fns::all_null::AllNull;
21use crate::aggregate_fn::fns::bounded_max::BoundedMax;
22use crate::aggregate_fn::fns::bounded_min::BoundedMin;
23use crate::aggregate_fn::fns::count::Count;
24use crate::aggregate_fn::fns::count::CountGroupedKernel;
25use crate::aggregate_fn::fns::first::First;
26use crate::aggregate_fn::fns::is_constant::IsConstant;
27use crate::aggregate_fn::fns::is_sorted::IsSorted;
28use crate::aggregate_fn::fns::last::Last;
29use crate::aggregate_fn::fns::max::Max;
30use crate::aggregate_fn::fns::min::Min;
31use crate::aggregate_fn::fns::min_max::MinMax;
32use crate::aggregate_fn::fns::nan_count::NanCount;
33use crate::aggregate_fn::fns::null_count::NullCount;
34use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel;
35use crate::aggregate_fn::fns::sum::Sum;
36use crate::aggregate_fn::fns::sum_v2::PrimitiveGroupedSumV2EncodingKernel;
37use crate::aggregate_fn::fns::sum_v2::SumV2;
38use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
39use crate::aggregate_fn::kernels::DynAggregateKernel;
40use crate::aggregate_fn::kernels::DynGroupedAggregateKernel;
41use crate::array::ArrayId;
42use crate::array::VTable;
43use crate::arrays::Chunked;
44use crate::arrays::Dict;
45use crate::arrays::Primitive;
46use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate;
47use crate::arrays::dict::compute::is_constant::DictIsConstantKernel;
48use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel;
49use crate::arrays::dict::compute::min_max::DictMinMaxKernel;
50use crate::dtype::DType;
51
52/// Session state for aggregate functions and encoding-specific aggregate kernels.
53///
54/// The default session registers the built-in aggregate functions and kernels. Additional
55/// aggregate functions and kernels may be registered by extensions when they are added to a
56/// [`VortexSession`](vortex_session::VortexSession).
57#[derive(Clone, Debug)]
58pub struct AggregateFnSession {
59    registry: AggregateFnRegistry,
60
61    kernels: AggregateKernelRegistry,
62    grouped_kernels: GroupedKernelRegistry,
63    grouped_encoding_kernels: GroupedEncodingKernelRegistry,
64}
65
66impl SessionVar for AggregateFnSession {
67    fn as_any(&self) -> &dyn Any {
68        self
69    }
70
71    fn as_any_mut(&mut self) -> &mut dyn Any {
72        self
73    }
74}
75
76type AggregateKernelKey = (ArrayId, Option<AggregateFnId>);
77type GroupedEncodingKernelKey = (ArrayId, AggregateFnId);
78
79/// Registry of aggregate function plugins, keyed by aggregate function id.
80type AggregateFnRegistry = ArcSwapMap<AggregateFnId, AggregateFnPluginRef>;
81/// Registry of aggregate kernels, keyed by encoding and optional aggregate function.
82type AggregateKernelRegistry = ArcSwapMap<AggregateKernelKey, &'static dyn DynAggregateKernel>;
83/// Registry of encoding-agnostic grouped aggregate kernels, keyed by aggregate function id.
84type GroupedKernelRegistry = ArcSwapMap<AggregateFnId, &'static dyn DynGroupedAggregateKernel>;
85/// Registry of grouped aggregate kernels, keyed by encoding and aggregate function.
86type GroupedEncodingKernelRegistry =
87    ArcSwapMap<GroupedEncodingKernelKey, &'static dyn DynGroupedAggregateKernel>;
88
89impl Default for AggregateFnSession {
90    fn default() -> Self {
91        let this = Self {
92            registry: AggregateFnRegistry::default(),
93            kernels: AggregateKernelRegistry::default(),
94            grouped_kernels: GroupedKernelRegistry::default(),
95            grouped_encoding_kernels: GroupedEncodingKernelRegistry::default(),
96        };
97
98        // Register the built-in aggregate functions
99        this.register(AllNonDistinct);
100        this.register(AllNonNan);
101        this.register(AllNonNull);
102        this.register(AllNan);
103        this.register(AllNull);
104        this.register(BoundedMax);
105        this.register(BoundedMin);
106        this.register(First);
107        this.register(IsConstant);
108        this.register(IsSorted);
109        this.register(Last);
110        this.register(Max);
111        this.register(Min);
112        this.register(MinMax);
113        this.register(NanCount);
114        this.register(NullCount);
115        this.register(Sum);
116        this.register(SumV2);
117        this.register(UncompressedSizeInBytes);
118
119        // Register the built-in aggregate kernels.
120        this.register_aggregate_kernel(Chunked.id(), None::<AggregateFnId>, &ChunkedArrayAggregate);
121        this.register_aggregate_kernel(Dict.id(), Some(MinMax.id()), &DictMinMaxKernel);
122        this.register_aggregate_kernel(Dict.id(), Some(IsConstant.id()), &DictIsConstantKernel);
123        this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel);
124
125        // Register the built-in grouped aggregate kernels.
126        this.register_grouped_kernel(Count.id(), &CountGroupedKernel);
127        this.register_grouped_encoding_kernel(
128            Primitive.id(),
129            Sum.id(),
130            &PrimitiveGroupedSumEncodingKernel,
131        );
132        this.register_grouped_encoding_kernel(
133            Primitive.id(),
134            SumV2.id(),
135            &PrimitiveGroupedSumV2EncodingKernel,
136        );
137
138        this
139    }
140}
141
142impl AggregateFnSession {
143    /// Returns the aggregate function plugin registered for `id`, if any.
144    pub fn find_plugin(&self, id: &AggregateFnId) -> Option<AggregateFnPluginRef> {
145        self.registry.get(id)
146    }
147
148    /// Register an aggregate function vtable in the session, replacing any existing vtable with
149    /// the same ID.
150    pub fn register<V: AggregateFnVTable>(&self, vtable: V) {
151        let id = vtable.id();
152        let pluginref = Arc::new(vtable) as AggregateFnPluginRef;
153        self.registry.insert(id, pluginref);
154    }
155
156    /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every
157    /// registered aggregate's `zone_stat_default`.
158    ///
159    /// Each call scans the whole plugin registry, so this is intended to be called once per
160    /// column when a zoned writer is opened, not per chunk or per row.
161    pub fn zone_stat_defaults(&self, input_dtype: &DType) -> Vec<AggregateFnRef> {
162        self.registry.read(|registry| {
163            let mut fns: Vec<AggregateFnRef> = registry
164                .values()
165                .filter_map(|plugin| plugin.zone_stat_default(input_dtype))
166                .collect();
167            fns.sort_by_key(|f| f.id());
168            fns
169        })
170    }
171
172    /// Returns the aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
173    ///
174    /// Lookup first checks for a kernel registered for the exact aggregate function, then falls
175    /// back to a kernel registered for all aggregate functions on the same array encoding.
176    pub fn find_aggregate_kernel(
177        &self,
178        array_id: impl Into<ArrayId>,
179        agg_fn_id: impl Into<AggregateFnId>,
180    ) -> Option<&'static dyn DynAggregateKernel> {
181        let id = array_id.into();
182        let fn_id = agg_fn_id.into();
183        self.kernels.read(|kernels| {
184            kernels
185                .get(&(id, Some(fn_id)))
186                .or_else(|| kernels.get(&(id, None)))
187                .copied()
188        })
189    }
190
191    /// Registers an aggregate kernel for an array encoding.
192    ///
193    /// When `agg_fn_id` is `Some`, the kernel is used only for that aggregate function. When
194    /// `agg_fn_id` is `None`, the kernel is used as the fallback for aggregate functions on the
195    /// array encoding that do not have a more specific kernel.
196    pub fn register_aggregate_kernel(
197        &self,
198        array_id: impl Into<ArrayId>,
199        agg_fn_id: Option<impl Into<AggregateFnId>>,
200        kernel: &'static dyn DynAggregateKernel,
201    ) {
202        let id = (array_id.into(), agg_fn_id.map(|id| id.into()));
203        self.kernels.insert(id, kernel);
204    }
205
206    /// Returns the grouped aggregate kernel registered for `agg_fn_id`, if any.
207    ///
208    /// These kernels are independent of the element encoding and are checked for each element
209    /// representation, after any kernel registered for the current element encoding.
210    pub fn find_grouped_kernel(
211        &self,
212        agg_fn_id: impl Into<AggregateFnId>,
213    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
214        let fn_id = agg_fn_id.into();
215        self.grouped_kernels
216            .read(|kernels| kernels.get(&fn_id).copied())
217    }
218
219    /// Registers a grouped aggregate kernel for an aggregate function.
220    pub fn register_grouped_kernel(
221        &self,
222        agg_fn_id: impl Into<AggregateFnId>,
223        kernel: &'static dyn DynGroupedAggregateKernel,
224    ) {
225        let fn_id = agg_fn_id.into();
226        self.grouped_kernels.insert(fn_id, kernel)
227    }
228
229    /// Returns the grouped aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
230    ///
231    /// These kernels are matched against each intermediate element encoding while the grouped
232    /// accumulator executes the element array.
233    pub fn find_grouped_encoding_kernel(
234        &self,
235        array_id: impl Into<ArrayId>,
236        agg_fn_id: impl Into<AggregateFnId>,
237    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
238        let id = array_id.into();
239        let fn_id = agg_fn_id.into();
240        self.grouped_encoding_kernels
241            .read(|kernels| kernels.get(&(id, fn_id)).copied())
242    }
243
244    /// Registers a grouped aggregate kernel for a specific aggregate function and array encoding.
245    pub fn register_grouped_encoding_kernel(
246        &self,
247        array_id: impl Into<ArrayId>,
248        agg_fn_id: impl Into<AggregateFnId>,
249        kernel: &'static dyn DynGroupedAggregateKernel,
250    ) {
251        let id = array_id.into();
252        let fn_id = agg_fn_id.into();
253        self.grouped_encoding_kernels.insert((id, fn_id), kernel)
254    }
255}
256
257/// Extension trait for accessing aggregate function session data.
258pub trait AggregateFnSessionExt: SessionExt {
259    /// Returns the aggregate function session data.
260    fn aggregate_fns(&self) -> SessionGuard<'_, AggregateFnSession> {
261        self.get::<AggregateFnSession>()
262    }
263}
264impl<S: SessionExt> AggregateFnSessionExt for S {}
vortex_array/aggregate_fn/session.rs

vortex_array/aggregate_fn/
session.rs
