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:

Coverage Stats

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

26.03

/vortex-array/src/arrays/chunked/array.rs

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

//! First-class chunked arrays.
//!
//! Vortex is a chunked array library that's able to

use std::fmt::Debug;

use futures_util::stream;
use itertools::Itertools;
use vortex_buffer::{Buffer, BufferMut};
use vortex_dtype::DType;
use vortex_error::{VortexExpect as _, VortexResult, VortexUnwrap, vortex_bail};
use vortex_mask::Mask;

use crate::arrays::ChunkedVTable;
use crate::iter::{ArrayIterator, ArrayIteratorAdapter};
use crate::search_sorted::{SearchSorted, SearchSortedSide};
use crate::stats::{ArrayStats, StatsSetRef};
use crate::stream::{ArrayStream, ArrayStreamAdapter};
use crate::vtable::{ArrayVTable, ValidityVTable};
use crate::{Array, ArrayRef, IntoArray};

#[derive(Clone, Debug)]
pub struct ChunkedArray {
    dtype: DType,
    len: usize,
    chunk_offsets: Buffer<u64>,
    chunks: Vec<ArrayRef>,
    stats_set: ArrayStats,
}

impl ChunkedArray {
    pub fn try_new(chunks: Vec<ArrayRef>, dtype: DType) -> VortexResult<Self> {
        for chunk in &chunks {
            if chunk.dtype() != &dtype {
                vortex_bail!(MismatchedTypes: dtype, chunk.dtype());
            }
        }

        Ok(Self::new_unchecked(chunks, dtype))
    }

    pub fn new_unchecked(chunks: Vec<ArrayRef>, dtype: DType) -> Self {
        let nchunks = chunks.len();

        let mut chunk_offsets = BufferMut::<u64>::with_capacity(nchunks + 1);
        unsafe { chunk_offsets.push_unchecked(0) }
        let mut curr_offset = 0;
        for c in &chunks {
            curr_offset += c.len() as u64;
            unsafe { chunk_offsets.push_unchecked(curr_offset) }
        }
        assert_eq!(chunk_offsets.len(), nchunks + 1);

        Self {
            dtype,
            len: curr_offset.try_into().vortex_unwrap(),
            chunk_offsets: chunk_offsets.freeze(),
            chunks,
            stats_set: Default::default(),
        }
    }

    #[inline]
    pub fn chunk(&self, idx: usize) -> &ArrayRef {
        assert!(idx < self.nchunks(), "chunk index {idx} out of bounds");

        &self.chunks[idx]
    }

    pub fn nchunks(&self) -> usize {
        self.chunks.len()
    }

    #[inline]
    pub fn chunk_offsets(&self) -> &Buffer<u64> {
        &self.chunk_offsets
    }

    pub(crate) fn find_chunk_idx(&self, index: usize) -> (usize, usize) {
        assert!(index <= self.len(), "Index out of bounds of the array");
        let index = index as u64;

        // Since there might be duplicate values in offsets because of empty chunks we want to search from right
        // and take the last chunk (we subtract 1 since there's a leading 0)
        let index_chunk = self
            .chunk_offsets()
            .search_sorted(&index, SearchSortedSide::Right)
            .to_ends_index(self.nchunks() + 1)
            .saturating_sub(1);
        let chunk_start = self.chunk_offsets()[index_chunk];

        let index_in_chunk =
            usize::try_from(index - chunk_start).vortex_expect("Index is too large for usize");
        (index_chunk, index_in_chunk)
    }

    pub fn chunks(&self) -> &[ArrayRef] {
        &self.chunks
    }

