16605969753

Committed 29 Jul 2025 07:50PM UTC coverage: 82.446% (-0.2%) from 82.684%

Build # 16605969753

Build Type

Pull #4057

github

Committed by

web-flow

Commit Message

Merge 1125a7099 into 6fb0f3e49

Pull Request Pull Request #4057: feat: `ArrayEquals` kernel

Run Details

0 of 158 new or added lines in 2 files covered. (0.0%)

52 existing lines in 1 file now uncovered.

45215 of 54842 relevant lines covered (82.45%)

184084.84 hits per line

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

0.0

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

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

use std::any::Any;
use std::sync::LazyLock;

use arcref::ArcRef;
use vortex_dtype::{DType, Nullability};
use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
use vortex_scalar::Scalar;

use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output, compare, Operator};
use crate::stats::{Precision, Stat, StatsProvider};
use crate::vtable::VTable;
use crate::Array;

pub fn array_equals(left: &dyn Array, right: &dyn Array) -> VortexResult<bool> {
    array_equals_opts(left, right, false)
}

pub fn array_equals_opts(left: &dyn Array, right: &dyn Array, ignore_nullability: bool) -> VortexResult<bool> {
    Ok(ARRAY_EQUALS_FN
        .invoke(&InvocationArgs {
            inputs: &[left.into(), right.into()],
            options: &ArrayEqualsOptions { ignore_nullability },
        })?
        .unwrap_scalar()?
        .as_bool()
        .value()
        .vortex_expect("non-nullable"))
}

#[derive(Clone, Copy)]
struct ArrayEqualsOptions {
    ignore_nullability: bool,
}

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

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

struct ArrayEquals;
impl ComputeFnVTable for ArrayEquals {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let ArrayEqualsArgs {
            left,
            right,
            ignore_nullability,
        } = ArrayEqualsArgs::try_from(args)?;

        if ignore_nullability && !left.dtype().eq_ignore_nullability(right.dtype()) {
            return Ok(Scalar::from(false).into());
        }

        if !ignore_nullability && !left.dtype().eq(right.dtype()) {
            return Ok(Scalar::from(false).into());
        }

        if left.len() != right.len() {
            return Ok(Scalar::from(false).into());
        }

        if let Some(l_scalar) = left.as_constant()
            && let Some(r_scalar) = right.as_constant()
        {
            return Ok(Scalar::from(l_scalar.eq(&r_scalar)).into());
        }

        if left.is_empty() && right.is_empty() {
            return Ok(Scalar::from(true).into());
        }

        for stat in [
            Stat::IsConstant,
            Stat::IsSorted,
            Stat::IsStrictSorted,
            Stat::Max, // todo: can we do that with e.g. float errors?
            Stat::Min,
            Stat::Sum,
            Stat::NullCount,
            Stat::NaNCount,
            // No Stat::UncompressedSizeInBytes because arrays may physically differ and has a different metric
        ] {
            let Some(Precision::Exact(left_v)) = left.statistics().get(stat) else {
                continue;
            };

            let Some(Precision::Exact(right_v)) = right.statistics().get(stat) else {
                continue;
            };

            if !left_v.eq(&right_v) {
                return Ok(Scalar::from(false).into());
            }
        }

        let args = InvocationArgs {
            inputs: &[left.into(), right.into()],
            options: &ArrayEqualsOptions { ignore_nullability },
        };

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

        if let Some(output) = left.invoke(&ARRAY_EQUALS_FN, &args)? {
            return Ok(output);
        }
        
        // Try swapping arguments
        let swapped_args = InvocationArgs {
            inputs: &[right.into(), left.into()],
            options: &ArrayEqualsOptions { ignore_nullability },
        };
        if let Some(output) = right.invoke(&ARRAY_EQUALS_FN, &swapped_args)? {
            return Ok(output);
        }

        // Try canonical arrays if not already canonical
        let canonical_equals = if !left.is_canonical() || !right.is_canonical() {
            let left_canonical = left.to_canonical()?;
            let right_canonical = right.to_canonical()?;
            
            array_equals_opts(left_canonical.as_ref(), right_canonical.as_ref(), ignore_nullability)?
        } else {
            // Fallback to chunked comparison
            const BATCH_SIZE: usize = 65536; // 64K elements per batch
            
            let mut offset = 0;
            while offset < left.len() {
                let end = (offset + BATCH_SIZE).min(left.len());
                
                let left_slice = left.slice(offset, end)?;
                let right_slice = right.slice(offset, end)?;
                
                let compare_result = compare(&left_slice, &right_slice, Operator::Eq)?;
                
                // Check if the result is a constant
                if let Some(constant_scalar) = compare_result.as_constant() {
                    // If constant is false, arrays are different
                    if !constant_scalar.is_valid() || constant_scalar.as_bool().value() == Some(false) {
                        return Ok(Scalar::from(false).into());
                    }
                    // If constant is true, continue to next batch
                } else {
                    // If not constant, arrays must be different
                    return Ok(Scalar::from(false).into());
                }
                
                offset = end;
            }
            
            true
        };
        
