16578751259

Committed 28 Jul 2025 07:46PM UTC coverage: 81.858% (+0.8%) from 81.087%

Build # 16578751259

Build Type

Pull #3992

github

Committed by

web-flow

Commit Message

Merge 0fe171459 into 1ae560509

Pull Request Pull Request #3992: feat: teach SparseArray to canonicalize lists

Run Details

222 of 241 new or added lines in 1 file covered. (92.12%)

495 existing lines in 40 files now uncovered.

43529 of 53176 relevant lines covered (81.86%)

169875.89 hits per line

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

77.78

/vortex-array/src/builders/primitive.rs

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

use std::any::Any;
use std::mem::MaybeUninit;
use std::ops::{Deref, DerefMut};

use vortex_buffer::BufferMut;
use vortex_dtype::{DType, NativePType, Nullability};
use vortex_error::{VortexResult, vortex_bail};
use vortex_mask::Mask;

use crate::arrays::PrimitiveArray;
use crate::builders::ArrayBuilder;
use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
use crate::{Array, ArrayRef, IntoArray, ToCanonical};

/// Builder for [`PrimitiveArray`].
pub struct PrimitiveBuilder<T> {
    values: BufferMut<T>,
    nulls: LazyNullBufferBuilder,
    dtype: DType,
}

impl<T: NativePType> PrimitiveBuilder<T> {
    pub fn new(nullability: Nullability) -> Self {
        Self::with_capacity(nullability, 1024) // Same as Arrow builders
    }

    pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
        Self {
            values: BufferMut::with_capacity(capacity),
            nulls: LazyNullBufferBuilder::new(capacity),
            dtype: DType::Primitive(T::PTYPE, nullability),
        }
    }

    /// Append a `Mask` to the null buffer.
    pub fn append_mask(&mut self, mask: Mask) {
        self.nulls.append_validity_mask(mask);
    }

    pub fn append_value(&mut self, value: T) {
        self.values.push(value);
        self.nulls.append(true);
    }

    pub fn append_option(&mut self, value: Option<T>) {
        match value {
            Some(value) => {
                self.values.push(value);
                self.nulls.append(true);
            }
            None => self.append_null(),
        }
    }

    pub fn values(&self) -> &[T] {
        self.values.as_ref()
    }

    /// Create a new handle to the next `len` uninitialized values in the builder.
    ///
    /// All reads/writes through the handle to the values buffer or the validity buffer will operate
    /// on indices relative to the start of the range.
    ///
    ///
    /// ## Example
    ///
    /// ```
    /// use std::mem::MaybeUninit;
    /// use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder};
    /// use vortex_dtype::Nullability;
    ///
    /// // Create a new builder.
    /// let mut builder: PrimitiveBuilder<i32> = PrimitiveBuilder::with_capacity(Nullability::NonNullable, 5);
    ///
    /// // Populate the values in reverse order.
    /// let mut range = builder.uninit_range(5);
    /// for i in [4, 3, 2, 1, 0] {
    ///     range[i] = MaybeUninit::new(i as i32);
    /// }
    /// range.finish();
    ///
    /// let built = builder.finish_into_primitive();
    ///
    /// assert_eq!(built.as_slice::<i32>(), &[0i32, 1, 2, 3, 4]);
    /// ```
    pub fn uninit_range(&mut self, len: usize) -> UninitRange<'_, T> {
        let offset = self.values.len();
        assert!(
            offset + len <= self.values.capacity(),
            "uninit_range of len {len} exceeds builder capacity {}",
            self.values.capacity()
        );

        UninitRange {
            offset,
            len,
            builder: self,
        }
    }

    pub fn finish_into_primitive(&mut self) -> PrimitiveArray {
        let validity = self
            .nulls
            .finish_with_nullability(self.dtype().nullability());

        PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity)
    }

    pub fn extend_with_iterator(&mut self, iter: impl IntoIterator<Item = T>, mask: Mask) {
        self.values.extend(iter);
        self.extend_with_validity_mask(mask)
    }

    fn extend_with_validity_mask(&mut self, validity_mask: Mask) {
        self.nulls.append_validity_mask(validity_mask);
    }
}

impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn dtype(&self) -> &DType {
        &self.dtype
    }

    fn len(&self) -> usize {
        self.values.len()
    }

    fn append_zeros(&mut self, n: usize) {
        self.values.push_n(T::default(), n);
        self.nulls.append_n_non_nulls(n);
    }

    fn append_nulls(&mut self, n: usize) {
        self.values.push_n(T::default(), n);
        self.nulls.append_n_nulls(n);
    }

    fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
        let array = array.to_primitive()?;
        if array.ptype() != T::PTYPE {
            vortex_bail!("Cannot extend from array with different ptype");
        }

        self.values.extend_from_slice(array.as_slice::<T>());

        self.extend_with_validity_mask(array.validity_mask()?);

        Ok(())
    }

    fn ensure_capacity(&mut self, capacity: usize) {
        if capacity > self.values.capacity() {
            self.values.reserve(capacity - self.values.len());
            self.nulls.ensure_capacity(capacity);
        }
    }

    fn set_validity(&mut self, validity: Mask) {
        self.nulls = LazyNullBufferBuilder::new(validity.len());
        self.nulls.append_validity_mask(validity);
    }

    fn finish(&mut self) -> ArrayRef {
        self.finish_into_primitive().into_array()
    }
}

pub struct UninitRange<'a, T> {
    offset: usize,
    len: usize,
    builder: &'a mut PrimitiveBuilder<T>,
}

impl<T> Deref for UninitRange<'_, T> {
    type Target = [MaybeUninit<T>];

    fn deref(&self) -> &[MaybeUninit<T>] {
        let start = self.builder.values.as_ptr();
        unsafe {
            // SAFETY: start + len is checked on construction to be in range.
            let dst = std::slice::from_raw_parts(start, self.len);

            // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
            let dst: &[MaybeUninit<T>] = std::mem::transmute(dst);

            dst
        }
    }
}

impl<T> DerefMut for UninitRange<'_, T> {
    fn deref_mut(&mut self) -> &mut [MaybeUninit<T>] {
        &mut self.builder.values.spare_capacity_mut()[..self.len]
    }
}