    pub fn non_empty_chunks(&self) -> impl Iterator<Item = &ArrayRef> + '_ {
        self.chunks().iter().filter(|c| !c.is_empty())
    }

    pub fn array_iterator(&self) -> impl ArrayIterator + '_ {
        ArrayIteratorAdapter::new(self.dtype().clone(), self.chunks().iter().cloned().map(Ok))
    }

    pub fn array_stream(&self) -> impl ArrayStream + '_ {
        ArrayStreamAdapter::new(
            self.dtype().clone(),
            stream::iter(self.chunks().iter().cloned().map(Ok)),
        )
    }

    pub fn rechunk(&self, target_bytesize: u64, target_rowsize: usize) -> VortexResult<Self> {
        let mut new_chunks = Vec::new();
        let mut chunks_to_combine = Vec::new();
        let mut new_chunk_n_bytes = 0;
        let mut new_chunk_n_elements = 0;
        for chunk in self.chunks() {
            let n_bytes = chunk.nbytes();
            let n_elements = chunk.len();

            if (new_chunk_n_bytes + n_bytes > target_bytesize
                || new_chunk_n_elements + n_elements > target_rowsize)
                && !chunks_to_combine.is_empty()
            {
                new_chunks.push(
                    ChunkedArray::new_unchecked(chunks_to_combine, self.dtype().clone())
                        .to_canonical()?
                        .into_array(),
                );

                new_chunk_n_bytes = 0;
                new_chunk_n_elements = 0;
                chunks_to_combine = Vec::new();
            }

            if n_bytes > target_bytesize || n_elements > target_rowsize {
                new_chunks.push(chunk.clone());
            } else {
                new_chunk_n_bytes += n_bytes;
                new_chunk_n_elements += n_elements;
                chunks_to_combine.push(chunk.clone());
            }
        }

        if !chunks_to_combine.is_empty() {
            new_chunks.push(
                ChunkedArray::new_unchecked(chunks_to_combine, self.dtype().clone())
                    .to_canonical()?
                    .into_array(),
            );
        }

        Ok(Self::new_unchecked(new_chunks, self.dtype().clone()))
    }
}

impl FromIterator<ArrayRef> for ChunkedArray {
    fn from_iter<T: IntoIterator<Item = ArrayRef>>(iter: T) -> Self {
        let chunks: Vec<ArrayRef> = iter.into_iter().collect();
        let dtype = chunks
            .first()
            .map(|c| c.dtype().clone())
            .vortex_expect("Cannot infer DType from an empty iterator");
        Self::try_new(chunks, dtype).vortex_expect("Failed to create chunked array from iterator")
    }
}

impl ArrayVTable<ChunkedVTable> for ChunkedVTable {
    fn len(array: &ChunkedArray) -> usize {
        array.len
    }

    fn dtype(array: &ChunkedArray) -> &DType {
        &array.dtype
    }

    fn stats(array: &ChunkedArray) -> StatsSetRef<'_> {
        array.stats_set.to_ref(array.as_ref())
    }
}

impl ValidityVTable<ChunkedVTable> for ChunkedVTable {
    fn is_valid(array: &ChunkedArray, index: usize) -> VortexResult<bool> {
        if !array.dtype.is_nullable() {
            return Ok(true);
        }
        let (chunk, offset_in_chunk) = array.find_chunk_idx(index);
        array.chunk(chunk).is_valid(offset_in_chunk)
    }

    fn all_valid(array: &ChunkedArray) -> VortexResult<bool> {
        if !array.dtype().is_nullable() {
            return Ok(true);
        }
        for chunk in array.non_empty_chunks() {
            if !chunk.all_valid()? {
                return Ok(false);
            }
        }
        Ok(true)
    }

    fn all_invalid(array: &ChunkedArray) -> VortexResult<bool> {
        if !array.dtype().is_nullable() {
            return Ok(false);
        }
        for chunk in array.non_empty_chunks() {
            if !chunk.all_invalid()? {
                return Ok(false);
            }
        }
        Ok(true)
    }

    fn validity_mask(array: &ChunkedArray) -> VortexResult<Mask> {
        array
            .chunks()
            .iter()
            .map(|a| a.validity_mask())
            .try_collect()
    }
}

#[cfg(test)]
mod test {
    use vortex_buffer::buffer;
    use vortex_dtype::{DType, Nullability, PType};
    use vortex_error::VortexResult;

    use crate::array::Array;
    use crate::arrays::chunked::ChunkedArray;
    use crate::arrays::{ChunkedVTable, PrimitiveArray};
    use crate::compute::sub_scalar;
    use crate::validity::Validity;
    use crate::{IntoArray, ToCanonical, assert_arrays_eq};

    fn chunked_array() -> ChunkedArray {
        ChunkedArray::try_new(
            vec![
                buffer![1u64, 2, 3].into_array(),
                buffer![4u64, 5, 6].into_array(),
                buffer![7u64, 8, 9].into_array(),
            ],
            DType::Primitive(PType::U64, Nullability::NonNullable),
        )
        .unwrap()
    }