        Ok(Scalar::from(canonical_equals).into())
    }

    fn return_dtype(&self, _args: &InvocationArgs) -> VortexResult<DType> {
        Ok(DType::Bool(Nullability::NonNullable))
    }

    fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {
        Ok(1)
    }

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

// todo: statistics
pub trait ArrayEqualsKernel: VTable {
    fn compare_array(
        &self,
        array: &Self::Array,
        other: &dyn Array,
        ignore_nullability: bool,
    ) -> VortexResult<Option<bool>>;
}

struct ArrayEqualsArgs<'a> {
    left: &'a dyn Array,
    right: &'a dyn Array,
    ignore_nullability: bool,
}

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

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 2 {
            vortex_bail!(
                "ArrayEquals function requires two arguments, got {}",
                value.inputs.len()
            );
        }
        let left = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("First argument must be an array"))?;

        let right = value.inputs[1]
            .array()
            .ok_or_else(|| vortex_err!("Second argument must be an array"))?;

        let options = value
            .options
            .as_any()
            .downcast_ref::<ArrayEqualsOptions>()
            .ok_or_else(|| vortex_err!("Invalid options type for array equals function"))?;

        Ok(ArrayEqualsArgs {
            left,
            right,
            ignore_nullability: options.ignore_nullability,
        })
    }
}

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

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

