16474110289

Committed 23 Jul 2025 02:52PM UTC coverage: 81.205% (+0.1%) from 81.11%

Build # 16474110289

Build Type

Pull #3992

github

Committed by

web-flow

Commit Message

Merge d7bd8890a into e793ee468

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

Run Details

235 of 246 new or added lines in 3 files covered. (95.53%)

8 existing lines in 4 files now uncovered.

42347 of 52148 relevant lines covered (81.21%)

173081.8 hits per line

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

77.52

/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, vortex_panic};
use vortex_mask::Mask;

use crate::arrays::{BoolArray, PrimitiveArray};
use crate::builders::ArrayBuilder;
use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
use crate::validity::Validity;
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 nulls = self.nulls.finish();

        if let Some(null_buf) = nulls.as_ref() {
            assert_eq!(
                null_buf.len(),
                self.values.len(),
                "null buffer length must equal value buffer length"
            );
        }

        let validity = match (nulls, self.dtype().nullability()) {
            (None, Nullability::NonNullable) => Validity::NonNullable,
            (Some(_), Nullability::NonNullable) => {
                vortex_panic!("Non-nullable builder has null values")
            }
            (None, Nullability::Nullable) => Validity::AllValid,
            (Some(nulls), Nullability::Nullable) => {
                if nulls.null_count() == nulls.len() {
                    Validity::AllInvalid
                } else {
                    Validity::Array(BoolArray::from(nulls.into_inner()).into_array())
                }
            }
        };

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

vortex-data / vortex / 16474110289

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