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_buffer::BufferAllocatorRef;
8use vortex_error::VortexResult;
9
10use crate::arrays::ListArray;
11use crate::builders::ArrayBuilder;
12use crate::builders::ListBuilder;
13use crate::dtype::DType;
14use crate::dtype::OffsetBuilderPType;
15use crate::scalar::Scalar;
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: OffsetBuilderPType, I: IntoIterator>(
22        iter: I,
23        dtype: Arc<DType>,
24    ) -> VortexResult<ListArray>
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_in(
31            Arc::clone(&dtype),
32            crate::dtype::Nullability::NonNullable,
33            2 * iter.size_hint().0,
34            iter.size_hint().0,
35            BufferAllocatorRef::static_ref(),
36        );
37
38        for v in iter {
39            let elem = Scalar::list(
40                Arc::clone(&dtype),
41                v.into_iter().map(|x| x.into()).collect_vec(),
42                dtype.nullability(),
43            );
44            builder.append_value(elem.as_list())?
45        }
46        Ok(builder.finish_into_list())
47    }
48
49    pub fn from_iter_opt_slow<O: OffsetBuilderPType, I: IntoIterator<Item = Option<T>>, T>(
50        iter: I,
51        dtype: Arc<DType>,
52    ) -> VortexResult<ListArray>
53    where
54        T: IntoIterator,
55        T::Item: Into<Scalar>,
56    {
57        let iter = iter.into_iter();
58        let mut builder = ListBuilder::<O>::with_capacity_in(
59            Arc::clone(&dtype),
60            crate::dtype::Nullability::Nullable,
61            2 * iter.size_hint().0,
62            iter.size_hint().0,
63            BufferAllocatorRef::static_ref(),
64        );
65
66        for v in iter {
67            if let Some(v) = v {
68                let elem = Scalar::list(
69                    Arc::clone(&dtype),
70                    v.into_iter().map(|x| x.into()).collect_vec(),
71                    dtype.nullability(),
72                );
73                builder.append_value(elem.as_list())?
74            } else {
75                builder.append_null()
76            }
77        }
78        Ok(builder.finish_into_list())
79    }
80}