vortex_array/aggregate_fn/
storage_pushdown.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Opt-in pushdown of an aggregate into the storage of one extension type.
5//!
6//! See [`ExtensionStoragePushdown`] for when such a pushdown is correct.
7
8use std::mem;
9
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_err;
13use vortex_session::VortexSession;
14
15use crate::ArrayRef;
16use crate::Canonical;
17use crate::Columnar;
18use crate::ExecutionCtx;
19use crate::IntoArray;
20use crate::aggregate_fn::Accumulator;
21use crate::aggregate_fn::AggregateArgs;
22use crate::aggregate_fn::AggregateDTypes;
23use crate::aggregate_fn::AggregateFnId;
24use crate::aggregate_fn::AggregateFnVTable;
25use crate::aggregate_fn::DynAccumulator;
26use crate::arrays::ConstantArray;
27use crate::arrays::Extension;
28use crate::arrays::extension::ExtensionArrayExt;
29use crate::dtype::DType;
30use crate::dtype::extension::ExtVTable;
31use crate::scalar::Scalar;
32
33/// Runs aggregate `A` on the storage values of extension type `E`.
34///
35/// The input must have extension type `E`, and `A` must support its storage dtype. Binding fails
36/// for any other input. The wrapper has a separate aggregate ID, `<A>[<E>]`, formed from the IDs
37/// of `A` and `E`. Registering it does not replace `A`.
38///
39/// Aggregates that already support extension dtypes directly do not need this wrapper.
40///
41/// # Correctness
42///
43/// A registration **must** satisfy both conditions:
44///
45/// - The aggregation is defined for logical values of `E`.
46/// - For every valid input, aggregating the storage values produces the storage representation
47///   of the logical aggregate result. This requires storage to preserve whichever properties `A`
48///   relies on, such as ordering, equality, arithmetic, or NaN classification.
49///
50/// For example, minimum is defined for timestamps, and their integer storage preserves their
51/// ordering, so both conditions hold. Summing timestamps fails the first condition even though
52/// their storage supports addition.
53///
54/// These conditions are not checked. An invalid registration can silently return incorrect results.
55///
56/// # Results and partials
57///
58/// Results and partial states are those of `A` over the storage dtype. Results are not converted
59/// back to `E`. For example, the minimum and maximum of a timestamp column are raw `i64` values
60/// in the column's time unit.
61///
62/// Partial states can only be combined when the aggregate, options, and full input dtype match.
63/// This includes extension metadata, so timestamp partials with different time units must not be
64/// combined even if their storage dtypes match.
65///
66/// # Example
67///
68/// Timestamps compare like their `i64` storage, so `is_sorted` can be pushed down. Register the
69/// pair once on the session, then bind it like any other aggregate:
70///
71/// ```
72/// # use vortex_array::IntoArray;
73/// # use vortex_array::VortexSessionExecute;
74/// # use vortex_array::aggregate_fn::Accumulator;
75/// # use vortex_array::aggregate_fn::DynAccumulator;
76/// # use vortex_array::aggregate_fn::ExtensionStoragePushdown;
77/// # use vortex_array::aggregate_fn::fns::is_sorted::IsSorted;
78/// # use vortex_array::aggregate_fn::fns::is_sorted::IsSortedOptions;
79/// # use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
80/// # use vortex_array::arrays::ExtensionArray;
81/// # use vortex_array::arrays::PrimitiveArray;
82/// # use vortex_array::dtype::DType;
83/// # use vortex_array::dtype::Nullability;
84/// # use vortex_array::extension::datetime::TimeUnit;
85/// # use vortex_array::extension::datetime::Timestamp;
86/// # fn main() -> vortex_error::VortexResult<()> {
87/// let session = vortex_array::array_session();
88///
89/// // Register `vortex.is_sorted[vortex.timestamp]`.
90/// let is_sorted = ExtensionStoragePushdown::new(IsSorted, Timestamp);
91/// session.aggregate_fns().register(is_sorted.clone());
92///
93/// // Aggregate a timestamp column through its storage.
94/// let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased();
95/// let column = ExtensionArray::new(
96///     ext_dtype.clone(),
97///     PrimitiveArray::from_iter([1_000i64, 2_000, 3_000]).into_array(),
98/// )
99/// .into_array();
100///
101/// let mut acc = Accumulator::try_new(
102///     is_sorted,
103///     IsSortedOptions { strict: true },
104///     DType::Extension(ext_dtype),
105/// )?;
106/// acc.accumulate(&column, &mut session.create_execution_ctx())?;
107/// assert!(bool::try_from(&acc.finish()?)?);
108/// # Ok(())
109/// # }
110/// ```
111#[derive(Clone, Debug)]
112pub struct ExtensionStoragePushdown<A: AggregateFnVTable, E: ExtVTable> {
113    aggregate: A,
114    ext: E,
115    id: AggregateFnId,
116}
117
118impl<A: AggregateFnVTable, E: ExtVTable> ExtensionStoragePushdown<A, E> {
119    /// Push `aggregate` down into the storage of extension type `ext`.
120    ///
121    /// Only do this when `aggregate` over the storage is correct for `ext`; see
122    /// [`ExtensionStoragePushdown`] for how to check.
123    pub fn new(aggregate: A, ext: E) -> Self {
124        #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
125        let id = AggregateFnId::new(&format!("{}[{}]", aggregate.id(), ext.id()));
126        Self { aggregate, ext, id }
127    }
128
129    /// The storage dtype of `dtype`, if it is an extension dtype of type `E`.
130    fn storage_dtype<'a>(&self, dtype: &'a DType) -> Option<&'a DType> {
131        dtype
132            .as_extension_opt()
133            .filter(|ext_dtype| ext_dtype.is::<E>())
134            .map(|ext_dtype| ext_dtype.storage_dtype())
135    }
136
137    /// The arguments of the storage aggregate `A`.
138    ///
139    /// `A` shares this aggregate's options, return and partial dtypes, and its input dtype is
140    /// borrowed from the extension dtype, so this allocates nothing.
141    fn storage_args<'a>(
142        &self,
143        args: AggregateArgs<'a, A::Options>,
144    ) -> VortexResult<AggregateArgs<'a, A::Options>> {
145        let dtype = self.storage_dtype(args.dtype).ok_or_else(|| {
146            vortex_err!(
147                "Aggregate function {} expects extension type {}, got {}",
148                self.id,
149                self.ext.id(),
150                args.dtype
151            )
152        })?;
153        Ok(AggregateArgs { dtype, ..args })
154    }
155
156    /// Accumulate a storage batch into `state` through an [`Accumulator`] of `A`, so that the
157    /// kernels and cached statistics of `A` on the storage encoding still apply.
158    fn accumulate_storage(
159        &self,
160        args: AggregateArgs<'_, A::Options>,
161        state: &mut A::Partial,
162        storage: &ArrayRef,
163        ctx: &mut ExecutionCtx,
164    ) -> VortexResult<()> {
165        let mut accumulator = Accumulator::from_dtypes(
166            self.aggregate.clone(),
167            args.options.clone(),
168            AggregateDTypes::new(
169                args.dtype.clone(),
170                args.return_dtype.clone(),
171                args.partial_dtype.clone(),
172            ),
173        );
174        accumulator.accumulate(storage, ctx)?;
175        let Some(batch) = accumulator.take_partial() else {
176            return Ok(());
177        };
178        let running = mem::replace(state, self.aggregate.empty_partial(args)?);
179        *state = self.aggregate.merge_partials(args, running, batch)?;
180        Ok(())
181    }
182}
183
184impl<A: AggregateFnVTable, E: ExtVTable> AggregateFnVTable for ExtensionStoragePushdown<A, E> {
185    type Options = A::Options;
186    type Partial = A::Partial;
187
188    fn id(&self) -> AggregateFnId {
189        self.id
190    }
191
192    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
193        self.aggregate.serialize(options)
194    }
195
196    fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult<Self::Options> {
197        self.aggregate.deserialize(metadata, session)
198    }
199
200    fn return_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
201        self.aggregate
202            .return_dtype(options, self.storage_dtype(input_dtype)?)
203    }
204
205    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
206        self.aggregate
207            .partial_dtype(options, self.storage_dtype(input_dtype)?)
208    }
209
210    fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult<Self::Partial> {
211        self.aggregate.empty_partial(self.storage_args(args)?)
212    }
213
214    fn partial_from_scalar(
215        &self,
216        args: AggregateArgs<'_, Self::Options>,
217        scalar: &Scalar,
218    ) -> VortexResult<Self::Partial> {
219        self.aggregate
220            .partial_from_scalar(self.storage_args(args)?, scalar)
221    }
222
223    fn merge_partials(
224        &self,
225        args: AggregateArgs<'_, Self::Options>,
226        first: Self::Partial,
227        second: Self::Partial,
228    ) -> VortexResult<Self::Partial> {
229        self.aggregate
230            .merge_partials(self.storage_args(args)?, first, second)
231    }
232
233    fn to_scalar(
234        &self,
235        args: AggregateArgs<'_, Self::Options>,
236        partial: &Self::Partial,
237    ) -> VortexResult<Scalar> {
238        self.aggregate.to_scalar(self.storage_args(args)?, partial)
239    }
240
241    fn is_saturated(
242        &self,
243        args: AggregateArgs<'_, Self::Options>,
244        partial: &Self::Partial,
245    ) -> bool {
246        self.storage_args(args)
247            .is_ok_and(|args| self.aggregate.is_saturated(args, partial))
248    }
249
250    /// Unwrap extension arrays to their storage before decoding, so the storage accumulator sees
251    /// the storage encoding.
252    fn try_accumulate(
253        &self,
254        args: AggregateArgs<'_, Self::Options>,
255        state: &mut Self::Partial,
256        batch: &ArrayRef,
257        ctx: &mut ExecutionCtx,
258    ) -> VortexResult<bool> {
259        let Some(array) = batch.as_opt::<Extension>() else {
260            return Ok(false);
261        };
262        self.accumulate_storage(self.storage_args(args)?, state, array.storage_array(), ctx)?;
263        Ok(true)
264    }
265
266    fn accumulate(
267        &self,
268        args: AggregateArgs<'_, Self::Options>,
269        state: &mut Self::Partial,
270        batch: &Columnar,
271        ctx: &mut ExecutionCtx,
272    ) -> VortexResult<()> {
273        let args = self.storage_args(args)?;
274        match batch {
275            Columnar::Canonical(Canonical::Extension(array)) => {
276                self.accumulate_storage(args, state, array.storage_array(), ctx)
277            }
278            Columnar::Constant(constant) => {
279                let storage = constant.scalar().as_extension().to_storage_scalar();
280                self.accumulate_storage(
281                    args,
282                    state,
283                    &ConstantArray::new(storage, constant.len()).into_array(),
284                    ctx,
285                )
286            }
287            Columnar::Canonical(_) => {
288                vortex_bail!(
289                    "Aggregate function {} expects an extension array, got {}",
290                    self.id,
291                    batch.dtype()
292                )
293            }
294        }
295    }
296
297    fn finalize(
298        &self,
299        args: AggregateArgs<'_, Self::Options>,
300        states: ArrayRef,
301    ) -> VortexResult<ArrayRef> {
302        self.aggregate.finalize(self.storage_args(args)?, states)
303    }
304
305    fn finalize_scalar(
306        &self,
307        args: AggregateArgs<'_, Self::Options>,
308        partial: &Self::Partial,
309    ) -> VortexResult<Scalar> {
310        self.aggregate
311            .finalize_scalar(self.storage_args(args)?, partial)
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use vortex_error::VortexResult;
318
319    use super::ExtensionStoragePushdown;
320    use crate::IntoArray;
321    use crate::VortexSessionExecute;
322    use crate::aggregate_fn::Accumulator;
323    use crate::aggregate_fn::AggregateFnVTable;
324    use crate::aggregate_fn::DynAccumulator;
325    use crate::aggregate_fn::NumericalAggregateOpts;
326    use crate::aggregate_fn::fns::min_max::MinMax;
327    use crate::aggregate_fn::session::AggregateFnSessionExt;
328    use crate::arrays::ExtensionArray;
329    use crate::arrays::PrimitiveArray;
330    use crate::dtype::DType;
331    use crate::dtype::Nullability;
332    use crate::extension::datetime::Date;
333    use crate::extension::datetime::TimeUnit;
334    use crate::extension::datetime::Timestamp;
335
336    #[test]
337    fn pushes_down_into_storage() -> VortexResult<()> {
338        let session = crate::array_session();
339        let aggregate = ExtensionStoragePushdown::new(MinMax, Timestamp);
340        session.aggregate_fns().register(aggregate.clone());
341        assert!(
342            session
343                .aggregate_fns()
344                .find_plugin(&aggregate.id())
345                .is_some()
346        );
347
348        let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased();
349        let storage =
350            PrimitiveArray::from_option_iter([Some(3_000i64), None, Some(1_000)]).into_array();
351        let array = ExtensionArray::new(ext_dtype.clone(), storage.clone()).into_array();
352
353        let mut ctx = session.create_execution_ctx();
354        let mut acc = Accumulator::try_new(
355            aggregate,
356            NumericalAggregateOpts::default(),
357            DType::Extension(ext_dtype),
358        )?;
359        acc.accumulate(&array, &mut ctx)?;
360
361        let mut expected = Accumulator::try_new(
362            MinMax,
363            NumericalAggregateOpts::default(),
364            storage.dtype().clone(),
365        )?;
366        expected.accumulate(&storage, &mut ctx)?;
367
368        assert_eq!(acc.finish()?, expected.finish()?);
369        Ok(())
370    }
371
372    #[test]
373    fn combines_flushed_partials() -> VortexResult<()> {
374        let session = crate::array_session();
375        let mut ctx = session.create_execution_ctx();
376        let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::Nullable).erased();
377        let batch = |values: [Option<i64>; 2]| {
378            ExtensionArray::new(
379                ext_dtype.clone(),
380                PrimitiveArray::from_option_iter(values).into_array(),
381            )
382            .into_array()
383        };
384        let accumulator = || {
385            Accumulator::try_new(
386                ExtensionStoragePushdown::new(MinMax, Timestamp),
387                NumericalAggregateOpts::default(),
388                DType::Extension(ext_dtype.clone()),
389            )
390        };
391
392        let mut first = accumulator()?;
393        first.accumulate(&batch([Some(3_000), None]), &mut ctx)?;
394        let mut second = accumulator()?;
395        second.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?;
396        first.combine_partials(&second.flush()?)?;
397
398        let mut expected = accumulator()?;
399        expected.accumulate(&batch([Some(3_000), None]), &mut ctx)?;
400        expected.accumulate(&batch([Some(1_000), Some(500)]), &mut ctx)?;
401
402        assert_eq!(first.finish()?, expected.finish()?);
403        Ok(())
404    }
405
406    #[test]
407    fn rejects_other_extension_types() {
408        let date = DType::Extension(Date::new(TimeUnit::Days, Nullability::Nullable).erased());
409        assert!(
410            ExtensionStoragePushdown::new(MinMax, Timestamp)
411                .return_dtype(&NumericalAggregateOpts::default(), &date)
412                .is_none()
413        );
414    }
415}
vortex_array/aggregate_fn/storage_pushdown.rs

vortex_array/aggregate_fn/
storage_pushdown.rs
