17050456606

Committed 18 Aug 2025 07:32PM UTC coverage: 47.596%. First build

Build Type

github

Committed by GitHub

Commit Message

Merge afbf04d54 into 7eb8ac9fa

Pull Request Pull Request #4177:

Run Details

555 of 1372 new or added lines in 154 files covered. (40.45%)

18648 of 39180 relevant lines covered (47.6%)

234900.69 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/vortex-array/src/iter.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! Iterator over slices of an array, and related utilities.

use std::sync::Arc;

use itertools::Itertools;
use vortex_dtype::DType;
use vortex_error::VortexResult;

use crate::arrays::{ChunkedArray, ChunkedVTable};
use crate::stream::{ArrayStream, ArrayStreamAdapter};
use crate::{Array, ArrayRef, IntoArray};

/// Iterator of array with a known [`DType`].
///
/// It's up to implementations to guarantee all arrays have the same [`DType`].
pub trait ArrayIterator: Iterator<Item = VortexResult<ArrayRef>> {
    fn dtype(&self) -> &DType;
}

impl ArrayIterator for Box<dyn ArrayIterator + Send> {
    fn dtype(&self) -> &DType {
        self.as_ref().dtype()
    }
}

pub struct ArrayIteratorAdapter<I> {
    dtype: DType,
    inner: I,
}

impl<I> ArrayIteratorAdapter<I> {
    pub fn new(dtype: DType, inner: I) -> Self {
        Self { dtype, inner }
    }
}

impl<I> Iterator for ArrayIteratorAdapter<I>
where
    I: Iterator<Item = VortexResult<ArrayRef>>,
{
    type Item = VortexResult<ArrayRef>;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

impl<I> ArrayIterator for ArrayIteratorAdapter<I>
where
    I: Iterator<Item = VortexResult<ArrayRef>>,
{
    fn dtype(&self) -> &DType {
        &self.dtype
    }
}

pub trait ArrayIteratorExt: ArrayIterator {
    fn into_array_stream(self) -> impl ArrayStream
    where
        Self: Sized,
    {
        ArrayStreamAdapter::new(self.dtype().clone(), futures_util::stream::iter(self))
    }

    /// Collect the iterator into a single `Array`.
    ///
    /// If the iterator yields multiple chunks, they will be returned as a [`ChunkedArray`].
    fn read_all(self) -> VortexResult<ArrayRef>
    where
        Self: Sized,
    {
        let dtype = self.dtype().clone();
        let mut chunks: Vec<ArrayRef> = self.try_collect()?;
        if chunks.len() == 1 {
            Ok(chunks.remove(0))
        } else {
            Ok(ChunkedArray::try_new(chunks, dtype)?.into_array())
        }
    }
}

impl<I: ArrayIterator> ArrayIteratorExt for I {}

impl dyn Array + '_ {
    /// Create an [`ArrayIterator`] over the array.
    pub fn to_array_iterator(&self) -> impl ArrayIterator + 'static {
        let dtype = self.dtype().clone();
        let iter = if let Some(chunked) = self.as_opt::<ChunkedVTable>() {
            ArrayChunkIterator::Chunked(Arc::new(chunked.clone()), 0)
        } else {
            ArrayChunkIterator::Single(Some(self.to_array()))
        };
        ArrayIteratorAdapter::new(dtype, iter)
    }
}

/// We define a single iterator that can handle both chunked and non-chunked arrays.
/// This avoids the need to create boxed static iterators for the two chunked and non-chunked cases.
enum ArrayChunkIterator {
    Single(Option<ArrayRef>),
    Chunked(Arc<ChunkedArray>, usize),
}

