16980200873

Committed 15 Aug 2025 12:53AM UTC coverage: 49.805%. First build

Build # 16980200873

Build Type

Pull #2456

github

Committed by

web-flow

Commit Message

Merge aff477380 into aaf3e36ad

Pull Request Pull Request #2456: feat: basic BoolBuffer / BoolBufferMut

Run Details

574 of 1074 new or added lines in 84 files covered. (53.45%)

20158 of 40474 relevant lines covered (49.8%)

238516.31 hits per line

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

0.0

/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_bit_buffer().into(), 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(\|\| {	×
20	let compute = ComputeFn::new("mask".into(), ArcRef::new_ref(&MaskFn));	×
21	for kernel in inventory::iter::<MaskKernelRef> {	×
22	compute.register_kernel(kernel.0.clone());	×
23	}	×
24	compute	×
25	});	×
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> {	×
57	MASK_FN	×
58	.invoke(&InvocationArgs {	×
59	inputs: &[array.into(), mask.into()],	×
60	options: &(),	×
61	})?	×
62	.unwrap_array()	×
63	}	×
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> {
77	pub const fn lift(&'static self) -> MaskKernelRef {	×
78	MaskKernelRef(ArcRef::new_ref(self))	×
79	}	×
80	}
81
82	impl<V: VTable + MaskKernel> Kernel for MaskKernelAdapter<V> {
83	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	×
84	let inputs = MaskArgs::try_from(args)?;	×
85	let Some(array) = inputs.array.as_opt::<V>() else {	×
86	return Ok(None);	×
87	};
88	Ok(Some(V::mask(&self.0, array, inputs.mask)?.into()))	×
89	}	×
90	}
91
92	struct MaskFn;
93
94	impl ComputeFnVTable for MaskFn {
95	fn invoke(	×
96	&self,	×
97	args: &InvocationArgs,	×
98	kernels: &[ArcRef<dyn Kernel>],	×
99	) -> VortexResult<Output> {	×
100	let MaskArgs { array, mask } = MaskArgs::try_from(args)?;	×
101
102	if matches!(mask, Mask::AllFalse(_)) {	×
103	// Fast-path for empty mask
104	return Ok(cast(array, &array.dtype().as_nullable())?.into());	×
105	}	×
106
107	if matches!(mask, Mask::AllTrue(_)) {	×
108	// Fast-path for full mask.
109	return Ok(ConstantArray::new(	×
110	Scalar::null(array.dtype().clone().as_nullable()),	×
111	array.len(),	×
112	)	×
113	.into_array()	×
114	.into());	×
115	}	×
116
117	for kernel in kernels {	×
118	if let Some(output) = kernel.invoke(args)? {	×
119	return Ok(output);	×
120	}	×
121	}
122	if let Some(output) = array.invoke(&MASK_FN, args)? {	×
123	return Ok(output);	×
124	}	×
125
126	// Fallback: implement using Arrow kernels.
127	log::debug!("No mask implementation found for {}", array.encoding_id());	×
128
129	let array_ref = array.to_array().into_arrow_preferred()?;	×
NEW 130	let mask = BooleanArray::new(mask.to_bit_buffer().into(), None);	×
131
132	let masked = arrow_select::nullif::nullif(array_ref.as_ref(), &mask)?;	×
133
134	Ok(ArrayRef::from_arrow(masked.as_ref(), true).into())	×
135	}	×
136
137	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	×
138	let MaskArgs { array, .. } = MaskArgs::try_from(args)?;	×
139	Ok(array.dtype().as_nullable())	×
140	}	×
141
142	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	×
143	let MaskArgs { array, mask } = MaskArgs::try_from(args)?;	×
144
145	if mask.len() != array.len() {	×
146	vortex_bail!(	×
147	"mask.len() is {}, does not equal array.len() of {}",	×
148	mask.len(),	×
149	array.len()	×
150	);
151	}	×
152
153	Ok(mask.len())	×
154	}	×
155
156	fn is_elementwise(&self) -> bool {	×
157	true	×
158	}	×
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> {	×
170	if value.inputs.len() != 2 {	×
171	vortex_bail!("Mask function requires 2 arguments");	×
172	}	×
173	let array = value.inputs[0]	×
174	.array()	×
175	.ok_or_else(\|\| vortex_err!("Expected input 0 to be an array"))?;	×
176	let mask = value.inputs[1]	×
177	.mask()	×
178	.ok_or_else(\|\| vortex_err!("Expected input 1 to be a mask"))?;	×
179
180	Ok(MaskArgs { array, mask })	×
181	}	×
182	}

vortex-data / vortex / 16980200873

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