16565982454

Committed 28 Jul 2025 09:57AM UTC coverage: 81.767% (-0.02%) from 81.789%

Build # 16565982454

Build Type

Pull #4031

github

Committed by

web-flow

Commit Message

Merge 73cac1802 into a0efcbe7a

Pull Request Pull Request #4031: zip compute function

Run Details

104 of 138 new or added lines in 1 file covered. (75.36%)

29 existing lines in 3 files now uncovered.

43326 of 52987 relevant lines covered (81.77%)

171359.09 hits per line

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

75.36

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

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

use std::sync::LazyLock;

use arcref::ArcRef;
use vortex_dtype::DType;
use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
use vortex_mask::{AllOr, Mask};

use super::{ComputeFnVTable, InvocationArgs, Output, cast};
use crate::builders::builder_with_capacity;
use crate::compute::{ComputeFn, Kernel};
use crate::vtable::VTable;
use crate::{Array, ArrayRef};

pub fn zip(mask: &Mask, if_true: &dyn Array, if_false: &dyn Array) -> VortexResult<ArrayRef> {
    ZIP_FN
        .invoke(&InvocationArgs {
            inputs: &[mask.into(), if_true.into(), if_false.into()],
            options: &(),
        })?
        .unwrap_array()
}

pub static ZIP_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("zip".into(), ArcRef::new_ref(&Zip));
    for kernel in inventory::iter::<ZipKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

pub struct Zip;

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

        if mask.all_true() {
            return Ok(cast(if_true, &zip_return_dtype(if_true, if_false))?.into());
        }

        if mask.all_false() {
            return Ok(cast(if_false, &zip_return_dtype(if_true, if_false))?.into());
        }

        if if_true.is_canonical() && if_false.is_canonical() {
            // skip kernel lookup if both arrays are canonical
            return Ok(zip_impl(mask, if_true, if_false)?.into());
        }

        // check if if_true supports zip directly
        for kernel in kernels {
            if let Some(output) = kernel.invoke(args)? {
                return Ok(output);
            }
        }

        if let Some(output) = if_true.invoke(&ZIP_FN, args)? {
            return Ok(output);
        }

        // TODO(os): add invert_mask opt and check if if_false has a kernel like:
        //           kernel.invoke(Args(if_false, if_true, mask, invert_mask = true))

        Ok(zip_impl(
            mask,
            if_true.to_canonical()?.as_ref(),
            if_false.to_canonical()?.as_ref(),
        )?
        .into())
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let ZipArgs {
            if_true, if_false, ..
        } = ZipArgs::try_from(args)?;

        if !if_true.dtype().eq_ignore_nullability(if_false.dtype()) {
            vortex_bail!("input arrays to zip must have the same dtype");
        }
        Ok(zip_return_dtype(if_true, if_false))
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        let ZipArgs { if_true, .. } = ZipArgs::try_from(args)?;
        // ComputeFn::invoke asserts if_true.len() == if_false.len(), because zip is elementwise
        Ok(if_true.len())
    }

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

struct ZipArgs<'a> {
    mask: &'a Mask,
    if_true: &'a dyn Array,
    if_false: &'a dyn Array,
}

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

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 3 {
            vortex_bail!("Expected 3 inputs for zip, found {}", value.inputs.len());
        }
        let mask = value.inputs[0]
            .mask()
            .ok_or_else(|| vortex_err!("Expected input 0 to be a mask"))?;

        let if_true = value.inputs[1]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 1 to be an array"))?;

        let if_false = value.inputs[2]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 2 to be an array"))?;
        Ok(Self {
            mask,
            if_true,
            if_false,
        })
    }
}

pub trait ZipKernel: VTable {
    fn zip(
        &self,
        mask: &Mask,
        if_true: &Self::Array,
        if_false: &dyn Array,
    ) -> VortexResult<ArrayRef>;
}

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

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

impl<V: VTable + ZipKernel> ZipKernelAdapter<V> {
    pub const fn lift(&'static self) -> ZipKernelRef {
        ZipKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + ZipKernel> Kernel for ZipKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let ZipArgs {
            mask,
            if_true,
            if_false,
        } = ZipArgs::try_from(args)?;
        let Some(if_true) = if_true.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(Some(V::zip(&self.0, mask, if_true, if_false)?.into()))
    }
}

