diff --git a/encodings/parquet-variant/src/compute/allnondistinct.rs b/encodings/parquet-variant/src/compute/allnondistinct.rs new file mode 100644 index 00000000000..4f5568fbcf6 --- /dev/null +++ b/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> { + if !aggregate_fn.is::() { + return Ok(None); + } + + let Some(batch) = batch.as_opt::() else { + return Ok(None); + }; + let lhs = batch.unmasked_field(0); + let rhs = batch.unmasked_field(1); + let (Some(lhs), Some(rhs)) = ( + lhs.as_opt::(), + rhs.as_opt::(), + ) 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 = 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>(values: impl IntoIterator) -> ArrayRef { + VarBinViewArray::from_iter_bin(values).into_array() + } + + fn parquet_variant( + len: usize, + value: Option, + typed_value: Option, + ) -> VortexResult { + 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(()) + } +} diff --git a/encodings/parquet-variant/src/compute/mod.rs b/encodings/parquet-variant/src/compute/mod.rs new file mode 100644 index 00000000000..f69f592b89d --- /dev/null +++ b/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::*; diff --git a/encodings/parquet-variant/src/kernel.rs b/encodings/parquet-variant/src/kernel.rs index 88a4b8931ab..bd7d0bfaa7c 100644 --- a/encodings/parquet-variant/src/kernel.rs +++ b/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)] diff --git a/encodings/parquet-variant/src/lib.rs b/encodings/parquet-variant/src/lib.rs index 05e93c5d08c..c71ea81e734 100644 --- a/encodings/parquet-variant/src/lib.rs +++ b/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; diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs index 445f4e17b6f..fcb9b18346d 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/mod.rs +++ b/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 {:?}", diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs index f11a9a2bc52..1ad0ca6d646 100644 --- a/vortex-array/src/aggregate_fn/fns/all_non_distinct/primitive.rs +++ b/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::

() == rhs.as_slice::

()) + Ok(lhs + .as_slice::

() + .iter() + .zip(rhs.as_slice::

()) + .all(|(l, r)| l.is_eq(*r))) }) } diff --git a/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs b/vortex-array/src/aggregate_fn/fns/all_non_distinct/variant.rs new file mode 100644 index 00000000000..1576b8c3d6a --- /dev/null +++ b/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 { + 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) +} diff --git a/vortex-array/src/arrays/assertions.rs b/vortex-array/src/arrays/assertions.rs index a445f26ec15..aff8d4e32cb 100644 --- a/vortex-array/src/arrays/assertions.rs +++ b/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>(indices: I) -> impl Display { indices.into_iter().format(",") @@ -117,10 +118,14 @@ 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!( @@ -128,6 +133,8 @@ pub fn assert_arrays_eq_impl(left: &ArrayRef, right: &ArrayRef, ctx: &mut Execut 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: {}", diff --git a/vortex-ipc/src/lib.rs b/vortex-ipc/src/lib.rs index 2a198387b6f..667f175cbc4 100644 --- a/vortex-ipc/src/lib.rs +++ b/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::() .with::() .with::() + .with::() .build() }); }