16595455362

Committed 29 Jul 2025 12:00PM UTC coverage: 82.182% (+0.4%) from 81.796%

Build # 16595455362

Build Type

Pull #4036

github

Committed by

web-flow

Commit Message

Merge 28b120788 into 261aabd6a

Pull Request Pull Request #4036: varbinview builder buffer deduplication

Run Details

147 of 155 new or added lines in 2 files covered. (94.84%)

404 existing lines in 34 files now uncovered.

44415 of 54045 relevant lines covered (82.18%)

167869.47 hits per line

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

89.16

/vortex-array/src/compute/mask.rs

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

use std::sync::LazyLock;

use arcref::ArcRef;
use arrow_array::BooleanArray;
use vortex_dtype::DType;
use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
use vortex_mask::Mask;
use vortex_scalar::Scalar;

use crate::arrays::ConstantArray;
use crate::arrow::{FromArrowArray, IntoArrowArray};
use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, cast};
use crate::vtable::VTable;
use crate::{Array, ArrayRef, IntoArray};

static MASK_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("mask".into(), ArcRef::new_ref(&MaskFn));
    for kernel in inventory::iter::<MaskKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

/// Replace values with null where the mask is true.
///
/// The returned array is nullable but otherwise has the same dtype and length as `array`.
///
/// # Examples
///
/// ```
/// use vortex_array::IntoArray;
/// use vortex_array::arrays::{BoolArray, PrimitiveArray};
/// use vortex_array::compute::{ mask};
/// use vortex_mask::Mask;
/// use vortex_scalar::Scalar;
///
/// let array =
///     PrimitiveArray::from_option_iter([Some(0i32), None, Some(1i32), None, Some(2i32)]);
/// let mask_array = Mask::try_from(
///     &BoolArray::from_iter([true, false, false, false, true]),
/// )
/// .unwrap();
///
/// let masked = mask(array.as_ref(), &mask_array).unwrap();
/// assert_eq!(masked.len(), 5);
/// assert!(!masked.is_valid(0).unwrap());
/// assert!(!masked.is_valid(1).unwrap());
/// assert_eq!(masked.scalar_at(2).unwrap(), Scalar::from(Some(1)));
/// assert!(!masked.is_valid(3).unwrap());
/// assert!(!masked.is_valid(4).unwrap());
/// ```
///
pub fn mask(array: &dyn Array, mask: &Mask) -> VortexResult<ArrayRef> {
    MASK_FN
        .invoke(&InvocationArgs {
            inputs: &[array.into(), mask.into()],
            options: &(),
        })?
        .unwrap_array()
}

pub struct MaskKernelRef(ArcRef<dyn Kernel>);
inventory::collect!(MaskKernelRef);

pub trait MaskKernel: VTable {
    /// Replace masked values with null in array.
    fn mask(&self, array: &Self::Array, mask: &Mask) -> VortexResult<ArrayRef>;
}

#[derive(Debug)]
pub struct MaskKernelAdapter<V: VTable>(pub V);

impl<V: VTable + MaskKernel> MaskKernelAdapter<V> {
    pub const fn lift(&'static self) -> MaskKernelRef {
        MaskKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + MaskKernel> Kernel for MaskKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let inputs = MaskArgs::try_from(args)?;
        let Some(array) = inputs.array.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(Some(V::mask(&self.0, array, inputs.mask)?.into()))
    }
}

struct MaskFn;

impl ComputeFnVTable for MaskFn {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let MaskArgs { array, mask } = MaskArgs::try_from(args)?;

        if matches!(mask, Mask::AllFalse(_)) {
            // Fast-path for empty mask
            return Ok(cast(array, &array.dtype().as_nullable())?.into());
        }

        if matches!(mask, Mask::AllTrue(_)) {
            // Fast-path for full mask.
            return Ok(ConstantArray::new(
                Scalar::null(array.dtype().clone().as_nullable()),
                array.len(),
            )
            .into_array()
            .into());
        }

        for kernel in kernels {
            if let Some(output) = kernel.invoke(args)? {
                return Ok(output);
            }
        }
        if let Some(output) = array.invoke(&MASK_FN, args)? {
            return Ok(output);
        }

        // Fallback: implement using Arrow kernels.
        log::debug!("No mask implementation found for {}", array.encoding_id());