fn zip_return_dtype(if_true: &dyn Array, if_false: &dyn Array) -> DType {
    if_true
        .dtype()
        .union_nullability(if_false.dtype().nullability())
}

fn zip_impl(mask: &Mask, if_true: &dyn Array, if_false: &dyn Array) -> VortexResult<ArrayRef> {
    // if_true.len() == if_false.len() from ComputeFn::invoke
    let mut builder = builder_with_capacity(&zip_return_dtype(if_true, if_false), if_true.len());

    match mask.slices() {
        AllOr::All => Ok(if_true.to_array()),
        AllOr::None => Ok(if_false.to_array()),
        AllOr::Some(slices) => {
            for (start, end) in slices {
                builder.extend_from_array(&if_false.slice(builder.len(), *start)?)?;
                builder.extend_from_array(&if_true.slice(*start, *end)?)?;
            }
            if builder.len() < if_false.len() {
                builder.extend_from_array(&if_false.slice(builder.len(), if_false.len())?)?;
            }
            Ok(builder.finish())
        }
    }
}

#[cfg(test)]
mod tests {
    use vortex_array::arrays::{BoolArray, PrimitiveArray};
    use vortex_array::compute::zip;
    use vortex_array::{IntoArray, ToCanonical};
    use vortex_mask::Mask;

    #[test]
    fn test_zip_basic() {
        let mask =
            Mask::try_from(&BoolArray::from_iter([true, false, false, true, false])).unwrap();
        let if_true = PrimitiveArray::from_iter([10, 20, 30, 40, 50]).into_array();
        let if_false = PrimitiveArray::from_iter([1, 2, 3, 4, 5]).into_array();

        let result = zip(&mask, &if_true, &if_false).unwrap();
        let expected = PrimitiveArray::from_iter([10, 2, 3, 40, 5]);

        assert_eq!(
            result.to_primitive().unwrap().as_slice::<i32>(),
            expected.as_slice::<i32>()
        );
    }

    #[test]
    fn test_zip_all_true() {
        let mask = Mask::new_true(4);
        let if_true = PrimitiveArray::from_iter([10, 20, 30, 40]).into_array();
        let if_false = PrimitiveArray::from_iter([1, 2, 3, 4]).into_array();

        let result = zip(&mask, &if_true, &if_false).unwrap();

        assert_eq!(
            result.to_primitive().unwrap().as_slice::<i32>(),
            if_true.to_primitive().unwrap().as_slice::<i32>()
        );
    }

