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_null_sensitive(&self, _options: &Self::Options) -> bool {
112        false
113    }
114
115    fn is_fallible(&self, _options: &Self::Options) -> bool {
116        false
117    }
118}
119
120fn scalar_list_length(scalar: &Scalar, nullability: Nullability) -> VortexResult<Scalar> {
121    if scalar.is_null() {
122        let dtype = DType::Primitive(PType::U64, Nullability::Nullable);
123        return Ok(Scalar::null(dtype));
124    }
125    let len: u64 = scalar.as_list().len().as_();
126    Ok(Scalar::primitive(len, nullability))
127}
128
129pub(crate) fn list_length(
130    array: &ArrayRef,
131    nullability: Nullability,
132    ctx: &mut ExecutionCtx,
133) -> VortexResult<ArrayRef> {
134    let any_list = array.clone().execute_until::<AnyList>(ctx)?;
135
136    let (lengths, validity) = if let Some(fsl) = any_list.as_opt::<FixedSizeList>() {
137        // The length of fixed-size list is constant, so just need to carry over validity
138        let size = fsl.list_size() as u64;
139        let lengths =
140            ConstantArray::new(Scalar::primitive(size, Nullability::NonNullable), fsl.len())
141                .into_array();
142        (lengths, fsl.validity()?)
143    } else if let Some(lv) = any_list.as_opt::<ListView>() {
144        // Length array is exactly the sizes child
145        (lv.sizes().clone(), lv.listview_validity())
146    } else if let Some(l) = any_list.as_opt::<List>() {
147        let lengths = list_length_from_offsets(l)?;
148        (lengths, l.list_validity())
149    } else {
150        let dtype = any_list.dtype();
151        vortex_bail!("list_length() requires List, ListView, or FixedSizeList but got {dtype}")
152    };
153
154    // Cast to `U64`
155    let len = lengths.len();
156    let lengths = lengths.cast(DType::Primitive(PType::U64, nullability))?;
157
158    // Carry over validity mask for nullable arrays
159    if matches!(nullability, Nullability::Nullable) {
160        lengths.mask(validity.to_array(len))
161    } else {
162        Ok(lengths)
163    }
164}
165
166/// Calculate the lengths of `ListArray` elements via the `offsets` child:
167/// `length[i] = offsets[i + 1] - offsets[i]`.
168fn list_length_from_offsets(list: ArrayView<'_, List>) -> VortexResult<ArrayRef> {
169    let offsets = list.offsets();
170    let n = offsets.len().saturating_sub(1);
171
172    offsets
173        .slice(1..offsets.len())?
174        .binary(offsets.slice(0..n)?, Operator::Sub)
175}
176
177/// Matches an `Array<List>`, `Array<ListView>`, or `Array<FixedSizeList>`
178struct AnyList;
179
180impl Matcher for AnyList {
181    type Match<'a> = ();
182
183    fn try_match(array: &ArrayRef) -> Option<Self::Match<'_>> {
184        (array.as_opt::<List>().is_some()
185            || array.as_opt::<ListView>().is_some()
186            || array.as_opt::<FixedSizeList>().is_some())
187        .then_some(())
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use std::sync::Arc;
194
195    use rstest::rstest;
196    use vortex_buffer::buffer;
197    use vortex_error::VortexResult;
198
199    use crate::ArrayRef;
200    use crate::IntoArray;
201    use crate::VortexSessionExecute;
202    use crate::array_session;
203    use crate::arrays::BoolArray;
204    use crate::arrays::ConstantArray;
205    use crate::arrays::FixedSizeListArray;
206    use crate::arrays::ListArray;
207    use crate::arrays::ListViewArray;
208    use crate::arrays::PrimitiveArray;
209    use crate::assert_arrays_eq;
210    use crate::dtype::DType;
211    use crate::dtype::Nullability;
212    use crate::dtype::PType;
213    use crate::expr::cast;
214    use crate::expr::list_length;
215    use crate::expr::root;
216    use crate::scalar::Scalar;
217    use crate::validity::Validity;
218
219    fn create_list_elements() -> ArrayRef {
220        PrimitiveArray::from_option_iter::<i32, _>([
221            Some(1),
222            Some(2),
223            Some(3),
224            Some(4),
225            Some(5),
226            Some(6),
227            None,
228        ])
229        .into_array()
230    }
231
232    #[rstest]
233    #[case(buffer![0u32, 2, 5, 5, 7].into_array())]
234    #[case(buffer![0u64, 2, 5, 5, 7].into_array())]
235    fn test_list_length(#[case] offsets: ArrayRef) -> VortexResult<()> {
236        let elements = create_list_elements();
237        let list = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array();
238        let result = list.apply(&list_length(root()))?;
239        let mut ctx = array_session().create_execution_ctx();
240        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx);
241        Ok(())
242    }
243
244    #[rstest]
245    #[case(buffer![0u32, 2, 5, 5, 7].into_array())]
246    #[case(buffer![0u64, 2, 5, 5, 7].into_array())]
247    fn test_nullable_list_length(#[case] offsets: ArrayRef) -> VortexResult<()> {
248        let elements = create_list_elements();
249        let list = ListArray::try_new(
250            elements,
251            offsets,
252            Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()),
253        )?
254        .into_array();
255        let result = list.apply(&list_length(root()))?;
256
257        let mut ctx = array_session().create_execution_ctx();
258        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
259
260        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(0), None]);
261
262        assert_arrays_eq!(result, expected, &mut ctx);
263
264        Ok(())
265    }
266
267    #[test]
268    fn test_null_scalar_list_length() -> VortexResult<()> {
269        let null_scalar = Scalar::null(DType::List(
270            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
271            Nullability::Nullable,
272        ));
273        let array = ConstantArray::new(null_scalar, 2).into_array();
274        let result = array.apply(&list_length(root()))?;
275
276        let mut ctx = array_session().create_execution_ctx();
277        assert!(!result.is_valid(0, &mut ctx)?);
278        assert!(!result.is_valid(1, &mut ctx)?);
279        Ok(())
280    }
281
282    #[test]
283    fn test_listview_length() -> VortexResult<()> {
284        let elements = create_list_elements();
285        let lv = ListViewArray::new(
286            elements,
287            buffer![5u32, 0, 4, 1].into_array(),
288            buffer![2u32, 3, 0, 2].into_array(),
289            Validity::NonNullable,
290        )
291        .into_array();
292        let result = lv.apply(&list_length(root()))?;
293        let mut ctx = array_session().create_execution_ctx();
294        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 3, 0, 2]), &mut ctx);
295        Ok(())
296    }
297
298    #[test]
299    fn test_listview_length_nullable() -> VortexResult<()> {
300        let elements = create_list_elements();
301        let lv = ListViewArray::new(
302            elements,
303            buffer![5u32, 0, 4, 1].into_array(),
304            buffer![2u32, 3, 0, 2].into_array(),
305            Validity::Array(BoolArray::from_iter([true, false, true, false]).into_array()),
306        )
307        .into_array();
308        let result = lv.apply(&list_length(root()))?;
309
310        let mut ctx = array_session().create_execution_ctx();
311        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
312
313        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(0), None]);
314        assert_arrays_eq!(result, expected, &mut ctx);
315        Ok(())
316    }
317
318    #[test]
319    fn test_list_length_take() -> VortexResult<()> {
320        let elements = create_list_elements();
321        let list = ListArray::try_new(
322            elements,
323            buffer![0u32, 2, 5, 5, 7].into_array(),
324            Validity::NonNullable,
325        )?
326        .into_array();
327        let taken = list.take(buffer![3u64, 0, 2].into_array())?;
328
329        let result = taken.apply(&list_length(root()))?;
330        let mut ctx = array_session().create_execution_ctx();
331        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 0]), &mut ctx);
332        Ok(())
333    }
334
335    fn create_fixed_size_list(validity: Validity) -> ArrayRef {
336        // 4 lists of size 2 over 8 primitive elements.
337        let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6, 7, 8]).into_array();
338        FixedSizeListArray::new(elements, 2, validity, 4).into_array()
339    }
340
341    #[test]
342    fn test_fixed_size_list_length() -> VortexResult<()> {
343        let fsl = create_fixed_size_list(Validity::NonNullable);
344        let result = fsl.apply(&list_length(root()))?;
345
346        let mut ctx = array_session().create_execution_ctx();
347        assert_arrays_eq!(result, PrimitiveArray::from_iter([2u64, 2, 2, 2]), &mut ctx);
348        Ok(())
349    }
350
351    #[test]
352    fn test_fixed_size_list_length_nullable() -> VortexResult<()> {
353        let fsl = create_fixed_size_list(Validity::Array(
354            BoolArray::from_iter([true, false, true, false]).into_array(),
355        ));
356        let result = fsl.apply(&list_length(root()))?;
357
358        let mut ctx = array_session().create_execution_ctx();
359        let result = result.execute::<PrimitiveArray>(&mut ctx)?;
360
361        let expected = PrimitiveArray::from_option_iter::<u64, _>([Some(2), None, Some(2), None]);
362        assert_arrays_eq!(result, expected, &mut ctx);
363        Ok(())
364    }
365
366    #[test]
367    fn test_fallible_child_expression_fails() -> VortexResult<()> {
368        let fsl = create_fixed_size_list(Validity::Array(
369            BoolArray::from_iter([true, false, true, false]).into_array(),
370        ));
371        let failing_cast_dtype = DType::FixedSizeList(
372            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
373            2,
374            Nullability::NonNullable,
375        );
376
377        let lengths = fsl.apply(&list_length(cast(root(), failing_cast_dtype)))?;
378
379        let mut ctx = array_session().create_execution_ctx();
380        let result = lengths.execute::<ArrayRef>(&mut ctx);
381
382        assert!(result.is_err());
383
384        let err_message = result.unwrap_err().to_string();
385
386        assert!(
387            err_message.contains("Cannot cast array with invalid values to non-nullable type.")
388        );
389
390        Ok(())
391    }
392
393    #[test]
394    fn test_display() {
395        let expr = list_length(root());
396        assert_eq!(expr.to_string(), "vortex.list.length($)");
397    }
398}
vortex_array/scalar_fn/fns/list_length.rs

vortex_array/scalar_fn/fns/
list_length.rs
