Skip to main content

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