impl Iterator for ArrayChunkIterator {
    type Item = VortexResult<ArrayRef>;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            ArrayChunkIterator::Single(array) => array.take().map(Ok),
            ArrayChunkIterator::Chunked(chunked, idx) => (*idx < chunked.nchunks()).then(|| {
                let chunk = chunked.chunk(*idx);
                *idx += 1;
                Ok(chunk.clone())
            }),
        }
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	//! Iterator over slices of an array, and related utilities.
5
6	use std::sync::Arc;
7
8	use itertools::Itertools;
9	use vortex_dtype::DType;
10	use vortex_error::VortexResult;
11
12	use crate::arrays::{ChunkedArray, ChunkedVTable};
13	use crate::stream::{ArrayStream, ArrayStreamAdapter};
14	use crate::{Array, ArrayRef, IntoArray};
15
16	/// Iterator of array with a known [`DType`].
17	///
18	/// It's up to implementations to guarantee all arrays have the same [`DType`].
19	pub trait ArrayIterator: Iterator<Item = VortexResult<ArrayRef>> {
20	fn dtype(&self) -> &DType;
21	}
22
23	impl ArrayIterator for Box<dyn ArrayIterator + Send> {
24	fn dtype(&self) -> &DType {	×
25	self.as_ref().dtype()	×
26	}	×
27	}
28
29	pub struct ArrayIteratorAdapter<I> {
30	dtype: DType,
31	inner: I,
32	}
33
34	impl<I> ArrayIteratorAdapter<I> {
35	pub fn new(dtype: DType, inner: I) -> Self {	×
36	Self { dtype, inner }	×
37	}	×
38	}
39
40	impl<I> Iterator for ArrayIteratorAdapter<I>
41	where
42	I: Iterator<Item = VortexResult<ArrayRef>>,
43	{
44	type Item = VortexResult<ArrayRef>;
45
46	fn next(&mut self) -> Option<Self::Item> {	×
47	self.inner.next()	×
48	}	×
49	}
50
51	impl<I> ArrayIterator for ArrayIteratorAdapter<I>
52	where
53	I: Iterator<Item = VortexResult<ArrayRef>>,
54	{
55	fn dtype(&self) -> &DType {	×
56	&self.dtype	×
57	}	×
58	}
59
60	pub trait ArrayIteratorExt: ArrayIterator {
61	fn into_array_stream(self) -> impl ArrayStream	×
62	where	×
63	Self: Sized,	×
64	{
65	ArrayStreamAdapter::new(self.dtype().clone(), futures_util::stream::iter(self))	×
66	}	×
67
68	/// Collect the iterator into a single `Array`.
69	///
70	/// If the iterator yields multiple chunks, they will be returned as a [`ChunkedArray`].
71	fn read_all(self) -> VortexResult<ArrayRef>	×
72	where	×
73	Self: Sized,	×
74	{
75	let dtype = self.dtype().clone();	×
76	let mut chunks: Vec<ArrayRef> = self.try_collect()?;	×
77	if chunks.len() == 1 {	×
78	Ok(chunks.remove(0))	×
79	} else {
80	Ok(ChunkedArray::try_new(chunks, dtype)?.into_array())	×
81	}
82	}	×
83	}
84
85	impl<I: ArrayIterator> ArrayIteratorExt for I {}
86
87	impl dyn Array + '_ {
88	/// Create an [`ArrayIterator`] over the array.
89	pub fn to_array_iterator(&self) -> impl ArrayIterator + 'static {	×
90	let dtype = self.dtype().clone();	×
91	let iter = if let Some(chunked) = self.as_opt::<ChunkedVTable>() {	×
92	ArrayChunkIterator::Chunked(Arc::new(chunked.clone()), 0)	×
93	} else {
94	ArrayChunkIterator::Single(Some(self.to_array()))	×
95	};
96	ArrayIteratorAdapter::new(dtype, iter)	×
97	}	×
98	}
99
100	/// We define a single iterator that can handle both chunked and non-chunked arrays.
101	/// This avoids the need to create boxed static iterators for the two chunked and non-chunked cases.
102	enum ArrayChunkIterator {
103	Single(Option<ArrayRef>),
104	Chunked(Arc<ChunkedArray>, usize),
105	}
106
107	impl Iterator for ArrayChunkIterator {
108	type Item = VortexResult<ArrayRef>;
109
110	fn next(&mut self) -> Option<Self::Item> {	×
111	match self {	×
112	ArrayChunkIterator::Single(array) => array.take().map(Ok),	×
113	ArrayChunkIterator::Chunked(chunked, idx) => (*idx < chunked.nchunks()).then(\|\| {	×
NEW 114	let chunk = chunked.chunk(*idx);	×
115	*idx += 1;	×
116	Ok(chunk.clone())	×
117	}),	×
118	}
119	}	×
120	}

vortex-data / vortex / 17050456606

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous