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

Coverage Stats

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

70.74

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

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

use core::fmt;
use std::any::Any;
use std::fmt::{Display, Formatter};
use std::sync::LazyLock;

use arcref::ArcRef;
use arrow_ord::cmp;
use vortex_buffer::BitBuffer;
use vortex_dtype::{DType, NativePType, Nullability};
use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
use vortex_scalar::Scalar;

use crate::arrays::ConstantArray;
use crate::arrow::{Datum, from_arrow_array_with_len};
use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
use crate::vtable::VTable;
use crate::{Array, ArrayRef, Canonical, IntoArray};

static COMPARE_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("compare".into(), ArcRef::new_ref(&Compare));
    for kernel in inventory::iter::<CompareKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

/// Compares two arrays and returns a new boolean array with the result of the comparison.
/// Or, returns None if comparison is not supported for these arrays.
pub fn compare(left: &dyn Array, right: &dyn Array, operator: Operator) -> VortexResult<ArrayRef> {
    COMPARE_FN
        .invoke(&InvocationArgs {
            inputs: &[left.into(), right.into()],
            options: &operator,
        })?
        .unwrap_array()
}

#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
pub enum Operator {
    /// Equality (`=`)
    Eq,
    /// Inequality (`!=`)
    NotEq,
    /// Greater than (`>`)
    Gt,
    /// Greater than or equal (`>=`)
    Gte,
    /// Less than (`<`)
    Lt,
    /// Less than or equal (`<=`)
    Lte,
}

impl Display for Operator {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let display = match &self {
            Operator::Eq => "=",
            Operator::NotEq => "!=",
            Operator::Gt => ">",
            Operator::Gte => ">=",
            Operator::Lt => "<",
            Operator::Lte => "<=",
        };
        Display::fmt(display, f)
    }
}

impl Operator {
    pub fn inverse(self) -> Self {
        match self {
            Operator::Eq => Operator::NotEq,
            Operator::NotEq => Operator::Eq,
            Operator::Gt => Operator::Lte,
            Operator::Gte => Operator::Lt,
            Operator::Lt => Operator::Gte,
            Operator::Lte => Operator::Gt,
        }
    }

    /// Change the sides of the operator, where changing lhs and rhs won't change the result of the operation
    pub fn swap(self) -> Self {
        match self {
            Operator::Eq => Operator::Eq,
            Operator::NotEq => Operator::NotEq,
            Operator::Gt => Operator::Lt,
            Operator::Gte => Operator::Lte,
            Operator::Lt => Operator::Gt,
            Operator::Lte => Operator::Gte,
        }
    }
}

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

pub trait CompareKernel: VTable {
    fn compare(
        &self,
        lhs: &Self::Array,
        rhs: &dyn Array,
        operator: Operator,
    ) -> VortexResult<Option<ArrayRef>>;
}

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

impl<V: VTable + CompareKernel> CompareKernelAdapter<V> {
    pub const fn lift(&'static self) -> CompareKernelRef {
        CompareKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + CompareKernel> Kernel for CompareKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let inputs = CompareArgs::try_from(args)?;
        let Some(array) = inputs.lhs.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(V::compare(&self.0, array, inputs.rhs, inputs.operator)?.map(|array| array.into()))
    }
}

struct Compare;

impl ComputeFnVTable for Compare {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let CompareArgs { lhs, rhs, operator } = CompareArgs::try_from(args)?;

        let return_dtype = self.return_dtype(args)?;

        if lhs.is_empty() {
            return Ok(Canonical::empty(&return_dtype).into_array().into());
        }

        let left_constant_null = lhs.as_constant().map(|l| l.is_null()).unwrap_or(false);
        let right_constant_null = rhs.as_constant().map(|r| r.is_null()).unwrap_or(false);
        if left_constant_null || right_constant_null {
            return Ok(ConstantArray::new(Scalar::null(return_dtype), lhs.len())
                .into_array()
                .into());
        }

