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;
11use vortex_utils::aliases::hash_map::HashMap;
12
13use crate::aggregate_fn::AggregateFnId;
14use crate::aggregate_fn::AggregateFnPluginRef;
15use crate::aggregate_fn::AggregateFnRef;
16use crate::aggregate_fn::AggregateFnVTable;
17use crate::aggregate_fn::ExtensionStoragePushdown;
18use crate::aggregate_fn::fns::all_nan::AllNan;
19use crate::aggregate_fn::fns::all_non_distinct::AllNonDistinct;
20use crate::aggregate_fn::fns::all_non_nan::AllNonNan;
21use crate::aggregate_fn::fns::all_non_null::AllNonNull;
22use crate::aggregate_fn::fns::all_null::AllNull;
23use crate::aggregate_fn::fns::bounded_max::BoundedMax;
24use crate::aggregate_fn::fns::bounded_min::BoundedMin;
25use crate::aggregate_fn::fns::count::Count;
26use crate::aggregate_fn::fns::count::CountGroupedKernel;
27use crate::aggregate_fn::fns::first::First;
28use crate::aggregate_fn::fns::is_constant::IsConstant;
29use crate::aggregate_fn::fns::is_sorted::IsSorted;
30use crate::aggregate_fn::fns::last::Last;
31use crate::aggregate_fn::fns::max::Max;
32use crate::aggregate_fn::fns::min::Min;
33use crate::aggregate_fn::fns::min_max::MinMax;
34use crate::aggregate_fn::fns::nan_count::NanCount;
35use crate::aggregate_fn::fns::null_count::NullCount;
36use crate::aggregate_fn::fns::sum::PrimitiveGroupedSumEncodingKernel;
37use crate::aggregate_fn::fns::sum::Sum;
38use crate::aggregate_fn::fns::sum_v2::PrimitiveGroupedSumV2EncodingKernel;
39use crate::aggregate_fn::fns::sum_v2::SumV2;
40use crate::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes;
41use crate::aggregate_fn::kernels::DynAggregateKernel;
42use crate::aggregate_fn::kernels::DynGroupedAggregateKernel;
43use crate::array::ArrayId;
44use crate::array::VTable;
45use crate::arrays::Chunked;
46use crate::arrays::Dict;
47use crate::arrays::Primitive;
48use crate::arrays::chunked::compute::aggregate::ChunkedArrayAggregate;
49use crate::arrays::dict::compute::is_constant::DictIsConstantKernel;
50use crate::arrays::dict::compute::is_sorted::DictIsSortedKernel;
51use crate::arrays::dict::compute::min_max::DictMinMaxKernel;
52use crate::dtype::DType;
53use crate::dtype::extension::ExtVTable;
54use crate::extension::datetime::Date;
55use crate::extension::datetime::Time;
56use crate::extension::datetime::Timestamp;
57
58/// Session state for aggregate functions and encoding-specific aggregate kernels.
59///
60/// The default session registers the built-in aggregate functions and kernels. Additional
61/// aggregate functions and kernels may be registered by extensions when they are added to a
62/// [`VortexSession`](vortex_session::VortexSession).
63#[derive(Clone, Debug)]
64pub struct AggregateFnSession {
65    registry: AggregateFnRegistry,
66
67    kernels: AggregateKernelRegistry,
68    grouped_kernels: GroupedKernelRegistry,
69    grouped_encoding_kernels: GroupedEncodingKernelRegistry,
70}
71
72impl SessionVar for AggregateFnSession {
73    fn as_any(&self) -> &dyn Any {
74        self
75    }
76
77    fn as_any_mut(&mut self) -> &mut dyn Any {
78        self
79    }
80}
81
82type AggregateKernelKey = (ArrayId, Option<AggregateFnId>);
83type GroupedEncodingKernelKey = (ArrayId, AggregateFnId);
84
85/// Registry of aggregate function plugins, keyed by aggregate function id.
86type AggregateFnRegistry = ArcSwapMap<AggregateFnId, AggregateFnPluginRef>;
87/// Registry of aggregate kernels, keyed by encoding and optional aggregate function.
88type AggregateKernelRegistry = ArcSwapMap<AggregateKernelKey, &'static dyn DynAggregateKernel>;
89/// Registry of encoding-agnostic grouped aggregate kernels, keyed by aggregate function id.
90type GroupedKernelRegistry = ArcSwapMap<AggregateFnId, &'static dyn DynGroupedAggregateKernel>;
91/// Registry of grouped aggregate kernels, keyed by encoding and aggregate function.
92type GroupedEncodingKernelRegistry =
93    ArcSwapMap<GroupedEncodingKernelKey, &'static dyn DynGroupedAggregateKernel>;
94
95impl Default for AggregateFnSession {
96    fn default() -> Self {
97        let this = Self {
98            registry: AggregateFnRegistry::default(),
99            kernels: AggregateKernelRegistry::default(),
100            grouped_kernels: GroupedKernelRegistry::default(),
101            grouped_encoding_kernels: GroupedEncodingKernelRegistry::default(),
102        };
103
104        // Register the built-in aggregate functions
105        this.register(AllNonDistinct);
106        this.register(AllNonNan);
107        this.register(AllNonNull);
108        this.register(AllNan);
109        this.register(AllNull);
110        this.register(BoundedMax);
111        this.register(BoundedMin);
112        this.register(First);
113        this.register(IsConstant);
114        this.register(IsSorted);
115        this.register(Last);
116        this.register(Max);
117        this.register(Min);
118        this.register(MinMax);
119        this.register(NanCount);
120        this.register(NullCount);
121        this.register(Sum);
122        this.register(SumV2);
123        this.register(UncompressedSizeInBytes);
124
125        // Temporal extension types compare like their integer storage, so comparison-based
126        // aggregates are pushed down into it.
127        this.register_storage_ordered(Date);
128        this.register_storage_ordered(Time);
129        this.register_storage_ordered(Timestamp);
130
131        // Register the built-in aggregate kernels.
132        this.register_aggregate_kernel(Chunked.id(), None::<AggregateFnId>, &ChunkedArrayAggregate);
133        this.register_aggregate_kernel(Dict.id(), Some(MinMax.id()), &DictMinMaxKernel);
134        this.register_aggregate_kernel(Dict.id(), Some(IsConstant.id()), &DictIsConstantKernel);
135        this.register_aggregate_kernel(Dict.id(), Some(IsSorted.id()), &DictIsSortedKernel);
136
137        // Register the built-in grouped aggregate kernels.
138        this.register_grouped_kernel(Count.id(), &CountGroupedKernel);
139        this.register_grouped_encoding_kernel(
140            Primitive.id(),
141            Sum.id(),
142            &PrimitiveGroupedSumEncodingKernel,
143        );
144        this.register_grouped_encoding_kernel(
145            Primitive.id(),
146            SumV2.id(),
147            &PrimitiveGroupedSumV2EncodingKernel,
148        );
149
150        this
151    }
152}
153
154impl AggregateFnSession {
155    /// Returns the aggregate function plugin registered for `id`, if any.
156    pub fn find_plugin(&self, id: &AggregateFnId) -> Option<AggregateFnPluginRef> {
157        self.registry.get(id)
158    }
159
160    /// Register an aggregate function vtable in the session, replacing any existing vtable with
161    /// the same ID.
162    pub fn register<V: AggregateFnVTable>(&self, vtable: V) {
163        let id = vtable.id();
164        let pluginref = Arc::new(vtable) as AggregateFnPluginRef;
165        self.registry.insert(id, pluginref);
166    }
167
168    /// Push the comparison-based aggregates `is_sorted` and `min_max` down into the storage of
169    /// extension type `ext`, through [`ExtensionStoragePushdown`].
170    ///
171    /// Only call this when comparing two values of `ext` is the same as comparing their storage,
172    /// as for the temporal types and for UUIDs, which are stored big-endian. See
173    /// [`ExtensionStoragePushdown`] for how to check.
174    pub fn register_storage_ordered<E: ExtVTable>(&self, ext: E) {
175        self.register(ExtensionStoragePushdown::new(IsSorted, ext.clone()));
176        self.register(ExtensionStoragePushdown::new(MinMax, ext));
177    }
178
179    /// The default per-chunk zone statistics for a column of `input_dtype`, collected from every
180    /// registered aggregate's `zone_stat_default`.
181    ///
182    /// Each call scans the whole plugin registry, so this is intended to be called once per
183    /// column when a zoned writer is opened, not per chunk or per row.
184    pub fn zone_stat_defaults(&self, input_dtype: &DType) -> Vec<AggregateFnRef> {
185        self.registry.read(|registry| {
186            let mut fns: Vec<AggregateFnRef> = registry
187                .values()
188                .filter_map(|plugin| plugin.zone_stat_default(input_dtype))
189                .collect();
190            fns.sort_by_key(|f| f.id());
191            fns
192        })
193    }
194
195    /// Returns the aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
196    ///
197    /// Lookup first checks for a kernel registered for the exact aggregate function, then falls
198    /// back to a kernel registered for all aggregate functions on the same array encoding.
199    pub fn find_aggregate_kernel(
200        &self,
201        array_id: impl Into<ArrayId>,
202        agg_fn_id: impl Into<AggregateFnId>,
203    ) -> Option<&'static dyn DynAggregateKernel> {
204        let id = array_id.into();
205        let fn_id = agg_fn_id.into();
206        self.kernels
207            .read(|kernels| lookup_aggregate_kernel(kernels, id, fn_id))
208    }
209
210    /// Registers an aggregate kernel for an array encoding.
211    ///
212    /// When `agg_fn_id` is `Some`, the kernel is used only for that aggregate function. When
213    /// `agg_fn_id` is `None`, the kernel is used as the fallback for aggregate functions on the
214    /// array encoding that do not have a more specific kernel.
215    pub fn register_aggregate_kernel(
216        &self,
217        array_id: impl Into<ArrayId>,
218        agg_fn_id: Option<impl Into<AggregateFnId>>,
219        kernel: &'static dyn DynAggregateKernel,
220    ) {
221        let id = (array_id.into(), agg_fn_id.map(|id| id.into()));
222        self.kernels.insert(id, kernel);
223    }
224
225    /// Returns the grouped aggregate kernel registered for `agg_fn_id`, if any.
226    ///
227    /// These kernels are independent of the element encoding and are checked for each element
228    /// representation, after any kernel registered for the current element encoding.
229    pub fn find_grouped_kernel(
230        &self,
231        agg_fn_id: impl Into<AggregateFnId>,
232    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
233        let fn_id = agg_fn_id.into();
234        self.grouped_kernels
235            .read(|kernels| kernels.get(&fn_id).copied())
236    }
237
238    /// Registers a grouped aggregate kernel for an aggregate function.
239    pub fn register_grouped_kernel(
240        &self,
241        agg_fn_id: impl Into<AggregateFnId>,
242        kernel: &'static dyn DynGroupedAggregateKernel,
243    ) {
244        let fn_id = agg_fn_id.into();
245        self.grouped_kernels.insert(fn_id, kernel)
246    }
247
248    /// Returns the grouped aggregate kernel registered for `array_id` and `agg_fn_id`, if any.
249    ///
250    /// These kernels are matched against each intermediate element encoding while the grouped
251    /// accumulator executes the element array.
252    pub fn find_grouped_encoding_kernel(
253        &self,
254        array_id: impl Into<ArrayId>,
255        agg_fn_id: impl Into<AggregateFnId>,
256    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
257        let id = array_id.into();
258        let fn_id = agg_fn_id.into();
259        self.grouped_encoding_kernels
260            .read(|kernels| kernels.get(&(id, fn_id)).copied())
261    }
262
263    /// Registers a grouped aggregate kernel for a specific aggregate function and array encoding.
264    pub fn register_grouped_encoding_kernel(
265        &self,
266        array_id: impl Into<ArrayId>,
267        agg_fn_id: impl Into<AggregateFnId>,
268        kernel: &'static dyn DynGroupedAggregateKernel,
269    ) {
270        let id = array_id.into();
271        let fn_id = agg_fn_id.into();
272        self.grouped_encoding_kernels.insert((id, fn_id), kernel)
273    }
274
275    /// Captures the currently registered aggregate kernels, so repeated lookups skip the session
276    /// and registry loads.
277    pub(crate) fn kernels_snapshot(&self) -> AggregateKernels {
278        AggregateKernels {
279            kernels: self.kernels.snapshot(),
280            grouped_kernels: self.grouped_kernels.snapshot(),
281            grouped_encoding_kernels: self.grouped_encoding_kernels.snapshot(),
282        }
283    }
284}
285
286/// A point-in-time snapshot of the aggregate kernel registries of an [`AggregateFnSession`].
287///
288/// Lookups follow the same rules as the corresponding `find_*` methods on
289/// [`AggregateFnSession`].
290#[derive(Clone, Debug)]
291pub(crate) struct AggregateKernels {
292    kernels: Arc<HashMap<AggregateKernelKey, &'static dyn DynAggregateKernel>>,
293    grouped_kernels: Arc<HashMap<AggregateFnId, &'static dyn DynGroupedAggregateKernel>>,
294    grouped_encoding_kernels:
295        Arc<HashMap<GroupedEncodingKernelKey, &'static dyn DynGroupedAggregateKernel>>,
296}
297
298impl AggregateKernels {
299    /// See [`AggregateFnSession::find_aggregate_kernel`].
300    pub(crate) fn find_aggregate_kernel(
301        &self,
302        array_id: ArrayId,
303        agg_fn_id: AggregateFnId,
304    ) -> Option<&'static dyn DynAggregateKernel> {
305        lookup_aggregate_kernel(&self.kernels, array_id, agg_fn_id)
306    }
307
308    /// See [`AggregateFnSession::find_grouped_kernel`].
309    pub(crate) fn find_grouped_kernel(
310        &self,
311        agg_fn_id: AggregateFnId,
312    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
313        self.grouped_kernels.get(&agg_fn_id).copied()
314    }
315
316    /// See [`AggregateFnSession::find_grouped_encoding_kernel`].
317    pub(crate) fn find_grouped_encoding_kernel(
318        &self,
319        array_id: ArrayId,
320        agg_fn_id: AggregateFnId,
321    ) -> Option<&'static dyn DynGroupedAggregateKernel> {
322        self.grouped_encoding_kernels
323            .get(&(array_id, agg_fn_id))
324            .copied()
325    }
326}
327
328fn lookup_aggregate_kernel(
329    kernels: &HashMap<AggregateKernelKey, &'static dyn DynAggregateKernel>,
330    array_id: ArrayId,
331    agg_fn_id: AggregateFnId,
332) -> Option<&'static dyn DynAggregateKernel> {
333    kernels
334        .get(&(array_id, Some(agg_fn_id)))
335        .or_else(|| kernels.get(&(array_id, None)))
336        .copied()
337}
338
339/// Extension trait for accessing aggregate function session data.
340pub trait AggregateFnSessionExt: SessionExt {
341    /// Returns the aggregate function session data.
342    fn aggregate_fns(&self) -> SessionGuard<'_, AggregateFnSession> {
343        self.get::<AggregateFnSession>()
344    }
345}
346impl<S: SessionExt> AggregateFnSessionExt for S {}
vortex_array/aggregate_fn/session.rs

vortex_array/aggregate_fn/
session.rs
