vortex_array/arrays/list/compute/
mod.rs

1mod mask;
2mod to_arrow;
3
4use std::sync::Arc;
5
6use itertools::Itertools;
7use vortex_error::VortexResult;
8use vortex_scalar::Scalar;
9
10use crate::arrays::{ListArray, ListEncoding};
11use crate::compute::{
12    IsConstantFn, IsConstantOpts, IsSortedFn, MinMaxFn, MinMaxResult, ScalarAtFn, SliceFn,
13    ToArrowFn, UncompressedSizeFn, scalar_at, slice, uncompressed_size,
14};
15use crate::vtable::ComputeVTable;
16use crate::{Array, ArrayRef};
17
18impl ComputeVTable for ListEncoding {
19    fn is_constant_fn(&self) -> Option<&dyn IsConstantFn<&dyn Array>> {
20        Some(self)
21    }
22
23    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<&dyn Array>> {
24        Some(self)
25    }
26
27    fn slice_fn(&self) -> Option<&dyn SliceFn<&dyn Array>> {
28        Some(self)
29    }
30
31    fn to_arrow_fn(&self) -> Option<&dyn ToArrowFn<&dyn Array>> {
32        Some(self)
33    }
34
35    fn min_max_fn(&self) -> Option<&dyn MinMaxFn<&dyn Array>> {
36        Some(self)
37    }
38
39    fn uncompressed_size_fn(&self) -> Option<&dyn UncompressedSizeFn<&dyn Array>> {
40        Some(self)
41    }
42
43    fn is_sorted_fn(&self) -> Option<&dyn IsSortedFn<&dyn Array>> {
44        Some(self)
45    }
46}
47
48impl ScalarAtFn<&ListArray> for ListEncoding {
49    fn scalar_at(&self, array: &ListArray, index: usize) -> VortexResult<Scalar> {
50        let elem = array.elements_at(index)?;
51        let scalars: Vec<Scalar> = (0..elem.len()).map(|i| scalar_at(&elem, i)).try_collect()?;
52
53        Ok(Scalar::list(
54            Arc::new(elem.dtype().clone()),
55            scalars,
56            array.dtype().nullability(),
57        ))
58    }
59}
60
61impl SliceFn<&ListArray> for ListEncoding {
62    fn slice(&self, array: &ListArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
63        Ok(ListArray::try_new(
64            array.elements().clone(),
65            slice(array.offsets(), start, stop + 1)?,
66            array.validity().slice(start, stop)?,
67        )?
68        .into_array())
69    }
70}
71
72impl MinMaxFn<&ListArray> for ListEncoding {
73    fn min_max(&self, _array: &ListArray) -> VortexResult<Option<MinMaxResult>> {
74        // TODO(joe): Implement list min max
75        Ok(None)
76    }
77}
78
79impl IsConstantFn<&ListArray> for ListEncoding {
80    fn is_constant(
81        &self,
82        _array: &ListArray,
83        _opts: &IsConstantOpts,
84    ) -> VortexResult<Option<bool>> {
85        // TODO(adam): Do we want to fallback to arrow here?
86        Ok(None)
87    }
88}
89
90impl UncompressedSizeFn<&ListArray> for ListEncoding {
91    fn uncompressed_size(&self, array: &ListArray) -> VortexResult<usize> {
92        let size = uncompressed_size(array.elements())? + uncompressed_size(array.offsets())?;
93        Ok(size + array.validity().uncompressed_size())
94    }
95}
96
97impl IsSortedFn<&ListArray> for ListEncoding {
98    fn is_sorted(&self, _array: &ListArray) -> VortexResult<bool> {
99        Ok(false)
100    }
101
102    fn is_strict_sorted(&self, _array: &ListArray) -> VortexResult<bool> {
103        Ok(false)
104    }
105}
106
107#[cfg(test)]
108mod test {
109    use crate::array::Array;
110    use crate::arrays::{ListArray, PrimitiveArray};
111    use crate::compute::conformance::mask::test_mask;
112    use crate::validity::Validity;
113
114    #[test]
115    fn test_mask_list() {
116        let elements = PrimitiveArray::from_iter(0..35);
117        let offsets = PrimitiveArray::from_iter([0, 5, 11, 18, 26, 35]);
118        let validity = Validity::AllValid;
119        let array =
120            ListArray::try_new(elements.into_array(), offsets.into_array(), validity).unwrap();
121
122        test_mask(&array);
123    }
124}
vortex_array/arrays/list/compute/mod.rs

vortex_array/arrays/list/compute/
mod.rs
