vortex_array/arrays/chunked/
array.rs

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

vortex_array/arrays/chunked/
array.rs
