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