vortex_array/aggregate_fn/fns/first/
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 first non-null value of an array.
21///
22/// See [`First`] for details.
23pub fn first(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Scalar> {
24    let mut acc = Accumulator::try_new(First, EmptyOptions, array.dtype().clone())?;
25    acc.accumulate(array, ctx)?;
26    acc.finish()
27}
28
29/// Return the first non-null value seen across all batches.
30#[derive(Clone, Debug)]
31pub struct First;
32
33/// Partial accumulator state for the [`First`] aggregate.
34pub struct FirstPartial {
35    /// The first 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 First {
40    type Options = EmptyOptions;
41    type Partial = FirstPartial;
42
43    fn id(&self) -> AggregateFnId {
44        static ID: CachedId = CachedId::new("vortex.first");
45        *ID
46    }
47
48    fn serialize(&self, _options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
49        vortex_bail!("First 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(FirstPartial { 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(FirstPartial {
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 earlier value wins.
85        Ok(FirstPartial {
86            value: first.value.or(second.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        partial.value.is_some()
108    }
109
110    fn try_accumulate(
111        &self,
112        _args: AggregateArgs<'_, Self::Options>,
113        partial: &mut Self::Partial,
114        batch: &ArrayRef,
115        ctx: &mut ExecutionCtx,
116    ) -> VortexResult<bool> {
117        if partial.value.is_some() {
118            return Ok(true);
119        }
120        if let Some(idx) = batch.validity()?.execute_mask(batch.len(), ctx)?.first() {
121            let scalar = batch.execute_scalar(idx, ctx)?;
122            partial.value = Some(scalar.into_nullable());
123        }
124        Ok(true)
125    }
126
127    fn accumulate(
128        &self,
129        _args: AggregateArgs<'_, Self::Options>,
130        _partial: &mut Self::Partial,
131        _batch: &Columnar,
132        _ctx: &mut ExecutionCtx,
133    ) -> VortexResult<()> {
134        unreachable!("First::try_accumulate handles all arrays")
135    }
136
137    fn finalize(
138        &self,
139        _args: AggregateArgs<'_, Self::Options>,
140        partials: ArrayRef,
141    ) -> VortexResult<ArrayRef> {
142        Ok(partials)
143    }
144
145    fn finalize_scalar(
146        &self,
147        args: AggregateArgs<'_, Self::Options>,
148        partial: &Self::Partial,
149    ) -> VortexResult<Scalar> {
150        self.to_scalar(args, partial)
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use vortex_buffer::buffer;
157    use vortex_error::VortexResult;
158
159    use crate::IntoArray;
160    use crate::VortexSessionExecute;
161    use crate::aggregate_fn::Accumulator;
162    use crate::aggregate_fn::AggregateDTypes;
163    use crate::aggregate_fn::AggregateFnVTable;
164    use crate::aggregate_fn::DynAccumulator;
165    use crate::aggregate_fn::EmptyOptions;
166    use crate::aggregate_fn::fns::first::First;
167    use crate::aggregate_fn::fns::first::first;
168    use crate::array_session;
169    use crate::arrays::ChunkedArray;
170    use crate::arrays::ConstantArray;
171    use crate::arrays::PrimitiveArray;
172    use crate::arrays::VarBinArray;
173    use crate::dtype::DType;
174    use crate::dtype::Nullability;
175    use crate::dtype::Nullability::Nullable;
176    use crate::dtype::PType;
177    use crate::scalar::Scalar;
178    use crate::validity::Validity;
179
180    #[test]
181    fn first_non_null() -> VortexResult<()> {
182        let array = PrimitiveArray::new(buffer![10i32, 20, 30], Validity::NonNullable).into_array();
183        let mut ctx = array_session().create_execution_ctx();
184        assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(10i32, Nullable));
185        Ok(())
186    }
187
188    #[test]
189    fn first_skips_leading_nulls() -> VortexResult<()> {
190        let array =
191            PrimitiveArray::from_option_iter([None, None, Some(7i32), Some(8)]).into_array();
192        let mut ctx = array_session().create_execution_ctx();
193        assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(7i32, Nullable));
194        Ok(())
195    }
196
197    #[test]
198    fn first_all_null() -> VortexResult<()> {
199        let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]).into_array();
200        let mut ctx = array_session().create_execution_ctx();
201        let dtype = DType::Primitive(PType::I32, Nullable);
202        assert_eq!(first(&array, &mut ctx)?, Scalar::null(dtype));
203        Ok(())
204    }
205
206    #[test]
207    fn first_empty() -> VortexResult<()> {
208        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
209        let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?;
210        let result = acc.finish()?;
211        assert_eq!(result, Scalar::null(DType::Primitive(PType::I32, Nullable)));
212        Ok(())
213    }
214
215    #[test]
216    fn first_constant() -> VortexResult<()> {
217        let array = ConstantArray::new(42i32, 10).into_array();
218        let mut ctx = array_session().create_execution_ctx();
219        assert_eq!(first(&array, &mut ctx)?, Scalar::primitive(42i32, Nullable));
220        Ok(())
221    }
222
223    #[test]
224    fn first_constant_null() -> VortexResult<()> {
225        let dtype = DType::Primitive(PType::I32, Nullable);
226        let array = ConstantArray::new(Scalar::null(dtype.clone()), 10).into_array();
227        let mut ctx = array_session().create_execution_ctx();
228        assert_eq!(first(&array, &mut ctx)?, Scalar::null(dtype));
229        Ok(())
230    }
231
232    #[test]
233    fn first_varbin() -> VortexResult<()> {
234        let array = VarBinArray::from_iter(
235            vec![None, Some("hello"), Some("world")],
236            DType::Utf8(Nullable),
237        )
238        .into_array();
239        let mut ctx = array_session().create_execution_ctx();
240        assert_eq!(first(&array, &mut ctx)?, Scalar::utf8("hello", Nullable));
241        Ok(())
242    }
243
244    #[test]
245    fn first_multi_batch_picks_earliest_non_null() -> VortexResult<()> {
246        let mut ctx = array_session().create_execution_ctx();
247        let dtype = DType::Primitive(PType::I32, Nullable);
248        let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?;
249
250        // First batch is all null - should not saturate.
251        let batch1 = PrimitiveArray::from_option_iter::<i32, _>([None, None]).into_array();
252        acc.accumulate(&batch1, &mut ctx)?;
253        assert!(!acc.is_saturated());
254
255        // Second batch contains the first non-null value.
256        let batch2 = PrimitiveArray::from_option_iter([None, Some(99i32), Some(100)]).into_array();
257        acc.accumulate(&batch2, &mut ctx)?;
258        assert!(acc.is_saturated());
259
260        // Third batch must be ignored - First is already saturated.
261        let batch3 = PrimitiveArray::from_option_iter([Some(1i32)]).into_array();
262        acc.accumulate(&batch3, &mut ctx)?;
263
264        let result = acc.finish()?;
265        assert_eq!(result, Scalar::primitive(99i32, Nullable));
266        Ok(())
267    }
268
269    #[test]
270    fn first_finish_resets_state() -> VortexResult<()> {
271        let mut ctx = array_session().create_execution_ctx();
272        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
273        let mut acc = Accumulator::try_new(First, EmptyOptions, dtype)?;
274
275        let batch1 = PrimitiveArray::new(buffer![10i32, 20], Validity::NonNullable).into_array();
276        acc.accumulate(&batch1, &mut ctx)?;
277        assert_eq!(acc.finish()?, Scalar::primitive(10i32, Nullable));
278
279        let batch2 = PrimitiveArray::new(buffer![3i32, 6, 9], Validity::NonNullable).into_array();
280        acc.accumulate(&batch2, &mut ctx)?;
281        assert_eq!(acc.finish()?, Scalar::primitive(3i32, Nullable));
282        Ok(())
283    }
284
285    #[test]
286    fn first_state_merge() -> VortexResult<()> {
287        let dtype = DType::Primitive(PType::I32, Nullability::NonNullable);
288        let owned = AggregateDTypes::try_new(&First, &EmptyOptions, dtype)?;
289        let args = owned.args(&EmptyOptions);
290        let partial_of =
291            |value: i32| First.partial_from_scalar(args, Scalar::primitive(value, Nullable));
292
293        let empty = First.empty_partial(args)?;
294        assert!(!First.is_saturated(args, &empty));
295
296        // The first non-empty partial wins; later ones are dropped.
297        let five = partial_of(5)?;
298        let seven = partial_of(7)?;
299        let merge = |first, second| First.merge_partials(args, first, second);
300        let state = merge(merge(empty, five)?, seven)?;
301        assert!(First.is_saturated(args, &state));
302        assert_eq!(
303            First.to_scalar(args, &state)?,
304            Scalar::primitive(5i32, Nullable)
305        );
306        Ok(())
307    }
308
309    #[test]
310    fn first_chunked() -> VortexResult<()> {
311        let chunk1 = PrimitiveArray::from_option_iter::<i32, _>([None, None]);
312        let chunk2 = PrimitiveArray::from_option_iter([None, Some(42i32), Some(100)]);
313        let dtype = chunk1.dtype().clone();
314        let chunked = ChunkedArray::try_new(vec![chunk1.into_array(), chunk2.into_array()], dtype)?;
315        let mut ctx = array_session().create_execution_ctx();
316        assert_eq!(
317            first(&chunked.into_array(), &mut ctx)?,
318            Scalar::primitive(42i32, Nullable)
319        );
320        Ok(())
321    }
322}
vortex_array/aggregate_fn/fns/first/mod.rs

vortex_array/aggregate_fn/fns/first/
mod.rs
