16606232494

Committed 29 Jul 2025 08:03PM UTC coverage: 82.715% (+0.03%) from 82.684%

Build # 16606232494

Build Type

Pull #4057

github

Committed by

web-flow

Commit Message

Merge eed79082a into 6fb0f3e49

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

Run Details

326 of 376 new or added lines in 2 files covered. (86.7%)

1 existing line in 1 file now uncovered.

45543 of 55060 relevant lines covered (82.72%)

185089.06 hits per line

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

88.22

/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::Array;
use crate::compute::{
    ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Options, Output, compare,
};
use crate::stats::{Precision, Stat, StatsProvider};
use crate::vtable::VTable;

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,
                batch_size: None,
            },
        })?
        .unwrap_scalar()?
        .as_bool()
        .value()
        .vortex_expect("non-nullable"))
}

#[derive(Clone, Copy)]
struct ArrayEqualsOptions {
    ignore_nullability: bool,
    batch_size: Option<usize>,
}

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,
            batch_size,
        } = 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());
        }

        // Handle constant null arrays - they are equal only if both are null constants
        let left_constant_null = left.as_constant().map(|l| l.is_null()).unwrap_or(false);
        let right_constant_null = right.as_constant().map(|r| r.is_null()).unwrap_or(false);
        
        if left_constant_null || right_constant_null {
            return Ok(Scalar::from(left_constant_null && right_constant_null).into());
        }

        // Check statistics for early exit
        // TODO(optimization): Add more sophisticated statistical comparisons for floating point arrays
        if !check_stats_equality(left, right) {
            return Ok(Scalar::from(false).into());
        }

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

        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,
                batch_size,
            },
        };
        if let Some(output) = right.invoke(&ARRAY_EQUALS_FN, &swapped_args)? {
            return Ok(output);
        }

        // Try canonical arrays if not already canonical
        if !left.is_canonical() || !right.is_canonical() {
            log::debug!(
                "Falling back to canonical array_equals for encodings {} and {}",
                left.encoding_id(),
                right.encoding_id()
            );
            
            let left_canonical = left.to_canonical()?;
            let right_canonical = right.to_canonical()?;

            return Ok(Scalar::from(array_equals_opts(
                left_canonical.as_ref(),
                right_canonical.as_ref(),
                ignore_nullability,
            )?)
            .into());
        }

        // Final fallback to chunked comparison for canonical arrays
        log::debug!(
            "Using chunked comparison fallback for canonical arrays {} and {}",
            left.encoding_id(),
            right.encoding_id()
        );
        
        let all_equal = compare_chunked(left, right, batch_size)?;
        Ok(Scalar::from(all_equal).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,
    batch_size: Option<usize>,
}

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,
            batch_size: options.batch_size,
        })
    }
}

#[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,
            batch_size: _, // Not used in kernel adapters
        } = 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()))
    }
}

/// Compare arrays in chunks to avoid loading entire arrays into memory
fn compare_chunked(
    left: &dyn Array,
    right: &dyn Array,
    batch_size: Option<usize>,
) -> VortexResult<bool> {
    const DEFAULT_BATCH_SIZE: usize = 65536; // 64K elements per batch
    let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);

    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)?;

        if !compare_batch(&left_slice, &right_slice)? {
            return Ok(false);
        }

        offset = end;
    }

    Ok(true)
}

/// Compare a single batch of arrays
fn compare_batch(left: &dyn Array, right: &dyn Array) -> VortexResult<bool> {
    let compare_result = compare(left, right, Operator::Eq)?;

    // Check if the comparison result indicates all equal
    if let Some(all_equal) = check_constant_result(&compare_result)? {
        return Ok(all_equal);
    }

    // Not constant - need to check each value
    check_non_constant_result(&compare_result, left, right)
}

/// Check if a constant comparison result indicates equality
fn check_constant_result(compare_result: &dyn Array) -> VortexResult<Option<bool>> {
    if let Some(constant_scalar) = compare_result.as_constant() {
        // If constant is true, all are equal
        Ok(Some(
            constant_scalar.is_valid() && constant_scalar.as_bool().value() == Some(true),
        ))
    } else {
        Ok(None)
    }
}

/// Check non-constant comparison results, handling null comparisons
fn check_non_constant_result(
    compare_result: &dyn Array,
    left: &dyn Array,
    right: &dyn Array,
) -> VortexResult<bool> {
    // First, check statistics for quick rejection
    if let Some(all_true) = check_comparison_stats(compare_result) {
        return Ok(all_true);
    }

    // Fallback to element-wise check
    for i in 0..compare_result.len() {
        let cmp_scalar = compare_result.scalar_at(i)?;

        // Check for definite inequality
        if cmp_scalar.is_valid() && cmp_scalar.as_bool().value() == Some(false) {
            return Ok(false);
        }

        // Handle null comparison results
        if cmp_scalar.is_null() && !check_null_equality(left, right, i)? {
            return Ok(false);
        }
    }

    Ok(true)
}