    #[test]
    fn test_scalar_subtract() {
        let chunked = chunked_array().into_array();
        let to_subtract = 1u64;
        let array = sub_scalar(&chunked, to_subtract.into()).unwrap();

        let chunked = array.as_::<ChunkedVTable>();
        let chunks_out = chunked.chunks();

        let results = chunks_out[0]
            .to_primitive()
            .unwrap()
            .as_slice::<u64>()
            .to_vec();
        assert_eq!(results, &[0u64, 1, 2]);
        let results = chunks_out[1]
            .to_primitive()
            .unwrap()
            .as_slice::<u64>()
            .to_vec();
        assert_eq!(results, &[3u64, 4, 5]);
        let results = chunks_out[2]
            .to_primitive()
            .unwrap()
            .as_slice::<u64>()
            .to_vec();
        assert_eq!(results, &[6u64, 7, 8]);
    }

    #[test]
    fn test_rechunk_one_chunk() {
        let chunked = ChunkedArray::try_new(
            vec![buffer![0u64].into_array()],
            DType::Primitive(PType::U64, Nullability::NonNullable),
        )
        .unwrap();

        let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap();

        assert_arrays_eq!(chunked, rechunked);
    }

    #[test]
    fn test_rechunk_two_chunks() {
        let chunked = ChunkedArray::try_new(
            vec![buffer![0u64].into_array(), buffer![5u64].into_array()],
            DType::Primitive(PType::U64, Nullability::NonNullable),
        )
        .unwrap();

        let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap();

        assert_eq!(rechunked.nchunks(), 1);
        assert_arrays_eq!(chunked, rechunked);
    }

    #[test]
    fn test_rechunk_tiny_target_chunks() {
        let chunked = ChunkedArray::try_new(
            vec![
                buffer![0u64, 1, 2, 3].into_array(),
                buffer![4u64, 5].into_array(),
            ],
            DType::Primitive(PType::U64, Nullability::NonNullable),
        )
        .unwrap();

        let rechunked = chunked.rechunk(1 << 16, 5).unwrap();

        assert_eq!(rechunked.nchunks(), 2);
        assert!(rechunked.chunks().iter().all(|c| c.len() < 5));
        assert_arrays_eq!(chunked, rechunked);
    }

    #[test]
    fn test_rechunk_with_too_big_chunk() {
        let chunked = ChunkedArray::try_new(
            vec![
                buffer![0u64, 1, 2].into_array(),
                buffer![42_u64; 6].into_array(),
                buffer![4u64, 5].into_array(),
                buffer![6u64, 7].into_array(),
                buffer![8u64, 9].into_array(),
            ],
            DType::Primitive(PType::U64, Nullability::NonNullable),
        )
        .unwrap();

        let rechunked = chunked.rechunk(1 << 16, 5).unwrap();
        // greedy so should be: [0, 1, 2] [42, 42, 42, 42, 42, 42] [4, 5, 6, 7] [8, 9]

        assert_eq!(rechunked.nchunks(), 4);
        assert_arrays_eq!(chunked, rechunked);
    }

    #[test]
    fn test_empty_chunks_all_valid() -> VortexResult<()> {
        // Create chunks where some are empty but all non-empty chunks have all valid values
        let chunks = vec![
            PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid).into_array(),
            PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk
            PrimitiveArray::new(buffer![4u64, 5], Validity::AllValid).into_array(),
            PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk
        ];

        let chunked =
            ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?;

        // Should be all_valid since all non-empty chunks are all_valid
        assert!(chunked.all_valid()?);
        assert!(!chunked.all_invalid()?);

