16935267080

Committed 13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%

Build # 16935267080

Build Type

Pull #4226

github

Committed by

web-flow

Commit Message

Merge 81b48c7fb into baa6ea202

Pull Request Pull Request #4226: Support converting TimestampTZ to and from duckdb

Run Details

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

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

0.0

/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::{ArrayBuilder, builder_with_capacity};
use crate::compute::{ComputeFn, Kernel};
use crate::vtable::VTable;
use crate::{Array, ArrayRef};

/// Performs element-wise conditional selection between two arrays based on a mask.
///
/// Returns a new array where `result[i] = if_true[i]` when `mask[i]` is true,
/// otherwise `result[i] = if_false[i]`.
pub fn zip(if_true: &dyn Array, if_false: &dyn Array, mask: &Mask) -> VortexResult<ArrayRef> {
    ZIP_FN
        .invoke(&InvocationArgs {
            inputs: &[if_true.into(), if_false.into(), mask.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
});

struct Zip;

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

        // 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(
            if_true.to_canonical()?.as_ref(),
            if_false.to_canonical()?.as_ref(),
            mask,
        )?
        .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, mask, .. } = ZipArgs::try_from(args)?;
        // ComputeFn::invoke asserts if_true.len() == if_false.len(), because zip is elementwise
        if if_true.len() != mask.len() {
            vortex_bail!("input arrays must have the same length as the mask");
        }
        Ok(if_true.len())
    }

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

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

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 if_true = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;

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

        let mask = value.inputs[2]
            .mask()
            .ok_or_else(|| vortex_err!("Expected input 2 to be a mask"))?;

        Ok(Self {
            if_true,
            if_false,
            mask,
        })
    }
}

pub trait ZipKernel: VTable {
    fn zip(
        &self,
        if_true: &Self::Array,
        if_false: &dyn Array,
        mask: &Mask,
    ) -> VortexResult<Option<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 {
            if_true,
            if_false,
            mask,
        } = ZipArgs::try_from(args)?;
        let Some(if_true) = if_true.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(V::zip(&self.0, if_true, if_false, mask)?.map(Into::into))
    }
}

pub(crate) 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(if_true: &dyn Array, if_false: &dyn Array, mask: &Mask) -> VortexResult<ArrayRef> {
    // if_true.len() == if_false.len() from ComputeFn::invoke
    let builder = builder_with_capacity(&zip_return_dtype(if_true, if_false), if_true.len());
    zip_impl_with_builder(if_true, if_false, mask, builder)
}

pub(crate) fn zip_impl_with_builder(
    if_true: &dyn Array,
    if_false: &dyn Array,
    mask: &Mask,
    mut builder: Box<dyn ArrayBuilder>,
) -> VortexResult<ArrayRef> {
    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(&if_true, &if_false, &mask).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_option_iter([Some(1), Some(2), Some(3), None]).into_array();

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

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

        // result must be nullable even if_true was not
        assert_eq!(result.dtype(), if_false.dtype())
    }

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

vortex-data / vortex / 16935267080

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