diff --git a/vortex-array/benches/compare.rs b/vortex-array/benches/compare.rs index 95702c9f8f5..aaeda5ab3b6 100644 --- a/vortex-array/benches/compare.rs +++ b/vortex-array/benches/compare.rs @@ -8,12 +8,21 @@ use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::DecimalArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinViewArray; use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DecimalDType; +use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::validity::Validity; use vortex_buffer::Buffer; fn main() { @@ -22,51 +31,190 @@ fn main() { const ARRAY_SIZE: usize = 65_536; -#[divan::bench] -fn compare_bool(bencher: Bencher) { - let mut rng = StdRng::seed_from_u64(0); - let range = Uniform::new(0u8, 1).unwrap(); - - let arr1 = BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(range) == 0)).into_array(); - let arr2 = BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(range) == 0)).into_array(); +fn bench_compare(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef, op: Operator) { let session = vortex_array::array_session(); - bencher - .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) + .with_inputs(|| (&lhs, &rhs, session.create_execution_ctx())) .bench_refs(|input| { input .0 .clone() - .binary(input.1.clone(), Operator::Gte) + .binary(input.1.clone(), op) .unwrap() .execute::(&mut input.2) }); } -#[divan::bench] -fn compare_int(bencher: Bencher) { - let mut rng = StdRng::seed_from_u64(0); - let range = Uniform::new(0i64, 100_000_000).unwrap(); +fn bool_array(rng: &mut StdRng) -> ArrayRef { + BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.5))).into_array() +} + +fn bool_array_nullable(rng: &mut StdRng) -> ArrayRef { + BoolArray::new( + (0..ARRAY_SIZE).map(|_| rng.random_bool(0.5)).collect(), + Validity::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.9))), + ) + .into_array() +} - let arr1 = (0..ARRAY_SIZE) +fn int_array(rng: &mut StdRng) -> ArrayRef { + let range = Uniform::new(0i64, 100_000_000).unwrap(); + (0..ARRAY_SIZE) .map(|_| rng.sample(range)) .collect::>() - .into_array(); + .into_array() +} - let arr2 = (0..ARRAY_SIZE) - .map(|_| rng.sample(range)) +fn int_array_nullable(rng: &mut StdRng) -> ArrayRef { + let range = Uniform::new(0i64, 100_000_000).unwrap(); + PrimitiveArray::new( + (0..ARRAY_SIZE) + .map(|_| rng.sample(range)) + .collect::>(), + Validity::from_iter((0..ARRAY_SIZE).map(|_| rng.random_bool(0.9))), + ) + .into_array() +} + +fn float_array(rng: &mut StdRng) -> ArrayRef { + (0..ARRAY_SIZE) + .map(|_| rng.random_range(0.0f64..1.0)) .collect::>() - .into_array(); - let session = vortex_array::array_session(); + .into_array() +} - bencher - .with_inputs(|| (&arr1, &arr2, session.create_execution_ctx())) - .bench_refs(|input| { - input - .0 - .clone() - .binary(input.1.clone(), Operator::Gte) - .unwrap() - .execute::(&mut input.2) - }); +fn string_array(rng: &mut StdRng) -> ArrayRef { + VarBinViewArray::from_iter_str((0..ARRAY_SIZE).map(|_| { + let len = rng.random_range(1usize..24); + (0..len) + .map(|_| char::from(rng.random_range(b'a'..=b'z'))) + .collect::() + })) + .into_array() +} + +fn decimal_array(rng: &mut StdRng) -> ArrayRef { + DecimalArray::from_iter::( + (0..ARRAY_SIZE).map(|_| rng.random_range(0i128..100_000_000)), + DecimalDType::new(38, 2), + ) + .into_array() +} + +#[divan::bench] +fn compare_bool(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = bool_array(&mut rng); + let arr2 = bool_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_bool_nullable(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = bool_array_nullable(&mut rng); + let arr2 = bool_array_nullable(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_bool_constant(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr = bool_array(&mut rng); + let constant = ConstantArray::new(true, ARRAY_SIZE).into_array(); + bench_compare(bencher, arr, constant, Operator::Eq); +} + +#[divan::bench] +fn compare_int(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = int_array(&mut rng); + let arr2 = int_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_int_nullable(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = int_array_nullable(&mut rng); + let arr2 = int_array_nullable(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_int_constant(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr = int_array(&mut rng); + let constant = ConstantArray::new(50_000_000i64, ARRAY_SIZE).into_array(); + bench_compare(bencher, arr, constant, Operator::Gte); +} + +#[divan::bench] +fn compare_int_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = int_array(&mut rng); + let arr2 = int_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_float(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = float_array(&mut rng); + let arr2 = float_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_decimal(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = decimal_array(&mut rng); + let arr2 = decimal_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Gte); +} + +#[divan::bench] +fn compare_string_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = string_array(&mut rng); + let arr2 = string_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); +} + +#[divan::bench] +fn compare_string_lt(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = string_array(&mut rng); + let arr2 = string_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Lt); +} + +#[divan::bench] +fn compare_string_constant(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr = string_array(&mut rng); + let constant = ConstantArray::new(Scalar::from("mmmmmmmmmmmm"), ARRAY_SIZE).into_array(); + bench_compare(bencher, arr, constant, Operator::Lt); +} + +fn struct_array(rng: &mut StdRng) -> ArrayRef { + StructArray::from_fields(&[("a", int_array(rng)), ("b", int_array(rng))]) + .unwrap() + .into_array() +} + +#[divan::bench] +fn compare_struct_lt(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = struct_array(&mut rng); + let arr2 = struct_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Lt); +} + +#[divan::bench] +fn compare_struct_eq(bencher: Bencher) { + let mut rng = StdRng::seed_from_u64(0); + let arr1 = struct_array(&mut rng); + let arr2 = struct_array(&mut rng); + bench_compare(bencher, arr1, arr2, Operator::Eq); } diff --git a/vortex-array/src/arrays/extension/compute/compare.rs b/vortex-array/src/arrays/extension/compute/compare.rs index ecd6bbd5614..b3f978ef503 100644 --- a/vortex-array/src/arrays/extension/compute/compare.rs +++ b/vortex-array/src/arrays/extension/compute/compare.rs @@ -22,6 +22,12 @@ impl CompareKernel for Extension { operator: CompareOperator, _ctx: &mut ExecutionCtx, ) -> VortexResult> { + // Storage values are only comparable when both sides share the same extension dtype + // (e.g. timestamps in different units must not compare their raw storage). + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + return Ok(None); + } + // If the RHS is a constant, we can extract the storage scalar. if let Some(const_ext) = rhs.as_constant() { let storage_scalar = const_ext.as_extension().to_storage_scalar(); diff --git a/vortex-array/src/scalar_fn/fns/binary/compare.rs b/vortex-array/src/scalar_fn/fns/binary/compare.rs deleted file mode 100644 index b7a53c21d62..00000000000 --- a/vortex-array/src/scalar_fn/fns/binary/compare.rs +++ /dev/null @@ -1,590 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::cmp::Ordering; - -use arrow_array::BooleanArray; -use arrow_buffer::NullBuffer; -use arrow_ord::cmp; -use arrow_ord::ord::make_comparator; -use arrow_schema::Field; -use arrow_schema::SortOptions; -use vortex_error::VortexResult; -use vortex_error::vortex_err; - -use crate::ArrayRef; -use crate::Canonical; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::array::VTable; -use crate::arrays::Constant; -use crate::arrays::ConstantArray; -use crate::arrays::ScalarFn; -use crate::arrays::scalar_fn::ExactScalarFn; -use crate::arrays::scalar_fn::ScalarFnArrayExt; -use crate::arrays::scalar_fn::ScalarFnArrayView; -use crate::arrow::ArrowSessionExt; -use crate::arrow::Datum; -use crate::arrow::from_arrow_columnar; -use crate::dtype::DType; -use crate::dtype::Nullability; -use crate::kernel::ExecuteParentKernel; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::Binary; -use crate::scalar_fn::fns::operators::CompareOperator; - -/// Trait for encoding-specific comparison kernels that operate in encoded space. -/// -/// Implementations can compare an encoded array against another array (typically a constant) -/// without first decompressing. The adaptor normalizes operand order so `array` is always -/// the left-hand side, swapping the operator when necessary. -pub trait CompareKernel: VTable { - fn compare( - lhs: ArrayView<'_, Self>, - rhs: &ArrayRef, - operator: CompareOperator, - ctx: &mut ExecutionCtx, - ) -> VortexResult>; -} - -/// Adaptor that bridges [`CompareKernel`] implementations to [`ExecuteParentKernel`]. -/// -/// When a `ScalarFnArray(Binary, cmp_op)` wraps a child that implements `CompareKernel`, -/// this adaptor extracts the comparison operator and other operand, normalizes operand order -/// (swapping the operator if the encoded array is on the RHS), and delegates to the kernel. -#[derive(Default, Debug)] -pub struct CompareExecuteAdaptor(pub V); - -impl ExecuteParentKernel for CompareExecuteAdaptor -where - V: CompareKernel, -{ - type Parent = ExactScalarFn; - - fn execute_parent( - &self, - array: ArrayView<'_, V>, - parent: ScalarFnArrayView<'_, Binary>, - child_idx: usize, - ctx: &mut ExecutionCtx, - ) -> VortexResult> { - // Only handle comparison operators - let Ok(cmp_op) = CompareOperator::try_from(*parent.options) else { - return Ok(None); - }; - - // Get the ScalarFnArray to access children - let Some(scalar_fn_array) = parent.as_opt::() else { - return Ok(None); - }; - // Normalize so `array` is always LHS, swapping the operator if needed - // TODO(joe): should be go this here or in the Rule/Kernel - let (cmp_op, other) = match child_idx { - 0 => (cmp_op, scalar_fn_array.get_child(1)), - 1 => (cmp_op.swap(), scalar_fn_array.get_child(0)), - _ => return Ok(None), - }; - - let len = array.len(); - let nullable = array.dtype().is_nullable() || other.dtype().is_nullable(); - - // Empty array → empty bool result - if len == 0 { - return Ok(Some( - Canonical::empty(&DType::Bool(nullable.into())).into_array(), - )); - } - - // Null constant on either side → all-null bool result - if other.as_constant().is_some_and(|s| s.is_null()) { - return Ok(Some( - ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), len).into_array(), - )); - } - - V::compare(array, other, cmp_op, ctx) - } -} - -/// Execute a compare operation between two arrays. -/// -/// This is the entry point for compare operations from the binary expression. -/// Handles empty, constant-null, and constant-constant directly, otherwise falls back to Arrow. -pub(crate) fn execute_compare( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: CompareOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); - - if lhs.is_empty() { - return Ok(Canonical::empty(&DType::Bool(nullable.into())).into_array()); - } - - 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(DType::Bool(nullable.into())), lhs.len()).into_array(), - ); - } - - // Constant-constant fast path - if let (Some(lhs_const), Some(rhs_const)) = (lhs.as_opt::(), rhs.as_opt::()) - { - let result = scalar_cmp(lhs_const.scalar(), rhs_const.scalar(), op)?; - return Ok(ConstantArray::new(result, lhs.len()).into_array()); - } - - arrow_compare_arrays(lhs, rhs, op, ctx) -} - -/// Fall back to Arrow for comparison. -fn arrow_compare_arrays( - left: &ArrayRef, - right: &ArrayRef, - operator: CompareOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - assert_eq!(left.len(), right.len()); - - let nullable = left.dtype().is_nullable() || right.dtype().is_nullable(); - - // Arrow's vectorized comparison kernels don't support nested types. - // For nested types, fall back to `make_comparator` which does element-wise comparison. - let arrow_array: BooleanArray = if left.dtype().is_nested() || right.dtype().is_nested() { - let session = ctx.session().clone(); - let lhs = session.arrow().execute_arrow(left.clone(), None, ctx)?; - let target_field = Field::new("", lhs.data_type().clone(), right.dtype().is_nullable()); - let rhs = session - .arrow() - .execute_arrow(right.clone(), Some(&target_field), ctx)?; - - compare_nested_arrow_arrays(lhs.as_ref(), rhs.as_ref(), operator)? - } else { - // Fast path: use vectorized kernels for primitive types. - let lhs = Datum::try_new(left, ctx)?; - let rhs = Datum::try_new_with_target_datatype(right, lhs.data_type(), ctx)?; - - match operator { - CompareOperator::Eq => cmp::eq(&lhs, &rhs)?, - CompareOperator::NotEq => cmp::neq(&lhs, &rhs)?, - CompareOperator::Gt => cmp::gt(&lhs, &rhs)?, - CompareOperator::Gte => cmp::gt_eq(&lhs, &rhs)?, - CompareOperator::Lt => cmp::lt(&lhs, &rhs)?, - CompareOperator::Lte => cmp::lt_eq(&lhs, &rhs)?, - } - }; - - from_arrow_columnar(&arrow_array, left.len(), nullable, ctx) -} - -pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> VortexResult { - if lhs.is_null() | rhs.is_null() { - return Ok(Scalar::null(DType::Bool(Nullability::Nullable))); - } - - let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); - - // We use `partial_cmp` to ensure we do not lose a type mismatch error. - let ordering = lhs.partial_cmp(rhs).ok_or_else(|| { - vortex_err!( - "Cannot compare scalars with incompatible types: {} and {}", - lhs.dtype(), - rhs.dtype() - ) - })?; - - let b = match operator { - CompareOperator::Eq => ordering.is_eq(), - CompareOperator::NotEq => ordering.is_ne(), - CompareOperator::Gt => ordering.is_gt(), - CompareOperator::Gte => ordering.is_ge(), - CompareOperator::Lt => ordering.is_lt(), - CompareOperator::Lte => ordering.is_le(), - }; - - Ok(Scalar::bool(b, nullability)) -} - -/// Compare two Arrow arrays element-wise using [`make_comparator`]. -/// -/// This function is required for nested types (Struct, List, FixedSizeList) because Arrow's -/// vectorized comparison kernels ([`cmp::eq`], [`cmp::neq`], etc.) do not support them. -/// -/// The vectorized kernels are faster but only work on primitive types, so for non-nested types, -/// prefer using the vectorized kernels directly for better performance. -pub fn compare_nested_arrow_arrays( - lhs: &dyn arrow_array::Array, - rhs: &dyn arrow_array::Array, - operator: CompareOperator, -) -> VortexResult { - let compare_arrays_at = make_comparator(lhs, rhs, SortOptions::default())?; - - let cmp_fn = match operator { - CompareOperator::Eq => Ordering::is_eq, - CompareOperator::NotEq => Ordering::is_ne, - CompareOperator::Gt => Ordering::is_gt, - CompareOperator::Gte => Ordering::is_ge, - CompareOperator::Lt => Ordering::is_lt, - CompareOperator::Lte => Ordering::is_le, - }; - - let values = (0..lhs.len()) - .map(|i| cmp_fn(compare_arrays_at(i, i))) - .collect(); - let nulls = NullBuffer::union(lhs.nulls(), rhs.nulls()); - - Ok(BooleanArray::new(values, nulls)) -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use rstest::rstest; - use vortex_buffer::BitBuffer; - use vortex_buffer::buffer; - use vortex_error::VortexExpect; - - use crate::ArrayRef; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::ListArray; - use crate::arrays::ListViewArray; - use crate::arrays::PrimitiveArray; - use crate::arrays::StructArray; - use crate::arrays::VarBinArray; - use crate::arrays::VarBinViewArray; - use crate::assert_arrays_eq; - use crate::builtins::ArrayBuiltins; - use crate::dtype::DType; - use crate::dtype::FieldName; - use crate::dtype::FieldNames; - use crate::dtype::Nullability; - use crate::dtype::PType; - use crate::extension::datetime::TimeUnit; - use crate::extension::datetime::Timestamp; - use crate::extension::datetime::TimestampOptions; - use crate::scalar::Scalar; - use crate::scalar_fn::fns::binary::compare::ConstantArray; - use crate::scalar_fn::fns::binary::scalar_cmp; - use crate::scalar_fn::fns::operators::CompareOperator; - use crate::scalar_fn::fns::operators::Operator; - use crate::test_harness::to_int_indices; - use crate::validity::Validity; - - #[test] - fn test_bool_basic_comparisons() { - let ctx = &mut array_session().create_execution_ctx(); - let arr = BoolArray::new( - BitBuffer::from_iter([true, true, false, true, false]), - Validity::from_iter([false, true, true, true, true]), - ); - - let matches = arr - .clone() - .into_array() - .binary(arr.clone().into_array(), Operator::Eq) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [1u64, 2, 3, 4]); - - let matches = arr - .clone() - .into_array() - .binary(arr.clone().into_array(), Operator::NotEq) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - let empty: [u64; 0] = []; - assert_eq!(to_int_indices(matches, ctx).unwrap(), empty); - - let other = BoolArray::new( - BitBuffer::from_iter([false, false, false, true, true]), - Validity::from_iter([false, true, true, true, true]), - ); - - let matches = arr - .clone() - .into_array() - .binary(other.clone().into_array(), Operator::Lte) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); - - let matches = arr - .clone() - .into_array() - .binary(other.clone().into_array(), Operator::Lt) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); - - let matches = other - .clone() - .into_array() - .binary(arr.clone().into_array(), Operator::Gte) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); - - let matches = other - .into_array() - .binary(arr.into_array(), Operator::Gt) - .unwrap() - .execute::(ctx) - .vortex_expect("must be a bool array"); - assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); - } - - #[test] - fn constant_compare() { - let left = ConstantArray::new(Scalar::from(2u32), 10); - let right = ConstantArray::new(Scalar::from(10u32), 10); - - let result = left - .into_array() - .binary(right.into_array(), Operator::Gt) - .unwrap(); - assert_eq!(result.len(), 10); - let scalar = result - .execute_scalar(0, &mut array_session().create_execution_ctx()) - .unwrap(); - assert_eq!(scalar.as_bool().value(), Some(false)); - } - - #[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 mut ctx = array_session().create_execution_ctx(); - let res = left.binary(right, Operator::Eq).unwrap(); - let expected = BoolArray::from_iter([true, true]); - assert_arrays_eq!(res, expected, &mut ctx); - } - - #[test] - fn test_list_array_comparison() { - let mut ctx = array_session().create_execution_ctx(); - let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); - let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); - let list1 = ListArray::try_new( - values1.into_array(), - offsets1.into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let values2 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 7, 8]); - let offsets2 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); - let list2 = ListArray::try_new( - values2.into_array(), - offsets2.into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let result = list1 - .clone() - .into_array() - .binary(list2.clone().into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut ctx); - - let result = list1 - .clone() - .into_array() - .binary(list2.clone().into_array(), Operator::NotEq) - .unwrap(); - let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected, &mut ctx); - - let result = list1 - .into_array() - .binary(list2.into_array(), Operator::Lt) - .unwrap(); - let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected, &mut ctx); - } - - #[test] - fn test_list_array_constant_comparison() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); - let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); - let list = ListArray::try_new( - values.into_array(), - offsets.into_array(), - Validity::NonNullable, - ) - .unwrap(); - - let list_scalar = Scalar::list( - Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), - vec![3i32.into(), 4i32.into()], - Nullability::NonNullable, - ); - let constant = ConstantArray::new(list_scalar, 3); - - let result = list - .into_array() - .binary(constant.into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([false, true, false]); - assert_arrays_eq!(result, expected, &mut ctx); - } - - #[test] - fn test_struct_array_comparison() { - let mut ctx = array_session().create_execution_ctx(); - let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); - let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); - - let bool_field2 = BoolArray::from_iter([Some(true), Some(false), Some(false)]); - let int_field2 = PrimitiveArray::from_iter([1i32, 2, 4]); - - let struct1 = StructArray::from_fields(&[ - ("bool_col", bool_field1.into_array()), - ("int_col", int_field1.into_array()), - ]) - .unwrap(); - - let struct2 = StructArray::from_fields(&[ - ("bool_col", bool_field2.into_array()), - ("int_col", int_field2.into_array()), - ]) - .unwrap(); - - let result = struct1 - .clone() - .into_array() - .binary(struct2.clone().into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut ctx); - - let result = struct1 - .into_array() - .binary(struct2.into_array(), Operator::Gt) - .unwrap(); - let expected = BoolArray::from_iter([false, false, true]); - assert_arrays_eq!(result, expected, &mut ctx); - } - - #[test] - fn test_empty_struct_compare() { - let mut ctx = array_session().create_execution_ctx(); - let empty1 = StructArray::try_new( - FieldNames::from(Vec::::new()), - Vec::new(), - 5, - Validity::NonNullable, - ) - .unwrap(); - - let empty2 = StructArray::try_new( - FieldNames::from(Vec::::new()), - Vec::new(), - 5, - Validity::NonNullable, - ) - .unwrap(); - - let result = empty1 - .into_array() - .binary(empty2.into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, true, true, true]); - assert_arrays_eq!(result, expected, &mut ctx); - } - - /// Regression test: comparing struct arrays where the same logical field is backed by - /// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. - #[test] - fn struct_compare_mixed_binary_encodings() { - let mut ctx = array_session().create_execution_ctx(); - // LHS: struct with a VarBinArray (offset-based) binary field - let bin_field1 = VarBinArray::from(vec![ - "apple".as_bytes(), - "banana".as_bytes(), - "cherry".as_bytes(), - ]); - let struct1 = StructArray::from_fields(&[("data", bin_field1.into_array())]).unwrap(); - - // RHS: struct with a VarBinViewArray (view-based) binary field — same logical DType - let bin_field2 = VarBinViewArray::from_iter_bin([ - "apple".as_bytes(), - "banana".as_bytes(), - "durian".as_bytes(), - ]); - let struct2 = StructArray::from_fields(&[("data", bin_field2.into_array())]).unwrap(); - - let result = struct1 - .into_array() - .binary(struct2.into_array(), Operator::Eq) - .unwrap(); - let expected = BoolArray::from_iter([true, true, false]); - assert_arrays_eq!(result, expected, &mut ctx); - } - - /// Regression test: `scalar_cmp` must error when comparing scalars with incompatible - /// extension types (e.g., timestamps with different time units) rather than silently - /// returning a wrong result. - #[test] - fn scalar_cmp_incompatible_extension_types_errors() { - let ms_scalar = Scalar::extension::( - TimestampOptions { - unit: TimeUnit::Milliseconds, - tz: None, - }, - Scalar::from(1704067200000i64), - ); - let s_scalar = Scalar::extension::( - TimestampOptions { - unit: TimeUnit::Seconds, - tz: None, - }, - Scalar::from(1704067200i64), - ); - - // Ordering comparisons must error on incompatible types. - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gt).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lt).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gte).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lte).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Eq).is_err()); - assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::NotEq).is_err()); - } - - #[test] - fn test_empty_list() { - let ctx = &mut array_session().create_execution_ctx(); - let list = ListViewArray::new( - BoolArray::from_iter(Vec::::new()).into_array(), - buffer![0i32, 0i32, 0i32].into_array(), - buffer![0i32, 0i32, 0i32].into_array(), - Validity::AllValid, - ); - - let result = list - .clone() - .into_array() - .binary(list.into_array(), Operator::Eq) - .unwrap(); - assert!(result.execute_scalar(0, ctx).unwrap().is_valid()); - assert!(result.execute_scalar(1, ctx).unwrap().is_valid()); - assert!(result.execute_scalar(2, ctx).unwrap().is_valid()); - } -} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs b/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs new file mode 100644 index 00000000000..15bd486d341 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/boolean.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of boolean arrays using word-wise bit operations. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::dtype::Nullability; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +enum BoolOperand { + Array { bits: BitBuffer, validity: Validity }, + Constant { value: bool, validity: Validity }, +} + +impl BoolOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + let value = constant + .scalar() + .as_bool_opt() + .ok_or_else(|| vortex_err!("expected boolean scalar"))? + .value() + .ok_or_else(|| vortex_err!("null constant handled by execute_compare"))?; + return Ok(Self::Constant { + value, + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }); + } + + let array = array.clone().execute::(ctx)?; + let validity = array.validity()?; + Ok(Self::Array { + bits: array.into_bit_buffer(), + validity, + }) + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + } + } +} + +/// Compare two boolean arrays. +/// +/// Values compare as `false < true`; every operator reduces to at most two word-wise passes over +/// the value bit buffers. +pub(super) fn compare_bool( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = BoolOperand::try_new(lhs, ctx)?; + let rhs = BoolOperand::try_new(rhs, ctx)?; + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (lhs, rhs) { + (BoolOperand::Array { bits: l, .. }, BoolOperand::Array { bits: r, .. }) => { + compare_bits(l, r, op) + } + (BoolOperand::Array { bits, .. }, BoolOperand::Constant { value, .. }) => { + compare_bits_constant(bits, value, op) + } + (BoolOperand::Constant { value, .. }, BoolOperand::Array { bits, .. }) => { + compare_bits_constant(bits, value, op.swap()) + } + (BoolOperand::Constant { value: l, .. }, BoolOperand::Constant { value: r, .. }) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + let result = super::ordering_predicate(op)(l.cmp(&r)); + BitBuffer::full(result, len) + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn compare_bits(lhs: BitBuffer, rhs: BitBuffer, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => !(lhs ^ &rhs), + CompareOperator::NotEq => lhs ^ &rhs, + // a < b ⟺ !a & b + CompareOperator::Lt => rhs.into_bitand_not(&lhs), + // a <= b ⟺ !(a & !b) + CompareOperator::Lte => !lhs.into_bitand_not(&rhs), + // a > b ⟺ a & !b + CompareOperator::Gt => lhs.into_bitand_not(&rhs), + // a >= b ⟺ !(!a & b) + CompareOperator::Gte => !rhs.into_bitand_not(&lhs), + } +} + +/// Compare array bits against a non-null constant: `bits value`. +fn compare_bits_constant(bits: BitBuffer, value: bool, op: CompareOperator) -> BitBuffer { + let len = bits.len(); + match (op, value) { + (CompareOperator::Eq, true) + | (CompareOperator::NotEq, false) + | (CompareOperator::Gt, false) + | (CompareOperator::Gte, true) => bits, + (CompareOperator::Eq, false) + | (CompareOperator::NotEq, true) + | (CompareOperator::Lt, true) + | (CompareOperator::Lte, false) => !bits, + (CompareOperator::Lt, false) | (CompareOperator::Gt, true) => BitBuffer::new_unset(len), + (CompareOperator::Lte, true) | (CompareOperator::Gte, false) => BitBuffer::new_set(len), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs new file mode 100644 index 00000000000..70c56411d08 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/bytes.rs @@ -0,0 +1,400 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of UTF-8 and binary arrays over canonical [`VarBinViewArray`]s. +//! +//! Equality first compares the leading 8 bytes of each view (length plus 4-byte prefix), which +//! answers most lanes without touching the data buffers. Ordering compares the inline 4-byte +//! prefixes first and only dereferences the full value on a prefix tie. UTF-8 values compare by +//! their byte representation, which matches code-point order. + +use std::cmp::Ordering; + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::VarBinViewArray; +use crate::arrays::varbinview::BinaryView; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::binary::compare::ordering_predicate; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +enum BytesOperand { + Array { + values: VarBinViewArray, + validity: Validity, + }, + Constant { + value: Vec, + validity: Validity, + }, +} + +impl BytesOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + let value = constant_bytes(constant.scalar())?; + return Ok(Self::Constant { + value, + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }); + } + + let values = array.clone().execute::(ctx)?; + let validity = values.validity()?; + Ok(Self::Array { values, validity }) + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + } + } +} + +fn constant_bytes(scalar: &Scalar) -> VortexResult> { + let value = match scalar.dtype() { + DType::Utf8(_) => scalar + .as_utf8() + .value() + .map(|s| s.as_str().as_bytes().to_vec()), + DType::Binary(_) => scalar.as_binary().value().map(|b| b.to_vec()), + _ => vortex_bail!("expected utf8 or binary scalar, got {}", scalar.dtype()), + }; + value.ok_or_else(|| vortex_err!("null constant handled by execute_compare")) +} + +/// A resolved view over a canonical [`VarBinViewArray`]: the view structs plus borrowed slices of +/// every data buffer, supporting cheap per-lane byte access. +struct ViewsSide<'a> { + views: &'a [BinaryView], + buffers: Vec<&'a [u8]>, +} + +impl<'a> ViewsSide<'a> { + fn new(array: &'a VarBinViewArray) -> Self { + Self { + views: array.views(), + buffers: (0..array.data_buffers().len()) + .map(|idx| array.buffer(idx).as_slice()) + .collect(), + } + } + + fn len(&self) -> usize { + self.views.len() + } + + /// The view at `index` without a bounds check. + /// + /// # Safety + /// + /// `index` must be strictly less than `self.len()`. + #[inline] + unsafe fn view_unchecked(&self, index: usize) -> &'a BinaryView { + // SAFETY: caller guarantees index < self.views.len(). + unsafe { self.views.get_unchecked(index) } + } + + /// The full bytes of `view`, which must belong to this side. + #[inline] + fn view_bytes(&self, view: &'a BinaryView) -> &'a [u8] { + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &self.buffers[view.buffer_index as usize][view.as_range()] + } + } +} + +/// The leading 8 bytes of a view: the `u32` length plus the first 4 bytes of the value +/// (zero-padded for values shorter than 4 bytes). +#[inline] +#[expect(clippy::cast_possible_truncation, reason = "intentional bit slicing")] +fn view_head(view: &BinaryView) -> u64 { + view.as_u128() as u64 +} + +/// The first 4 value bytes of a view as a big-endian `u32` (zero-padded for values shorter than +/// 4 bytes), so that numeric comparison matches lexicographic byte order. +/// +/// Zero padding preserves lexicographic order: a padded position differs from a real byte only +/// when one value is a strict prefix of the other within the first 4 bytes, and the shorter +/// value orders first exactly as `0` orders before any later real byte. A padded tie falls +/// through to the tail or full byte comparison. +#[inline] +#[expect(clippy::cast_possible_truncation, reason = "intentional bit slicing")] +fn view_prefix(view: &BinaryView) -> u32 { + ((view.as_u128() >> 32) as u32).swap_bytes() +} + +/// Value bytes 4..12 of an *inlined* view as a big-endian `u64` (zero-padded past the value's +/// length). Only meaningful for inlined views: reference views store the buffer index and offset +/// in these bytes. +#[inline] +fn view_tail(view: &BinaryView) -> u64 { + ((view.as_u128() >> 64) as u64).swap_bytes() +} + +/// Compare `lhs_view` from `lhs` against `rhs_view` from `rhs` for equality. +#[inline] +fn view_eq( + lhs: &ViewsSide<'_>, + lhs_view: &BinaryView, + rhs: &ViewsSide<'_>, + rhs_view: &BinaryView, +) -> bool { + if view_head(lhs_view) != view_head(rhs_view) { + return false; + } + if lhs_view.is_inlined() { + // Lengths are equal and at most 12: the whole value lives in the view. + return lhs_view.as_u128() == rhs_view.as_u128(); + } + // Equal lengths above 12 and equal prefixes: compare the out-of-line suffixes. + lhs.view_bytes(lhs_view)[4..] == rhs.view_bytes(rhs_view)[4..] +} + +/// Compare `lhs_view` from `lhs` against `rhs_view` from `rhs`. +#[inline] +fn view_cmp( + lhs: &ViewsSide<'_>, + lhs_view: &BinaryView, + rhs: &ViewsSide<'_>, + rhs_view: &BinaryView, +) -> Ordering { + let lhs_prefix = view_prefix(lhs_view); + let rhs_prefix = view_prefix(rhs_view); + if lhs_prefix != rhs_prefix { + return lhs_prefix.cmp(&rhs_prefix); + } + if lhs_view.is_inlined() && rhs_view.is_inlined() { + // Both values live entirely in their views: compare the remaining 8 (zero-padded) + // value bytes, then lengths. A tie on padded windows means the shorter value is a + // prefix of the longer one. + let lhs_tail = view_tail(lhs_view); + let rhs_tail = view_tail(rhs_view); + if lhs_tail != rhs_tail { + return lhs_tail.cmp(&rhs_tail); + } + return lhs_view.len().cmp(&rhs_view.len()); + } + lhs.view_bytes(lhs_view).cmp(rhs.view_bytes(rhs_view)) +} + +/// Compare two UTF-8 or binary arrays. +pub(super) fn compare_bytes( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = BytesOperand::try_new(lhs, ctx)?; + let rhs = BytesOperand::try_new(rhs, ctx)?; + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (&lhs, &rhs) { + (BytesOperand::Array { values: l, .. }, BytesOperand::Array { values: r, .. }) => { + compare_views(&ViewsSide::new(l), &ViewsSide::new(r), op) + } + (BytesOperand::Array { values, .. }, BytesOperand::Constant { value, .. }) => { + compare_views_constant(&ViewsSide::new(values), value, op) + } + (BytesOperand::Constant { value, .. }, BytesOperand::Array { values, .. }) => { + compare_views_constant(&ViewsSide::new(values), value, op.swap()) + } + (BytesOperand::Constant { value: l, .. }, BytesOperand::Constant { value: r, .. }) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + BitBuffer::full(ordering_predicate(op)(l.as_slice().cmp(r.as_slice())), len) + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn compare_views(lhs: &ViewsSide<'_>, rhs: &ViewsSide<'_>, op: CompareOperator) -> BitBuffer { + let len = lhs.len(); + // The unchecked view accesses below index both sides with `i < len`, so this must hold even + // in release builds. + assert_eq!(len, rhs.len(), "compared views must have equal lengths"); + // Dispatch the operator outside the lane loop so each predicate inlines into its own loop; + // a shared `fn(Ordering) -> bool` pointer would cost an indirect call per lane. + match op { + CompareOperator::Eq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + unsafe { view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } + }), + CompareOperator::NotEq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + unsafe { !view_eq(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) } + }), + CompareOperator::Gt => collect_ordering_bits(lhs, rhs, Ordering::is_gt), + CompareOperator::Gte => collect_ordering_bits(lhs, rhs, Ordering::is_ge), + CompareOperator::Lt => collect_ordering_bits(lhs, rhs, Ordering::is_lt), + CompareOperator::Lte => collect_ordering_bits(lhs, rhs, Ordering::is_le), + } +} + +/// Bit-pack `predicate(view_cmp(lhs[i], rhs[i]))` over two equal-length view sides. +fn collect_ordering_bits( + lhs: &ViewsSide<'_>, + rhs: &ViewsSide<'_>, + predicate: impl Fn(Ordering) -> bool, +) -> BitBuffer { + let len = lhs.len(); + assert_eq!(len, rhs.len(), "compared views must have equal lengths"); + BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len() for both sides. + predicate(unsafe { view_cmp(lhs, lhs.view_unchecked(i), rhs, rhs.view_unchecked(i)) }) + }) +} + +fn compare_views_constant(lhs: &ViewsSide<'_>, constant: &[u8], op: CompareOperator) -> BitBuffer { + let len = lhs.len(); + // The same head/prefix/tail words a view stores, precomputed once for the constant. + let mut prefix_bytes = [0u8; 4]; + let prefix_len = constant.len().min(4); + prefix_bytes[..prefix_len].copy_from_slice(&constant[..prefix_len]); + let mut tail_bytes = [0u8; 8]; + if constant.len() > 4 { + let tail_len = (constant.len() - 4).min(8); + tail_bytes[..tail_len].copy_from_slice(&constant[4..4 + tail_len]); + } + let constant_head = + (constant.len() as u64) | (u64::from(u32::from_le_bytes(prefix_bytes)) << 32); + let constant_prefix = u32::from_be_bytes(prefix_bytes); + let constant_tail = u64::from_be_bytes(tail_bytes); + // The full 16-byte word an inlined view holding `constant` would carry; only meaningful when + // the constant is short enough to inline, and only reached in that case (a longer constant + // never head-matches an inlined view). + let constant_inlined = + u128::from(constant_head) | (u128::from(u64::from_le_bytes(tail_bytes)) << 64); + + match op { + CompareOperator::Eq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + constant_eq(lhs, view, constant, constant_head, constant_inlined) + }), + CompareOperator::NotEq => BitBuffer::collect_bool(len, |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + !constant_eq(lhs, view, constant, constant_head, constant_inlined) + }), + CompareOperator::Gt => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_gt, + ), + CompareOperator::Gte => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_ge, + ), + CompareOperator::Lt => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_lt, + ), + CompareOperator::Lte => collect_constant_ordering_bits( + lhs, + constant, + constant_prefix, + constant_tail, + Ordering::is_le, + ), + } +} + +/// Bit-pack `predicate(constant_cmp(lhs[i], constant))` over one view side. +fn collect_constant_ordering_bits( + lhs: &ViewsSide<'_>, + constant: &[u8], + constant_prefix: u32, + constant_tail: u64, + predicate: impl Fn(Ordering) -> bool, +) -> BitBuffer { + BitBuffer::collect_bool(lhs.len(), |i| { + // SAFETY: `collect_bool` yields i < len == views.len(). + let view = unsafe { lhs.view_unchecked(i) }; + predicate(constant_cmp( + lhs, + view, + constant, + constant_prefix, + constant_tail, + )) + }) +} + +/// Compare a view against a constant for equality using the constant's precomputed head and +/// inline words. +#[inline] +fn constant_eq( + side: &ViewsSide<'_>, + view: &BinaryView, + constant: &[u8], + constant_head: u64, + constant_inlined: u128, +) -> bool { + if view_head(view) != constant_head { + return false; + } + if view.is_inlined() { + // An equal head implies equal lengths, so the constant is also at most 12 bytes and + // `constant_inlined` holds its exact inlined representation. + return view.as_u128() == constant_inlined; + } + side.view_bytes(view)[4..] == constant[4..] +} + +/// Compare a view against a constant using the constant's precomputed prefix and tail words. +#[inline] +fn constant_cmp( + side: &ViewsSide<'_>, + view: &BinaryView, + constant: &[u8], + constant_prefix: u32, + constant_tail: u64, +) -> Ordering { + let prefix = view_prefix(view); + if prefix != constant_prefix { + return prefix.cmp(&constant_prefix); + } + if view.is_inlined() { + // The view's whole value lives in its first 12 (zero-padded) bytes, and the constant's + // first 12 (zero-padded) bytes are precomputed: compare tails, then lengths. A padded + // tie means one value is a prefix of the other within the first 12 bytes. + let tail = view_tail(view); + if tail != constant_tail { + return tail.cmp(&constant_tail); + } + return (view.len() as usize).cmp(&constant.len()); + } + side.view_bytes(view).cmp(constant) +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs new file mode 100644 index 00000000000..0825d0674af --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/decimal.rs @@ -0,0 +1,188 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of decimal arrays. +//! +//! Both operands share a logical [`DecimalDType`] (equal precision and scale), so comparing the +//! unscaled integer values is sufficient. The physical storage width may differ per operand; when +//! it does, both sides are widened once to the larger storage type before the lane loop. +//! +//! [`DecimalDType`]: crate::dtype::DecimalDType + +use vortex_buffer::BitBuffer; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +use crate::arrays::DecimalArray; +use crate::dtype::NativeDecimalType; +use crate::dtype::Nullability; +use crate::dtype::i256; +use crate::match_each_decimal_value_type; +use crate::scalar::DecimalValue; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +enum DecimalOperand { + Array { + values: DecimalArray, + validity: Validity, + }, + Constant { + value: DecimalValue, + validity: Validity, + }, +} + +impl DecimalOperand { + fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if let Some(constant) = array.as_opt::() { + let value = constant + .scalar() + .as_decimal() + .decimal_value() + .ok_or_else(|| vortex_err!("null constant handled by execute_compare"))?; + return Ok(Self::Constant { + value, + validity: if constant.scalar().dtype().is_nullable() { + Validity::AllValid + } else { + Validity::NonNullable + }, + }); + } + + let values = array.clone().execute::(ctx)?; + let validity = values.validity()?; + Ok(Self::Array { values, validity }) + } + + fn validity(&self) -> Validity { + match self { + Self::Array { validity, .. } | Self::Constant { validity, .. } => validity.clone(), + } + } +} + +/// Compare two decimal arrays with the same logical decimal dtype. +pub(super) fn compare_decimal( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = DecimalOperand::try_new(lhs, ctx)?; + let rhs = DecimalOperand::try_new(rhs, ctx)?; + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (lhs, rhs) { + (DecimalOperand::Array { values: l, .. }, DecimalOperand::Array { values: r, .. }) => { + compare_decimal_values(&l, &r, op) + } + (DecimalOperand::Array { values, .. }, DecimalOperand::Constant { value, .. }) => { + compare_decimal_constant(&values, value, op) + } + (DecimalOperand::Constant { value, .. }, DecimalOperand::Array { values, .. }) => { + compare_decimal_constant(&values, value, op.swap()) + } + (DecimalOperand::Constant { value: l, .. }, DecimalOperand::Constant { value: r, .. }) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + let ordering = l.as_i256().cmp(&r.as_i256()); + BitBuffer::full(super::ordering_predicate(op)(ordering), len) + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +fn compare_decimal_values( + lhs: &DecimalArray, + rhs: &DecimalArray, + op: CompareOperator, +) -> BitBuffer { + let common = lhs.values_type().max(rhs.values_type()); + match_each_decimal_value_type!(common, |W| { + let lhs = widened_buffer::(lhs); + let rhs = widened_buffer::(rhs); + compare_slices::(&lhs, &rhs, op) + }) +} + +/// Return the array's unscaled values widened to `W`, which must be at least as wide as the +/// array's storage type. +pub(super) fn widened_buffer(array: &DecimalArray) -> Buffer { + if array.values_type() == W::DECIMAL_TYPE { + return array.buffer::(); + } + match_each_decimal_value_type!(array.values_type(), |T| { + array + .buffer::() + .iter() + .map(|v| W::from(*v).vortex_expect("widening decimal cast must succeed")) + .collect() + }) +} + +fn compare_decimal_constant( + array: &DecimalArray, + constant: DecimalValue, + op: CompareOperator, +) -> BitBuffer { + match_each_decimal_value_type!(array.values_type(), |T| { + match constant.cast::() { + Some(value) => compare_slice_constant::(&array.buffer::(), value, op), + None => { + // The constant does not fit the array's storage type, so it is either greater + // than every possible array value or less than every possible array value; the + // sign tells us which. + let constant_greater = constant.as_i256() > i256::ZERO; + let result = match op { + CompareOperator::Eq => false, + CompareOperator::NotEq => true, + // array constant + CompareOperator::Lt | CompareOperator::Lte => constant_greater, + CompareOperator::Gt | CompareOperator::Gte => !constant_greater, + }; + BitBuffer::full(result, array.len()) + } + } + }) +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a == b), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| a != b), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, |a: T, b: T| a > b), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, |a: T, b: T| a >= b), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, |a: T, b: T| a < b), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, |a: T, b: T| a <= b), + } +} + +fn compare_slice_constant( + values: &[T], + constant: T, + op: CompareOperator, +) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(values, |a: T| a == constant), + CompareOperator::NotEq => collect_bits(values, |a: T| a != constant), + CompareOperator::Gt => collect_bits(values, |a: T| a > constant), + CompareOperator::Gte => collect_bits(values, |a: T| a >= constant), + CompareOperator::Lt => collect_bits(values, |a: T| a < constant), + CompareOperator::Lte => collect_bits(values, |a: T| a <= constant), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs new file mode 100644 index 00000000000..c7d6c79c7ae --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/mod.rs @@ -0,0 +1,320 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison kernels. +//! +//! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every +//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from +//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise +//! comparator for nested types. There is no Arrow fallback. +//! +//! Floating point values compare with Vortex's total ordering (`NaN` is the largest value, +//! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics. + +use std::cmp::Ordering; + +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_compute::lane_kernels::LaneZip; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::Canonical; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::ScalarFn; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::scalar_fn::ExactScalarFn; +use crate::arrays::scalar_fn::ScalarFnArrayExt; +use crate::arrays::scalar_fn::ScalarFnArrayView; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::kernel::ExecuteParentKernel; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::validity::Validity; + +mod boolean; +mod bytes; +mod decimal; +mod nested; +mod primitive; +#[cfg(test)] +mod tests; + +/// Trait for encoding-specific comparison kernels that operate in encoded space. +/// +/// Implementations can compare an encoded array against another array (typically a constant) +/// without first decompressing. The adaptor normalizes operand order so `array` is always +/// the left-hand side, swapping the operator when necessary. +pub trait CompareKernel: VTable { + fn compare( + lhs: ArrayView<'_, Self>, + rhs: &ArrayRef, + operator: CompareOperator, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; +} + +/// Adaptor that bridges [`CompareKernel`] implementations to [`ExecuteParentKernel`]. +/// +/// When a `ScalarFnArray(Binary, cmp_op)` wraps a child that implements `CompareKernel`, +/// this adaptor extracts the comparison operator and other operand, normalizes operand order +/// (swapping the operator if the encoded array is on the RHS), and delegates to the kernel. +#[derive(Default, Debug)] +pub struct CompareExecuteAdaptor(pub V); + +impl ExecuteParentKernel for CompareExecuteAdaptor +where + V: CompareKernel, +{ + type Parent = ExactScalarFn; + + fn execute_parent( + &self, + array: ArrayView<'_, V>, + parent: ScalarFnArrayView<'_, Binary>, + child_idx: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + // Only handle comparison operators + let Ok(cmp_op) = CompareOperator::try_from(*parent.options) else { + return Ok(None); + }; + + // Get the ScalarFnArray to access children + let Some(scalar_fn_array) = parent.as_opt::() else { + return Ok(None); + }; + // Normalize so `array` is always LHS, swapping the operator if needed + // TODO(joe): should be go this here or in the Rule/Kernel + let (cmp_op, other) = match child_idx { + 0 => (cmp_op, scalar_fn_array.get_child(1)), + 1 => (cmp_op.swap(), scalar_fn_array.get_child(0)), + _ => return Ok(None), + }; + + let len = array.len(); + let nullable = array.dtype().is_nullable() || other.dtype().is_nullable(); + + // Empty array → empty bool result + if len == 0 { + return Ok(Some( + Canonical::empty(&DType::Bool(nullable.into())).into_array(), + )); + } + + // Null constant on either side → all-null bool result + if other.as_constant().is_some_and(|s| s.is_null()) { + return Ok(Some( + ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), len).into_array(), + )); + } + + V::compare(array, other, cmp_op, ctx) + } +} + +/// Execute a compare operation between two arrays. +/// +/// This is the entry point for compare operations from the binary expression. +/// Handles empty, constant-null, and constant-constant directly, otherwise dispatches to a +/// native per-dtype kernel. +pub(crate) fn execute_compare( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); + + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + if lhs.is_empty() { + return Ok(Canonical::empty(&DType::Bool(nullable.into())).into_array()); + } + + 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(DType::Bool(nullable.into())), lhs.len()).into_array(), + ); + } + + // Constant-constant fast path + if let (Some(lhs_const), Some(rhs_const)) = (lhs.as_opt::(), rhs.as_opt::()) + { + let result = scalar_cmp(lhs_const.scalar(), rhs_const.scalar(), op)?; + return Ok(ConstantArray::new(result, lhs.len()).into_array()); + } + + compare_arrays(lhs, rhs, op, ctx) +} + +/// Dispatch a comparison to the native kernel for the operands' logical dtype. +fn compare_arrays( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable()); + + // Extension arrays compare through their storage. When both sides are extensions they must + // agree on the full extension dtype (a timestamp in milliseconds must not compare its raw + // storage against one in seconds); a lone extension side may compare against its raw storage. + if lhs.dtype().is_extension() || rhs.dtype().is_extension() { + if lhs.dtype().is_extension() + && rhs.dtype().is_extension() + && !lhs.dtype().eq_ignore_nullability(rhs.dtype()) + { + vortex_bail!( + "Cannot compare extension dtypes {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + let lhs = extension_storage(lhs, ctx)?; + let rhs = extension_storage(rhs, ctx)?; + return compare_arrays(&lhs, &rhs, op, ctx); + } + + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + vortex_bail!( + "Cannot compare different DTypes {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + + match lhs.dtype() { + // Every value of the null dtype is null, so every comparison result is null. + DType::Null => Ok(ConstantArray::new( + Scalar::null(DType::Bool(Nullability::Nullable)), + lhs.len(), + ) + .into_array()), + DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx), + DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx), + DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx), + DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx), + DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) => { + nested::compare_nested(lhs, rhs, op, nullability, ctx) + } + DType::Union(_) | DType::Variant(_) | DType::Extension(_) => { + vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) + } + } +} + +/// Replace an extension array (or extension constant) with its storage. Non-extension arrays are +/// returned unchanged. +fn extension_storage(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + if !array.dtype().is_extension() { + return Ok(array.clone()); + } + if let Some(constant) = array.as_opt::() { + return Ok(ConstantArray::new( + constant.scalar().as_extension().to_storage_scalar(), + constant.len(), + ) + .into_array()); + } + let ext = array.clone().execute::(ctx)?; + Ok(ext.storage_array().clone()) +} + +pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> VortexResult { + if lhs.is_null() | rhs.is_null() { + return Ok(Scalar::null(DType::Bool(Nullability::Nullable))); + } + + let nullability = lhs.dtype().nullability() | rhs.dtype().nullability(); + + // We use `partial_cmp` to ensure we do not lose a type mismatch error. + let ordering = lhs.partial_cmp(rhs).ok_or_else(|| { + vortex_err!( + "Cannot compare scalars with incompatible types: {} and {}", + lhs.dtype(), + rhs.dtype() + ) + })?; + + let b = match operator { + CompareOperator::Eq => ordering.is_eq(), + CompareOperator::NotEq => ordering.is_ne(), + CompareOperator::Gt => ordering.is_gt(), + CompareOperator::Gte => ordering.is_ge(), + CompareOperator::Lt => ordering.is_lt(), + CompareOperator::Lte => ordering.is_le(), + }; + + Ok(Scalar::bool(b, nullability)) +} + +/// Returns the predicate `Ordering -> bool` for a comparison operator. +#[inline] +pub(super) fn ordering_predicate(op: CompareOperator) -> fn(Ordering) -> bool { + match op { + CompareOperator::Eq => Ordering::is_eq, + CompareOperator::NotEq => Ordering::is_ne, + CompareOperator::Gt => Ordering::is_gt, + CompareOperator::Gte => Ordering::is_ge, + CompareOperator::Lt => Ordering::is_lt, + CompareOperator::Lte => Ordering::is_le, + } +} + +/// Freeze `len` bits packed into `words` (LSB-first, 64 lanes per word) into a [`BitBuffer`]. +pub(super) fn bit_buffer_from_words(words: BufferMut, len: usize) -> BitBuffer { + debug_assert!(words.len() * 64 >= len); + let mut bytes = words.into_byte_buffer(); + bytes.truncate(len.div_ceil(8)); + BitBuffer::new(bytes.freeze(), len) +} + +/// Bit-pack the predicate `f(lhs[i], rhs[i])` over two equal-length slices into a [`BitBuffer`]. +pub(super) fn collect_zip_bits( + lhs: &[T], + rhs: &[T], + mut f: impl FnMut(T, T) -> bool, +) -> BitBuffer { + let len = lhs.len(); + let mut words = BufferMut::::zeroed(len.div_ceil(64)); + LaneZip::new(lhs, rhs).map_bits_into(words.as_mut_slice(), |(a, b)| f(a, b)); + bit_buffer_from_words(words, len) +} + +/// Bit-pack the predicate `f(values[i])` over a slice into a [`BitBuffer`]. +pub(super) fn collect_bits(values: &[T], f: impl FnMut(T) -> bool) -> BitBuffer { + let len = values.len(); + let mut words = BufferMut::::zeroed(len.div_ceil(64)); + values.map_bits_into(words.as_mut_slice(), f); + bit_buffer_from_words(words, len) +} + +/// Combine operand validities into the validity of a comparison result with the given result +/// nullability. +pub(super) fn compare_validity( + lhs: Validity, + rhs: Validity, + nullability: Nullability, +) -> VortexResult { + Ok(lhs.and(rhs)?.union_nullability(nullability)) +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs new file mode 100644 index 00000000000..cbc38d6b23e --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/nested.rs @@ -0,0 +1,256 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-wise comparison of nested (struct, list, fixed-size list) arrays. +//! +//! Nested comparisons canonicalize both operands recursively and build a tree of row +//! comparators, one per nested level. The ordering semantics match [`Scalar`] comparison: +//! structs compare field-by-field in declaration order, lists compare element-by-element and +//! then by length, and a null value (at any nesting depth) orders before every non-null value. +//! Only top-level nulls make the comparison result null. +//! +//! [`Scalar`]: crate::scalar::Scalar + +use std::cmp::Ordering; + +use num_traits::AsPrimitive; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::RecursiveCanonical; +use crate::arrays::BoolArray; +use crate::arrays::DecimalArray; +use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinViewArray; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::fixed_size_list::FixedSizeListArrayExt; +use crate::arrays::listview::ListViewArrayExt; +use crate::arrays::struct_::StructArrayExt; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::match_each_integer_ptype; +use crate::match_each_native_ptype; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// A row comparator: compares row `i` of the left operand against row `j` of the right operand. +type RowComparator = Box Ordering>; + +/// Compare two nested arrays row by row. +pub(super) fn compare_nested( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = lhs + .clone() + .execute::(ctx)? + .0 + .into_array(); + let rhs = rhs + .clone() + .execute::(ctx)? + .0 + .into_array(); + + let validity = compare_validity(lhs.validity()?, rhs.validity()?, nullability)?; + let comparator = build_comparator(&lhs, &rhs, ctx)?; + // Dispatch the operator outside the row loop so the predicate inlines into each loop; the + // comparator call itself stays virtual. + let bits = match op { + CompareOperator::Eq => collect_ordering_bits(len, &comparator, Ordering::is_eq), + CompareOperator::NotEq => collect_ordering_bits(len, &comparator, Ordering::is_ne), + CompareOperator::Gt => collect_ordering_bits(len, &comparator, Ordering::is_gt), + CompareOperator::Gte => collect_ordering_bits(len, &comparator, Ordering::is_ge), + CompareOperator::Lt => collect_ordering_bits(len, &comparator, Ordering::is_lt), + CompareOperator::Lte => collect_ordering_bits(len, &comparator, Ordering::is_le), + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +/// Bit-pack `predicate(comparator(i, i))` over `len` rows. +fn collect_ordering_bits( + len: usize, + comparator: &RowComparator, + predicate: impl Fn(Ordering) -> bool, +) -> BitBuffer { + BitBuffer::collect_bool(len, |i| predicate(comparator(i, i))) +} + +/// The validity mask of a recursively canonical array. +fn validity_mask(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + array.validity()?.execute_mask(array.len(), ctx) +} + +/// Build a row comparator over two recursively canonical arrays of the same logical dtype +/// (ignoring nullability). Null values order before all non-null values at every level. +fn build_comparator( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let lhs_mask = validity_mask(lhs, ctx)?; + let rhs_mask = validity_mask(rhs, ctx)?; + let values = build_values_comparator(lhs, rhs, ctx)?; + + if lhs_mask.all_true() && rhs_mask.all_true() { + return Ok(values); + } + + Ok(Box::new(move |i, j| { + match (lhs_mask.value(i), rhs_mask.value(j)) { + (true, true) => values(i, j), + (true, false) => Ordering::Greater, + (false, true) => Ordering::Less, + (false, false) => Ordering::Equal, + } + })) +} + +/// Build a comparator over the non-null values of two recursively canonical arrays. +fn build_values_comparator( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) { + vortex_bail!( + "Cannot compare different DTypes {} and {}", + lhs.dtype(), + rhs.dtype() + ); + } + + Ok(match lhs.dtype() { + // Null values compare through the validity wrapper; the value comparator is trivial. + DType::Null => Box::new(|_, _| Ordering::Equal), + DType::Bool(_) => { + let lhs = lhs.clone().execute::(ctx)?.into_bit_buffer(); + let rhs = rhs.clone().execute::(ctx)?.into_bit_buffer(); + Box::new(move |i, j| lhs.value(i).cmp(&rhs.value(j))) + } + DType::Primitive(ptype, _) => { + match_each_native_ptype!(*ptype, |T| { primitive_comparator::(lhs, rhs, ctx)? }) + } + DType::Decimal(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let common = lhs.values_type().max(rhs.values_type()); + crate::match_each_decimal_value_type!(common, |W| { + let lhs = super::decimal::widened_buffer::(&lhs); + let rhs = super::decimal::widened_buffer::(&rhs); + Box::new(move |i: usize, j: usize| lhs[i].cmp(&rhs[j])) as RowComparator + }) + } + DType::Utf8(_) | DType::Binary(_) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + Box::new(move |i, j| view_bytes(&lhs, i).cmp(view_bytes(&rhs, j))) + } + DType::Struct(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let fields = lhs + .iter_unmasked_fields() + .zip(rhs.iter_unmasked_fields()) + .map(|(lhs_field, rhs_field)| build_comparator(lhs_field, rhs_field, ctx)) + .collect::>>()?; + Box::new(move |i, j| { + fields + .iter() + .map(|cmp| cmp(i, j)) + .find(|ordering| ordering.is_ne()) + .unwrap_or(Ordering::Equal) + }) + } + DType::List(..) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let elements = build_comparator(lhs.elements(), rhs.elements(), ctx)?; + // Materialize offsets and sizes once instead of paying `offset_at`/`size_at`'s + // per-call encoding dispatch inside the row loop. + let lhs_offsets = usize_values(lhs.offsets(), ctx)?; + let lhs_sizes = usize_values(lhs.sizes(), ctx)?; + let rhs_offsets = usize_values(rhs.offsets(), ctx)?; + let rhs_sizes = usize_values(rhs.sizes(), ctx)?; + Box::new(move |i, j| { + let lhs_len = lhs_sizes[i]; + let rhs_len = rhs_sizes[j]; + (0..lhs_len.min(rhs_len)) + .map(|el| elements(lhs_offsets[i] + el, rhs_offsets[j] + el)) + .find(|ordering| ordering.is_ne()) + .unwrap_or_else(|| lhs_len.cmp(&rhs_len)) + }) + } + DType::FixedSizeList(_, size, _) => { + let size = *size as usize; + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + let elements = build_comparator(lhs.elements(), rhs.elements(), ctx)?; + Box::new(move |i, j| { + (0..size) + .map(|el| elements(i * size + el, j * size + el)) + .find(|ordering| ordering.is_ne()) + .unwrap_or(Ordering::Equal) + }) + } + DType::Extension(_) => { + let lhs = lhs.clone().execute::(ctx)?; + let rhs = rhs.clone().execute::(ctx)?; + build_comparator(lhs.storage_array(), rhs.storage_array(), ctx)? + } + DType::Union(_) | DType::Variant(_) => { + vortex_bail!("compare is not supported for dtype {}", lhs.dtype()) + } + }) +} + +fn primitive_comparator( + lhs: &ArrayRef, + rhs: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let lhs = lhs + .clone() + .execute::(ctx)? + .into_buffer::(); + let rhs = rhs + .clone() + .execute::(ctx)? + .into_buffer::(); + Ok(Box::new(move |i, j| lhs[i].total_compare(rhs[j]))) +} + +/// Materialize an integer array as `usize` values for `O(1)` row access. +fn usize_values(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + let primitive = array.clone().execute::(ctx)?; + match_each_integer_ptype!(primitive.ptype(), |P| { + Ok(primitive.as_slice::

().iter().map(|v| v.as_()).collect()) + }) +} + +/// The bytes of row `index`, resolved directly from the view without cloning buffer handles the +/// way `VarBinViewArray::bytes_at` does. +fn view_bytes(array: &VarBinViewArray, index: usize) -> &[u8] { + let view = &array.views()[index]; + if view.is_inlined() { + view.as_inlined().value() + } else { + let view = view.as_view(); + &array.buffer(view.buffer_index as usize).as_slice()[view.as_range()] + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs new file mode 100644 index 00000000000..3bfb11a266e --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/primitive.rs @@ -0,0 +1,130 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native comparison of primitive arrays via bit-packing lane kernels. + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::PrimitiveOperand; +use crate::scalar_fn::fns::binary::compare::collect_bits; +use crate::scalar_fn::fns::binary::compare::collect_zip_bits; +use crate::scalar_fn::fns::binary::compare::compare_validity; +use crate::scalar_fn::fns::operators::CompareOperator; + +/// Compare two primitive arrays of the same [`PType`]. +/// +/// Floats compare with Vortex's total ordering: `NaN` is the largest value, `-0.0 < +0.0`, and +/// equality is bitwise. +pub(super) fn compare_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let ptype = PType::try_from(lhs.dtype())?; + match_each_native_ptype!(ptype, |T| { + compare_primitive_typed::(lhs, rhs, op, nullability, ctx) + }) +} + +fn compare_primitive_typed( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: CompareOperator, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = lhs.len(); + let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; + let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; + if lhs.len() != rhs.len() { + vortex_bail!( + "compare operator requires equal lengths, got {} and {}", + lhs.len(), + rhs.len() + ); + } + + let validity = compare_validity(lhs.validity(), rhs.validity(), nullability)?; + + let bits = match (&lhs, &rhs) { + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slices(lhs, rhs, op), + ( + PrimitiveOperand::Array { values: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => compare_slice_constant(lhs, *rhs, op), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Array { values: rhs, .. }, + ) => compare_slice_constant(rhs, *lhs, op.swap()), + ( + PrimitiveOperand::Constant { value: lhs, .. }, + PrimitiveOperand::Constant { value: rhs, .. }, + ) => { + // Unreachable through `execute_compare` (constant-constant is folded there), but + // cheap to answer anyway. + BitBuffer::full(apply_op(*lhs, *rhs, op), len) + } + (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => { + return Ok( + ConstantArray::new(Scalar::null(DType::Bool(Nullability::Nullable)), len) + .into_array(), + ); + } + }; + + Ok(BoolArray::try_new(bits, validity)?.into_array()) +} + +#[inline(always)] +fn apply_op(lhs: T, rhs: T, op: CompareOperator) -> bool { + match op { + CompareOperator::Eq => lhs.is_eq(rhs), + CompareOperator::NotEq => !lhs.is_eq(rhs), + CompareOperator::Gt => lhs.is_gt(rhs), + CompareOperator::Gte => lhs.is_ge(rhs), + CompareOperator::Lt => lhs.is_lt(rhs), + CompareOperator::Lte => lhs.is_le(rhs), + } +} + +fn compare_slices(lhs: &[T], rhs: &[T], op: CompareOperator) -> BitBuffer { + // Dispatch the operator outside the lane loop so each instantiation vectorizes a single + // branch-free predicate. + match op { + CompareOperator::Eq => collect_zip_bits(lhs, rhs, |a: T, b: T| a.is_eq(b)), + CompareOperator::NotEq => collect_zip_bits(lhs, rhs, |a: T, b: T| !a.is_eq(b)), + CompareOperator::Gt => collect_zip_bits(lhs, rhs, T::is_gt), + CompareOperator::Gte => collect_zip_bits(lhs, rhs, T::is_ge), + CompareOperator::Lt => collect_zip_bits(lhs, rhs, T::is_lt), + CompareOperator::Lte => collect_zip_bits(lhs, rhs, T::is_le), + } +} + +fn compare_slice_constant(lhs: &[T], rhs: T, op: CompareOperator) -> BitBuffer { + match op { + CompareOperator::Eq => collect_bits(lhs, |a: T| a.is_eq(rhs)), + CompareOperator::NotEq => collect_bits(lhs, |a: T| !a.is_eq(rhs)), + CompareOperator::Gt => collect_bits(lhs, |a: T| a.is_gt(rhs)), + CompareOperator::Gte => collect_bits(lhs, |a: T| a.is_ge(rhs)), + CompareOperator::Lt => collect_bits(lhs, |a: T| a.is_lt(rhs)), + CompareOperator::Lte => collect_bits(lhs, |a: T| a.is_le(rhs)), + } +} diff --git a/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs new file mode 100644 index 00000000000..c2af83561a4 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/compare/tests.rs @@ -0,0 +1,673 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use rstest::rstest; +use vortex_buffer::BitBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexExpect; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::ExtensionArray; +use crate::arrays::FixedSizeListArray; +use crate::arrays::ListArray; +use crate::arrays::ListViewArray; +use crate::arrays::PrimitiveArray; +use crate::arrays::StructArray; +use crate::arrays::VarBinArray; +use crate::arrays::VarBinViewArray; +use crate::assert_arrays_eq; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::DecimalDType; +use crate::dtype::FieldName; +use crate::dtype::FieldNames; +use crate::dtype::Nullability; +use crate::dtype::PType; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; +use crate::extension::datetime::TimestampOptions; +use crate::scalar::DecimalValue; +use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::scalar_cmp; +use crate::scalar_fn::fns::operators::CompareOperator; +use crate::scalar_fn::fns::operators::Operator; +use crate::test_harness::to_int_indices; +use crate::validity::Validity; + +#[test] +fn test_bool_basic_comparisons() { + let ctx = &mut array_session().create_execution_ctx(); + let arr = BoolArray::new( + BitBuffer::from_iter([true, true, false, true, false]), + Validity::from_iter([false, true, true, true, true]), + ); + + let matches = arr + .clone() + .into_array() + .binary(arr.clone().into_array(), Operator::Eq) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [1u64, 2, 3, 4]); + + let matches = arr + .clone() + .into_array() + .binary(arr.clone().into_array(), Operator::NotEq) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + let empty: [u64; 0] = []; + assert_eq!(to_int_indices(matches, ctx).unwrap(), empty); + + let other = BoolArray::new( + BitBuffer::from_iter([false, false, false, true, true]), + Validity::from_iter([false, true, true, true, true]), + ); + + let matches = arr + .clone() + .into_array() + .binary(other.clone().into_array(), Operator::Lte) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); + + let matches = arr + .clone() + .into_array() + .binary(other.clone().into_array(), Operator::Lt) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); + + let matches = other + .clone() + .into_array() + .binary(arr.clone().into_array(), Operator::Gte) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [2u64, 3, 4]); + + let matches = other + .into_array() + .binary(arr.into_array(), Operator::Gt) + .unwrap() + .execute::(ctx) + .vortex_expect("must be a bool array"); + assert_eq!(to_int_indices(matches, ctx).unwrap(), [4u64]); +} + +#[test] +fn constant_compare() { + let left = ConstantArray::new(Scalar::from(2u32), 10); + let right = ConstantArray::new(Scalar::from(10u32), 10); + + let result = left + .into_array() + .binary(right.into_array(), Operator::Gt) + .unwrap(); + assert_eq!(result.len(), 10); + let scalar = result + .execute_scalar(0, &mut array_session().create_execution_ctx()) + .unwrap(); + assert_eq!(scalar.as_bool().value(), Some(false)); +} + +#[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 compare_different_encodings(#[case] left: ArrayRef, #[case] right: ArrayRef) { + let mut ctx = array_session().create_execution_ctx(); + let res = left.binary(right, Operator::Eq).unwrap(); + let expected = BoolArray::from_iter([true, true]); + assert_arrays_eq!(res, expected, &mut ctx); +} + +#[test] +fn test_list_array_comparison() { + let mut ctx = array_session().create_execution_ctx(); + let values1 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); + let offsets1 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); + let list1 = ListArray::try_new( + values1.into_array(), + offsets1.into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let values2 = PrimitiveArray::from_iter([1i32, 2, 3, 4, 7, 8]); + let offsets2 = PrimitiveArray::from_iter([0i32, 2, 4, 6]); + let list2 = ListArray::try_new( + values2.into_array(), + offsets2.into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let result = list1 + .clone() + .into_array() + .binary(list2.clone().into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, false]); + assert_arrays_eq!(result, expected, &mut ctx); + + let result = list1 + .clone() + .into_array() + .binary(list2.clone().into_array(), Operator::NotEq) + .unwrap(); + let expected = BoolArray::from_iter([false, false, true]); + assert_arrays_eq!(result, expected, &mut ctx); + + let result = list1 + .into_array() + .binary(list2.into_array(), Operator::Lt) + .unwrap(); + let expected = BoolArray::from_iter([false, false, true]); + assert_arrays_eq!(result, expected, &mut ctx); +} + +#[test] +fn test_list_array_constant_comparison() { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]); + let offsets = PrimitiveArray::from_iter([0i32, 2, 4, 6]); + let list = ListArray::try_new( + values.into_array(), + offsets.into_array(), + Validity::NonNullable, + ) + .unwrap(); + + let list_scalar = Scalar::list( + Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)), + vec![3i32.into(), 4i32.into()], + Nullability::NonNullable, + ); + let constant = ConstantArray::new(list_scalar, 3); + + let result = list + .into_array() + .binary(constant.into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([false, true, false]); + assert_arrays_eq!(result, expected, &mut ctx); +} + +#[test] +fn test_struct_array_comparison() { + let mut ctx = array_session().create_execution_ctx(); + let bool_field1 = BoolArray::from_iter([Some(true), Some(false), Some(true)]); + let int_field1 = PrimitiveArray::from_iter([1i32, 2, 3]); + + let bool_field2 = BoolArray::from_iter([Some(true), Some(false), Some(false)]); + let int_field2 = PrimitiveArray::from_iter([1i32, 2, 4]); + + let struct1 = StructArray::from_fields(&[ + ("bool_col", bool_field1.into_array()), + ("int_col", int_field1.into_array()), + ]) + .unwrap(); + + let struct2 = StructArray::from_fields(&[ + ("bool_col", bool_field2.into_array()), + ("int_col", int_field2.into_array()), + ]) + .unwrap(); + + let result = struct1 + .clone() + .into_array() + .binary(struct2.clone().into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, false]); + assert_arrays_eq!(result, expected, &mut ctx); + + let result = struct1 + .into_array() + .binary(struct2.into_array(), Operator::Gt) + .unwrap(); + let expected = BoolArray::from_iter([false, false, true]); + assert_arrays_eq!(result, expected, &mut ctx); +} + +#[test] +fn test_empty_struct_compare() { + let mut ctx = array_session().create_execution_ctx(); + let empty1 = StructArray::try_new( + FieldNames::from(Vec::::new()), + Vec::new(), + 5, + Validity::NonNullable, + ) + .unwrap(); + + let empty2 = StructArray::try_new( + FieldNames::from(Vec::::new()), + Vec::new(), + 5, + Validity::NonNullable, + ) + .unwrap(); + + let result = empty1 + .into_array() + .binary(empty2.into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, true, true, true]); + assert_arrays_eq!(result, expected, &mut ctx); +} + +/// Regression test: comparing struct arrays where the same logical field is backed by +/// different Vortex encodings (VarBinArray vs VarBinViewArray) must not panic. +#[test] +fn struct_compare_mixed_binary_encodings() { + let mut ctx = array_session().create_execution_ctx(); + // LHS: struct with a VarBinArray (offset-based) binary field + let bin_field1 = VarBinArray::from(vec![ + "apple".as_bytes(), + "banana".as_bytes(), + "cherry".as_bytes(), + ]); + let struct1 = StructArray::from_fields(&[("data", bin_field1.into_array())]).unwrap(); + + // RHS: struct with a VarBinViewArray (view-based) binary field — same logical DType + let bin_field2 = VarBinViewArray::from_iter_bin([ + "apple".as_bytes(), + "banana".as_bytes(), + "durian".as_bytes(), + ]); + let struct2 = StructArray::from_fields(&[("data", bin_field2.into_array())]).unwrap(); + + let result = struct1 + .into_array() + .binary(struct2.into_array(), Operator::Eq) + .unwrap(); + let expected = BoolArray::from_iter([true, true, false]); + assert_arrays_eq!(result, expected, &mut ctx); +} + +/// Regression test: `scalar_cmp` must error when comparing scalars with incompatible +/// extension types (e.g., timestamps with different time units) rather than silently +/// returning a wrong result. +#[test] +fn scalar_cmp_incompatible_extension_types_errors() { + let ms_scalar = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Milliseconds, + tz: None, + }, + Scalar::from(1704067200000i64), + ); + let s_scalar = Scalar::extension::( + TimestampOptions { + unit: TimeUnit::Seconds, + tz: None, + }, + Scalar::from(1704067200i64), + ); + + // Ordering comparisons must error on incompatible types. + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gt).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lt).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Gte).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Lte).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::Eq).is_err()); + assert!(scalar_cmp(&ms_scalar, &s_scalar, CompareOperator::NotEq).is_err()); +} + +#[test] +fn test_empty_list() { + let ctx = &mut array_session().create_execution_ctx(); + let list = ListViewArray::new( + BoolArray::from_iter(Vec::::new()).into_array(), + buffer![0i32, 0i32, 0i32].into_array(), + buffer![0i32, 0i32, 0i32].into_array(), + Validity::AllValid, + ); + + let result = list + .clone() + .into_array() + .binary(list.into_array(), Operator::Eq) + .unwrap(); + assert!(result.execute_scalar(0, ctx).unwrap().is_valid()); + assert!(result.execute_scalar(1, ctx).unwrap().is_valid()); + assert!(result.execute_scalar(2, ctx).unwrap().is_valid()); +} + +fn execute_compare_test(lhs: ArrayRef, rhs: ArrayRef, op: Operator) -> ArrayRef { + lhs.binary(rhs, op).unwrap() +} + +#[rstest] +#[case(Operator::Eq, [false, true, false, false])] +#[case(Operator::NotEq, [true, false, true, true])] +#[case(Operator::Lt, [true, false, false, true])] +#[case(Operator::Lte, [true, true, false, true])] +#[case(Operator::Gt, [false, false, true, false])] +#[case(Operator::Gte, [false, true, true, false])] +fn int_all_operators(#[case] op: Operator, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![1i32, 5, 9, 2].into_array(); + let rhs = buffer![3i32, 5, 7, 4].into_array(); + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); +} + +#[rstest] +#[case(Operator::Lt, [Some(true), None, Some(false), None])] +#[case(Operator::Eq, [Some(false), None, Some(false), None])] +fn int_nullable(#[case] op: Operator, #[case] expected: [Option; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let lhs = PrimitiveArray::from_option_iter([Some(1i64), None, Some(9), Some(2)]).into_array(); + let rhs = PrimitiveArray::from_option_iter([Some(3i64), Some(5), Some(7), None]).into_array(); + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); +} + +#[rstest] +#[case(Operator::Gt, [false, false, true, false])] +#[case(Operator::Lte, [true, true, false, true])] +fn int_constant_lhs_and_rhs(#[case] op: Operator, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let array = buffer![1i32, 5, 9, 2].into_array(); + let constant = ConstantArray::new(5i32, 4).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); + + // The swapped form must produce the swapped result. + let swapped = execute_compare_test(constant, array, op); + let expected_swapped: Vec = match op { + Operator::Gt => vec![true, false, false, true], + Operator::Lte => vec![false, true, true, false], + _ => unreachable!(), + }; + assert_arrays_eq!(swapped, BoolArray::from_iter(expected_swapped), &mut ctx); +} + +/// Floats compare with Vortex's total ordering: NaN is the largest value, equality is bitwise, +/// and -0.0 < +0.0. This matches `Scalar` comparison semantics. +#[test] +fn float_total_order() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = buffer![f64::NAN, f64::NAN, -0.0f64, 1.0].into_array(); + let rhs = buffer![f64::NAN, f64::INFINITY, 0.0f64, f64::NAN].into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Eq); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, false, false]), + &mut ctx + ); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, true, true]), + &mut ctx + ); +} + +#[rstest] +#[case(Operator::Eq, [true, false, true, true])] +#[case(Operator::Lt, [false, true, false, false])] +#[case(Operator::Gte, [true, false, true, true])] +fn bool_constant(#[case] op: Operator, #[case] expected: [bool; 4]) { + let mut ctx = array_session().create_execution_ctx(); + let array = BoolArray::from_iter([true, false, true, true]).into_array(); + let constant = ConstantArray::new(true, 4).into_array(); + let result = execute_compare_test(array, constant, op); + assert_arrays_eq!(result, BoolArray::from_iter(expected), &mut ctx); +} + +#[rstest] +// Inlined vs inlined, decided by prefix. +#[case("bad", "bat", Operator::Lt, true)] +// Inlined prefix tie decided by the tail bytes. +#[case("abcdefgh", "abcdefgi", Operator::Lt, true)] +// Inlined prefix and tail tie decided by length. +#[case("abc", "abcd", Operator::Lt, true)] +// Out-of-line values with equal prefixes. +#[case("aaaaaaaaaaaaaaaaaaaab", "aaaaaaaaaaaaaaaaaaaac", Operator::Lt, true)] +// Inlined vs out-of-line where one is a prefix of the other. +#[case("aaaa", "aaaaaaaaaaaaaaaaaaaa", Operator::Lt, true)] +// Equality across the inlined/out-of-line boundary. +#[case("aaaaaaaaaaaaaaaaaaaa", "aaaaaaaaaaaaaaaaaaaa", Operator::Eq, true)] +#[case("aaaaaaaaaaaaaaaaaaaa", "aaaaaaaaaaaaaaaaaaab", Operator::Eq, false)] +// Embedded NUL: "a\0" > "a" even though the padded prefixes tie. +#[case("a\0", "a", Operator::Gt, true)] +fn string_compare_cases( + #[case] lhs: &str, + #[case] rhs: &str, + #[case] op: Operator, + #[case] expected: bool, +) { + let mut ctx = array_session().create_execution_ctx(); + let lhs = VarBinViewArray::from_iter_str([lhs]).into_array(); + let rhs = VarBinViewArray::from_iter_str([rhs]).into_array(); + let result = execute_compare_test(lhs, rhs, op); + assert_arrays_eq!(result, BoolArray::from_iter([expected]), &mut ctx); +} + +#[test] +fn string_constant_compare() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str([ + "apple", + "banana", + "banan", + "bananarama-bananarama", + "cherry", + ]) + .into_array(); + let constant = ConstantArray::new(Scalar::from("banana"), 5).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Eq); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, true, false, false, false]), + &mut ctx + ); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([true, false, true, false, false]), + &mut ctx + ); + + let result = execute_compare_test(constant, array, Operator::Lt); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false, true, true]), + &mut ctx + ); +} + +#[test] +fn string_constant_longer_than_inline() { + let mut ctx = array_session().create_execution_ctx(); + let array = VarBinViewArray::from_iter_str(["short", "averyveryverylongstring", "averyvery"]) + .into_array(); + let constant = ConstantArray::new(Scalar::from("averyveryverylongstring"), 3).into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, false]), &mut ctx); + + let result = execute_compare_test(array, constant, Operator::Lte); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, true]), &mut ctx); +} + +#[test] +fn decimal_compare() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100, 250, 300], dtype).into_array(); + let rhs = DecimalArray::from_iter::([200, 250, 100], dtype).into_array(); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); +} + +/// Two decimal arrays with the same logical dtype but different storage widths compare through +/// the widened common storage type. +#[test] +fn decimal_compare_mixed_storage_widths() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(10, 2); + let lhs = DecimalArray::from_iter::([100, 250, 300], dtype).into_array(); + let rhs = DecimalArray::from_iter::([200, 250, 100], dtype).into_array(); + + let result = execute_compare_test(lhs, rhs, Operator::Lte); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, false]), &mut ctx); +} + +/// A decimal constant that does not fit the array's narrow storage type still compares +/// correctly: it is greater than every representable array value. +#[test] +fn decimal_constant_out_of_storage_range() { + let mut ctx = array_session().create_execution_ctx(); + let dtype = DecimalDType::new(20, 2); + let array = DecimalArray::from_iter::([1, 50, 100], dtype).into_array(); + let constant = ConstantArray::new( + Scalar::decimal( + DecimalValue::from(10_000_000i64), + dtype, + Nullability::NonNullable, + ), + 3, + ) + .into_array(); + + let result = execute_compare_test(array.clone(), constant.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, true, true]), &mut ctx); + + let result = execute_compare_test(array, constant, Operator::Gte); + assert_arrays_eq!( + result, + BoolArray::from_iter([false, false, false]), + &mut ctx + ); +} + +/// Extension arrays compare through their storage values. +#[test] +fn extension_timestamp_compare() { + let mut ctx = array_session().create_execution_ctx(); + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let lhs = ExtensionArray::new(ext_dtype.clone(), buffer![1000i64, 2000, 3000].into_array()) + .into_array(); + let rhs = + ExtensionArray::new(ext_dtype, buffer![1500i64, 2000, 2500].into_array()).into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, false]), &mut ctx); +} + +/// Comparing extension arrays with different extension dtypes (e.g. timestamps in different +/// units) must error rather than silently comparing raw storage values. +#[test] +fn extension_mismatched_units_errors() { + let mut ctx = array_session().create_execution_ctx(); + let ms = ExtensionArray::new( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + buffer![1000i64, 2000].into_array(), + ) + .into_array(); + let secs = ExtensionArray::new( + Timestamp::new(TimeUnit::Seconds, Nullability::NonNullable).erased(), + buffer![1i64, 3].into_array(), + ) + .into_array(); + + let result = ms + .binary(secs, Operator::Eq) + .and_then(|a| a.execute::(&mut ctx)); + assert!(result.is_err()); +} + +/// Struct fields containing nulls order null-first, matching `Scalar` comparison semantics; +/// only top-level nulls make the result null. +#[test] +fn struct_compare_null_fields_order_first() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = StructArray::from_fields(&[( + "a", + PrimitiveArray::from_option_iter([None, Some(5i32), None]).into_array(), + )]) + .unwrap() + .into_array(); + let rhs = StructArray::from_fields(&[( + "a", + PrimitiveArray::from_option_iter([Some(1i32), Some(5), None]).into_array(), + )]) + .unwrap() + .into_array(); + + // null field < non-null field; null == null. + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, true]), &mut ctx); +} + +#[test] +fn fixed_size_list_compare() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = FixedSizeListArray::new( + buffer![1i32, 2, 3, 4, 5, 6].into_array(), + 2, + Validity::NonNullable, + 3, + ) + .into_array(); + let rhs = FixedSizeListArray::new( + buffer![1i32, 2, 3, 5, 4, 6].into_array(), + 2, + Validity::NonNullable, + 3, + ) + .into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([false, true, false]), &mut ctx); +} + +/// Binary (non-utf8) comparison over VarBinView arrays. +#[test] +fn binary_compare() { + let mut ctx = array_session().create_execution_ctx(); + let lhs = VarBinViewArray::from_iter_bin([b"bad".as_slice(), b"\xff\x00", b""]).into_array(); + let rhs = VarBinViewArray::from_iter_bin([b"bat".as_slice(), b"\xff", b""]).into_array(); + + let result = execute_compare_test(lhs.clone(), rhs.clone(), Operator::Lt); + assert_arrays_eq!(result, BoolArray::from_iter([true, false, false]), &mut ctx); + + let result = execute_compare_test(lhs, rhs, Operator::Eq); + assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index e95627155e9..ac76d2e4934 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -212,7 +212,9 @@ where } } -enum PrimitiveOperand { +/// A primitive binary-operator operand: a materialized buffer, a non-null constant, or an +/// all-null constant. +pub(super) enum PrimitiveOperand { Array { values: Buffer, validity: Validity, @@ -226,7 +228,7 @@ enum PrimitiveOperand { } impl PrimitiveOperand { - fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + pub(super) fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { if let Some(constant) = array.as_opt::() { return Ok( match constant.scalar().as_primitive().try_typed_value::()? { @@ -250,14 +252,14 @@ impl PrimitiveOperand { Ok(Self::Array { values, validity }) } - fn len(&self) -> usize { + pub(super) fn len(&self) -> usize { match self { Self::Array { values, .. } => values.len(), Self::Constant { len, .. } | Self::Null(len) => *len, } } - fn validity(&self) -> Validity { + pub(super) fn validity(&self) -> Validity { match self { Self::Array { validity, .. } => validity.clone(), Self::Constant { validity, .. } => validity.clone(), diff --git a/vortex-compute/src/lane_kernels/map_into.rs b/vortex-compute/src/lane_kernels/map_into.rs index d20b261c274..258913e9fc7 100644 --- a/vortex-compute/src/lane_kernels/map_into.rs +++ b/vortex-compute/src/lane_kernels/map_into.rs @@ -156,6 +156,67 @@ pub trait IndexedSourceExt: IndexedSource + Sized { } } + /// Apply the predicate `f(value)` lane-by-lane and bit-pack the results into + /// `words`, LSB-first, 64 lanes per `u64`. + /// + /// This is the kernel shape behind comparison operators: each lane read is an + /// independent indexed load (drive two columns via [`LaneZip`]) and the 64 + /// per-lane booleans of a chunk reduce into a single word with `OR + shift`, + /// which the autovectorizer lowers to a vector compare plus movemask. + /// + /// Words are written with `=` (not `|=`), so `words` need not be + /// zero-initialised. Bits at positions `>= self.len()` in the last word are + /// written as zero. + /// + /// Like [`map_into`], this kernel has no validity awareness; pair the packed + /// bits with a separately computed validity mask. + /// + /// [`LaneZip`]: crate::lane_kernels::LaneZip + /// [`map_into`]: IndexedSourceExt::map_into + /// + /// # Panics + /// + /// Panics if `words.len() < self.len().div_ceil(64)`. + #[inline] + fn map_bits_into(self, words: &mut [u64], mut f: F) + where + F: FnMut(Self::Item) -> bool, + { + #[inline(always)] + fn chunk(values: &S, f: &mut F, base: usize, count: usize) -> u64 + where + S: IndexedSource, + F: FnMut(S::Item) -> bool, + { + let mut packed: u64 = 0; + for bit_idx in 0..count { + // SAFETY: caller guarantees base + count <= len. + let val = unsafe { values.get_unchecked(base + bit_idx) }; + packed |= (f(val) as u64) << bit_idx; + } + packed + } + + let values = self; + let len = values.len(); + let num_words = len.div_ceil(64); + assert!( + words.len() >= num_words, + "words slice has {} entries, need at least {num_words}", + words.len(), + ); + + let full = len / 64; + let remainder = len % 64; + + for word_idx in 0..full { + words[word_idx] = chunk(&values, &mut f, word_idx * 64, 64); + } + if remainder != 0 { + words[full] = chunk(&values, &mut f, full * 64, remainder); + } + } + /// Fallible map with **no validity awareness at all** — every `None` returned /// by the closure is treated as a failure, even at null lanes. /// @@ -485,6 +546,43 @@ mod tests { assert!(res.is_ok(), "null lane should bypass the range check"); } + #[test] + fn map_bits_into_packs_full_and_remainder_words() { + let values: Vec = (0..130).collect(); + let mut words = vec![u64::MAX; 3]; + values.as_slice().map_bits_into(&mut words, |v| v % 2 == 0); + + for idx in 0..130 { + let bit = (words[idx / 64] >> (idx % 64)) & 1; + assert_eq!(bit == 1, idx % 2 == 0, "lane {idx}"); + } + // Bits past `len` in the remainder word must be written as zero. + assert_eq!(words[2] >> 2, 0); + } + + #[test] + fn map_bits_into_lane_zip_compare() { + use crate::lane_kernels::LaneZip; + + let lhs: Vec = (0..100).collect(); + let rhs: Vec = (0..100).rev().collect(); + let mut words = vec![0u64; 2]; + LaneZip::new(lhs.as_slice(), rhs.as_slice()).map_bits_into(&mut words, |(a, b)| a >= b); + + for idx in 0..100 { + let bit = (words[idx / 64] >> (idx % 64)) & 1; + assert_eq!(bit == 1, lhs[idx] >= rhs[idx], "lane {idx}"); + } + } + + #[test] + #[should_panic(expected = "words slice has 1 entries")] + fn map_bits_into_words_too_short_panics() { + let values: Vec = (0..65).collect(); + let mut words = vec![0u64; 1]; + values.as_slice().map_bits_into(&mut words, |v| v > 0); + } + #[test] fn try_map_masked_into_overflow_in_remainder() { let mut values: Vec = (0..130).collect();