Use AllNonDistinct in assert_arrays_eq and implement it for variant types (#8546)

AllNonDistinct is exactly the semantics we want, it should be faster and
in case
of failure we run the slow method to produce precise error message

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

encodings/parquet-variant/src/compute/allnondistinct.rs

@@ -0,0 +1,161 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt::Debug;
+
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::aggregate_fn::AggregateFnRef;
+use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct;
+use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct;
+use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
+use vortex_array::arrays::Struct;
+use vortex_array::arrays::struct_::StructArrayExt;
+use vortex_array::dtype::Nullability;
+use vortex_array::scalar::Scalar;
+use vortex_error::VortexResult;
+
+use crate::ParquetVariant;
+use crate::ParquetVariantArrayExt;
+
+/// Lets `AllNonDistinct` compare two `ParquetVariant` arrays without canonicalizing them.
+///
+/// `AllNonDistinct` accumulates over a `Struct{lhs, rhs}` batch, so this kernel is registered for
+/// the struct encoding and inspects the two children. When both are `ParquetVariant`, we compare
+/// the typed (`typed_value`) arrays if both sides are shredded, and fall back to the raw `value`
+/// arrays otherwise. Comparing these child arrays directly avoids re-canonicalizing the variant
+/// (which would recurse through the `Variant` canonical form).
+#[derive(Debug)]
+pub struct AllNonDistinctParquetVariant;
+
+impl DynAggregateKernel for AllNonDistinctParquetVariant {
+    fn aggregate(
+        &self,
+        aggregate_fn: &AggregateFnRef,
+        batch: &ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<Scalar>> {
+        if !aggregate_fn.is::<AllNonDistinct>() {
+            return Ok(None);
+        }
+
+        let Some(batch) = batch.as_opt::<Struct>() else {
+            return Ok(None);
+        };
+        let lhs = batch.unmasked_field(0);
+        let rhs = batch.unmasked_field(1);
+        let (Some(lhs), Some(rhs)) = (
+            lhs.as_opt::<ParquetVariant>(),
+            rhs.as_opt::<ParquetVariant>(),
+        ) else {
+            return Ok(None);
+        };
+
+        let typed_identical = match (lhs.typed_value_array(), rhs.typed_value_array()) {
+            (Some(lhs_typed), Some(rhs_typed)) => {
+                if lhs_typed.dtype().eq_ignore_nullability(rhs_typed.dtype()) {
+                    all_non_distinct(lhs_typed, rhs_typed, ctx)?
+                } else {
+                    return Ok(None);
+                }
+            }
+            _ => true,
+        };
+
+        if typed_identical {
+            let values_identical = match (lhs.value_array(), rhs.value_array()) {
+                (Some(lhs_value), Some(rhs_value)) => all_non_distinct(lhs_value, rhs_value, ctx)?,
+                (None, None) => true,
+                // Mixed shredding layouts: let the generic canonical path handle it.
+                _ => return Ok(None),
+            };
+            Ok(Some(Scalar::bool(
+                values_identical,
+                Nullability::NonNullable,
+            )))
+        } else {
+            Ok(Some(Scalar::bool(false, Nullability::NonNullable)))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::LazyLock;
+
+    use vortex_array::ArrayRef;
+    use vortex_array::IntoArray;
+    use vortex_array::VortexSessionExecute;
+    use vortex_array::aggregate_fn::fns::all_non_distinct::all_non_distinct;
+    use vortex_array::arrays::VarBinViewArray;
+    use vortex_array::validity::Validity;
+    use vortex_buffer::buffer;
+    use vortex_error::VortexResult;
+    use vortex_session::VortexSession;
+
+    use crate::ParquetVariant;
+
+    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+        let session = vortex_array::array_session();
+        crate::initialize(&session);
+        session
+    });
+
+    /// Non-nullable, minimally-valid metadata column of `len` rows.
+    fn metadata(len: usize) -> ArrayRef {
+        VarBinViewArray::from_iter_bin(vec![b"\x01\x00"; len]).into_array()
+    }
+
+    /// Non-nullable binary `value` column.
+    fn binary<T: AsRef<[u8]>>(values: impl IntoIterator<Item = T>) -> ArrayRef {
+        VarBinViewArray::from_iter_bin(values).into_array()
+    }
+
+    fn parquet_variant(
+        len: usize,
+        value: Option<ArrayRef>,
+        typed_value: Option<ArrayRef>,
+    ) -> VortexResult<ArrayRef> {
+        Ok(
+            ParquetVariant::try_new(Validity::NonNullable, metadata(len), value, typed_value)?
+                .into_array(),
+        )
+    }
+
+    #[test]
+    fn all_non_distinct_matches_equal_unshredded() -> VortexResult<()> {
+        let lhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), None)?;
+        let rhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), None)?;
+        let mut ctx = SESSION.create_execution_ctx();
+        assert!(all_non_distinct(&lhs, &rhs, &mut ctx)?);
+        Ok(())
+    }
+
+    #[test]
+    fn all_non_distinct_detects_distinct_unshredded() -> VortexResult<()> {
+        let lhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), None)?;
+        let rhs = parquet_variant(2, Some(binary([b"\x10", b"\x12"])), None)?;
+        let mut ctx = SESSION.create_execution_ctx();
+        assert!(!all_non_distinct(&lhs, &rhs, &mut ctx)?);
+        Ok(())
+    }
+
+    #[test]
+    fn all_non_distinct_matches_equal_value_and_typed() -> VortexResult<()> {
+        let typed = || buffer![1i32, 2].into_array();
+        let lhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), Some(typed()))?;
+        let rhs = parquet_variant(2, Some(binary([b"\x10", b"\x11"])), Some(typed()))?;
+        let mut ctx = SESSION.create_execution_ctx();
+        assert!(all_non_distinct(&lhs, &rhs, &mut ctx)?);
+        Ok(())
+    }
+
+    #[test]
+    fn all_non_distinct_empty_is_true() -> VortexResult<()> {
+        let lhs = parquet_variant(0, Some(binary(Vec::<&[u8]>::new())), None)?;
+        let rhs = parquet_variant(0, Some(binary(Vec::<&[u8]>::new())), None)?;
+        let mut ctx = SESSION.create_execution_ctx();
+        assert!(all_non_distinct(&lhs, &rhs, &mut ctx)?);
+        Ok(())
+    }
+}