        Ok(())
    }

    #[test]
    fn test_empty_chunks_all_invalid() -> VortexResult<()> {
        // Create chunks where some are empty but all non-empty chunks have all invalid values
        let chunks = vec![
            PrimitiveArray::new(buffer![1u64, 2], Validity::AllInvalid).into_array(),
            PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk
            PrimitiveArray::new(buffer![3u64, 4, 5], Validity::AllInvalid).into_array(),
            PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk
        ];

        let chunked =
            ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?;

        // Should be all_invalid since all non-empty chunks are all_invalid
        assert!(!chunked.all_valid()?);
        assert!(chunked.all_invalid()?);

        Ok(())
    }

    #[test]
    fn test_empty_chunks_mixed_validity() -> VortexResult<()> {
        // Create chunks with mixed validity including empty chunks
        let chunks = vec![
            PrimitiveArray::new(buffer![1u64, 2], Validity::AllValid).into_array(),
            PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk
            PrimitiveArray::new(buffer![3u64, 4], Validity::AllInvalid).into_array(),
            PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk
        ];

        let chunked =
            ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?;

        // Should be neither all_valid nor all_invalid
        assert!(!chunked.all_valid()?);
        assert!(!chunked.all_invalid()?);

        Ok(())
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	//! First-class chunked arrays.
5	//!
6	//! Vortex is a chunked array library that's able to
7
8	use std::fmt::Debug;
9
10	use futures_util::stream;
11	use itertools::Itertools;
12	use vortex_buffer::{Buffer, BufferMut};
13	use vortex_dtype::DType;
14	use vortex_error::{VortexExpect as _, VortexResult, VortexUnwrap, vortex_bail};
15	use vortex_mask::Mask;
16
17	use crate::arrays::ChunkedVTable;
18	use crate::iter::{ArrayIterator, ArrayIteratorAdapter};
19	use crate::search_sorted::{SearchSorted, SearchSortedSide};
20	use crate::stats::{ArrayStats, StatsSetRef};
21	use crate::stream::{ArrayStream, ArrayStreamAdapter};
22	use crate::vtable::{ArrayVTable, ValidityVTable};
23	use crate::{Array, ArrayRef, IntoArray};
24
25	#[derive(Clone, Debug)]
26	pub struct ChunkedArray {
27	dtype: DType,
28	len: usize,
29	chunk_offsets: Buffer<u64>,
30	chunks: Vec<ArrayRef>,
31	stats_set: ArrayStats,
32	}
33
34	impl ChunkedArray {
35	pub fn try_new(chunks: Vec<ArrayRef>, dtype: DType) -> VortexResult<Self> {	2,001✔
36	for chunk in &chunks {	24,531✔
37	if chunk.dtype() != &dtype {	22,530✔
38	vortex_bail!(MismatchedTypes: dtype, chunk.dtype());	×
39	}	22,530✔
40	}
41
42	Ok(Self::new_unchecked(chunks, dtype))	2,001✔
43	}	2,001✔
44
45	pub fn new_unchecked(chunks: Vec<ArrayRef>, dtype: DType) -> Self {	7,125✔
46	let nchunks = chunks.len();	7,125✔
47
48	let mut chunk_offsets = BufferMut::<u64>::with_capacity(nchunks + 1);	7,125✔
49	unsafe { chunk_offsets.push_unchecked(0) }	7,125✔
50	let mut curr_offset = 0;	7,125✔
51	for c in &chunks {	38,427✔
52	curr_offset += c.len() as u64;	31,302✔
53	unsafe { chunk_offsets.push_unchecked(curr_offset) }	31,302✔
54	}
55	assert_eq!(chunk_offsets.len(), nchunks + 1);	7,125✔
56
57	Self {	7,125✔
58	dtype,	7,125✔
59	len: curr_offset.try_into().vortex_unwrap(),	7,125✔
60	chunk_offsets: chunk_offsets.freeze(),	7,125✔
61	chunks,	7,125✔
62	stats_set: Default::default(),	7,125✔
63	}	7,125✔
64	}	7,125✔
65
66	#[inline]
NEW 67	pub fn chunk(&self, idx: usize) -> &ArrayRef {	×
NEW 68	assert!(idx < self.nchunks(), "chunk index {idx} out of bounds");	×
69
NEW 70	&self.chunks[idx]	×
71	}	×
72
73	pub fn nchunks(&self) -> usize {	14,250✔
74	self.chunks.len()	14,250✔
75	}	14,250✔
76
77	#[inline]
78	pub fn chunk_offsets(&self) -> &Buffer<u64> {	×
79	&self.chunk_offsets	×
80	}	×
81
82	pub(crate) fn find_chunk_idx(&self, index: usize) -> (usize, usize) {	×
83	assert!(index <= self.len(), "Index out of bounds of the array");	×
84	let index = index as u64;	×
85
86	// Since there might be duplicate values in offsets because of empty chunks we want to search from right
87	// and take the last chunk (we subtract 1 since there's a leading 0)
88	let index_chunk = self	×
89	.chunk_offsets()	×
90	.search_sorted(&index, SearchSortedSide::Right)	×
91	.to_ends_index(self.nchunks() + 1)	×
92	.saturating_sub(1);	×
93	let chunk_start = self.chunk_offsets()[index_chunk];	×
94
95	let index_in_chunk =	×
96	usize::try_from(index - chunk_start).vortex_expect("Index is too large for usize");	×
97	(index_chunk, index_in_chunk)	×
98	}	×
99
100	pub fn chunks(&self) -> &[ArrayRef] {	7,125✔
101	&self.chunks	7,125✔
102	}	7,125✔
103
104	pub fn non_empty_chunks(&self) -> impl Iterator<Item = &ArrayRef> + '_ {	×
105	self.chunks().iter().filter(\|c\| !c.is_empty())	×
106	}	×
107
108	pub fn array_iterator(&self) -> impl ArrayIterator + '_ {	×
109	ArrayIteratorAdapter::new(self.dtype().clone(), self.chunks().iter().cloned().map(Ok))	×
110	}	×
111
112	pub fn array_stream(&self) -> impl ArrayStream + '_ {	×
113	ArrayStreamAdapter::new(	×
114	self.dtype().clone(),	×
115	stream::iter(self.chunks().iter().cloned().map(Ok)),	×
116	)
117	}	×
118
119	pub fn rechunk(&self, target_bytesize: u64, target_rowsize: usize) -> VortexResult<Self> {	×
120	let mut new_chunks = Vec::new();	×
121	let mut chunks_to_combine = Vec::new();	×
122	let mut new_chunk_n_bytes = 0;	×
123	let mut new_chunk_n_elements = 0;	×
124	for chunk in self.chunks() {	×
125	let n_bytes = chunk.nbytes();	×
126	let n_elements = chunk.len();	×
127
128	if (new_chunk_n_bytes + n_bytes > target_bytesize	×
129	\|\| new_chunk_n_elements + n_elements > target_rowsize)	×
130	&& !chunks_to_combine.is_empty()	×
131	{
132	new_chunks.push(	×
133	ChunkedArray::new_unchecked(chunks_to_combine, self.dtype().clone())	×
134	.to_canonical()?	×
135	.into_array(),	×
136	);
137
138	new_chunk_n_bytes = 0;	×
139	new_chunk_n_elements = 0;	×
140	chunks_to_combine = Vec::new();	×
141	}	×
142
143	if n_bytes > target_bytesize \|\| n_elements > target_rowsize {	×
144	new_chunks.push(chunk.clone());	×
145	} else {	×
146	new_chunk_n_bytes += n_bytes;	×
147	new_chunk_n_elements += n_elements;	×
148	chunks_to_combine.push(chunk.clone());	×
149	}	×
150	}
151
152	if !chunks_to_combine.is_empty() {	×
153	new_chunks.push(	×
154	ChunkedArray::new_unchecked(chunks_to_combine, self.dtype().clone())	×
155	.to_canonical()?	×
156	.into_array(),	×
157	);
158	}	×
159
160	Ok(Self::new_unchecked(new_chunks, self.dtype().clone()))	×
161	}	×
162	}
163
164	impl FromIterator<ArrayRef> for ChunkedArray {
165	fn from_iter<T: IntoIterator<Item = ArrayRef>>(iter: T) -> Self {	×
166	let chunks: Vec<ArrayRef> = iter.into_iter().collect();	×
167	let dtype = chunks	×
168	.first()	×
169	.map(\|c\| c.dtype().clone())	×
170	.vortex_expect("Cannot infer DType from an empty iterator");	×
171	Self::try_new(chunks, dtype).vortex_expect("Failed to create chunked array from iterator")	×
172	}	×
173	}
174
175	impl ArrayVTable<ChunkedVTable> for ChunkedVTable {
176	fn len(array: &ChunkedArray) -> usize {	16,893✔
177	array.len	16,893✔
178	}	16,893✔
179
180	fn dtype(array: &ChunkedArray) -> &DType {	14,091✔
181	&array.dtype	14,091✔
182	}	14,091✔
183
184	fn stats(array: &ChunkedArray) -> StatsSetRef<'_> {	7,125✔
185	array.stats_set.to_ref(array.as_ref())	7,125✔
186	}	7,125✔
187	}
188
189	impl ValidityVTable<ChunkedVTable> for ChunkedVTable {
190	fn is_valid(array: &ChunkedArray, index: usize) -> VortexResult<bool> {	×
191	if !array.dtype.is_nullable() {	×
192	return Ok(true);	×
193	}	×
194	let (chunk, offset_in_chunk) = array.find_chunk_idx(index);	×
NEW 195	array.chunk(chunk).is_valid(offset_in_chunk)	×
196	}	×
197
198	fn all_valid(array: &ChunkedArray) -> VortexResult<bool> {	×
199	if !array.dtype().is_nullable() {	×
200	return Ok(true);	×
201	}	×
202	for chunk in array.non_empty_chunks() {	×
203	if !chunk.all_valid()? {	×
204	return Ok(false);	×
205	}	×
206	}
207	Ok(true)	×
208	}	×
209
210	fn all_invalid(array: &ChunkedArray) -> VortexResult<bool> {	×
211	if !array.dtype().is_nullable() {	×
212	return Ok(false);	×
213	}	×
214	for chunk in array.non_empty_chunks() {	×
215	if !chunk.all_invalid()? {	×
216	return Ok(false);	×
217	}	×
218	}
219	Ok(true)	×
220	}	×
221
222	fn validity_mask(array: &ChunkedArray) -> VortexResult<Mask> {	×
223	array	×
224	.chunks()	×
225	.iter()	×
226	.map(\|a\| a.validity_mask())	×
227	.try_collect()	×
228	}	×
229	}
230
231	#[cfg(test)]
232	mod test {
233	use vortex_buffer::buffer;
234	use vortex_dtype::{DType, Nullability, PType};
235	use vortex_error::VortexResult;
236
237	use crate::array::Array;
238	use crate::arrays::chunked::ChunkedArray;
239	use crate::arrays::{ChunkedVTable, PrimitiveArray};
240	use crate::compute::sub_scalar;
241	use crate::validity::Validity;
242	use crate::{IntoArray, ToCanonical, assert_arrays_eq};
243
244	fn chunked_array() -> ChunkedArray {
245	ChunkedArray::try_new(
246	vec![
247	buffer![1u64, 2, 3].into_array(),
248	buffer![4u64, 5, 6].into_array(),
249	buffer![7u64, 8, 9].into_array(),
250	],
251	DType::Primitive(PType::U64, Nullability::NonNullable),
252	)
253	.unwrap()
254	}
255
256	#[test]
257	fn test_scalar_subtract() {
258	let chunked = chunked_array().into_array();
259	let to_subtract = 1u64;
260	let array = sub_scalar(&chunked, to_subtract.into()).unwrap();
261
262	let chunked = array.as_::<ChunkedVTable>();
263	let chunks_out = chunked.chunks();
264
265	let results = chunks_out[0]
266	.to_primitive()
267	.unwrap()
268	.as_slice::<u64>()
269	.to_vec();
270	assert_eq!(results, &[0u64, 1, 2]);
271	let results = chunks_out[1]
272	.to_primitive()
273	.unwrap()
274	.as_slice::<u64>()
275	.to_vec();
276	assert_eq!(results, &[3u64, 4, 5]);
277	let results = chunks_out[2]
278	.to_primitive()
279	.unwrap()
280	.as_slice::<u64>()
281	.to_vec();
282	assert_eq!(results, &[6u64, 7, 8]);
283	}
284
285	#[test]
286	fn test_rechunk_one_chunk() {
287	let chunked = ChunkedArray::try_new(
288	vec![buffer![0u64].into_array()],
289	DType::Primitive(PType::U64, Nullability::NonNullable),
290	)
291	.unwrap();
292
293	let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap();
294
295	assert_arrays_eq!(chunked, rechunked);
296	}
297
298	#[test]
299	fn test_rechunk_two_chunks() {
300	let chunked = ChunkedArray::try_new(
301	vec![buffer![0u64].into_array(), buffer![5u64].into_array()],
302	DType::Primitive(PType::U64, Nullability::NonNullable),
303	)
304	.unwrap();
305
306	let rechunked = chunked.rechunk(1 << 16, 1 << 16).unwrap();
307
308	assert_eq!(rechunked.nchunks(), 1);
309	assert_arrays_eq!(chunked, rechunked);
310	}
311
312	#[test]
313	fn test_rechunk_tiny_target_chunks() {
314	let chunked = ChunkedArray::try_new(
315	vec![
316	buffer![0u64, 1, 2, 3].into_array(),
317	buffer![4u64, 5].into_array(),
318	],
319	DType::Primitive(PType::U64, Nullability::NonNullable),
320	)
321	.unwrap();
322
323	let rechunked = chunked.rechunk(1 << 16, 5).unwrap();
324
325	assert_eq!(rechunked.nchunks(), 2);
326	assert!(rechunked.chunks().iter().all(\|c\| c.len() < 5));
327	assert_arrays_eq!(chunked, rechunked);
328	}
329
330	#[test]
331	fn test_rechunk_with_too_big_chunk() {
332	let chunked = ChunkedArray::try_new(
333	vec![
334	buffer![0u64, 1, 2].into_array(),
335	buffer![42_u64; 6].into_array(),
336	buffer![4u64, 5].into_array(),
337	buffer![6u64, 7].into_array(),
338	buffer![8u64, 9].into_array(),
339	],
340	DType::Primitive(PType::U64, Nullability::NonNullable),
341	)
342	.unwrap();
343
344	let rechunked = chunked.rechunk(1 << 16, 5).unwrap();
345	// greedy so should be: [0, 1, 2] [42, 42, 42, 42, 42, 42] [4, 5, 6, 7] [8, 9]
346
347	assert_eq!(rechunked.nchunks(), 4);
348	assert_arrays_eq!(chunked, rechunked);
349	}
350
351	#[test]
352	fn test_empty_chunks_all_valid() -> VortexResult<()> {
353	// Create chunks where some are empty but all non-empty chunks have all valid values
354	let chunks = vec![
355	PrimitiveArray::new(buffer![1u64, 2, 3], Validity::AllValid).into_array(),
356	PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk
357	PrimitiveArray::new(buffer![4u64, 5], Validity::AllValid).into_array(),
358	PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk
359	];
360
361	let chunked =
362	ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?;
363
364	// Should be all_valid since all non-empty chunks are all_valid
365	assert!(chunked.all_valid()?);
366	assert!(!chunked.all_invalid()?);
367
368	Ok(())
369	}
370
371	#[test]
372	fn test_empty_chunks_all_invalid() -> VortexResult<()> {
373	// Create chunks where some are empty but all non-empty chunks have all invalid values
374	let chunks = vec![
375	PrimitiveArray::new(buffer![1u64, 2], Validity::AllInvalid).into_array(),
376	PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk
377	PrimitiveArray::new(buffer![3u64, 4, 5], Validity::AllInvalid).into_array(),
378	PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk
379	];
380
381	let chunked =
382	ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?;
383
384	// Should be all_invalid since all non-empty chunks are all_invalid
385	assert!(!chunked.all_valid()?);
386	assert!(chunked.all_invalid()?);
387
388	Ok(())
389	}
390
391	#[test]
392	fn test_empty_chunks_mixed_validity() -> VortexResult<()> {
393	// Create chunks with mixed validity including empty chunks
394	let chunks = vec![
395	PrimitiveArray::new(buffer![1u64, 2], Validity::AllValid).into_array(),
396	PrimitiveArray::new(buffer![0u64; 0], Validity::AllValid).into_array(), // empty chunk
397	PrimitiveArray::new(buffer![3u64, 4], Validity::AllInvalid).into_array(),
398	PrimitiveArray::new(buffer![0u64; 0], Validity::AllInvalid).into_array(), // empty chunk
399	];
400
401	let chunked =
402	ChunkedArray::try_new(chunks, DType::Primitive(PType::U64, Nullability::Nullable))?;
403
404	// Should be neither all_valid nor all_invalid
405	assert!(!chunked.all_valid()?);
406	assert!(!chunked.all_invalid()?);
407
408	Ok(())
409	}
410	}

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