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

vortex_array/aggregate_fn/
accumulator_grouped.rs