        let right_is_constant = rhs.is_constant();

        // Always try to put constants on the right-hand side so encodings can optimise themselves.
        if lhs.is_constant() && !right_is_constant {
            return Ok(compare(rhs, lhs, operator.swap())?.into());
        }

        // First try lhs op rhs, then invert and try again.
        for kernel in kernels {
            if let Some(output) = kernel.invoke(args)? {
                return Ok(output);
            }
        }
        if let Some(output) = lhs.invoke(&COMPARE_FN, args)? {
            return Ok(output);
        }

        // Try inverting the operator and swapping the arguments
        let inverted_args = InvocationArgs {
            inputs: &[rhs.into(), lhs.into()],
            options: &operator.swap(),
        };
        for kernel in kernels {
            if let Some(output) = kernel.invoke(&inverted_args)? {
                return Ok(output);
            }
        }
        if let Some(output) = rhs.invoke(&COMPARE_FN, &inverted_args)? {
            return Ok(output);
        }

        // Only log missing compare implementation if there's possibly better one than arrow,
        // i.e. lhs isn't arrow or rhs isn't arrow or constant
        if !(lhs.is_arrow() && (rhs.is_arrow() || right_is_constant)) {
            log::debug!(
                "No compare implementation found for LHS {}, RHS {}, and operator {} (or inverse)",
                lhs.encoding_id(),
                rhs.encoding_id(),
                operator,
            );
        }

        // Fallback to arrow on canonical types
        Ok(arrow_compare(lhs, rhs, operator)?.into())
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let CompareArgs { lhs, rhs, .. } = CompareArgs::try_from(args)?;

        if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) {
            vortex_bail!(
                "Cannot compare different DTypes {} and {}",
                lhs.dtype(),
                rhs.dtype()
            );
        }

        // TODO(ngates): no reason why not
        if lhs.dtype().is_struct() {
            vortex_bail!(
                "Compare does not support arrays with Struct DType, got: {} and {}",
                lhs.dtype(),
                rhs.dtype()
            )
        }

        Ok(DType::Bool(
            lhs.dtype().nullability() | rhs.dtype().nullability(),
        ))
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        let CompareArgs { lhs, rhs, .. } = CompareArgs::try_from(args)?;
        if lhs.len() != rhs.len() {
            vortex_bail!(
                "Compare operations only support arrays of the same length, got {} and {}",
                lhs.len(),
                rhs.len()
            );
        }
        Ok(lhs.len())
    }

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

struct CompareArgs<'a> {
    lhs: &'a dyn Array,
    rhs: &'a dyn Array,
    operator: Operator,
}

impl Options for Operator {
    fn as_any(&self) -> &dyn Any {
        self
    }
}

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

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 2 {
            vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());
        }
        let lhs = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected first input to be an array"))?;
        let rhs = value.inputs[1]
            .array()
            .ok_or_else(|| vortex_err!("Expected second input to be an array"))?;
        let operator = *value
            .options
            .as_any()
            .downcast_ref::<Operator>()
            .vortex_expect("Expected options to be an operator");

        Ok(CompareArgs { lhs, rhs, operator })
    }
}

/// Helper function to compare empty values with arrays that have external value length information
/// like `VarBin`.
pub fn compare_lengths_to_empty<P, I>(lengths: I, op: Operator) -> BitBuffer
where
    P: NativePType,
    I: Iterator<Item = P>,
{
    // All comparison can be expressed in terms of equality. "" is the absolute min of possible value.
    let cmp_fn = match op {
        Operator::Eq | Operator::Lte => |v| v == P::zero(),
        Operator::NotEq | Operator::Gt => |v| v != P::zero(),
        Operator::Gte => |_| true,
        Operator::Lt => |_| false,
    };

    lengths.map(cmp_fn).collect()
}