    #[test]
    #[should_panic]
    fn test_invalid_lengths() {
        let mask = Mask::new_false(4);
        let if_true = PrimitiveArray::from_iter([10, 20, 30]).into_array();
        let if_false = PrimitiveArray::from_iter([1, 2, 3, 4]).into_array();

        zip(&mask, &if_true, &if_false).unwrap();
    }
}

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 vortex_dtype::DType;
8	use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
9	use vortex_mask::{AllOr, Mask};
10
11	use super::{ComputeFnVTable, InvocationArgs, Output, cast};
12	use crate::builders::builder_with_capacity;
13	use crate::compute::{ComputeFn, Kernel};
14	use crate::vtable::VTable;
15	use crate::{Array, ArrayRef};
16
17	pub fn zip(mask: &Mask, if_true: &dyn Array, if_false: &dyn Array) -> VortexResult<ArrayRef> {	114✔
18	ZIP_FN	114✔
19	.invoke(&InvocationArgs {	114✔
20	inputs: &[mask.into(), if_true.into(), if_false.into()],	114✔
21	options: &(),	114✔
22	})?	114✔
23	.unwrap_array()	76✔
24	}	114✔
25
26	pub static ZIP_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	114✔
27	let compute = ComputeFn::new("zip".into(), ArcRef::new_ref(&Zip));	114✔
28	for kernel in inventory::iter::<ZipKernelRef> {	114✔
NEW 29	compute.register_kernel(kernel.0.clone());	×
NEW 30	}	×
31	compute	114✔
32	});	114✔
33
34	pub struct Zip;
35
36	impl ComputeFnVTable for Zip {
37	fn invoke(	76✔
38	&self,	76✔
39	args: &InvocationArgs,	76✔
40	kernels: &[ArcRef<dyn Kernel>],	76✔
41	) -> VortexResult<Output> {	76✔
42	let ZipArgs {
43	mask,	76✔
44	if_true,	76✔
45	if_false,	76✔
46	} = ZipArgs::try_from(args)?;	76✔
47
48	if mask.all_true() {	76✔
49	return Ok(cast(if_true, &zip_return_dtype(if_true, if_false))?.into());	38✔
50	}	38✔
51
52	if mask.all_false() {	38✔
NEW 53	return Ok(cast(if_false, &zip_return_dtype(if_true, if_false))?.into());	×
54	}	38✔
55
56	if if_true.is_canonical() && if_false.is_canonical() {	38✔
57	// skip kernel lookup if both arrays are canonical
58	return Ok(zip_impl(mask, if_true, if_false)?.into());	38✔
NEW 59	}	×
60
61	// check if if_true supports zip directly
NEW 62	for kernel in kernels {	×
NEW 63	if let Some(output) = kernel.invoke(args)? {	×
NEW 64	return Ok(output);	×
NEW 65	}	×
66	}
67
NEW 68	if let Some(output) = if_true.invoke(&ZIP_FN, args)? {	×
NEW 69	return Ok(output);	×
NEW 70	}	×
71
72	// TODO(os): add invert_mask opt and check if if_false has a kernel like:
73	// kernel.invoke(Args(if_false, if_true, mask, invert_mask = true))
74
NEW 75	Ok(zip_impl(	×
NEW 76	mask,	×
NEW 77	if_true.to_canonical()?.as_ref(),	×
NEW 78	if_false.to_canonical()?.as_ref(),	×
NEW 79	)?	×
NEW 80	.into())	×
81	}	76✔
82
83	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	76✔
84	let ZipArgs {
85	if_true, if_false, ..	76✔
86	} = ZipArgs::try_from(args)?;	76✔
87
88	if !if_true.dtype().eq_ignore_nullability(if_false.dtype()) {	76✔
NEW 89	vortex_bail!("input arrays to zip must have the same dtype");	×
90	}	76✔
91	Ok(zip_return_dtype(if_true, if_false))	76✔
92	}	76✔
93
94	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	76✔
95	let ZipArgs { if_true, .. } = ZipArgs::try_from(args)?;	76✔
96	// ComputeFn::invoke asserts if_true.len() == if_false.len(), because zip is elementwise
97	Ok(if_true.len())	76✔
98	}	76✔
99
100	fn is_elementwise(&self) -> bool {	114✔
101	true	114✔
102	}	114✔
103	}
104
105	struct ZipArgs<'a> {
106	mask: &'a Mask,
107	if_true: &'a dyn Array,
108	if_false: &'a dyn Array,
109	}
110
111	impl<'a> TryFrom<&InvocationArgs<'a>> for ZipArgs<'a> {
112	type Error = VortexError;
113
114	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	228✔
115	if value.inputs.len() != 3 {	228✔
NEW 116	vortex_bail!("Expected 3 inputs for zip, found {}", value.inputs.len());	×
117	}	228✔
118	let mask = value.inputs[0]	228✔
119	.mask()	228✔
120	.ok_or_else(\|\| vortex_err!("Expected input 0 to be a mask"))?;	228✔
121
122	let if_true = value.inputs[1]	228✔
123	.array()	228✔
124	.ok_or_else(\|\| vortex_err!("Expected input 1 to be an array"))?;	228✔
125
126	let if_false = value.inputs[2]	228✔
127	.array()	228✔
128	.ok_or_else(\|\| vortex_err!("Expected input 2 to be an array"))?;	228✔
129	Ok(Self {	228✔
130	mask,	228✔
131	if_true,	228✔
132	if_false,	228✔
133	})	228✔
134	}	228✔
135	}
136
137	pub trait ZipKernel: VTable {
138	fn zip(
139	&self,
140	mask: &Mask,
141	if_true: &Self::Array,
142	if_false: &dyn Array,
143	) -> VortexResult<ArrayRef>;
144	}
145
146	pub struct ZipKernelRef(pub ArcRef<dyn Kernel>);
147	inventory::collect!(ZipKernelRef);
148
149	#[derive(Debug)]
150	pub struct ZipKernelAdapter<V: VTable>(pub V);
151
152	impl<V: VTable + ZipKernel> ZipKernelAdapter<V> {
NEW 153	pub const fn lift(&'static self) -> ZipKernelRef {	×
NEW 154	ZipKernelRef(ArcRef::new_ref(self))	×
NEW 155	}	×
156	}
157
158	impl<V: VTable + ZipKernel> Kernel for ZipKernelAdapter<V> {
NEW 159	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	×
160	let ZipArgs {
NEW 161	mask,	×
NEW 162	if_true,	×
NEW 163	if_false,	×
NEW 164	} = ZipArgs::try_from(args)?;	×
NEW 165	let Some(if_true) = if_true.as_opt::<V>() else {	×
NEW 166	return Ok(None);	×
167	};
NEW 168	Ok(Some(V::zip(&self.0, mask, if_true, if_false)?.into()))	×
NEW 169	}	×
170	}
171
172	fn zip_return_dtype(if_true: &dyn Array, if_false: &dyn Array) -> DType {	152✔
173	if_true	152✔
174	.dtype()	152✔
175	.union_nullability(if_false.dtype().nullability())	152✔
176	}	152✔
177
178	fn zip_impl(mask: &Mask, if_true: &dyn Array, if_false: &dyn Array) -> VortexResult<ArrayRef> {	38✔
179	// if_true.len() == if_false.len() from ComputeFn::invoke
180	let mut builder = builder_with_capacity(&zip_return_dtype(if_true, if_false), if_true.len());	38✔
181
182	match mask.slices() {	38✔
NEW 183	AllOr::All => Ok(if_true.to_array()),	×
NEW 184	AllOr::None => Ok(if_false.to_array()),	×
185	AllOr::Some(slices) => {	38✔
186	for (start, end) in slices {	114✔
187	builder.extend_from_array(&if_false.slice(builder.len(), *start)?)?;	76✔
188	builder.extend_from_array(&if_true.slice(start, end)?)?;	76✔
189	}
190	if builder.len() < if_false.len() {	38✔
191	builder.extend_from_array(&if_false.slice(builder.len(), if_false.len())?)?;	38✔
NEW 192	}	×
193	Ok(builder.finish())	38✔
194	}
195	}
196	}	38✔
197
198	#[cfg(test)]
199	mod tests {
200	use vortex_array::arrays::{BoolArray, PrimitiveArray};
201	use vortex_array::compute::zip;
202	use vortex_array::{IntoArray, ToCanonical};
203	use vortex_mask::Mask;
204
205	#[test]
206	fn test_zip_basic() {	1✔
207	let mask =	1✔
208	Mask::try_from(&BoolArray::from_iter([true, false, false, true, false])).unwrap();	1✔
209	let if_true = PrimitiveArray::from_iter([10, 20, 30, 40, 50]).into_array();	1✔
210	let if_false = PrimitiveArray::from_iter([1, 2, 3, 4, 5]).into_array();	1✔
211
212	let result = zip(&mask, &if_true, &if_false).unwrap();	1✔
213	let expected = PrimitiveArray::from_iter([10, 2, 3, 40, 5]);	1✔
214
215	assert_eq!(	1✔
216	result.to_primitive().unwrap().as_slice::<i32>(),	1✔
217	expected.as_slice::<i32>()	1✔
218	);
219	}	1✔
220
221	#[test]
222	fn test_zip_all_true() {	1✔
223	let mask = Mask::new_true(4);	1✔
224	let if_true = PrimitiveArray::from_iter([10, 20, 30, 40]).into_array();	1✔
225	let if_false = PrimitiveArray::from_iter([1, 2, 3, 4]).into_array();	1✔
226
227	let result = zip(&mask, &if_true, &if_false).unwrap();	1✔
228
229	assert_eq!(	1✔
230	result.to_primitive().unwrap().as_slice::<i32>(),	1✔
231	if_true.to_primitive().unwrap().as_slice::<i32>()	1✔
232	);
233	}	1✔
234
235	#[test]
236	#[should_panic]
237	fn test_invalid_lengths() {	1✔
238	let mask = Mask::new_false(4);	1✔
239	let if_true = PrimitiveArray::from_iter([10, 20, 30]).into_array();	1✔
240	let if_false = PrimitiveArray::from_iter([1, 2, 3, 4]).into_array();	1✔
241
242	zip(&mask, &if_true, &if_false).unwrap();	1✔
243	}	1✔
244	}

vortex-data / vortex / 16565982454

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