vortex_array/aggregate_fn/
accumulator_grouped.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_buffer::ArrowNativeType;
5use vortex_buffer::Buffer;
6use vortex_error::VortexExpect;
7use vortex_error::VortexResult;
8use vortex_error::vortex_bail;
9use vortex_error::vortex_ensure;
10use vortex_error::vortex_err;
11use vortex_error::vortex_panic;
12use vortex_mask::Mask;
13
14use crate::ArrayRef;
15use crate::Canonical;
16use crate::Columnar;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::aggregate_fn::Accumulator;
20use crate::aggregate_fn::AggregateFn;
21use crate::aggregate_fn::AggregateFnRef;
22use crate::aggregate_fn::AggregateFnVTable;
23use crate::aggregate_fn::DynAccumulator;
24use crate::aggregate_fn::session::AggregateFnSessionExt;
25use crate::arrays::ChunkedArray;
26use crate::arrays::FixedSizeListArray;
27use crate::arrays::ListViewArray;
28use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
29use crate::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
30use crate::arrays::listview::ListViewArraySlotsExt;
31use crate::builders::builder_with_capacity;
32use crate::builtins::ArrayBuiltins;
33use crate::columnar::AnyColumnar;
34use crate::dtype::DType;
35use crate::executor::max_iterations;
36use crate::match_each_integer_ptype;
37
38/// Reference-counted type-erased grouped accumulator.
39pub type GroupedAccumulatorRef = Box<dyn DynGroupedAccumulator>;
40
41/// A batch of grouped values to aggregate.
42///
43/// Each outer list value is one group, and the inner element array is shared by all groups.
44/// Aggregate implementations can inspect the concrete grouped representation directly, or ask for
45/// derived ranges when their algorithm is expressed in terms of `(offset, size)` pairs.
46pub enum GroupedArray {
47    /// Groups represented as a list-view array with per-group offsets and sizes.
48    ListView(ListViewArray),
49    /// Groups represented as a fixed-size list array.
50    FixedSizeList(FixedSizeListArray),
51}
52
53impl From<ListViewArray> for GroupedArray {
54    fn from(groups: ListViewArray) -> Self {
55        Self::ListView(groups)
56    }
57}
58
59impl From<FixedSizeListArray> for GroupedArray {
60    fn from(groups: FixedSizeListArray) -> Self {
61        Self::FixedSizeList(groups)
62    }
63}
64
65impl GroupedArray {
66    /// The inner element array shared by all groups.
67    pub fn elements(&self) -> &ArrayRef {
68        match self {
69            Self::ListView(groups) => groups.elements(),
70            Self::FixedSizeList(groups) => groups.elements(),
71        }
72    }
73
74    /// Return the `(offset, size)` ranges describing each group in `elements`.
75    pub fn group_ranges(&self, ctx: &mut ExecutionCtx) -> VortexResult<GroupRanges> {
76        match self {
77            Self::ListView(groups) => list_view_group_ranges(groups, ctx),
78            Self::FixedSizeList(groups) => Ok(fixed_size_list_group_ranges(groups)),
79        }
80    }
81
82    /// Return the per-group validity mask.
83    pub fn group_validity(&self, ctx: &mut ExecutionCtx) -> VortexResult<Mask> {
84        match self {
85            Self::ListView(groups) => groups.validity()?.execute_mask(groups.len(), ctx),
86            Self::FixedSizeList(groups) => groups.validity()?.execute_mask(groups.len(), ctx),
87        }
88    }
89
90    /// The number of groups in this batch.
91    pub fn len(&self) -> usize {
92        match self {
93            Self::ListView(groups) => groups.len(),
94            Self::FixedSizeList(groups) => groups.len(),
95        }
96    }
97
98    /// Returns true when this batch contains no groups.
99    pub fn is_empty(&self) -> bool {
100        self.len() == 0
101    }
102
103    /// Returns true when every group is valid.
104    pub fn all_groups_valid(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
105        Ok(self.group_validity(ctx)?.all_true())
106    }
107
108    unsafe fn with_elements_unchecked(&self, elements: ArrayRef) -> VortexResult<Self> {
109        Ok(match self {
110            Self::ListView(groups) => unsafe {
111                ListViewArray::new_unchecked(
112                    elements,
113                    groups.offsets().clone(),
114                    groups.sizes().clone(),
115                    groups.validity()?,
116                )
117            }
118            .into(),
119            Self::FixedSizeList(groups) => unsafe {
120                FixedSizeListArray::new_unchecked(
121                    elements,
122                    groups.list_size(),
123                    groups.validity()?,
124                    groups.len(),
125                )
126            }
127            .into(),
128        })
129    }
130}
131
132/// The physical ranges of a grouped array.
133pub enum GroupRanges {
134    /// Explicit ranges extracted from a list-view array.
135    ListView {
136        /// The `(offset, size)` ranges.
137        ranges: Vec<(usize, usize)>,
138    },
139    /// Uniform ranges derived from a fixed-size list array.
140    FixedSizeList {
141        /// The number of groups.
142        len: usize,
143        /// The number of elements in each group.
144        size: usize,
145    },
146}
147
148impl GroupRanges {
149    /// The number of groups described by these ranges.
150    pub fn len(&self) -> usize {
151        match self {
152            Self::ListView { ranges } => ranges.len(),
153            Self::FixedSizeList { len, .. } => *len,
154        }
155    }
156
157    /// Returns true when there are no groups.
158    pub fn is_empty(&self) -> bool {
159        self.len() == 0
160    }
161
162    /// Return the `(offset, size)` range for the group at `index`.
163    fn range(&self, index: usize) -> (usize, usize) {
164        match self {
165            Self::ListView { ranges } => ranges[index],
166            Self::FixedSizeList { len, size } => {
167                assert!(index < *len, "range index out of bounds");
168                (index * size, *size)
169            }
170        }
171    }
172
173    /// Iterate over all `(offset, size)` group ranges.
174    pub fn iter(&self) -> impl Iterator<Item = (usize, usize)> + '_ {
175        (0..self.len()).map(|index| self.range(index))
176    }
177}
178
179/// An accumulator used for computing grouped aggregates.
180///
181/// Note that the groups must be processed in order, and the accumulator does not support random
182/// access to groups.
183pub struct GroupedAccumulator<V: AggregateFnVTable> {
184    /// The vtable of the aggregate function.
185    vtable: V,
186    /// The options of the aggregate function.
187    options: V::Options,
188    /// Type-erased aggregate function used for kernel dispatch.
189    aggregate_fn: AggregateFnRef,
190    /// The DType of the input.
191    dtype: DType,
192    /// The DType of the aggregate.
193    return_dtype: DType,
194    /// The DType of the partial accumulator state.
195    partial_dtype: DType,
196    /// The accumulated state for prior batches of groups.
197    partials: Vec<ArrayRef>,
198}
199
200impl<V: AggregateFnVTable> GroupedAccumulator<V> {
201    pub fn try_new(vtable: V, options: V::Options, dtype: DType) -> VortexResult<Self> {
202        let aggregate_fn = AggregateFn::new(vtable.clone(), options.clone()).erased();
203        let return_dtype = vtable.return_dtype(&options, &dtype).ok_or_else(|| {
204            vortex_err!(
205                "Aggregate function {} cannot be applied to dtype {}",
206                vtable.id(),
207                dtype
208            )
209        })?;
210        let partial_dtype = vtable.partial_dtype(&options, &dtype).ok_or_else(|| {
211            vortex_err!(
212                "Aggregate function {} cannot be applied to dtype {}",
213                vtable.id(),
214                dtype
215            )
216        })?;
217
218        Ok(Self {
219            vtable,
220            options,
221            aggregate_fn,
222            dtype,
223            return_dtype,
224            partial_dtype,
225            partials: vec![],
226        })
227    }
228}
229
230/// A trait object for type-erased grouped accumulators, used for dynamic dispatch when the aggregate
231/// function is not known at compile time.
232pub trait DynGroupedAccumulator: 'static + Send {
233    /// Accumulate a list of groups into the accumulator.
234    fn accumulate_list(&mut self, groups: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()>;
235
236    /// Finish the accumulation and return the partial aggregate results for all groups.
237    /// Resets the accumulator state for the next round of accumulation.
238    fn flush(&mut self) -> VortexResult<ArrayRef>;
239
240    /// Finish the accumulation and return the final aggregate results for all groups.
241    /// Resets the accumulator state for the next round of accumulation.
242    fn finish(&mut self) -> VortexResult<ArrayRef>;
243}
244
245impl<V: AggregateFnVTable> DynGroupedAccumulator for GroupedAccumulator<V> {
246    fn accumulate_list(&mut self, groups: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> {
247        let elements_dtype = match groups.dtype() {
248            DType::List(elem, _) => elem,
249            DType::FixedSizeList(elem, ..) => elem,
250            _ => vortex_bail!(
251                "Input DType mismatch: expected List or FixedSizeList, got {}",
252                groups.dtype()
253            ),
254        };
255        vortex_ensure!(
256            elements_dtype.as_ref() == &self.dtype,
257            "Input DType mismatch: expected {}, got {}",
258            self.dtype,
259            elements_dtype
260        );
261
262        // We first execute the groups until it is a ListView or FixedSizeList, since we only
263        // dispatch the aggregate kernel over the elements of these arrays.
264        let canonical = match groups.clone().execute::<Columnar>(ctx)? {
265            Columnar::Canonical(c) => c,
266            Columnar::Constant(c) => c.into_array().execute::<Canonical>(ctx)?,
267        };
268        match canonical {
269            Canonical::List(groups) => self.accumulate_grouped_array(groups.into(), ctx),
270            Canonical::FixedSizeList(groups) => self.accumulate_grouped_array(groups.into(), ctx),
271            _ => vortex_panic!("We checked the DType above, so this should never happen"),
272        }
273    }
274
275    fn flush(&mut self) -> VortexResult<ArrayRef> {
276        let states = std::mem::take(&mut self.partials);
277        Ok(ChunkedArray::try_new(states, self.partial_dtype.clone())?.into_array())
278    }
279
280    fn finish(&mut self) -> VortexResult<ArrayRef> {
281        let states = self.flush()?;
282        let results = self.vtable.finalize(states)?;
283
284        vortex_ensure!(
285            results.dtype() == &self.return_dtype,
286            "Return DType mismatch: expected {}, got {}",
287            self.return_dtype,
288            results.dtype()
289        );
290
291        Ok(results)
292    }
293}
294
295impl<V: AggregateFnVTable> GroupedAccumulator<V> {
296    fn accumulate_grouped_array(
297        &mut self,
298        groups: GroupedArray,
299        ctx: &mut ExecutionCtx,
300    ) -> VortexResult<()> {
301        let mut elements = groups.elements().clone();
302        let session = ctx.session().clone();
303
304        for _ in 0..max_iterations() {
305            // Try a registered grouped kernel for the current element encoding.
306            if let Some(kernel) = session
307                .aggregate_fns()
308                .find_grouped_encoding_kernel(elements.encoding_id(), self.aggregate_fn.id())
309            {
310                // SAFETY: we assume that elements execution is safe
311                let kernel_groups = unsafe { groups.with_elements_unchecked(elements.clone())? };
312                if let Some(result) =
313                    kernel.grouped_aggregate(&self.aggregate_fn, &kernel_groups, ctx)?
314                {
315                    return self.push_result(result);
316                }
317            }
318
319            // Try a grouped kernel for the current aggregate regardless of element encoding.
320            if let Some(kernel) = session
321                .aggregate_fns()
322                .find_grouped_kernel(self.aggregate_fn.id())
323            {
324                // SAFETY: we preserve the grouped shape and validity while replacing the
325                // elements with another representation of the same logical array.
326                let kernel_groups = unsafe { groups.with_elements_unchecked(elements.clone())? };
327                if let Some(result) =
328                    kernel.grouped_aggregate(&self.aggregate_fn, &kernel_groups, ctx)?
329                {
330                    return self.push_result(result);
331                }
332            }
333
334            if elements.is::<AnyColumnar>() {
335                break;
336            }
337
338            // Execute one step and try again
339            elements = elements.execute(ctx)?;
340        }
341
342        let elements = elements.execute::<Columnar>(ctx)?.into_array();
343        // SAFETY: we preserve the grouped shape and validity while replacing the elements with an
344        // executed form of the same logical array.
345        let grouped = unsafe { groups.with_elements_unchecked(elements)? };
346
347        // Otherwise, we iterate the offsets and sizes and accumulate each group one by one.
348        self.accumulate_grouped_fallback(&grouped, ctx)
349    }
350
351    fn accumulate_grouped_fallback(
352        &mut self,
353        grouped: &GroupedArray,
354        ctx: &mut ExecutionCtx,
355    ) -> VortexResult<()> {
356        let mut accumulator = Accumulator::try_new(
357            self.vtable.clone(),
358            self.options.clone(),
359            self.dtype.clone(),
360        )?;
361        let mut states = builder_with_capacity(&self.partial_dtype, grouped.len());
362        let group_ranges = grouped.group_ranges(ctx)?;
363        let group_validity = grouped.group_validity(ctx)?;
364
365        for ((offset, size), valid) in group_ranges.iter().zip(group_validity.iter()) {
366            if valid {
367                let group = grouped.elements().slice(offset..offset + size)?;
368                accumulator.accumulate(&group, ctx)?;
369                states.append_scalar(&accumulator.flush()?)?;
370            } else {
371                states.append_null()
372            }
373        }
374
375        self.push_result(states.finish())
376    }
377
378    fn push_result(&mut self, state: ArrayRef) -> VortexResult<()> {
379        vortex_ensure!(
380            state.dtype() == &self.partial_dtype,
381            "State DType mismatch: expected {}, got {}",
382            self.partial_dtype,
383            state.dtype()
384        );
385        self.partials.push(state);
386        Ok(())
387    }
388}
389fn list_view_group_ranges(
390    groups: &ListViewArray,
391    ctx: &mut ExecutionCtx,
392) -> VortexResult<GroupRanges> {
393    let offsets = groups.offsets();
394    let sizes = groups.sizes().cast(offsets.dtype().clone())?;
395
396    let ranges = match_each_integer_ptype!(offsets.dtype().as_ptype(), |O| {
397        let offsets = offsets.clone().execute::<Buffer<O>>(ctx)?;
398        let sizes = sizes.execute::<Buffer<O>>(ctx)?;
399        offsets
400            .as_ref()
401            .iter()
402            .zip(sizes.as_ref().iter())
403            .map(|(offset, size)| {
404                (
405                    offset.to_usize().vortex_expect("Offset value is not usize"),
406                    size.to_usize().vortex_expect("Size value is not usize"),
407                )
408            })
409            .collect::<Vec<_>>()
410    });
411
412    Ok(GroupRanges::ListView { ranges })
413}
414
415fn fixed_size_list_group_ranges(groups: &FixedSizeListArray) -> GroupRanges {
416    GroupRanges::FixedSizeList {
417        len: groups.len(),
418        size: groups.list_size() as usize,
419    }
420}
vortex_array/aggregate_fn/accumulator_grouped.rs

vortex_array/aggregate_fn/
accumulator_grouped.rs
