vortex_array/aggregate_fn/fns/last/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_session::registry::CachedId;
7
8use crate::ArrayRef;
9use crate::Columnar;
10use crate::ExecutionCtx;
11use crate::aggregate_fn::Accumulator;
12use crate::aggregate_fn::AggregateArgs;
13use crate::aggregate_fn::AggregateFnId;
14use crate::aggregate_fn::AggregateFnVTable;
15use crate::aggregate_fn::DynAccumulator;
16use crate::aggregate_fn::EmptyOptions;
17use crate::dtype::DType;
18use crate::scalar::Scalar;
19
20/// Return the last non-null value of an array.
21///
22/// See [`Last`] for details.
23pub fn last(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
24    let mut acc = Accumulator::try_new(Last, EmptyOptions, array.dtype().clone())?;
25    acc.accumulate(array, ctx)?;
26    acc.finish()
27}
28
29/// Return the last non-null value seen across all batches.
30#[derive(Clone, Debug)]
31pub struct Last;
32
33/// Partial accumulator state for the [`Last`] aggregate.
34pub struct LastPartial {
35    /// The last non-null value seen so far, or `None` if no non-null value has been observed.
36    value: Option<Scalar>,
37}
38
39impl AggregateFnVTable for Last {
40    type Options = EmptyOptions;
41    type Partial = LastPartial;
42
43    fn id(&self) -> AggregateFnId {
44        static ID: CachedId = CachedId::new("vortex.last");
45        *ID
46    }
47
48    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
49        vortex_bail!("Last is not yet serializable");
50    }
51
52    fn return_dtype(&self, _options: &Self::Options, input_dtype: &DType) -> Option<DType> {
53        Some(input_dtype.as_nullable())
54    }
55
56    fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option<DType> {
57        self.return_dtype(options, input_dtype)
58    }
59
60    fn empty_partial(
61        &self,
62        _args: AggregateArgs<'_, Self::Options>,
63    ) -> VortexResult<Self::Partial> {
64        Ok(LastPartial { value: None })
65    }
66
67    fn partial_from_scalar(
68        &self,
69        _args: AggregateArgs<'_, Self::Options>,
70        scalar: Scalar,
71    ) -> VortexResult<Self::Partial> {
72        // A null partial means the producing accumulator saw nothing valid.
73        Ok(LastPartial {
74            value: (!scalar.is_null()).then_some(scalar),
75        })
76    }
77
78    fn merge_partials(
79        &self,
80        _args: AggregateArgs<'_, Self::Options>,
81        first: Self::Partial,
82        second: Self::Partial,
83    ) -> VortexResult<Self::Partial> {
84        // The later value wins.
85        Ok(LastPartial {
86            value: second.value.or(first.value),
87        })
88    }
89
90    fn to_scalar(
91        &self,
92        args: AggregateArgs<'_, Self::Options>,
93        partial: &Self::Partial,
94    ) -> VortexResult<Scalar> {
95        Ok(match &partial.value {
96            Some(v) => v.clone(),
97            None => Scalar::null(args.return_dtype.clone()),
98        })
99    }
100
101    #[inline]
102    fn is_saturated(
103        &self,
104        _args: AggregateArgs<'_, Self::Options>,
105        _partial: &Self::Partial,
106    ) -> bool {
107        // Last can never short-circuit: a later batch can always supersede the current value.
108        false
109    }
110
111    fn try_accumulate(
112        &self,
113        _args: AggregateArgs<'_, Self::Options>,
114        partial: &mut Self::Partial,
115        batch: &ArrayRef,
116        ctx: &mut ExecutionCtx,
117    ) -> VortexResult<bool> {
118        if let Some(idx) = batch.validity()?.execute_mask(batch.len(), ctx)?.last() {
119            let scalar = batch.execute_scalar(idx, ctx)?;
120            partial.value = Some(scalar.into_nullable());
121        }
122        Ok(true)
123    }
124
125    fn accumulate(
126        &self,
127        _args: AggregateArgs<'_, Self::Options>,
128        _partial: &mut Self::Partial,
129        _batch: &Columnar,
130        _ctx: &mut ExecutionCtx,
131    ) -> VortexResult<()> {
132        unreachable!("Last::try_accumulate handles all arrays")
133    }
134
135    fn finalize(
136        &self,
137        _args: AggregateArgs<'_, Self::Options>,
138        partials: ArrayRef,
139    ) -> VortexResult<ArrayRef> {
140        Ok(partials)
141    }
142
143    fn finalize_scalar(
144        &self,
145        args: AggregateArgs<'_, Self::Options>,
146        partial: &Self::Partial,
147    ) -> VortexResult<Scalar> {
148        self.to_scalar(args, partial)
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use vortex_buffer::buffer;
155    use vortex_error::VortexResult;
156
157    use crate::IntoArray;
158    use crate::VortexSessionExecute;
159    use crate::aggregate_fn::Accumulator;
160    use crate::aggregate_fn::AggregateDTypes;
161    use crate::aggregate_fn::AggregateFnVTable;
162    use crate::aggregate_fn::DynAccumulator;
163    use crate::aggregate_fn::EmptyOptions;
164    use crate::aggregate_fn::fns::last::Last;
165    use crate::aggregate_fn::fns::last::last;
166    use crate::array_session;
167    use crate::arrays::ChunkedArray;
168    use crate::arrays::ConstantArray;
169    use crate::arrays::PrimitiveArray;
170    use crate::arrays::VarBinArray;
171    use crate::dtype::DType;
172    use crate::dtype::Nullability;
173    use crate::dtype::Nullability::Nullable;
174    use crate::dtype::PType;
175    use crate::scalar::Scalar;
176    use crate::validity::Validity;
177
178    #[test]
179    fn last_non_null() -> VortexResult<()> {
180        let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array();
181        let mut ctx = array_session().create_execution_ctx();
182        assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(30i32, Nullable));
183        Ok(())
184    }
185
186    #[test]
187    fn last_skips_trailing_nulls() -> VortexResult<()> {
188        let array =
189            PrimitiveArray::from_option_iter([Some(7i32), Some(8), None, None]).into_array();
190        let mut ctx = array_session().create_execution_ctx();
191        assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(8i32, Nullable));
192        Ok(())
193    }
194
195    #[test]
196    fn last_all_null() -> VortexResult<()> {
197        let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]).into_array();
198        let mut ctx = array_session().create_execution_ctx();
199        let dtype = DType::Primitive(PType::I32, Nullable);
200        assert_eq!(last(&array, &mut ctx)?, Scalar::null(dtype));
201        Ok(())
202    }
203
204    #[test]
205    fn last_empty() -> VortexResult<()> {
206        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
207        let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?;
208        let result = acc.finish()?;
209        assert_eq!(result, Scalar::null(DType::Primitive(PType::I32, Nullable)));
210        Ok(())
211    }
212
213    #[test]
214    fn last_constant() -> VortexResult<()> {
215        let array = ConstantArray::new(42i32, 10).into_array();
216        let mut ctx = array_session().create_execution_ctx();
217        assert_eq!(last(&array, &mut ctx)?, Scalar::primitive(42i32, Nullable));
218        Ok(())
219    }
220
221    #[test]
222    fn last_constant_null() -> VortexResult<()> {
223        let dtype = DType::Primitive(PType::I32, Nullable);
224        let array = ConstantArray::new(Scalar::null(dtype.clone()), 10).into_array();
225        let mut ctx = array_session().create_execution_ctx();
226        assert_eq!(last(&array, &mut ctx)?, Scalar::null(dtype));
227        Ok(())
228    }
229
230    #[test]
231    fn last_varbin() -> VortexResult<()> {
232        let array = VarBinArray::from_iter(
233            vec![Some("hello"), Some("world"), None],
234            DType::Utf8(Nullable),
235        )
236        .into_array();
237        let mut ctx = array_session().create_execution_ctx();
238        assert_eq!(last(&array, &mut ctx)?, Scalar::utf8("world", Nullable));
239        Ok(())
240    }
241
242    #[test]
243    fn last_multi_batch_picks_latest_non_null() -> VortexResult<()> {
244        let mut ctx = array_session().create_execution_ctx();
245        let dtype = DType::Primitive(PType::I32, Nullable);
246        let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?;
247
248        let batch1 = PrimitiveArray::from_option_iter([Some(1i32), Some(2)]).into_array();
249        acc.accumulate(&batch1, &mut ctx)?;
250
251        // All-null batch must not clobber the previously-stored value.
252        let batch2 = PrimitiveArray::from_option_iter::<i32, _>([None, None]).into_array();
253        acc.accumulate(&batch2, &mut ctx)?;
254
255        let batch3 = PrimitiveArray::from_option_iter([Some(99i32), None]).into_array();
256        acc.accumulate(&batch3, &mut ctx)?;
257
258        // Last is never saturated; later batches keep updating it.
259        assert!(!acc.is_saturated());
260
261        let result = acc.finish()?;
262        assert_eq!(result, Scalar::primitive(99i32, Nullable));
263        Ok(())
264    }
265
266    #[test]
267    fn last_finish_resets_state() -> VortexResult<()> {
268        let mut ctx = array_session().create_execution_ctx();
269        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
270        let mut acc = Accumulator::try_new(Last, EmptyOptions, dtype)?;
271
272        let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array();
273        acc.accumulate(&batch1, &mut ctx)?;
274        assert_eq!(acc.finish()?, Scalar::primitive(20i32, Nullable));
275
276        let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array();
277        acc.accumulate(&batch2, &mut ctx)?;
278        assert_eq!(acc.finish()?, Scalar::primitive(9i32, Nullable));
279        Ok(())
280    }
281
282    #[test]
283    fn last_state_merge() -> VortexResult<()> {
284        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
285        let owned = AggregateDTypes::try_new(&Last, &EmptyOptions, dtype)?;
286        let args = owned.args(&EmptyOptions);
287        let partial_of =
288            |value: i32| Last.partial_from_scalar(args, Scalar::primitive(value, Nullable));
289
290        let five = partial_of(5)?;
291        let seven = partial_of(7)?;
292        let empty = Last.empty_partial(args)?;
293
294        // The last non-empty partial wins, and a trailing empty one changes nothing.
295        let merge = |first, second| Last.merge_partials(args, first, second);
296        let state = merge(merge(five, seven)?, empty)?;
297        assert_eq!(
298            Last.to_scalar(args, &state)?,
299            Scalar::primitive(7i32, Nullable)
300        );
301        Ok(())
302    }
303
304    #[test]
305    fn last_chunked() -> VortexResult<()> {
306        let chunk1 = PrimitiveArray::from_option_iter([Some(42i32), Some(100)]);
307        let chunk2 = PrimitiveArray::from_option_iter::<i32, _>([None, None]);
308        let dtype = chunk1.dtype().clone();
309        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
310        let mut ctx = array_session().create_execution_ctx();
311        assert_eq!(
312            last(&chunked.into_array(), &mut ctx)?,
313            Scalar::primitive(100i32, Nullable)
314        );
315        Ok(())
316    }
317}
vortex_array/aggregate_fn/fns/last/mod.rs

vortex_array/aggregate_fn/fns/last/
mod.rs