impl<T> UninitRange<'_, T> {
    /// Set a validity bit at the given index. The index is relative to the start of this range
    /// of the builder.
    pub fn set_bit(&mut self, index: usize, v: bool) {
        self.builder.nulls.set_bit(self.offset + index, v);
    }

    /// Set values from an initialized range.
    pub fn copy_from_init(&mut self, offset: usize, len: usize, src: &[T])
    where
        T: Copy,
    {
        // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
        let uninit_src: &[MaybeUninit<T>] = unsafe { std::mem::transmute(src) };

        let dst = &mut self[offset..][..len];
        dst.copy_from_slice(uninit_src);
    }

    /// Finish building this range, marking it as initialized and advancing the length of the
    /// underlying values buffer.
    pub fn finish(self) {
        // SAFETY: constructor enforces that offset + len does not exceed the capacity of the array.
        unsafe { self.builder.values.set_len(self.offset + self.len) };
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::any::Any;
5	use std::mem::MaybeUninit;
6	use std::ops::{Deref, DerefMut};
7
8	use vortex_buffer::BufferMut;
9	use vortex_dtype::{DType, NativePType, Nullability};
10	use vortex_error::{VortexResult, vortex_bail};
11	use vortex_mask::Mask;
12
13	use crate::arrays::PrimitiveArray;
14	use crate::builders::ArrayBuilder;
15	use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
16	use crate::{Array, ArrayRef, IntoArray, ToCanonical};
17
18	/// Builder for [`PrimitiveArray`].
19	pub struct PrimitiveBuilder<T> {
20	values: BufferMut<T>,
21	nulls: LazyNullBufferBuilder,
22	dtype: DType,
23	}
24
25	impl<T: NativePType> PrimitiveBuilder<T> {
UNCOV 26	pub fn new(nullability: Nullability) -> Self {	×
27	Self::with_capacity(nullability, 1024) // Same as Arrow builders	×
28	}	×
29
30	pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {	26,425✔
31	Self {	26,425✔
32	values: BufferMut::with_capacity(capacity),	26,425✔
33	nulls: LazyNullBufferBuilder::new(capacity),	26,425✔
34	dtype: DType::Primitive(T::PTYPE, nullability),	26,425✔
35	}	26,425✔
36	}	26,425✔
37
38	/// Append a `Mask` to the null buffer.
39	pub fn append_mask(&mut self, mask: Mask) {	7,560✔
40	self.nulls.append_validity_mask(mask);	7,560✔
41	}	7,560✔
42
43	pub fn append_value(&mut self, value: T) {	19✔
44	self.values.push(value);	19✔
45	self.nulls.append(true);	19✔
46	}	19✔
47
48	pub fn append_option(&mut self, value: Option<T>) {	33,762✔
49	match value {	33,762✔
50	Some(value) => {	33,762✔
51	self.values.push(value);	33,762✔
52	self.nulls.append(true);	33,762✔
53	}	33,762✔
UNCOV 54	None => self.append_null(),	×
55	}
56	}	33,762✔
57
UNCOV 58	pub fn values(&self) -> &[T] {	×
59	self.values.as_ref()	×
60	}	×
61
62	/// Create a new handle to the next `len` uninitialized values in the builder.
63	///
64	/// All reads/writes through the handle to the values buffer or the validity buffer will operate
65	/// on indices relative to the start of the range.
66	///
67	///
68	/// ## Example
69	///
70	/// ```
71	/// use std::mem::MaybeUninit;
72	/// use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder};
73	/// use vortex_dtype::Nullability;
74	///
75	/// // Create a new builder.
76	/// let mut builder: PrimitiveBuilder<i32> = PrimitiveBuilder::with_capacity(Nullability::NonNullable, 5);
77	///
78	/// // Populate the values in reverse order.
79	/// let mut range = builder.uninit_range(5);
80	/// for i in [4, 3, 2, 1, 0] {
81	/// range[i] = MaybeUninit::new(i as i32);
82	/// }
83	/// range.finish();
84	///
85	/// let built = builder.finish_into_primitive();
86	///
87	/// assert_eq!(built.as_slice::<i32>(), &[0i32, 1, 2, 3, 4]);
88	/// ```
89	pub fn uninit_range(&mut self, len: usize) -> UninitRange<'_, T> {	7,560✔
90	let offset = self.values.len();	7,560✔
91	assert!(	7,560✔
92	offset + len <= self.values.capacity(),	7,560✔
UNCOV 93	"uninit_range of len {len} exceeds builder capacity {}",	×
94	self.values.capacity()	×
95	);
96
97	UninitRange {	7,560✔
98	offset,	7,560✔
99	len,	7,560✔
100	builder: self,	7,560✔
101	}	7,560✔
102	}	7,560✔
103
104	pub fn finish_into_primitive(&mut self) -> PrimitiveArray {	26,423✔
105	let validity = self	26,423✔
106	.nulls	26,423✔
107	.finish_with_nullability(self.dtype().nullability());	26,423✔
108
109	PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity)	26,423✔
110	}	26,423✔
111
112	pub fn extend_with_iterator(&mut self, iter: impl IntoIterator<Item = T>, mask: Mask) {	×
UNCOV 113	self.values.extend(iter);	×
UNCOV 114	self.extend_with_validity_mask(mask)	×
UNCOV 115	}	×
116
117	fn extend_with_validity_mask(&mut self, validity_mask: Mask) {	33,000✔
118	self.nulls.append_validity_mask(validity_mask);	33,000✔
119	}	33,000✔
120	}
121
122	impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
UNCOV 123	fn as_any(&self) -> &dyn Any {	×
124	self	×
UNCOV 125	}	×
126
127	fn as_any_mut(&mut self) -> &mut dyn Any {	34,872✔
128	self	34,872✔
129	}	34,872✔
130
131	fn dtype(&self) -> &DType {	118,151✔
132	&self.dtype	118,151✔
133	}	118,151✔
134
135	fn len(&self) -> usize {	82,695✔
136	self.values.len()	82,695✔
137	}	82,695✔
138
139	fn append_zeros(&mut self, n: usize) {	490✔
140	self.values.push_n(T::default(), n);	490✔
141	self.nulls.append_n_non_nulls(n);	490✔
142	}	490✔
143
UNCOV 144	fn append_nulls(&mut self, n: usize) {	×
145	self.values.push_n(T::default(), n);	×
146	self.nulls.append_n_nulls(n);	×
147	}	×
148
149	fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {	33,000✔
150	let array = array.to_primitive()?;	33,000✔
151	if array.ptype() != T::PTYPE {	33,000✔
UNCOV 152	vortex_bail!("Cannot extend from array with different ptype");	×
153	}	33,000✔
154
155	self.values.extend_from_slice(array.as_slice::<T>());	33,000✔
156
157	self.extend_with_validity_mask(array.validity_mask()?);	33,000✔
158
159	Ok(())	33,000✔
160	}	33,000✔
161
162	fn ensure_capacity(&mut self, capacity: usize) {	439✔
163	if capacity > self.values.capacity() {	439✔
UNCOV 164	self.values.reserve(capacity - self.values.len());	×
UNCOV 165	self.nulls.ensure_capacity(capacity);	×
166	}	439✔
167	}	439✔
168
169	fn set_validity(&mut self, validity: Mask) {	1✔
170	self.nulls = LazyNullBufferBuilder::new(validity.len());	1✔
171	self.nulls.append_validity_mask(validity);	1✔
172	}	1✔
173
174	fn finish(&mut self) -> ArrayRef {	17,621✔
175	self.finish_into_primitive().into_array()	17,621✔
176	}	17,621✔
177	}
178
179	pub struct UninitRange<'a, T> {
180	offset: usize,
181	len: usize,
182	builder: &'a mut PrimitiveBuilder<T>,
183	}
184
185	impl<T> Deref for UninitRange<'_, T> {
186	type Target = [MaybeUninit<T>];
187
188	fn deref(&self) -> &[MaybeUninit<T>] {	633✔
189	let start = self.builder.values.as_ptr();	633✔
190	unsafe {
191	// SAFETY: start + len is checked on construction to be in range.
192	let dst = std::slice::from_raw_parts(start, self.len);	633✔
193
194	// SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
195	let dst: &[MaybeUninit<T>] = std::mem::transmute(dst);	633✔
196
197	dst	633✔
198	}
199	}	633✔
200	}
201
202	impl<T> DerefMut for UninitRange<'_, T> {
203	fn deref_mut(&mut self) -> &mut [MaybeUninit<T>] {	614,282✔
204	&mut self.builder.values.spare_capacity_mut()[..self.len]	614,282✔
205	}	614,282✔
206	}
207
208	impl<T> UninitRange<'_, T> {
209	/// Set a validity bit at the given index. The index is relative to the start of this range
210	/// of the builder.
UNCOV 211	pub fn set_bit(&mut self, index: usize, v: bool) {	×
UNCOV 212	self.builder.nulls.set_bit(self.offset + index, v);	×
UNCOV 213	}	×
214
215	/// Set values from an initialized range.
216	pub fn copy_from_init(&mut self, offset: usize, len: usize, src: &[T])	4,681✔
217	where	4,681✔
218	T: Copy,	4,681✔
219	{
220	// SAFETY: &[T] and &[MaybeUninit<T>] have the same layout
221	let uninit_src: &[MaybeUninit<T>] = unsafe { std::mem::transmute(src) };	4,681✔
222
223	let dst = &mut self[offset..][..len];	4,681✔
224	dst.copy_from_slice(uninit_src);	4,681✔
225	}	4,681✔
226
227	/// Finish building this range, marking it as initialized and advancing the length of the
228	/// underlying values buffer.
229	pub fn finish(self) {	7,560✔
230	// SAFETY: constructor enforces that offset + len does not exceed the capacity of the array.
231	unsafe { self.builder.values.set_len(self.offset + self.len) };	7,560✔
232	}	7,560✔
233	}

vortex-data / vortex / 16578751259

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