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

vortex_array/aggregate_fn/
session.rs
