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

vortex_array/aggregate_fn/
session.rs