        let array_ref = array.to_array().into_arrow_preferred()?;
        let mask = BooleanArray::new(mask.to_boolean_buffer(), None);

        let masked = arrow_select::nullif::nullif(array_ref.as_ref(), &mask)?;

        Ok(ArrayRef::from_arrow(masked.as_ref(), true).into())
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let MaskArgs { array, .. } = MaskArgs::try_from(args)?;
        Ok(array.dtype().as_nullable())
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        let MaskArgs { array, mask } = MaskArgs::try_from(args)?;

        if mask.len() != array.len() {
            vortex_bail!(
                "mask.len() is {}, does not equal array.len() of {}",
                mask.len(),
                array.len()
            );
        }

        Ok(mask.len())
    }

    fn is_elementwise(&self) -> bool {
        true
    }
}

struct MaskArgs<'a> {
    array: &'a dyn Array,
    mask: &'a Mask,
}

impl<'a> TryFrom<&InvocationArgs<'a>> for MaskArgs<'a> {
    type Error = VortexError;

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 2 {
            vortex_bail!("Mask function requires 2 arguments");
        }
        let array = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
        let mask = value.inputs[1]
            .mask()
            .ok_or_else(|| vortex_err!("Expected input 1 to be a mask"))?;

        Ok(MaskArgs { array, mask })
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::sync::LazyLock;
5
6	use arcref::ArcRef;
7	use arrow_array::BooleanArray;
8	use vortex_dtype::DType;
9	use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
10	use vortex_mask::Mask;
11	use vortex_scalar::Scalar;
12
13	use crate::arrays::ConstantArray;
14	use crate::arrow::{FromArrowArray, IntoArrowArray};
15	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, cast};
16	use crate::vtable::VTable;
17	use crate::{Array, ArrayRef, IntoArray};
18
19	static MASK_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	354✔
20	let compute = ComputeFn::new("mask".into(), ArcRef::new_ref(&MaskFn));	354✔
21	for kernel in inventory::iter::<MaskKernelRef> {	3,642✔
22	compute.register_kernel(kernel.0.clone());	3,288✔
23	}	3,288✔
24	compute	354✔
25	});	354✔
26
27	/// Replace values with null where the mask is true.
28	///
29	/// The returned array is nullable but otherwise has the same dtype and length as `array`.
30	///
31	/// # Examples
32	///
33	/// ```
34	/// use vortex_array::IntoArray;
35	/// use vortex_array::arrays::{BoolArray, PrimitiveArray};
36	/// use vortex_array::compute::{ mask};
37	/// use vortex_mask::Mask;
38	/// use vortex_scalar::Scalar;
39	///
40	/// let array =
41	/// PrimitiveArray::from_option_iter([Some(0i32), None, Some(1i32), None, Some(2i32)]);
42	/// let mask_array = Mask::try_from(
43	/// &BoolArray::from_iter([true, false, false, false, true]),
44	/// )
45	/// .unwrap();
46	///
47	/// let masked = mask(array.as_ref(), &mask_array).unwrap();
48	/// assert_eq!(masked.len(), 5);
49	/// assert!(!masked.is_valid(0).unwrap());
50	/// assert!(!masked.is_valid(1).unwrap());
51	/// assert_eq!(masked.scalar_at(2).unwrap(), Scalar::from(Some(1)));
52	/// assert!(!masked.is_valid(3).unwrap());
53	/// assert!(!masked.is_valid(4).unwrap());
54	/// ```
55	///
56	pub fn mask(array: &dyn Array, mask: &Mask) -> VortexResult<ArrayRef> {	3,873✔
57	MASK_FN	3,873✔
58	.invoke(&InvocationArgs {	3,873✔
59	inputs: &[array.into(), mask.into()],	3,873✔
60	options: &(),	3,873✔
61	})?	3,873✔
62	.unwrap_array()	3,873✔
63	}	3,873✔
64
65	pub struct MaskKernelRef(ArcRef<dyn Kernel>);
66	inventory::collect!(MaskKernelRef);
67
68	pub trait MaskKernel: VTable {
69	/// Replace masked values with null in array.
70	fn mask(&self, array: &Self::Array, mask: &Mask) -> VortexResult<ArrayRef>;
71	}
72
73	#[derive(Debug)]
74	pub struct MaskKernelAdapter<V: VTable>(pub V);
75
76	impl<V: VTable + MaskKernel> MaskKernelAdapter<V> {
UNCOV 77	pub const fn lift(&'static self) -> MaskKernelRef {	×
UNCOV 78	MaskKernelRef(ArcRef::new_ref(self))	×
UNCOV 79	}	×
80	}
81
82	impl<V: VTable + MaskKernel> Kernel for MaskKernelAdapter<V> {
83	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	20,340✔
84	let inputs = MaskArgs::try_from(args)?;	20,340✔
85	let Some(array) = inputs.array.as_opt::<V>() else {	20,340✔
86	return Ok(None);	19,454✔
87	};
88	Ok(Some(V::mask(&self.0, array, inputs.mask)?.into()))	886✔
89	}	20,340✔
90	}
91
92	struct MaskFn;
93
94	impl ComputeFnVTable for MaskFn {
95	fn invoke(	3,873✔
96	&self,	3,873✔
97	args: &InvocationArgs,	3,873✔
98	kernels: &[ArcRef<dyn Kernel>],	3,873✔
99	) -> VortexResult<Output> {	3,873✔
100	let MaskArgs { array, mask } = MaskArgs::try_from(args)?;	3,873✔
101
102	if matches!(mask, Mask::AllFalse(_)) {	3,873✔
103	// Fast-path for empty mask
104	return Ok(cast(array, &array.dtype().as_nullable())?.into());	591✔
105	}	3,282✔
106
107	if matches!(mask, Mask::AllTrue(_)) {	3,282✔
108	// Fast-path for full mask.
109	return Ok(ConstantArray::new(	593✔
110	Scalar::null(array.dtype().clone().as_nullable()),	593✔
111	array.len(),	593✔
112	)	593✔
113	.into_array()	593✔
114	.into());	593✔
115	}	2,689✔
116
117	for kernel in kernels {	22,957✔
118	if let Some(output) = kernel.invoke(args)? {	22,264✔
119	return Ok(output);	1,996✔
120	}	20,268✔
121	}
122	if let Some(output) = array.invoke(&MASK_FN, args)? {	693✔
123	return Ok(output);	×
124	}	693✔
125
126	// Fallback: implement using Arrow kernels.
127	log::debug!("No mask implementation found for {}", array.encoding_id());	693✔
128
129	let array_ref = array.to_array().into_arrow_preferred()?;	693✔
130	let mask = BooleanArray::new(mask.to_boolean_buffer(), None);	693✔
131
132	let masked = arrow_select::nullif::nullif(array_ref.as_ref(), &mask)?;	693✔
133
134	Ok(ArrayRef::from_arrow(masked.as_ref(), true).into())	693✔
135	}	3,873✔
136
137	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	3,873✔
138	let MaskArgs { array, .. } = MaskArgs::try_from(args)?;	3,873✔
139	Ok(array.dtype().as_nullable())	3,873✔
140	}	3,873✔
141
142	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	3,873✔
143	let MaskArgs { array, mask } = MaskArgs::try_from(args)?;	3,873✔
144
145	if mask.len() != array.len() {	3,873✔
146	vortex_bail!(	×
147	"mask.len() is {}, does not equal array.len() of {}",	×
148	mask.len(),	×
149	array.len()	×
150	);
151	}	3,873✔
152
153	Ok(mask.len())	3,873✔
154	}	3,873✔
155
156	fn is_elementwise(&self) -> bool {	3,873✔
157	true	3,873✔
158	}	3,873✔
159	}
160
161	struct MaskArgs<'a> {
162	array: &'a dyn Array,
163	mask: &'a Mask,
164	}
165
166	impl<'a> TryFrom<&InvocationArgs<'a>> for MaskArgs<'a> {
167	type Error = VortexError;
168
169	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	33,883✔
170	if value.inputs.len() != 2 {	33,883✔
171	vortex_bail!("Mask function requires 2 arguments");	×
172	}	33,883✔
173	let array = value.inputs[0]	33,883✔
174	.array()	33,883✔
175	.ok_or_else(\|\| vortex_err!("Expected input 0 to be an array"))?;	33,883✔
176	let mask = value.inputs[1]	33,883✔
177	.mask()	33,883✔
178	.ok_or_else(\|\| vortex_err!("Expected input 1 to be a mask"))?;	33,883✔
179
180	Ok(MaskArgs { array, mask })	33,883✔
181	}	33,883✔
182	}

vortex-data / vortex / 16595455362

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