encodings/parquet-variant/src/compute/mod.rs

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod allnondistinct;
+
+pub use allnondistinct::*;

encodings/parquet-variant/src/kernel.rs

@@ -20,10 +20,14 @@ use vortex_array::ArrayVTable;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
 use vortex_array::IntoArray;
+use vortex_array::aggregate_fn::AggregateFnVTable;
+use vortex_array::aggregate_fn::fns::all_non_distinct::AllNonDistinct;
+use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
 use vortex_array::arrays::Dict;
 use vortex_array::arrays::Extension;
 use vortex_array::arrays::Filter;
 use vortex_array::arrays::Slice;
+use vortex_array::arrays::Struct;
 use vortex_array::arrays::dict::TakeExecute;
 use vortex_array::arrays::dict::TakeExecuteAdaptor;
 use vortex_array::arrays::extension::ExtensionArrayExt;
@@ -52,6 +56,7 @@ use vortex_session::VortexSession;
 
 use crate::ParquetVariant;
 use crate::ParquetVariantArrayExt;
+use crate::compute::AllNonDistinctParquetVariant;
 
 pub(crate) fn initialize(session: &VortexSession) {
     let kernels = session.kernels();
@@ -76,6 +81,12 @@ pub(crate) fn initialize(session: &VortexSession) {
         Extension,
         JsonExtensionToVariantKernel,
     );
+    let aggregates = session.aggregate_fns();
+    aggregates.register_aggregate_kernel(
+        Struct.id(),
+        Some(AllNonDistinct.id()),
+        &AllNonDistinctParquetVariant,
+    );
 }
 
 #[derive(Default, Debug)]

encodings/parquet-variant/src/lib.rs

@@ -26,6 +26,7 @@
 
 mod array;
 mod arrow;
+mod compute;
 #[cfg(test)]
 mod json_to_variant_tests;
 mod kernel;

vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs

@@ -12,6 +12,7 @@ mod struct_;
 #[cfg(test)]
 mod tests;
 mod varbin;
+mod variant;
 
 use std::sync::LazyLock;
 
@@ -38,6 +39,7 @@ use crate::aggregate_fn::AggregateFnId;
 use crate::aggregate_fn::AggregateFnVTable;
 use crate::aggregate_fn::DynAccumulator;
 use crate::aggregate_fn::EmptyOptions;
+use crate::aggregate_fn::fns::all_non_distinct::variant::check_variant_identical;
 use crate::arrays::StructArray;
 use crate::arrays::struct_::StructArrayExt;
 use crate::dtype::DType;