impl<V: VTable + ArrayEqualsKernel> ArrayEqualsKernelAdapter<V> {
    pub const fn lift(&'static self) -> ArrayEqualsKernelRef {
        ArrayEqualsKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + ArrayEqualsKernel> Kernel for ArrayEqualsKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let ArrayEqualsArgs {
            left,
            right,
            ignore_nullability,
        } = ArrayEqualsArgs::try_from(args)?;

        let Some(left) = left.as_opt::<V>() else {
            return Ok(None);
        };

        let is_equal = V::compare_array(&self.0, left, right, ignore_nullability)?;
        Ok(is_equal.map(|b| Scalar::from(b).into()))
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::any::Any;
5	use std::sync::LazyLock;
6
7	use arcref::ArcRef;
8	use vortex_dtype::{DType, Nullability};
9	use vortex_error::{VortexError, VortexExpect, VortexResult, vortex_bail, vortex_err};
10	use vortex_scalar::Scalar;
11
12	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Options, Output, compare, Operator};
13	use crate::stats::{Precision, Stat, StatsProvider};
14	use crate::vtable::VTable;
15	use crate::Array;
16
NEW 17	pub fn array_equals(left: &dyn Array, right: &dyn Array) -> VortexResult<bool> {	×
NEW 18	array_equals_opts(left, right, false)	×
NEW 19	}	×
20
NEW 21	pub fn array_equals_opts(left: &dyn Array, right: &dyn Array, ignore_nullability: bool) -> VortexResult<bool> {	×
NEW 22	Ok(ARRAY_EQUALS_FN	×
NEW 23	.invoke(&InvocationArgs {	×
NEW 24	inputs: &[left.into(), right.into()],	×
NEW 25	options: &ArrayEqualsOptions { ignore_nullability },	×
NEW 26	})?	×
NEW 27	.unwrap_scalar()?	×
NEW 28	.as_bool()	×
NEW 29	.value()	×
NEW 30	.vortex_expect("non-nullable"))	×
NEW 31	}	×
32
33	#[derive(Clone, Copy)]
34	struct ArrayEqualsOptions {
35	ignore_nullability: bool,
36	}
37
38	impl Options for ArrayEqualsOptions {
NEW 39	fn as_any(&self) -> &dyn Any {	×
NEW 40	self	×
NEW 41	}	×
42	}
43
NEW 44	pub static ARRAY_EQUALS_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	×
NEW 45	let compute = ComputeFn::new("array_equals".into(), ArcRef::new_ref(&ArrayEquals));	×
NEW 46	for kernel in inventory::iter::<ArrayEqualsKernelRef> {	×
NEW 47	compute.register_kernel(kernel.0.clone());	×
NEW 48	}	×
NEW 49	compute	×
NEW 50	});	×
51
52	struct ArrayEquals;
53	impl ComputeFnVTable for ArrayEquals {
NEW 54	fn invoke(	×
NEW 55	&self,	×
NEW 56	args: &InvocationArgs,	×
NEW 57	kernels: &[ArcRef<dyn Kernel>],	×
NEW 58	) -> VortexResult<Output> {	×
59	let ArrayEqualsArgs {
NEW 60	left,	×
NEW 61	right,	×
NEW 62	ignore_nullability,	×
NEW 63	} = ArrayEqualsArgs::try_from(args)?;	×
64
NEW 65	if ignore_nullability && !left.dtype().eq_ignore_nullability(right.dtype()) {	×
NEW 66	return Ok(Scalar::from(false).into());	×
NEW 67	}	×
68
NEW 69	if !ignore_nullability && !left.dtype().eq(right.dtype()) {	×
NEW 70	return Ok(Scalar::from(false).into());	×
NEW 71	}	×
72
NEW 73	if left.len() != right.len() {	×
NEW 74	return Ok(Scalar::from(false).into());	×
NEW 75	}	×
76
NEW 77	if let Some(l_scalar) = left.as_constant()	×
NEW 78	&& let Some(r_scalar) = right.as_constant()	×
79	{
NEW 80	return Ok(Scalar::from(l_scalar.eq(&r_scalar)).into());	×
NEW 81	}	×
82
NEW 83	if left.is_empty() && right.is_empty() {	×
NEW 84	return Ok(Scalar::from(true).into());	×
NEW 85	}	×
86
NEW 87	for stat in [	×
NEW 88	Stat::IsConstant,	×
NEW 89	Stat::IsSorted,	×
NEW 90	Stat::IsStrictSorted,	×
NEW 91	Stat::Max, // todo: can we do that with e.g. float errors?	×
NEW 92	Stat::Min,	×
NEW 93	Stat::Sum,	×
NEW 94	Stat::NullCount,	×
NEW 95	Stat::NaNCount,	×
96	// No Stat::UncompressedSizeInBytes because arrays may physically differ and has a different metric
97	] {
NEW 98	let Some(Precision::Exact(left_v)) = left.statistics().get(stat) else {	×
NEW 99	continue;	×
100	};
101
NEW 102	let Some(Precision::Exact(right_v)) = right.statistics().get(stat) else {	×
NEW 103	continue;	×
104	};
105
NEW 106	if !left_v.eq(&right_v) {	×
NEW 107	return Ok(Scalar::from(false).into());	×
NEW 108	}	×
109	}
110
NEW 111	let args = InvocationArgs {	×
NEW 112	inputs: &[left.into(), right.into()],	×
NEW 113	options: &ArrayEqualsOptions { ignore_nullability },	×
NEW 114	};	×
115
NEW 116	for kernel in kernels {	×
NEW 117	if let Some(output) = kernel.invoke(&args)? {	×
NEW 118	return Ok(output);	×
NEW 119	}	×
120	}
121
NEW 122	if let Some(output) = left.invoke(&ARRAY_EQUALS_FN, &args)? {	×
NEW 123	return Ok(output);	×
NEW 124	}	×
125
126	// Try swapping arguments
NEW 127	let swapped_args = InvocationArgs {	×
NEW 128	inputs: &[right.into(), left.into()],	×
NEW 129	options: &ArrayEqualsOptions { ignore_nullability },	×
NEW 130	};	×
NEW 131	if let Some(output) = right.invoke(&ARRAY_EQUALS_FN, &swapped_args)? {	×
NEW 132	return Ok(output);	×
NEW 133	}	×
134
135	// Try canonical arrays if not already canonical
NEW 136	let canonical_equals = if !left.is_canonical() \|\| !right.is_canonical() {	×
NEW 137	let left_canonical = left.to_canonical()?;	×
NEW 138	let right_canonical = right.to_canonical()?;	×
139
NEW 140	array_equals_opts(left_canonical.as_ref(), right_canonical.as_ref(), ignore_nullability)?	×
141	} else {
142	// Fallback to chunked comparison
143	const BATCH_SIZE: usize = 65536; // 64K elements per batch
144
NEW 145	let mut offset = 0;	×
NEW 146	while offset < left.len() {	×
NEW 147	let end = (offset + BATCH_SIZE).min(left.len());	×
148
NEW 149	let left_slice = left.slice(offset, end)?;	×
NEW 150	let right_slice = right.slice(offset, end)?;	×
151
NEW 152	let compare_result = compare(&left_slice, &right_slice, Operator::Eq)?;	×
153
154	// Check if the result is a constant
NEW 155	if let Some(constant_scalar) = compare_result.as_constant() {	×
156	// If constant is false, arrays are different
NEW 157	if !constant_scalar.is_valid() \|\| constant_scalar.as_bool().value() == Some(false) {	×
NEW 158	return Ok(Scalar::from(false).into());	×
NEW 159	}	×
160	// If constant is true, continue to next batch
161	} else {
162	// If not constant, arrays must be different
NEW 163	return Ok(Scalar::from(false).into());	×
164	}
165
NEW 166	offset = end;	×
167	}
168
NEW 169	true	×
170	};
171
NEW 172	Ok(Scalar::from(canonical_equals).into())	×
NEW 173	}	×
174
NEW 175	fn return_dtype(&self, _args: &InvocationArgs) -> VortexResult<DType> {	×
NEW 176	Ok(DType::Bool(Nullability::NonNullable))	×
NEW 177	}	×
178
NEW 179	fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {	×
NEW 180	Ok(1)	×
NEW 181	}	×
182
NEW 183	fn is_elementwise(&self) -> bool {	×
NEW 184	false	×
NEW 185	}	×
186	}
187
188	// todo: statistics
189	pub trait ArrayEqualsKernel: VTable {
190	fn compare_array(
191	&self,
192	array: &Self::Array,
193	other: &dyn Array,
194	ignore_nullability: bool,
195	) -> VortexResult<Option<bool>>;
196	}
197
198	struct ArrayEqualsArgs<'a> {
199	left: &'a dyn Array,
200	right: &'a dyn Array,
201	ignore_nullability: bool,
202	}
203
204	impl<'a> TryFrom<&InvocationArgs<'a>> for ArrayEqualsArgs<'a> {
205	type Error = VortexError;
206
NEW 207	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	×
NEW 208	if value.inputs.len() != 2 {	×
NEW 209	vortex_bail!(	×
NEW 210	"ArrayEquals function requires two arguments, got {}",	×
NEW 211	value.inputs.len()	×
212	);
NEW 213	}	×
NEW 214	let left = value.inputs[0]	×
NEW 215	.array()	×
NEW 216	.ok_or_else(\|\| vortex_err!("First argument must be an array"))?;	×
217
NEW 218	let right = value.inputs[1]	×
NEW 219	.array()	×
NEW 220	.ok_or_else(\|\| vortex_err!("Second argument must be an array"))?;	×
221
NEW 222	let options = value	×
NEW 223	.options	×
NEW 224	.as_any()	×
NEW 225	.downcast_ref::<ArrayEqualsOptions>()	×
NEW 226	.ok_or_else(\|\| vortex_err!("Invalid options type for array equals function"))?;	×
227
NEW 228	Ok(ArrayEqualsArgs {	×
NEW 229	left,	×
NEW 230	right,	×
NEW 231	ignore_nullability: options.ignore_nullability,	×
NEW 232	})	×
NEW 233	}	×
234	}
235
236	#[derive(Debug)]
237	pub struct ArrayEqualsKernelAdapter<V: VTable>(pub V);
238
239	pub struct ArrayEqualsKernelRef(ArcRef<dyn Kernel>);
240	inventory::collect!(ArrayEqualsKernelRef);
241
242	impl<V: VTable + ArrayEqualsKernel> ArrayEqualsKernelAdapter<V> {
NEW 243	pub const fn lift(&'static self) -> ArrayEqualsKernelRef {	×
NEW 244	ArrayEqualsKernelRef(ArcRef::new_ref(self))	×
NEW 245	}	×
246	}
247
248	impl<V: VTable + ArrayEqualsKernel> Kernel for ArrayEqualsKernelAdapter<V> {
NEW 249	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	×
250	let ArrayEqualsArgs {
NEW 251	left,	×
NEW 252	right,	×
NEW 253	ignore_nullability,	×
NEW 254	} = ArrayEqualsArgs::try_from(args)?;	×
255
NEW 256	let Some(left) = left.as_opt::<V>() else {	×
NEW 257	return Ok(None);	×
258	};
259
NEW 260	let is_equal = V::compare_array(&self.0, left, right, ignore_nullability)?;	×
NEW 261	Ok(is_equal.map(\|b\| Scalar::from(b).into()))	×
NEW 262	}	×
263	}

vortex-data / vortex / 16605969753

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