/// Check comparison statistics for quick determination
fn check_comparison_stats(compare_result: &dyn Array) -> Option<bool> {
    // If min is false, we have at least one false
    if let Some(Precision::Exact(min)) = compare_result.statistics().get(Stat::Min) {
        if min.as_bool().ok()? == Some(false) {
            return Some(false);
        }
    }

    // If both min and max are true, all are true
    if let Some(Precision::Exact(min)) = compare_result.statistics().get(Stat::Min) {
        if let Some(Precision::Exact(max)) = compare_result.statistics().get(Stat::Max) {
            if min.as_bool().ok()? == Some(true) && max.as_bool().ok()? == Some(true) {
                return Some(true);
            }
        }
    }

    None
}

/// Check if two potentially null values at a given index are equal
fn check_null_equality(left: &dyn Array, right: &dyn Array, index: usize) -> VortexResult<bool> {
    let left_val = left.scalar_at(index)?;
    let right_val = right.scalar_at(index)?;

    // Both null or both non-null means they could be equal
    // (if both non-null, the comparison would have returned true/false, not null)
    Ok(left_val.is_null() == right_val.is_null())
}

/// Check statistics equality for early exit
fn check_stats_equality(left: &dyn Array, right: &dyn Array) -> bool {
    let stats_to_check = [
        Stat::IsConstant,
        Stat::IsSorted,
        Stat::IsStrictSorted,
        Stat::Max,
        Stat::Min,
        Stat::Sum,
        Stat::NullCount,
        Stat::NaNCount,
    ];

    for stat in stats_to_check {
        match (left.statistics().get(stat), right.statistics().get(stat)) {
            (Some(Precision::Exact(left_v)), Some(Precision::Exact(right_v))) => {
                if !left_v.eq(&right_v) {
                    return false;
                }
            }
            _ => continue,
        }
    }

    true
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::IntoArray;
    use crate::arrays::{BoolArray, ChunkedArray, ConstantArray, PrimitiveArray, VarBinArray};
    use crate::validity::Validity;
    use vortex_dtype::{DType, Nullability, PType};

    #[test]
    fn test_simple_equals() {
        let arr1 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]);
        let arr2 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]);
        let arr3 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 6]);

        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
        assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());
    }

    #[test]
    fn test_stats_comparison() {
        // Arrays with different stats should be detected as different early
        let arr1 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]);
        let arr2 = PrimitiveArray::from_iter(vec![10i32, 20, 30, 40, 50]);

        assert!(!array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
    }

    #[test]
    fn test_constant_arrays() {
        let const1 = ConstantArray::new(Scalar::from(42i32), 100);
        let const2 = ConstantArray::new(Scalar::from(42i32), 100);
        let const3 = ConstantArray::new(Scalar::from(43i32), 100);

        assert!(array_equals(const1.as_ref(), const2.as_ref()).unwrap());
        assert!(!array_equals(const1.as_ref(), const3.as_ref()).unwrap());
    }

    #[test]
    fn test_different_types() {
        let int_arr = PrimitiveArray::from_iter(vec![1i32, 2, 3]);
        let float_arr = PrimitiveArray::from_iter(vec![1.0f32, 2.0, 3.0]);

        assert!(!array_equals(int_arr.as_ref(), float_arr.as_ref()).unwrap());
    }

    #[test]
    fn test_with_nulls() {
        let arr1 = PrimitiveArray::from_option_iter(vec![Some(1i32), None, Some(3), Some(4)]);
        let arr2 = PrimitiveArray::from_option_iter(vec![Some(1i32), None, Some(3), Some(4)]);
        let arr3 = PrimitiveArray::from_option_iter(vec![Some(1i32), Some(2), Some(3), Some(4)]);

        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
        assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());
    }

    #[test]
    fn test_null_arrays() {
        let arr1 = PrimitiveArray::from_option_iter(vec![None::<i32>, None, None]);
        let arr2 = PrimitiveArray::from_option_iter(vec![None::<i32>, None, None]);

        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
    }

    #[test]
    fn test_bool_arrays() {
        use arrow_buffer::BooleanBuffer;

        let arr1 = BoolArray::new(
            BooleanBuffer::from_iter([true, false, true, false]),
            Validity::AllValid,
        );
        let arr2 = BoolArray::new(
            BooleanBuffer::from_iter([true, false, true, false]),
            Validity::AllValid,
        );
        let arr3 = BoolArray::new(
            BooleanBuffer::from_iter([true, false, false, false]),
            Validity::AllValid,
        );

        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
        assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());
    }

    #[test]
    fn test_empty_arrays() {
        let empty1 = PrimitiveArray::from_iter(Vec::<i32>::new());
        let empty2 = PrimitiveArray::from_iter(Vec::<i32>::new());

        assert!(array_equals(empty1.as_ref(), empty2.as_ref()).unwrap());
    }

    #[test]
    fn test_different_lengths() {
        let arr1 = PrimitiveArray::from_iter(vec![1i32, 2, 3]);
        let arr2 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4]);

        assert!(!array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
    }

    #[test]
    fn test_large_arrays() {
        // Test arrays larger than BATCH_SIZE
        let data1: Vec<i64> = (0..100_000).collect();
        let data2: Vec<i64> = (0..100_000).collect();
        let mut data3 = data1.clone();
        data3[99_999] = 999_999;

        let arr1 = PrimitiveArray::from_iter(data1);
        let arr2 = PrimitiveArray::from_iter(data2);
        let arr3 = PrimitiveArray::from_iter(data3);

        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
        assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());
    }

    #[test]
    fn test_non_canonical_arrays() {
        let varbin1 = VarBinArray::from_vec(
            vec!["hello".as_bytes(), "world".as_bytes()],
            DType::Utf8(Nullability::NonNullable),
        );
        let varbin2 = VarBinArray::from_vec(
            vec!["hello".as_bytes(), "world".as_bytes()],
            DType::Utf8(Nullability::NonNullable),
        );
        let varbin3 = VarBinArray::from_vec(
            vec!["hello".as_bytes(), "earth".as_bytes()],
            DType::Utf8(Nullability::NonNullable),
        );

        assert!(array_equals(varbin1.as_ref(), varbin2.as_ref()).unwrap());
        assert!(!array_equals(varbin1.as_ref(), varbin3.as_ref()).unwrap());
    }

    #[test]
    fn test_float_precision() {
        // Test if statistics-based comparison can handle float precision issues
        let arr1 = PrimitiveArray::from_iter(vec![1.0f64, 2.0, 3.0, 4.0, 5.0]);
        let arr2 = PrimitiveArray::from_iter(vec![1.0f64, 2.0, 3.0, 4.0, 5.0]);

        // Arrays with exact same values should be equal
        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());

        // Arrays with slightly different values should not be equal
        let arr3 = PrimitiveArray::from_iter(vec![1.0f64, 2.0, 3.0, 4.0, 5.0000000001]);
        assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());
    }

    #[test]
    fn test_batch_size_functionality() {
        // Test arrays larger than default batch size with different batch sizes
        let data1: Vec<i32> = (0..150_000).collect();
        let data2: Vec<i32> = (0..150_000).collect();

        let arr1 = PrimitiveArray::from_iter(data1);
        let arr2 = PrimitiveArray::from_iter(data2);

        // Test with different batch sizes (though we can't pass batch_size directly in public API)
        assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());
    }

    #[test]
    fn test_primitive_vs_dict_array() {
        // Test comparing primitive array with dictionary-encoded array containing same values

        let primitive_arr = PrimitiveArray::from_iter(vec![1i32, 2, 1, 3, 2, 1]);

        // Create a chunked array as a proxy for non-canonical encoding
        let chunk1 = PrimitiveArray::from_iter(vec![1i32, 2, 1]);
        let chunk2 = PrimitiveArray::from_iter(vec![3i32, 2, 1]);
        let chunked_arr = ChunkedArray::try_new(
            vec![chunk1.into_array(), chunk2.into_array()],
            primitive_arr.dtype().clone(),
        )
        .unwrap();

        // Should be equal as they contain the same logical values
        assert!(array_equals(primitive_arr.as_ref(), chunked_arr.as_ref()).unwrap());

        // Test with different values
        let chunk1_copy = PrimitiveArray::from_iter(vec![1i32, 2, 1]);
        let different_chunk2 = PrimitiveArray::from_iter(vec![3i32, 2, 4]);
        let different_chunked = ChunkedArray::try_new(
            vec![chunk1_copy.into_array(), different_chunk2.into_array()],
            primitive_arr.dtype().clone(),
        )
        .unwrap();

        assert!(!array_equals(primitive_arr.as_ref(), different_chunked.as_ref()).unwrap());
    }

    #[test]
    fn test_constant_null_arrays() {
        // Test constant null arrays - should be equal to each other but not to non-null constants
        let null_const1 = ConstantArray::new(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), 5);
        let null_const2 = ConstantArray::new(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), 5);
        let non_null_const = ConstantArray::new(Scalar::from(42i32), 5);
        
        // Both null constants should be equal
        assert!(array_equals(null_const1.as_ref(), null_const2.as_ref()).unwrap());
        
        // Null constant should not equal non-null constant
        assert!(!array_equals(null_const1.as_ref(), non_null_const.as_ref()).unwrap());
        assert!(!array_equals(non_null_const.as_ref(), null_const1.as_ref()).unwrap());
    }
}

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::Array;
13	use crate::compute::{
14	ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Options, Output, compare,
15	};
16	use crate::stats::{Precision, Stat, StatsProvider};
17	use crate::vtable::VTable;
18
19	pub fn array_equals(left: &dyn Array, right: &dyn Array) -> VortexResult<bool> {	25✔
20	array_equals_opts(left, right, false)	25✔
21	}	25✔
22
23	pub fn array_equals_opts(	29✔
24	left: &dyn Array,	29✔
25	right: &dyn Array,	29✔
26	ignore_nullability: bool,	29✔
27	) -> VortexResult<bool> {	29✔
28	Ok(ARRAY_EQUALS_FN	29✔
29	.invoke(&InvocationArgs {	29✔
30	inputs: &[left.into(), right.into()],	29✔
31	options: &ArrayEqualsOptions {	29✔
32	ignore_nullability,	29✔
33	batch_size: None,	29✔
34	},	29✔
35	})?	29✔
36	.unwrap_scalar()?	29✔
37	.as_bool()	29✔
38	.value()	29✔
39	.vortex_expect("non-nullable"))	29✔
40	}	29✔
41
42	#[derive(Clone, Copy)]
43	struct ArrayEqualsOptions {
44	ignore_nullability: bool,
45	batch_size: Option<usize>,
46	}
47
48	impl Options for ArrayEqualsOptions {
49	fn as_any(&self) -> &dyn Any {	29✔
50	self	29✔
51	}	29✔
52	}
53
54	pub static ARRAY_EQUALS_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	15✔
55	let compute = ComputeFn::new("array_equals".into(), ArcRef::new_ref(&ArrayEquals));	15✔
56	for kernel in inventory::iter::<ArrayEqualsKernelRef> {	15✔
NEW 57	compute.register_kernel(kernel.0.clone());	×
NEW 58	}	×
59	compute	15✔
60	});	15✔
61
62	struct ArrayEquals;
63	impl ComputeFnVTable for ArrayEquals {
64	fn invoke(	29✔
65	&self,	29✔
66	args: &InvocationArgs,	29✔
67	kernels: &[ArcRef<dyn Kernel>],	29✔
68	) -> VortexResult<Output> {	29✔
69	let ArrayEqualsArgs {
70	left,	29✔
71	right,	29✔
72	ignore_nullability,	29✔
73	batch_size,	29✔
74	} = ArrayEqualsArgs::try_from(args)?;	29✔
75
76	if ignore_nullability && !left.dtype().eq_ignore_nullability(right.dtype()) {	29✔
NEW 77	return Ok(Scalar::from(false).into());	×
78	}	29✔
79
80	if !ignore_nullability && !left.dtype().eq(right.dtype()) {	29✔
81	return Ok(Scalar::from(false).into());	3✔
82	}	26✔
83
84	if left.len() != right.len() {	26✔
85	return Ok(Scalar::from(false).into());	1✔
86	}	25✔
87
88	if let Some(l_scalar) = left.as_constant()	25✔
89	&& let Some(r_scalar) = right.as_constant()	4✔
90	{
91	return Ok(Scalar::from(l_scalar.eq(&r_scalar)).into());	4✔
92	}	21✔
93
94	if left.is_empty() && right.is_empty() {	21✔
95	return Ok(Scalar::from(true).into());	1✔
96	}	20✔
97
98	// Handle constant null arrays - they are equal only if both are null constants
99	let left_constant_null = left.as_constant().map(\|l\| l.is_null()).unwrap_or(false);	20✔
100	let right_constant_null = right.as_constant().map(\|r\| r.is_null()).unwrap_or(false);	20✔
101
102	if left_constant_null \|\| right_constant_null {	20✔
NEW 103	return Ok(Scalar::from(left_constant_null && right_constant_null).into());	×
104	}	20✔
105
106	// Check statistics for early exit
107	// TODO(optimization): Add more sophisticated statistical comparisons for floating point arrays
108	if !check_stats_equality(left, right) {	20✔
NEW 109	return Ok(Scalar::from(false).into());	×
110	}	20✔
111
112	let args = InvocationArgs {	20✔
113	inputs: &[left.into(), right.into()],	20✔
114	options: &ArrayEqualsOptions {	20✔
115	ignore_nullability,	20✔
116	batch_size,	20✔
117	},	20✔
118	};	20✔
119
120	for kernel in kernels {	20✔
NEW 121	if let Some(output) = kernel.invoke(&args)? {	×
NEW 122	return Ok(output);	×
NEW 123	}	×
124	}
125
126	if let Some(output) = left.invoke(&ARRAY_EQUALS_FN, &args)? {	20✔
NEW 127	return Ok(output);	×
128	}	20✔
129
130	// Try swapping arguments
131	let swapped_args = InvocationArgs {	20✔
132	inputs: &[right.into(), left.into()],	20✔
133	options: &ArrayEqualsOptions {	20✔
134	ignore_nullability,	20✔
135	batch_size,	20✔
136	},	20✔
137	};	20✔
138	if let Some(output) = right.invoke(&ARRAY_EQUALS_FN, &swapped_args)? {	20✔
NEW 139	return Ok(output);	×
140	}	20✔
141
142	// Try canonical arrays if not already canonical
143	if !left.is_canonical() \|\| !right.is_canonical() {	20✔
144	log::debug!(	4✔
NEW 145	"Falling back to canonical array_equals for encodings {} and {}",	×
NEW 146	left.encoding_id(),	×
NEW 147	right.encoding_id()	×
148	);
149
150	let left_canonical = left.to_canonical()?;	4✔
151	let right_canonical = right.to_canonical()?;	4✔
152
153	return Ok(Scalar::from(array_equals_opts(	4✔
154	left_canonical.as_ref(),	4✔
155	right_canonical.as_ref(),	4✔
156	ignore_nullability,	4✔
NEW 157	)?)	×
158	.into());	4✔
159	}	16✔
160
161	// Final fallback to chunked comparison for canonical arrays
162	log::debug!(	16✔
NEW 163	"Using chunked comparison fallback for canonical arrays {} and {}",	×
NEW 164	left.encoding_id(),	×
NEW 165	right.encoding_id()	×
166	);
167
168	let all_equal = compare_chunked(left, right, batch_size)?;	16✔
169	Ok(Scalar::from(all_equal).into())	16✔
170	}	29✔
171
172	fn return_dtype(&self, _args: &InvocationArgs) -> VortexResult<DType> {	29✔
173	Ok(DType::Bool(Nullability::NonNullable))	29✔
174	}	29✔
175
176	fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {	29✔
177	Ok(1)	29✔
178	}	29✔
179
180	fn is_elementwise(&self) -> bool {	31✔
181	false	31✔
182	}	31✔
183	}
184
185	// todo: statistics
186	pub trait ArrayEqualsKernel: VTable {
187	fn compare_array(
188	&self,
189	array: &Self::Array,
190	other: &dyn Array,
191	ignore_nullability: bool,
192	) -> VortexResult<Option<bool>>;
193	}
194
195	struct ArrayEqualsArgs<'a> {
196	left: &'a dyn Array,
197	right: &'a dyn Array,
198	ignore_nullability: bool,
199	batch_size: Option<usize>,
200	}
201
202	impl<'a> TryFrom<&InvocationArgs<'a>> for ArrayEqualsArgs<'a> {
203	type Error = VortexError;
204
205	fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {	29✔
206	if value.inputs.len() != 2 {	29✔
NEW 207	vortex_bail!(	×
NEW 208	"ArrayEquals function requires two arguments, got {}",	×
NEW 209	value.inputs.len()	×
210	);
211	}	29✔
212	let left = value.inputs[0]	29✔
213	.array()	29✔
214	.ok_or_else(\|\| vortex_err!("First argument must be an array"))?;	29✔
215
216	let right = value.inputs[1]	29✔
217	.array()	29✔
218	.ok_or_else(\|\| vortex_err!("Second argument must be an array"))?;	29✔
219
220	let options = value	29✔
221	.options	29✔
222	.as_any()	29✔
223	.downcast_ref::<ArrayEqualsOptions>()	29✔
224	.ok_or_else(\|\| vortex_err!("Invalid options type for array equals function"))?;	29✔
225
226	Ok(ArrayEqualsArgs {	29✔
227	left,	29✔
228	right,	29✔
229	ignore_nullability: options.ignore_nullability,	29✔
230	batch_size: options.batch_size,	29✔
231	})	29✔
232	}	29✔
233	}
234
235	#[derive(Debug)]
236	pub struct ArrayEqualsKernelAdapter<V: VTable>(pub V);
237
238	pub struct ArrayEqualsKernelRef(ArcRef<dyn Kernel>);
239	inventory::collect!(ArrayEqualsKernelRef);
240
241	impl<V: VTable + ArrayEqualsKernel> ArrayEqualsKernelAdapter<V> {
NEW 242	pub const fn lift(&'static self) -> ArrayEqualsKernelRef {	×
NEW 243	ArrayEqualsKernelRef(ArcRef::new_ref(self))	×
NEW 244	}	×
245	}
246
247	impl<V: VTable + ArrayEqualsKernel> Kernel for ArrayEqualsKernelAdapter<V> {
NEW 248	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	×
249	let ArrayEqualsArgs {
NEW 250	left,	×
NEW 251	right,	×
NEW 252	ignore_nullability,	×
253	batch_size: _, // Not used in kernel adapters
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	}
264
265	/// Compare arrays in chunks to avoid loading entire arrays into memory
266	fn compare_chunked(	16✔
267	left: &dyn Array,	16✔
268	right: &dyn Array,	16✔
269	batch_size: Option<usize>,	16✔
270	) -> VortexResult<bool> {	16✔
271	const DEFAULT_BATCH_SIZE: usize = 65536; // 64K elements per batch
272	let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);	16✔
273
274	let mut offset = 0;	16✔
275	while offset < left.len() {	28✔
276	let end = (offset + batch_size).min(left.len());	20✔
277
278	let left_slice = left.slice(offset, end)?;	20✔
279	let right_slice = right.slice(offset, end)?;	20✔
280
281	if !compare_batch(&left_slice, &right_slice)? {	20✔
282	return Ok(false);	8✔
283	}	12✔
284
285	offset = end;	12✔
286	}
287
288	Ok(true)	8✔
289	}	16✔
290
291	/// Compare a single batch of arrays
292	fn compare_batch(left: &dyn Array, right: &dyn Array) -> VortexResult<bool> {	20✔
293	let compare_result = compare(left, right, Operator::Eq)?;	20✔
294
295	// Check if the comparison result indicates all equal
296	if let Some(all_equal) = check_constant_result(&compare_result)? {	20✔
297	return Ok(all_equal);	12✔
298	}	8✔
299
300	// Not constant - need to check each value
301	check_non_constant_result(&compare_result, left, right)	8✔
302	}	20✔
303
304	/// Check if a constant comparison result indicates equality
305	fn check_constant_result(compare_result: &dyn Array) -> VortexResult<Option<bool>> {	20✔
306	if let Some(constant_scalar) = compare_result.as_constant() {	20✔
307	// If constant is true, all are equal
308	Ok(Some(
309	constant_scalar.is_valid() && constant_scalar.as_bool().value() == Some(true),	12✔
310	))
311	} else {
312	Ok(None)	8✔
313	}
314	}	20✔
315
316	/// Check non-constant comparison results, handling null comparisons
317	fn check_non_constant_result(	8✔
318	compare_result: &dyn Array,	8✔
319	left: &dyn Array,	8✔
320	right: &dyn Array,	8✔
321	) -> VortexResult<bool> {	8✔
322	// First, check statistics for quick rejection
323	if let Some(all_true) = check_comparison_stats(compare_result) {	8✔
NEW 324	return Ok(all_true);	×
325	}	8✔
326
327	// Fallback to element-wise check
328	for i in 0..compare_result.len() {	34,491✔
329	let cmp_scalar = compare_result.scalar_at(i)?;	34,491✔
330
331	// Check for definite inequality
332	if cmp_scalar.is_valid() && cmp_scalar.as_bool().value() == Some(false) {	34,491✔
333	return Ok(false);	6✔
334	}	34,485✔
335
336	// Handle null comparison results
337	if cmp_scalar.is_null() && !check_null_equality(left, right, i)? {	34,485✔
338	return Ok(false);	1✔
339	}	34,484✔
340	}
341
342	Ok(true)	1✔
343	}	8✔
344
345	/// Check comparison statistics for quick determination
346	fn check_comparison_stats(compare_result: &dyn Array) -> Option<bool> {	8✔
347	// If min is false, we have at least one false
348	if let Some(Precision::Exact(min)) = compare_result.statistics().get(Stat::Min) {	8✔
NEW 349	if min.as_bool().ok()? == Some(false) {	×
NEW 350	return Some(false);	×
NEW 351	}	×
352	}	8✔
353
354	// If both min and max are true, all are true
355	if let Some(Precision::Exact(min)) = compare_result.statistics().get(Stat::Min) {	8✔
NEW 356	if let Some(Precision::Exact(max)) = compare_result.statistics().get(Stat::Max) {	×
NEW 357	if min.as_bool().ok()? == Some(true) && max.as_bool().ok()? == Some(true) {	×
NEW 358	return Some(true);	×
NEW 359	}	×
NEW 360	}	×
361	}	8✔
362
363	None	8✔
364	}	8✔
365
366	/// Check if two potentially null values at a given index are equal
367	fn check_null_equality(left: &dyn Array, right: &dyn Array, index: usize) -> VortexResult<bool> {	2✔
368	let left_val = left.scalar_at(index)?;	2✔
369	let right_val = right.scalar_at(index)?;	2✔
370
371	// Both null or both non-null means they could be equal
372	// (if both non-null, the comparison would have returned true/false, not null)
373	Ok(left_val.is_null() == right_val.is_null())	2✔
374	}	2✔
375
376	/// Check statistics equality for early exit
377	fn check_stats_equality(left: &dyn Array, right: &dyn Array) -> bool {	20✔
378	let stats_to_check = [	20✔
379	Stat::IsConstant,	20✔
380	Stat::IsSorted,	20✔
381	Stat::IsStrictSorted,	20✔
382	Stat::Max,	20✔
383	Stat::Min,	20✔
384	Stat::Sum,	20✔
385	Stat::NullCount,	20✔
386	Stat::NaNCount,	20✔
387	];	20✔
388
389	for stat in stats_to_check {	180✔
390	match (left.statistics().get(stat), right.statistics().get(stat)) {	160✔
391	(Some(Precision::Exact(left_v)), Some(Precision::Exact(right_v))) => {	20✔
392	if !left_v.eq(&right_v) {	20✔
NEW 393	return false;	×
394	}	20✔
395	}
396	_ => continue,	140✔
397	}
398	}
399
400	true	20✔
401	}	20✔
402
403	#[cfg(test)]
404	mod tests {
405	use super::*;
406	use crate::IntoArray;
407	use crate::arrays::{BoolArray, ChunkedArray, ConstantArray, PrimitiveArray, VarBinArray};
408	use crate::validity::Validity;
409	use vortex_dtype::{DType, Nullability, PType};
410
411	#[test]
412	fn test_simple_equals() {	1✔
413	let arr1 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]);	1✔
414	let arr2 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]);	1✔
415	let arr3 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 6]);	1✔
416
417	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
418	assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());	1✔
419	}	1✔
420
421	#[test]
422	fn test_stats_comparison() {	1✔
423	// Arrays with different stats should be detected as different early
424	let arr1 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]);	1✔
425	let arr2 = PrimitiveArray::from_iter(vec![10i32, 20, 30, 40, 50]);	1✔
426
427	assert!(!array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
428	}	1✔
429
430	#[test]
431	fn test_constant_arrays() {	1✔
432	let const1 = ConstantArray::new(Scalar::from(42i32), 100);	1✔
433	let const2 = ConstantArray::new(Scalar::from(42i32), 100);	1✔
434	let const3 = ConstantArray::new(Scalar::from(43i32), 100);	1✔
435
436	assert!(array_equals(const1.as_ref(), const2.as_ref()).unwrap());	1✔
437	assert!(!array_equals(const1.as_ref(), const3.as_ref()).unwrap());	1✔
438	}	1✔
439
440	#[test]
441	fn test_different_types() {	1✔
442	let int_arr = PrimitiveArray::from_iter(vec![1i32, 2, 3]);	1✔
443	let float_arr = PrimitiveArray::from_iter(vec![1.0f32, 2.0, 3.0]);	1✔
444
445	assert!(!array_equals(int_arr.as_ref(), float_arr.as_ref()).unwrap());	1✔
446	}	1✔
447
448	#[test]
449	fn test_with_nulls() {	1✔
450	let arr1 = PrimitiveArray::from_option_iter(vec![Some(1i32), None, Some(3), Some(4)]);	1✔
451	let arr2 = PrimitiveArray::from_option_iter(vec![Some(1i32), None, Some(3), Some(4)]);	1✔
452	let arr3 = PrimitiveArray::from_option_iter(vec![Some(1i32), Some(2), Some(3), Some(4)]);	1✔
453
454	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
455	assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());	1✔
456	}	1✔
457
458	#[test]
459	fn test_null_arrays() {	1✔
460	let arr1 = PrimitiveArray::from_option_iter(vec![None::<i32>, None, None]);	1✔
461	let arr2 = PrimitiveArray::from_option_iter(vec![None::<i32>, None, None]);	1✔
462
463	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
464	}	1✔
465
466	#[test]
467	fn test_bool_arrays() {	1✔
468	use arrow_buffer::BooleanBuffer;
469
470	let arr1 = BoolArray::new(	1✔
471	BooleanBuffer::from_iter([true, false, true, false]),	1✔
472	Validity::AllValid,	1✔
473	);
474	let arr2 = BoolArray::new(	1✔
475	BooleanBuffer::from_iter([true, false, true, false]),	1✔
476	Validity::AllValid,	1✔
477	);
478	let arr3 = BoolArray::new(	1✔
479	BooleanBuffer::from_iter([true, false, false, false]),	1✔
480	Validity::AllValid,	1✔
481	);
482
483	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
484	assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());	1✔
485	}	1✔
486
487	#[test]
488	fn test_empty_arrays() {	1✔
489	let empty1 = PrimitiveArray::from_iter(Vec::<i32>::new());	1✔
490	let empty2 = PrimitiveArray::from_iter(Vec::<i32>::new());	1✔
491
492	assert!(array_equals(empty1.as_ref(), empty2.as_ref()).unwrap());	1✔
493	}	1✔
494
495	#[test]
496	fn test_different_lengths() {	1✔
497	let arr1 = PrimitiveArray::from_iter(vec![1i32, 2, 3]);	1✔
498	let arr2 = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4]);	1✔
499
500	assert!(!array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
501	}	1✔
502
503	#[test]
504	fn test_large_arrays() {	1✔
505	// Test arrays larger than BATCH_SIZE
506	let data1: Vec<i64> = (0..100_000).collect();	1✔
507	let data2: Vec<i64> = (0..100_000).collect();	1✔
508	let mut data3 = data1.clone();	1✔
509	data3[99_999] = 999_999;	1✔
510
511	let arr1 = PrimitiveArray::from_iter(data1);	1✔
512	let arr2 = PrimitiveArray::from_iter(data2);	1✔
513	let arr3 = PrimitiveArray::from_iter(data3);	1✔
514
515	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
516	assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());	1✔
517	}	1✔
518
519	#[test]
520	fn test_non_canonical_arrays() {	1✔
521	let varbin1 = VarBinArray::from_vec(	1✔
522	vec!["hello".as_bytes(), "world".as_bytes()],	1✔
523	DType::Utf8(Nullability::NonNullable),	1✔
524	);
525	let varbin2 = VarBinArray::from_vec(	1✔
526	vec!["hello".as_bytes(), "world".as_bytes()],	1✔
527	DType::Utf8(Nullability::NonNullable),	1✔
528	);
529	let varbin3 = VarBinArray::from_vec(	1✔
530	vec!["hello".as_bytes(), "earth".as_bytes()],	1✔
531	DType::Utf8(Nullability::NonNullable),	1✔
532	);
533
534	assert!(array_equals(varbin1.as_ref(), varbin2.as_ref()).unwrap());	1✔
535	assert!(!array_equals(varbin1.as_ref(), varbin3.as_ref()).unwrap());	1✔
536	}	1✔
537
538	#[test]
539	fn test_float_precision() {	1✔
540	// Test if statistics-based comparison can handle float precision issues
541	let arr1 = PrimitiveArray::from_iter(vec![1.0f64, 2.0, 3.0, 4.0, 5.0]);	1✔
542	let arr2 = PrimitiveArray::from_iter(vec![1.0f64, 2.0, 3.0, 4.0, 5.0]);	1✔
543
544	// Arrays with exact same values should be equal
545	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
546
547	// Arrays with slightly different values should not be equal
548	let arr3 = PrimitiveArray::from_iter(vec![1.0f64, 2.0, 3.0, 4.0, 5.0000000001]);	1✔
549	assert!(!array_equals(arr1.as_ref(), arr3.as_ref()).unwrap());	1✔
550	}	1✔
551
552	#[test]
553	fn test_batch_size_functionality() {	1✔
554	// Test arrays larger than default batch size with different batch sizes
555	let data1: Vec<i32> = (0..150_000).collect();	1✔
556	let data2: Vec<i32> = (0..150_000).collect();	1✔
557
558	let arr1 = PrimitiveArray::from_iter(data1);	1✔
559	let arr2 = PrimitiveArray::from_iter(data2);	1✔
560
561	// Test with different batch sizes (though we can't pass batch_size directly in public API)
562	assert!(array_equals(arr1.as_ref(), arr2.as_ref()).unwrap());	1✔
563	}	1✔
564
565	#[test]
566	fn test_primitive_vs_dict_array() {	1✔
567	// Test comparing primitive array with dictionary-encoded array containing same values
568
569	let primitive_arr = PrimitiveArray::from_iter(vec![1i32, 2, 1, 3, 2, 1]);	1✔
570
571	// Create a chunked array as a proxy for non-canonical encoding
572	let chunk1 = PrimitiveArray::from_iter(vec![1i32, 2, 1]);	1✔
573	let chunk2 = PrimitiveArray::from_iter(vec![3i32, 2, 1]);	1✔
574	let chunked_arr = ChunkedArray::try_new(	1✔
575	vec![chunk1.into_array(), chunk2.into_array()],	1✔
576	primitive_arr.dtype().clone(),	1✔
577	)
578	.unwrap();	1✔
579
580	// Should be equal as they contain the same logical values
581	assert!(array_equals(primitive_arr.as_ref(), chunked_arr.as_ref()).unwrap());	1✔
582
583	// Test with different values
584	let chunk1_copy = PrimitiveArray::from_iter(vec![1i32, 2, 1]);	1✔
585	let different_chunk2 = PrimitiveArray::from_iter(vec![3i32, 2, 4]);	1✔
586	let different_chunked = ChunkedArray::try_new(	1✔
587	vec![chunk1_copy.into_array(), different_chunk2.into_array()],	1✔
588	primitive_arr.dtype().clone(),	1✔
589	)
590	.unwrap();	1✔
591
592	assert!(!array_equals(primitive_arr.as_ref(), different_chunked.as_ref()).unwrap());	1✔
593	}	1✔
594
595	#[test]
596	fn test_constant_null_arrays() {	1✔
597	// Test constant null arrays - should be equal to each other but not to non-null constants
598	let null_const1 = ConstantArray::new(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), 5);	1✔
599	let null_const2 = ConstantArray::new(Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)), 5);	1✔
600	let non_null_const = ConstantArray::new(Scalar::from(42i32), 5);	1✔
601
602	// Both null constants should be equal
603	assert!(array_equals(null_const1.as_ref(), null_const2.as_ref()).unwrap());	1✔
604
605	// Null constant should not equal non-null constant
606	assert!(!array_equals(null_const1.as_ref(), non_null_const.as_ref()).unwrap());	1✔
607	assert!(!array_equals(non_null_const.as_ref(), null_const1.as_ref()).unwrap());	1✔
608	}	1✔
609	}

vortex-data / vortex / 16606232494

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