vortex_array/scalar_fn/fns/
list_length.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use num_traits::AsPrimitive;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7use vortex_session::VortexSession;
8use vortex_session::registry::CachedId;
9
10use crate::ArrayRef;
11use crate::ExecutionCtx;
12use crate::IntoArray;
13use crate::array::ArrayView;
14use crate::arrays::ConstantArray;
15use crate::arrays::FixedSizeList;
16use crate::arrays::List;
17use crate::arrays::ListView;
18use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
19use crate::arrays::list::ListArrayExt;
20use crate::arrays::list::ListArraySlotsExt;
21use crate::arrays::listview::ListViewArrayExt;
22use crate::arrays::listview::ListViewArraySlotsExt;
23use crate::builtins::ArrayBuiltins;
24use crate::dtype::DType;
25use crate::dtype::Nullability;
26use crate::dtype::PType;
27use crate::expr::Expression;
28use crate::matcher::Matcher;
29use crate::scalar::Scalar;
30use crate::scalar_fn::Arity;
31use crate::scalar_fn::ChildName;
32use crate::scalar_fn::EmptyOptions;
33use crate::scalar_fn::ExecutionArgs;
34use crate::scalar_fn::ScalarFnId;
35use crate::scalar_fn::ScalarFnVTable;
36use crate::scalar_fn::fns::operators::Operator;
37
38/// Number of elements in each list of a `List` or `FixedSizeList` typed array.
39///
40/// This is computed purely from the list's offsets (`ListArray`), sizes (`ListViewArray`), or
41/// dtype (`FixedSizeListArray`) without reading the element *values*. Validity is carried over
42/// from the original array.
43#[derive(Clone)]
44pub struct ListLength;
45
46impl ScalarFnVTable for ListLength {
47    type Options = EmptyOptions;
48
49    fn id(&self) -> ScalarFnId {
50        static ID: CachedId = CachedId::new("vortex.list.length");
51        *ID
52    }
53
54    fn serialize(&self, _instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
55        Ok(Some(vec![]))
56    }
57
58    fn deserialize(
59        &self,
60        _metadata: &[u8],
61        _session: &VortexSession,
62    ) -> VortexResult<Self::Options> {
63        Ok(EmptyOptions)
64    }
65
66    fn arity(&self, _options: &Self::Options) -> Arity {
67        Arity::Exact(1)
68    }
69
70    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
71        match child_idx {
72            0 => ChildName::from("input"),
73            _ => unreachable!("Invalid child index {child_idx} for list_length()"),
74        }
75    }
76
77    fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
78        match &arg_dtypes[0] {
79            DType::List(_, nullable) | DType::FixedSizeList(_, _, nullable) => {
80                Ok(DType::Primitive(PType::U64, *nullable))
81            }
82            other => vortex_bail!("list_length() requires List or FixedSizeList, got {other}"),
83        }
84    }
85
86    fn execute(
87        &self,
88        _options: &Self::Options,
89        args: &dyn ExecutionArgs,
90        ctx: &mut ExecutionCtx,
91    ) -> VortexResult<ArrayRef> {
92        let input = args.get(0)?;
93        let nullability = input.dtype().nullability();
94
95        if let Some(scalar) = input.as_constant() {
96            let len_scalar = scalar_list_length(&scalar, nullability)?;
97            return Ok(ConstantArray::new(len_scalar, args.row_count()).into_array());
98        }
99
100        list_length(&input, nullability, ctx)
101    }
102
103    fn validity(
104        &self,
105        _: &Self::Options,
106        expression: &Expression,
107    ) -> VortexResult<Option<Expression>> {
108        Ok(Some(expression.child(0).validity()?))
109    }
110
111    fn is_strict(&self, _options: &Self::Options) -> bool {
112        // A null list has a null length, and the length of a valid list is a non-null value
113        // determined by that list alone, with `return_dtype` carrying over the input nullability.
114        true
115    }
116
117    fn is_fallible(&self, _options: &Self::Options) -> bool {
118        false
119    }
120}
121
122fn scalar_list_length(scalar: &Scalar, nullability: Nullability) -> VortexResult<Scalar> {
123    if scalar.is_null() {
124        let dtype = DType::Primitive(PType::U64, Nullability::Nullable);
125        return Ok(Scalar::null(dtype));
126    }
127    let len: u64 = scalar.as_list().len().as_();
128    Ok(Scalar::primitive(len, nullability))
129}
130
131pub(crate) fn list_length(
132    array: &ArrayRef,
133    nullability: Nullability,
134    ctx: &mut ExecutionCtx,
135) -> VortexResult<ArrayRef> {
136    let any_list = array.clone().execute_until::<AnyList>(ctx)?;
137
138    let (lengths, validity) = if let Some(fsl) = any_list.as_opt::<FixedSizeList>() {
139        // The length of fixed-size list is constant, so just need to carry over validity
140        let size = fsl.list_size() as u64;
141        let lengths =
142            ConstantArray::new(Scalar::primitive(size, Nullability::NonNullable), fsl.len())
143                .into_array();
144        (lengths, fsl.validity()?)
145    } else if let Some(lv) = any_list.as_opt::<ListView>() {
146        // Length array is exactly the sizes child
147        (lv.sizes().clone(), lv.listview_validity())
148    } else if let Some(l) = any_list.as_opt::<List>() {
149        let lengths = list_length_from_offsets(l)?;
150        (lengths, l.list_validity())
151    } else {
152        let dtype = any_list.dtype();
153        vortex_bail!("list_length() requires List, ListView, or FixedSizeList but got {dtype}")
154    };
155
156    // Cast to `U64`
157    let len = lengths.len();
158    let lengths = lengths.cast(DType::Primitive(PType::U64, nullability))?;
159
160    // Carry over validity mask for nullable arrays
161    if matches!(nullability, Nullability::Nullable) {
162        lengths.mask(validity.to_array(len))
163    } else {
164        Ok(lengths)
165    }
166}
167
168/// Calculate the lengths of `ListArray` elements via the `offsets` child:
169/// `length[i] = offsets[i + 1] - offsets[i]`.
170fn list_length_from_offsets(list: ArrayView<'_, List>) -> VortexResult<ArrayRef> {
171    let offsets = list.offsets();
172    let n = offsets.len().saturating_sub(1);
173
174    offsets
175        .slice(1..offsets.len())?
176        .binary(offsets.slice(0..n)?, Operator::Sub)
177}
178
179/// Matches an `Array<List>`, `Array<ListView>`, or `Array<FixedSizeList>`
180struct AnyList;
181
182impl Matcher for AnyList {
183    type Match<'a> = ();
184
185    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
186        (array.as_opt::<List>().is_some()
187            || array.as_opt::<ListView>().is_some()
188            || array.as_opt::<FixedSizeList>().is_some())
189        .then_some(())
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use std::sync::Arc;
196
197    use rstest::rstest;
198    use vortex_buffer::buffer;
199    use vortex_error::VortexResult;
200
201    use crate::ArrayRef;
202    use crate::IntoArray;
203    use crate::VortexSessionExecute;
204    use crate::array_session;
205    use crate::arrays::BoolArray;
206    use crate::arrays::ConstantArray;
207    use crate::arrays::FixedSizeListArray;
208    use crate::arrays::ListArray;
209    use crate::arrays::ListViewArray;
210    use crate::arrays::PrimitiveArray;
211    use crate::assert_arrays_eq;
212    use crate::dtype::DType;
213    use crate::dtype::Nullability;
214    use crate::dtype::PType;
215    use crate::expr::cast;
216    use crate::expr::list_length;
217    use crate::expr::root;
218    use crate::scalar::Scalar;
219    use crate::validity::Validity;
220
221    fn create_list_elements() -> ArrayRef {
222        PrimitiveArray::from_option_iter::<i32, _>([
223            Some(1),
224            Some(2),
225            Some(3),
226            Some(4),
227            Some(5),
228            Some(6),
229            None,
230        ])
231        .into_array()
232    }
233
234    #[rstest]
235    #[case(buffer![0u32, 2, 5, 5, 7].into_array())]
236    #[case(buffer![0u64, 2, 5, 5, 7].into_array())]
237    fn test_list_length(#[case] offsets: ArrayRef) -> VortexResult<()> {
238        let elements = create_list_elements();
239        let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array();
240        let result = list.apply(&list_length(root()))?;
241        let mut ctx = array_session().create_execution_ctx();
242        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx);
243        Ok(())
244    }
245
246    #[rstest]
247    #[case(buffer![0u32, 2, 5, 5, 7].into_array())]
248    #[case(buffer![0u64, 2, 5, 5, 7].into_array())]
249    fn test_nullable_list_length(#[case] offsets: ArrayRef) -> VortexResult<()> {
250        let elements = create_list_elements();
251        let list = ListArray::try_new(
252            elements,
253            offsets,
254            Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()),
255        )?
256        .into_array();
257        let result = list.apply(&list_length(root()))?;
258
259        let mut ctx = array_session().create_execution_ctx();
260        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
261
262        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(0), None]);
263
264        assert_arrays_eq!(result, expected, &mut ctx);
265
266        Ok(())
267    }
268
269    #[test]
270    fn test_null_scalar_list_length() -> VortexResult<()> {
271        let null_scalar = Scalar::null(DType::List(
272            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
273            Nullability::Nullable,
274        ));
275        let array = ConstantArray::new(null_scalar, 2).into_array();
276        let result = array.apply(&list_length(root()))?;
277
278        let mut ctx = array_session().create_execution_ctx();
279        assert!(!result.is_valid(0, &mut ctx)?);
280        assert!(!result.is_valid(1, &mut ctx)?);
281        Ok(())
282    }
283
284    #[test]
285    fn test_listview_length() -> VortexResult<()> {
286        let elements = create_list_elements();
287        let lv = ListViewArray::new(
288            elements,
289            buffer![5u32, 0, 4, 1].into_array(),
290            buffer![2u32, 3, 0, 2].into_array(),
291            Validity::NonNullable,
292        )
293        .into_array();
294        let result = lv.apply(&list_length(root()))?;
295        let mut ctx = array_session().create_execution_ctx();
296        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx);
297        Ok(())
298    }
299
300    #[test]
301    fn test_listview_length_nullable() -> VortexResult<()> {
302        let elements = create_list_elements();
303        let lv = ListViewArray::new(
304            elements,
305            buffer![5u32, 0, 4, 1].into_array(),
306            buffer![2u32, 3, 0, 2].into_array(),
307            Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()),
308        )
309        .into_array();
310        let result = lv.apply(&list_length(root()))?;
311
312        let mut ctx = array_session().create_execution_ctx();
313        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
314
315        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(0), None]);
316        assert_arrays_eq!(result, expected, &mut ctx);
317        Ok(())
318    }
319
320    #[test]
321    fn test_list_length_take() -> VortexResult<()> {
322        let elements = create_list_elements();
323        let list = ListArray::try_new(
324            elements,
325            buffer![0u32, 2, 5, 5, 7].into_array(),
326            Validity::NonNullable,
327        )?
328        .into_array();
329        let taken = list.take(buffer![3u64, 0, 2].into_array())?;
330
331        let result = taken.apply(&list_length(root()))?;
332        let mut ctx = array_session().create_execution_ctx();
333        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 0]), &mut ctx);
334        Ok(())
335    }
336
337    fn create_fixed_size_list(validity: Validity) -> ArrayRef {
338        // 4 lists of size 2 over 8 primitive elements.
339        let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array();
340        FixedSizeListArray::new(elements, 2, validity, 4).into_array()
341    }
342
343    #[test]
344    fn test_fixed_size_list_length() -> VortexResult<()> {
345        let fsl = create_fixed_size_list(Validity::NonNullable);
346        let result = fsl.apply(&list_length(root()))?;
347
348        let mut ctx = array_session().create_execution_ctx();
349        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 2, 2]), &mut ctx);
350        Ok(())
351    }
352
353    #[test]
354    fn test_fixed_size_list_length_nullable() -> VortexResult<()> {
355        let fsl = create_fixed_size_list(Validity::Array(
356            BoolArray::from_iter([true, false, true, false]).into_array(),
357        ));
358        let result = fsl.apply(&list_length(root()))?;
359
360        let mut ctx = array_session().create_execution_ctx();
361        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
362
363        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(2), None]);
364        assert_arrays_eq!(result, expected, &mut ctx);
365        Ok(())
366    }
367
368    #[test]
369    fn test_fallible_child_expression_fails() -> VortexResult<()> {
370        let fsl = create_fixed_size_list(Validity::Array(
371            BoolArray::from_iter([true, false, true, false]).into_array(),
372        ));
373        let failing_cast_dtype = DType::FixedSizeList(
374            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
375            2,
376            Nullability::NonNullable,
377        );
378
379        let lengths = fsl.apply(&list_length(cast(root(), failing_cast_dtype)))?;
380
381        let mut ctx = array_session().create_execution_ctx();
382        let result = lengths.execute::<ArrayRef>(&mut ctx);
383
384        assert!(result.is_err());
385
386        let err_message = result.unwrap_err().to_string();
387
388        assert!(
389            err_message.contains("Cannot cast array with invalid values to non-nullable type.")
390        );
391
392        Ok(())
393    }
394
395    #[test]
396    fn test_display() {
397        let expr = list_length(root());
398        assert_eq!(expr.to_string(), "vortex.list.length($)");
399    }
400}
vortex_array/scalar_fn/fns/list_length.rs

vortex_array/scalar_fn/fns/
list_length.rs