/// Implementation of `CompareFn` using the Arrow crate.
fn arrow_compare(
    left: &dyn Array,
    right: &dyn Array,
    operator: Operator,
) -> VortexResult<ArrayRef> {
    let nullable = left.dtype().is_nullable() || right.dtype().is_nullable();
    let lhs = Datum::try_new(left)?;
    let rhs = Datum::try_new(right)?;

    let array = match operator {
        Operator::Eq => cmp::eq(&lhs, &rhs)?,
        Operator::NotEq => cmp::neq(&lhs, &rhs)?,
        Operator::Gt => cmp::gt(&lhs, &rhs)?,
        Operator::Gte => cmp::gt_eq(&lhs, &rhs)?,
        Operator::Lt => cmp::lt(&lhs, &rhs)?,
        Operator::Lte => cmp::lt_eq(&lhs, &rhs)?,
    };
    from_arrow_array_with_len(&array, left.len(), nullable)
}

pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: Operator) -> Scalar {
    if lhs.is_null() | rhs.is_null() {
        Scalar::null(DType::Bool(Nullability::Nullable))
    } else {
        let b = match operator {
            Operator::Eq => lhs == rhs,
            Operator::NotEq => lhs != rhs,
            Operator::Gt => lhs > rhs,
            Operator::Gte => lhs >= rhs,
            Operator::Lt => lhs < rhs,
            Operator::Lte => lhs <= rhs,
        };

        Scalar::bool(b, lhs.dtype().nullability() | rhs.dtype().nullability())
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;

    use super::*;
    use crate::ToCanonical;
    use crate::arrays::{BoolArray, ConstantArray, VarBinArray, VarBinViewArray};
    use crate::test_harness::to_int_indices;
    use crate::validity::Validity;

    #[test]
    fn test_bool_basic_comparisons() {
        let arr = BoolArray::new(
            BitBuffer::from_iter([true, true, false, true, false]),
            Validity::from_iter([false, true, true, true, true]),
        );

        let matches = compare(arr.as_ref(), arr.as_ref(), Operator::Eq)
            .unwrap()
            .to_bool()
            .unwrap();

        assert_eq!(to_int_indices(matches).unwrap(), [1u64, 2, 3, 4]);

        let matches = compare(arr.as_ref(), arr.as_ref(), Operator::NotEq)
            .unwrap()
            .to_bool()
            .unwrap();
        let empty: [u64; 0] = [];
        assert_eq!(to_int_indices(matches).unwrap(), empty);

        let other = BoolArray::new(
            BitBuffer::from_iter([false, false, false, true, true]),
            Validity::from_iter([false, true, true, true, true]),
        );

        let matches = compare(arr.as_ref(), other.as_ref(), Operator::Lte)
            .unwrap()
            .to_bool()
            .unwrap();
        assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]);

        let matches = compare(arr.as_ref(), other.as_ref(), Operator::Lt)
            .unwrap()
            .to_bool()
            .unwrap();
        assert_eq!(to_int_indices(matches).unwrap(), [4u64]);

        let matches = compare(other.as_ref(), arr.as_ref(), Operator::Gte)
            .unwrap()
            .to_bool()
            .unwrap();
        assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]);

        let matches = compare(other.as_ref(), arr.as_ref(), Operator::Gt)
            .unwrap()
            .to_bool()
            .unwrap();
        assert_eq!(to_int_indices(matches).unwrap(), [4u64]);
    }

    #[test]
    fn constant_compare() {
        let left = ConstantArray::new(Scalar::from(2u32), 10);
        let right = ConstantArray::new(Scalar::from(10u32), 10);

        let compare = compare(left.as_ref(), right.as_ref(), Operator::Gt).unwrap();
        let res = compare.as_constant().unwrap();
        assert_eq!(res.as_bool().value(), Some(false));
        assert_eq!(compare.len(), 10);

        let compare = arrow_compare(&left.into_array(), &right.into_array(), Operator::Gt).unwrap();
        let res = compare.as_constant().unwrap();
        assert_eq!(res.as_bool().value(), Some(false));
        assert_eq!(compare.len(), 10);
    }

    #[rstest]
    #[case(Operator::Eq, vec![false, false, false, true])]
    #[case(Operator::NotEq, vec![true, true, true, false])]
    #[case(Operator::Gt, vec![true, true, true, false])]
    #[case(Operator::Gte, vec![true, true, true, true])]
    #[case(Operator::Lt, vec![false, false, false, false])]
    #[case(Operator::Lte, vec![false, false, false, true])]
    fn test_cmp_to_empty(#[case] op: Operator, #[case] expected: Vec<bool>) {
        let lengths: Vec<i32> = vec![1, 5, 7, 0];

        let output = compare_lengths_to_empty(lengths.iter().copied(), op);
        assert_eq!(Vec::from_iter(output.iter()), expected);
    }

    #[rstest]
    #[case(VarBinArray::from(vec!["a", "b"]).into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array())]
    #[case(VarBinViewArray::from_iter_str(["a", "b"]).into_array(), VarBinArray::from(vec!["a", "b"]).into_array())]
    #[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())]
    #[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())]
    fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) {
        let res = compare(&left, &right, Operator::Eq).unwrap();
        assert_eq!(res.to_bool().unwrap().bit_buffer().true_count(), left.len());
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use core::fmt;
5	use std::any::Any;
6	use std::fmt::{Display, Formatter};
7	use std::sync::LazyLock;
8
9	use arcref::ArcRef;
10	use arrow_ord::cmp;
11	use vortex_buffer::BitBuffer;
12	use vortex_dtype::{DType, NativePType, Nullability};
13	use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
14	use vortex_scalar::Scalar;
15
16	use crate::arrays::ConstantArray;
17	use crate::arrow::{Datum, from_arrow_array_with_len};
18	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output};
19	use crate::vtable::VTable;
20	use crate::{Array, ArrayRef, Canonical, IntoArray};
21
22	static COMPARE_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	3✔
23	let compute = ComputeFn::new("compare".into(), ArcRef::new_ref(&Compare));	3✔
24	for kernel in inventory::iter::<CompareKernelRef> {	36✔
25	compute.register_kernel(kernel.0.clone());	33✔
26	}	33✔
27	compute	3✔
28	});	3✔
29
30	/// Compares two arrays and returns a new boolean array with the result of the comparison.
31	/// Or, returns None if comparison is not supported for these arrays.
32	pub fn compare(left: &dyn Array, right: &dyn Array, operator: Operator) -> VortexResult<ArrayRef> {	14,391✔
33	COMPARE_FN	14,391✔
34	.invoke(&InvocationArgs {	14,391✔
35	inputs: &[left.into(), right.into()],	14,391✔
36	options: &operator,	14,391✔
37	})?	14,391✔
38	.unwrap_array()	14,391✔
39	}	14,391✔
40
41	#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Hash)]
42	pub enum Operator {
43	/// Equality (`=`)
44	Eq,
45	/// Inequality (`!=`)
46	NotEq,
47	/// Greater than (`>`)
48	Gt,
49	/// Greater than or equal (`>=`)
50	Gte,
51	/// Less than (`<`)
52	Lt,
53	/// Less than or equal (`<=`)
54	Lte,
55	}
56
57	impl Display for Operator {
58	fn fmt(&self, f: &mut Formatter) -> fmt::Result {	×
59	let display = match &self {	×
60	Operator::Eq => "=",	×
61	Operator::NotEq => "!=",	×
62	Operator::Gt => ">",	×
63	Operator::Gte => ">=",	×
64	Operator::Lt => "<",	×
65	Operator::Lte => "<=",	×
66	};
67	Display::fmt(display, f)	×
68	}	×
69	}
70
71	impl Operator {
72	pub fn inverse(self) -> Self {	×
73	match self {	×
74	Operator::Eq => Operator::NotEq,	×
75	Operator::NotEq => Operator::Eq,	×
76	Operator::Gt => Operator::Lte,	×
77	Operator::Gte => Operator::Lt,	×
78	Operator::Lt => Operator::Gte,	×
79	Operator::Lte => Operator::Gt,	×
80	}
81	}	×
82
83	/// Change the sides of the operator, where changing lhs and rhs won't change the result of the operation
84	pub fn swap(self) -> Self {	8,424✔
85	match self {	8,424✔
86	Operator::Eq => Operator::Eq,	3,324✔
87	Operator::NotEq => Operator::NotEq,	6✔
88	Operator::Gt => Operator::Lt,	690✔
89	Operator::Gte => Operator::Lte,	1,302✔
90	Operator::Lt => Operator::Gt,	894✔
91	Operator::Lte => Operator::Gte,	2,208✔
92	}
93	}	8,424✔
94	}
95
96	pub struct CompareKernelRef(ArcRef<dyn Kernel>);
97	inventory::collect!(CompareKernelRef);
98
99	pub trait CompareKernel: VTable {
100	fn compare(
101	&self,
102	lhs: &Self::Array,
103	rhs: &dyn Array,
104	operator: Operator,
105	) -> VortexResult<Option<ArrayRef>>;
106	}
107
108	#[derive(Debug)]
109	pub struct CompareKernelAdapter<V: VTable>(pub V);
110
111	impl<V: VTable + CompareKernel> CompareKernelAdapter<V> {
112	pub const fn lift(&'static self) -> CompareKernelRef {	×
113	CompareKernelRef(ArcRef::new_ref(self))	×
114	}	×
115	}
116
117	impl<V: VTable + CompareKernel> Kernel for CompareKernelAdapter<V> {
118	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	135,927✔
119	let inputs = CompareArgs::try_from(args)?;	135,927✔
120	let Some(array) = inputs.lhs.as_opt::<V>() else {	135,927✔
121	return Ok(None);	121,356✔
122	};
123	Ok(V::compare(&self.0, array, inputs.rhs, inputs.operator)?.map(\|array\| array.into()))	14,571✔
124	}	135,927✔
125	}
126
127	struct Compare;
128
129	impl ComputeFnVTable for Compare {
130	fn invoke(	14,391✔
131	&self,	14,391✔
132	args: &InvocationArgs,	14,391✔
133	kernels: &[ArcRef<dyn Kernel>],	14,391✔
134	) -> VortexResult<Output> {	14,391✔
135	let CompareArgs { lhs, rhs, operator } = CompareArgs::try_from(args)?;	14,391✔
136
137	let return_dtype = self.return_dtype(args)?;	14,391✔
138
139	if lhs.is_empty() {	14,391✔
140	return Ok(Canonical::empty(&return_dtype).into_array().into());	×
141	}	14,391✔
142
143	let left_constant_null = lhs.as_constant().map(\|l\| l.is_null()).unwrap_or(false);	14,391✔
144	let right_constant_null = rhs.as_constant().map(\|r\| r.is_null()).unwrap_or(false);	14,391✔
145	if left_constant_null \|\| right_constant_null {	14,391✔
146	return Ok(ConstantArray::new(Scalar::null(return_dtype), lhs.len())	×
147	.into_array()	×
148	.into());	×
149	}	14,391✔
150
151	let right_is_constant = rhs.is_constant();	14,391✔
152
153	// Always try to put constants on the right-hand side so encodings can optimise themselves.
154	if lhs.is_constant() && !right_is_constant {	14,391✔
155	return Ok(compare(rhs, lhs, operator.swap())?.into());	2,877✔
156	}	11,514✔
157
158	// First try lhs op rhs, then invert and try again.
159	for kernel in kernels {	83,049✔
160	if let Some(output) = kernel.invoke(args)? {	77,502✔
161	return Ok(output);	5,967✔
162	}	71,535✔
163	}
164	if let Some(output) = lhs.invoke(&COMPARE_FN, args)? {	5,547✔
165	return Ok(output);	×
166	}	5,547✔
167
168	// Try inverting the operator and swapping the arguments
169	let inverted_args = InvocationArgs {	5,547✔
170	inputs: &[rhs.into(), lhs.into()],	5,547✔
171	options: &operator.swap(),	5,547✔
172	};	5,547✔
173	for kernel in kernels {	63,684✔
174	if let Some(output) = kernel.invoke(&inverted_args)? {	58,425✔
175	return Ok(output);	288✔
176	}	58,137✔
177	}
178	if let Some(output) = rhs.invoke(&COMPARE_FN, &inverted_args)? {	5,259✔
179	return Ok(output);	×
180	}	5,259✔
181
182	// Only log missing compare implementation if there's possibly better one than arrow,
183	// i.e. lhs isn't arrow or rhs isn't arrow or constant
184	if !(lhs.is_arrow() && (rhs.is_arrow() \|\| right_is_constant)) {	5,259✔
185	log::debug!(	4,737✔
186	"No compare implementation found for LHS {}, RHS {}, and operator {} (or inverse)",	×
187	lhs.encoding_id(),	×
188	rhs.encoding_id(),	×
189	operator,
190	);
191	}	522✔
192
193	// Fallback to arrow on canonical types
194	Ok(arrow_compare(lhs, rhs, operator)?.into())	5,259✔
195	}	14,391✔
196
197	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	28,782✔
198	let CompareArgs { lhs, rhs, .. } = CompareArgs::try_from(args)?;	28,782✔
199
200	if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) {	28,782✔
201	vortex_bail!(	×
202	"Cannot compare different DTypes {} and {}",	×
203	lhs.dtype(),	×
204	rhs.dtype()	×
205	);
206	}	28,782✔
207
208	// TODO(ngates): no reason why not
209	if lhs.dtype().is_struct() {	28,782✔
210	vortex_bail!(	×
211	"Compare does not support arrays with Struct DType, got: {} and {}",	×
212	lhs.dtype(),	×
213	rhs.dtype()	×
214	)
215	}	28,782✔
216
217	Ok(DType::Bool(	28,782✔
218	lhs.dtype().nullability() \| rhs.dtype().nullability(),	28,782✔
219	))	28,782✔
220	}	28,782✔
221
222	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	14,391✔
223	let CompareArgs { lhs, rhs, .. } = CompareArgs::try_from(args)?;	14,391✔
224	if lhs.len() != rhs.len() {	14,391✔
225	vortex_bail!(	×
226	"Compare operations only support arrays of the same length, got {} and {}",	×
227	lhs.len(),	×
228	rhs.len()	×
229	);
230	}	14,391✔
231	Ok(lhs.len())	14,391✔
232	}	14,391✔
233
234	fn is_elementwise(&self) -> bool {	14,391✔
235	true	14,391✔
236	}	14,391✔
237	}
238
239	struct CompareArgs<'a> {
240	lhs: &'a dyn Array,
241	rhs: &'a dyn Array,
242	operator: Operator,
243	}
244
245	impl Options for Operator {
246	fn as_any(&self) -> &dyn Any {	193,491✔
247	self	193,491✔
248	}	193,491✔
249	}
250
251	impl<'a> TryFrom<&InvocationArgs<'a>> for CompareArgs<'a> {
252	type Error = VortexError;
253
254	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	193,491✔
255	if value.inputs.len() != 2 {	193,491✔
256	vortex_bail!("Expected 2 inputs, found {}", value.inputs.len());	×
257	}	193,491✔
258	let lhs = value.inputs[0]	193,491✔
259	.array()	193,491✔
260	.ok_or_else(\|\| vortex_err!("Expected first input to be an array"))?;	193,491✔
261	let rhs = value.inputs[1]	193,491✔
262	.array()	193,491✔
263	.ok_or_else(\|\| vortex_err!("Expected second input to be an array"))?;	193,491✔
264	let operator = *value	193,491✔
265	.options	193,491✔
266	.as_any()	193,491✔
267	.downcast_ref::<Operator>()	193,491✔
268	.vortex_expect("Expected options to be an operator");	193,491✔
269
270	Ok(CompareArgs { lhs, rhs, operator })	193,491✔
271	}	193,491✔
272	}
273
274	/// Helper function to compare empty values with arrays that have external value length information
275	/// like `VarBin`.
NEW 276	pub fn compare_lengths_to_empty<P, I>(lengths: I, op: Operator) -> BitBuffer	×
277	where	×
278	P: NativePType,	×
279	I: Iterator<Item = P>,	×
280	{
281	// All comparison can be expressed in terms of equality. "" is the absolute min of possible value.
282	let cmp_fn = match op {	×
283	Operator::Eq \| Operator::Lte => \|v\| v == P::zero(),	×
284	Operator::NotEq \| Operator::Gt => \|v\| v != P::zero(),	×
285	Operator::Gte => \|_\| true,
286	Operator::Lt => \|_\| false,
287	};
288
NEW 289	lengths.map(cmp_fn).collect()	×
290	}	×
291
292	/// Implementation of `CompareFn` using the Arrow crate.
293	fn arrow_compare(	5,259✔
294	left: &dyn Array,	5,259✔
295	right: &dyn Array,	5,259✔
296	operator: Operator,	5,259✔
297	) -> VortexResult<ArrayRef> {	5,259✔
298	let nullable = left.dtype().is_nullable() \|\| right.dtype().is_nullable();	5,259✔
299	let lhs = Datum::try_new(left)?;	5,259✔
300	let rhs = Datum::try_new(right)?;	5,259✔
301
302	let array = match operator {	5,259✔
303	Operator::Eq => cmp::eq(&lhs, &rhs)?,	1,722✔
304	Operator::NotEq => cmp::neq(&lhs, &rhs)?,	6✔
305	Operator::Gt => cmp::gt(&lhs, &rhs)?,	480✔
306	Operator::Gte => cmp::gt_eq(&lhs, &rhs)?,	1,302✔
307	Operator::Lt => cmp::lt(&lhs, &rhs)?,	837✔
308	Operator::Lte => cmp::lt_eq(&lhs, &rhs)?,	912✔
309	};
310	from_arrow_array_with_len(&array, left.len(), nullable)	5,259✔
311	}	5,259✔
312
313	pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: Operator) -> Scalar {	2,877✔
314	if lhs.is_null() \| rhs.is_null() {	2,877✔
315	Scalar::null(DType::Bool(Nullability::Nullable))	×
316	} else {
317	let b = match operator {	2,877✔
318	Operator::Eq => lhs == rhs,	18✔
319	Operator::NotEq => lhs != rhs,	×
320	Operator::Gt => lhs > rhs,	1,491✔
321	Operator::Gte => lhs >= rhs,	33✔
322	Operator::Lt => lhs < rhs,	1,314✔
323	Operator::Lte => lhs <= rhs,	21✔
324	};
325
326	Scalar::bool(b, lhs.dtype().nullability() \| rhs.dtype().nullability())	2,877✔
327	}
328	}	2,877✔
329
330	#[cfg(test)]
331	mod tests {
332	use rstest::rstest;
333
334	use super::*;
335	use crate::ToCanonical;
336	use crate::arrays::{BoolArray, ConstantArray, VarBinArray, VarBinViewArray};
337	use crate::test_harness::to_int_indices;
338	use crate::validity::Validity;
339
340	#[test]
341	fn test_bool_basic_comparisons() {
342	let arr = BoolArray::new(
343	BitBuffer::from_iter([true, true, false, true, false]),
344	Validity::from_iter([false, true, true, true, true]),
345	);
346
347	let matches = compare(arr.as_ref(), arr.as_ref(), Operator::Eq)
348	.unwrap()
349	.to_bool()
350	.unwrap();
351
352	assert_eq!(to_int_indices(matches).unwrap(), [1u64, 2, 3, 4]);
353
354	let matches = compare(arr.as_ref(), arr.as_ref(), Operator::NotEq)
355	.unwrap()
356	.to_bool()
357	.unwrap();
358	let empty: [u64; 0] = [];
359	assert_eq!(to_int_indices(matches).unwrap(), empty);
360
361	let other = BoolArray::new(
362	BitBuffer::from_iter([false, false, false, true, true]),
363	Validity::from_iter([false, true, true, true, true]),
364	);
365
366	let matches = compare(arr.as_ref(), other.as_ref(), Operator::Lte)
367	.unwrap()
368	.to_bool()
369	.unwrap();
370	assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]);
371
372	let matches = compare(arr.as_ref(), other.as_ref(), Operator::Lt)
373	.unwrap()
374	.to_bool()
375	.unwrap();
376	assert_eq!(to_int_indices(matches).unwrap(), [4u64]);
377
378	let matches = compare(other.as_ref(), arr.as_ref(), Operator::Gte)
379	.unwrap()
380	.to_bool()
381	.unwrap();
382	assert_eq!(to_int_indices(matches).unwrap(), [2u64, 3, 4]);
383
384	let matches = compare(other.as_ref(), arr.as_ref(), Operator::Gt)
385	.unwrap()
386	.to_bool()
387	.unwrap();
388	assert_eq!(to_int_indices(matches).unwrap(), [4u64]);
389	}
390
391	#[test]
392	fn constant_compare() {
393	let left = ConstantArray::new(Scalar::from(2u32), 10);
394	let right = ConstantArray::new(Scalar::from(10u32), 10);
395
396	let compare = compare(left.as_ref(), right.as_ref(), Operator::Gt).unwrap();
397	let res = compare.as_constant().unwrap();
398	assert_eq!(res.as_bool().value(), Some(false));
399	assert_eq!(compare.len(), 10);
400
401	let compare = arrow_compare(&left.into_array(), &right.into_array(), Operator::Gt).unwrap();
402	let res = compare.as_constant().unwrap();
403	assert_eq!(res.as_bool().value(), Some(false));
404	assert_eq!(compare.len(), 10);
405	}
406
407	#[rstest]
408	#[case(Operator::Eq, vec![false, false, false, true])]
409	#[case(Operator::NotEq, vec![true, true, true, false])]
410	#[case(Operator::Gt, vec![true, true, true, false])]
411	#[case(Operator::Gte, vec![true, true, true, true])]
412	#[case(Operator::Lt, vec![false, false, false, false])]
413	#[case(Operator::Lte, vec![false, false, false, true])]
414	fn test_cmp_to_empty(#[case] op: Operator, #[case] expected: Vec<bool>) {
415	let lengths: Vec<i32> = vec![1, 5, 7, 0];
416
417	let output = compare_lengths_to_empty(lengths.iter().copied(), op);
418	assert_eq!(Vec::from_iter(output.iter()), expected);
419	}
420
421	#[rstest]
422	#[case(VarBinArray::from(vec!["a", "b"]).into_array(), VarBinViewArray::from_iter_str(["a", "b"]).into_array())]
423	#[case(VarBinViewArray::from_iter_str(["a", "b"]).into_array(), VarBinArray::from(vec!["a", "b"]).into_array())]
424	#[case(VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array())]
425	#[case(VarBinViewArray::from_iter_bin(["a".as_bytes(), "b".as_bytes()]).into_array(), VarBinArray::from(vec!["a".as_bytes(), "b".as_bytes()]).into_array())]
426	fn arrow_compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) {
427	let res = compare(&left, &right, Operator::Eq).unwrap();
428	assert_eq!(res.to_bool().unwrap().bit_buffer().true_count(), left.len());
429	}
430	}

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