@@ -264,8 +266,8 @@ fn check_canonical_identical(
         (Canonical::Extension(lhs), Canonical::Extension(rhs)) => {
             check_extension_identical(lhs, rhs, ctx)
         }
-        (Canonical::Variant(_), _) | (_, Canonical::Variant(_)) => {
-            vortex_bail!("Variant arrays don't support AllNonDistinct")
+        (Canonical::Variant(lhs), Canonical::Variant(rhs)) => {
+            check_variant_identical(lhs, rhs, ctx)
         }
         _ => Err(vortex_err!(
             "Canonical type mismatch in AllNonDistinct: {:?} vs {:?}",

vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_array::dtype::NativePType;
 use vortex_error::VortexResult;
 
 use crate::arrays::primitive::PrimitiveArrayExt;
@@ -12,6 +13,10 @@ where
     R: PrimitiveArrayExt,
 {
     match_each_native_ptype!(lhs.ptype(), |P| {
-        Ok(lhs.as_slice::<P>() == rhs.as_slice::<P>())
+        Ok(lhs
+            .as_slice::<P>()
+            .iter()
+            .zip(rhs.as_slice::<P>())
+            .all(|(l, r)| l.is_eq(*r)))
     })
 }

vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexResult;
+
+use crate::ExecutionCtx;
+use crate::arrays::VariantArray;
+
+/// Checks whether two canonical variant arrays are element-wise non-distinct.
+///
+/// Variant values cannot be routed back through [`all_non_distinct`]: canonicalizing a variant
+/// value array yields another canonical variant (with no shredded tree), which would recurse
+/// forever. The generic fallback therefore compares logical variant scalars row-by-row. Encodings
+/// that can compare their typed/value children more cheaply (e.g. `ParquetVariant`) register an
+/// aggregate kernel that intercepts the comparison before it reaches this fallback.
+///
+/// [`all_non_distinct`]: super::all_non_distinct
+pub(super) fn check_variant_identical(
+    lhs: &VariantArray,
+    rhs: &VariantArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<bool> {
+    if lhs.len() != rhs.len() {
+        return Ok(false);
+    }
+    for idx in 0..lhs.len() {
+        if lhs.execute_scalar(idx, ctx)? != rhs.execute_scalar(idx, ctx)? {
+            return Ok(false);
+        }
+    }
+    Ok(true)
+}

vortex-array/src/arrays/assertions.rs

@@ -10,6 +10,7 @@ use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::RecursiveCanonical;
+use crate::aggregate_fn::fns::all_non_distinct::all_non_distinct;
 
 fn format_indices<I: IntoIterator<Item = usize>>(indices: I) -> impl Display {
     indices.into_iter().format(",")
@@ -117,17 +118,23 @@ macro_rules! assert_arrays_eq {
 pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef, ctx: &mut ExecutionCtx) {
     let executed = execute_to_canonical(left.clone(), ctx);
 
-    let left_right = find_mismatched_indices(left, right, ctx);
-    let executed_right = find_mismatched_indices(&executed, right, ctx);
+    let left_right_the_same =
+        all_non_distinct(left, right, ctx).vortex_expect("failed to compare left and right");
+    let executed_right_the_same = all_non_distinct(&executed, right, ctx)
+        .vortex_expect("failed to compare executed left and right");
+
+    if !left_right_the_same || !executed_right_the_same {
+        let left_right = find_mismatched_indices(left, right, ctx);
 
-    if !left_right.is_empty() || !executed_right.is_empty() {
         let mut msg = String::new();
         if !left_right.is_empty() {
             msg.push_str(&format!(
                 "\n  left != right at indices: {}",
                 format_indices(left_right)
             ));
         }
+
+        let executed_right = find_mismatched_indices(&executed, right, ctx);
         if !executed_right.is_empty() {
             msg.push_str(&format!(
                 "\n  executed != right at indices: {}",

vortex-ipc/src/lib.rs

@@ -19,6 +19,7 @@ pub mod stream;
 mod test {
     use std::sync::LazyLock;
 
+    use vortex_array::aggregate_fn::session::AggregateFnSession;
     use vortex_array::dtype::session::DTypeSession;
     use vortex_array::optimizer::kernels::KernelSession;
     use vortex_array::session::ArraySession;
@@ -29,6 +30,7 @@ mod test {
             .with::<DTypeSession>()
             .with::<ArraySession>()
             .with::<KernelSession>()
+            .with::<AggregateFnSession>()
             .build()
     });
 }