vortex_array/arrays/list/
test_harness.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use itertools::Itertools;
7use vortex_dtype::{DType, IntegerPType};
8use vortex_error::VortexResult;
9use vortex_scalar::Scalar;
10
11use crate::ArrayRef;
12use crate::arrays::ListArray;
13use crate::builders::{ArrayBuilder, ListBuilder};
14
15impl ListArray {
16    /// This is a convenience method to create a list array from an iterator of iterators.
17    /// This method is slow however since each element is first converted to a scalar and then
18    /// appended to the array.
19    pub fn from_iter_slow<O: IntegerPType, I: IntoIterator>(
20        iter: I,
21        dtype: Arc<DType>,
22    ) -> VortexResult<ArrayRef>
23    where
24        I::Item: IntoIterator,
25        <I::Item as IntoIterator>::Item: Into<Scalar>,
26    {
27        let iter = iter.into_iter();
28        let mut builder = ListBuilder::<O>::with_capacity(
29            dtype.clone(),
30            vortex_dtype::Nullability::NonNullable,
31            2 * iter.size_hint().0,
32            iter.size_hint().0,
33        );
34
35        for v in iter {
36            let elem = Scalar::list(
37                dtype.clone(),
38                v.into_iter().map(|x| x.into()).collect_vec(),
39                dtype.nullability(),
40            );
41            builder.append_value(elem.as_list())?
42        }
43        Ok(builder.finish())
44    }
45
46    pub fn from_iter_opt_slow<O: IntegerPType, I: IntoIterator<Item = Option<T>>, T>(
47        iter: I,
48        dtype: Arc<DType>,
49    ) -> VortexResult<ArrayRef>
50    where
51        T: IntoIterator,
52        T::Item: Into<Scalar>,
53    {
54        let iter = iter.into_iter();
55        let mut builder = ListBuilder::<O>::with_capacity(
56            dtype.clone(),
57            vortex_dtype::Nullability::Nullable,
58            2 * iter.size_hint().0,
59            iter.size_hint().0,
60        );
61
62        for v in iter {
63            if let Some(v) = v {
64                let elem = Scalar::list(
65                    dtype.clone(),
66                    v.into_iter().map(|x| x.into()).collect_vec(),
67                    dtype.nullability(),
68                );
69                builder.append_value(elem.as_list())?
70            } else {
71                builder.append_null()
72            }
73        }
74        Ok(builder.finish())
75    }
76}