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

vortex_array/aggregate_fn/
session.rs
