vortex_array/scalar_fn/fns/binary/compare/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Native comparison kernels.
5//!
6//! [`execute_compare`] dispatches on the logical [`DType`] of its operands and evaluates every
7//! comparison directly over Vortex canonical arrays — bit buffers for booleans, lane kernels from
8//! `vortex-compute` for primitives and decimals, binary views for strings/bytes, and a row-wise
9//! comparator for nested types. There is no Arrow fallback.
10//!
11//! Floating point values compare with Vortex's total ordering (`NaN` is the largest value,
12//! `-0.0 < +0.0`, and equality is bitwise), matching [`Scalar`] comparison semantics.
13
14use std::cmp::Ordering;
15
16use vortex_error::VortexResult;
17use vortex_error::vortex_bail;
18use vortex_error::vortex_err;
19
20use crate::ArrayRef;
21use crate::Canonical;
22use crate::ExecutionCtx;
23use crate::IntoArray;
24use crate::array::ArrayView;
25use crate::array::VTable;
26use crate::arrays::Constant;
27use crate::arrays::ConstantArray;
28use crate::arrays::ExtensionArray;
29use crate::arrays::ScalarFn;
30use crate::arrays::extension::ExtensionArrayExt;
31use crate::arrays::scalar_fn::ExactScalarFn;
32use crate::arrays::scalar_fn::ScalarFnArrayExt;
33use crate::arrays::scalar_fn::ScalarFnArrayView;
34use crate::dtype::DType;
35use crate::dtype::Nullability;
36use crate::kernel::ExecuteParentKernel;
37use crate::scalar::Scalar;
38use crate::scalar_fn::fns::binary::Binary;
39use crate::scalar_fn::fns::operators::CompareOperator;
40use crate::validity::Validity;
41
42mod boolean;
43mod bytes;
44mod decimal;
45mod nested;
46pub(crate) use nested::build_comparator as build_row_comparator;
47mod primitive;
48#[cfg(test)]
49mod tests;
50
51/// Trait for encoding-specific comparison kernels that operate in encoded space.
52///
53/// Implementations can compare an encoded array against another array (typically a constant)
54/// without first decompressing. The adaptor normalizes operand order so `array` is always
55/// the left-hand side, swapping the operator when necessary.
56pub trait CompareKernel: VTable {
57    fn compare(
58        lhs: ArrayView<'_, Self>,
59        rhs: &ArrayRef,
60        operator: CompareOperator,
61        ctx: &mut ExecutionCtx,
62    ) -> VortexResult<Option<ArrayRef>>;
63}
64
65/// Adaptor that bridges [`CompareKernel`] implementations to [`ExecuteParentKernel`].
66///
67/// When a `ScalarFnArray(Binary, cmp_op)` wraps a child that implements `CompareKernel`,
68/// this adaptor extracts the comparison operator and other operand, normalizes operand order
69/// (swapping the operator if the encoded array is on the RHS), and delegates to the kernel.
70#[derive(Default, Debug)]
71pub struct CompareExecuteAdaptor<V>(pub V);
72
73impl<V> ExecuteParentKernel<V> for CompareExecuteAdaptor<V>
74where
75    V: CompareKernel,
76{
77    type Parent = ExactScalarFn<Binary>;
78
79    fn execute_parent(
80        &self,
81        array: ArrayView<'_, V>,
82        parent: ScalarFnArrayView<'_, Binary>,
83        child_idx: usize,
84        ctx: &mut ExecutionCtx,
85    ) -> VortexResult<Option<ArrayRef>> {
86        // Non-comparison operators are normal kernel misses, so decline them without constructing
87        // an error (which would capture a backtrace, or panic under `VORTEX_PANIC_ON_ERR`).
88        let Some(cmp_op) = parent.options.to_compare_operator() else {
89            return Ok(None);
90        };
91
92        // Get the ScalarFnArray to access children
93        let Some(scalar_fn_array) = parent.as_opt::<ScalarFn>() else {
94            return Ok(None);
95        };
96        // Normalize so `array` is always LHS, swapping the operator if needed
97        // TODO(joe): should be go this here or in the Rule/Kernel
98        let (cmp_op, other) = match child_idx {
99            0 => (cmp_op, scalar_fn_array.get_child(1)),
100            1 => (cmp_op.swap(), scalar_fn_array.get_child(0)),
101            _ => return Ok(None),
102        };
103
104        let len = array.len();
105        let nullable = array.dtype().is_nullable() || other.dtype().is_nullable();
106
107        // Empty array → empty bool result
108        if len == 0 {
109            return Ok(Some(
110                Canonical::empty(&DType::Bool(nullable.into())).into_array(),
111            ));
112        }
113
114        // Null constant on either side → all-null bool result
115        if other.as_constant().is_some_and(|s| s.is_null()) {
116            return Ok(Some(
117                ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), len).into_array(),
118            ));
119        }
120
121        V::compare(array, other, cmp_op, ctx)
122    }
123}
124
125/// Execute a compare operation between two arrays.
126///
127/// This is the entry point for compare operations from the binary expression.
128/// Handles empty, constant-null, and constant-constant directly, otherwise dispatches to a
129/// native per-dtype kernel.
130pub(crate) fn execute_compare(
131    lhs: &ArrayRef,
132    rhs: &ArrayRef,
133    op: CompareOperator,
134    ctx: &mut ExecutionCtx,
135) -> VortexResult<ArrayRef> {
136    let nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable();
137
138    if lhs.len() != rhs.len() {
139        vortex_bail!(
140            "compare operator requires equal lengths, got {} and {}",
141            lhs.len(),
142            rhs.len()
143        );
144    }
145
146    if lhs.is_empty() {
147        return Ok(Canonical::empty(&DType::Bool(nullable.into())).into_array());
148    }
149
150    let left_constant_null = lhs.as_constant().map(|l| l.is_null()).unwrap_or(false);
151    let right_constant_null = rhs.as_constant().map(|r| r.is_null()).unwrap_or(false);
152    if left_constant_null || right_constant_null {
153        return Ok(
154            ConstantArray::new(Scalar::null(DType::Bool(nullable.into())), lhs.len()).into_array(),
155        );
156    }
157
158    // Constant-constant fast path
159    if let (Some(lhs_const), Some(rhs_const)) = (lhs.as_opt::<Constant>(), rhs.as_opt::<Constant>())
160    {
161        let result = scalar_cmp(lhs_const.scalar(), rhs_const.scalar(), op)?;
162        return Ok(ConstantArray::new(result, lhs.len()).into_array());
163    }
164
165    compare_arrays(lhs, rhs, op, ctx)
166}
167
168/// Dispatch a comparison to the native kernel for the operands' logical dtype.
169fn compare_arrays(
170    lhs: &ArrayRef,
171    rhs: &ArrayRef,
172    op: CompareOperator,
173    ctx: &mut ExecutionCtx,
174) -> VortexResult<ArrayRef> {
175    let nullability = Nullability::from(lhs.dtype().is_nullable() || rhs.dtype().is_nullable());
176
177    // Extension arrays compare through their storage. When both sides are extensions they must
178    // agree on the full extension dtype (a timestamp in milliseconds must not compare its raw
179    // storage against one in seconds); a lone extension side may compare against its raw storage.
180    if lhs.dtype().is_extension() || rhs.dtype().is_extension() {
181        if lhs.dtype().is_extension()
182            && rhs.dtype().is_extension()
183            && !lhs.dtype().eq_ignore_nullability(rhs.dtype())
184        {
185            vortex_bail!(
186                "Cannot compare extension dtypes {} and {}",
187                lhs.dtype(),
188                rhs.dtype()
189            );
190        }
191        let lhs = extension_storage(lhs, ctx)?;
192        let rhs = extension_storage(rhs, ctx)?;
193        return compare_arrays(&lhs, &rhs, op, ctx);
194    }
195
196    if !lhs.dtype().eq_ignore_nullability(rhs.dtype()) {
197        vortex_bail!(
198            "Cannot compare different DTypes {} and {}",
199            lhs.dtype(),
200            rhs.dtype()
201        );
202    }
203
204    match lhs.dtype() {
205        // Every value of the null dtype is null, so every comparison result is null.
206        DType::Null => Ok(ConstantArray::new(
207            Scalar::null(DType::Bool(Nullability::Nullable)),
208            lhs.len(),
209        )
210        .into_array()),
211        DType::Bool(_) => boolean::compare_bool(lhs, rhs, op, nullability, ctx),
212        DType::Primitive(..) => primitive::compare_primitive(lhs, rhs, op, nullability, ctx),
213        DType::Decimal(..) => decimal::compare_decimal(lhs, rhs, op, nullability, ctx),
214        DType::Utf8(_) | DType::Binary(_) => bytes::compare_bytes(lhs, rhs, op, nullability, ctx),
215        DType::Struct(..) | DType::List(..) | DType::FixedSizeList(..) | DType::Map(..) => {
216            nested::compare_nested(lhs, rhs, op, nullability, ctx)
217        }
218        DType::Union(..) | DType::Variant(_) | DType::Extension(_) => {
219            vortex_bail!("compare is not supported for dtype {}", lhs.dtype())
220        }
221    }
222}
223
224/// Replace an extension array (or extension constant) with its storage. Non-extension arrays are
225/// returned unchanged.
226fn extension_storage(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
227    if !array.dtype().is_extension() {
228        return Ok(array.clone());
229    }
230    if let Some(constant) = array.as_opt::<Constant>() {
231        return Ok(ConstantArray::new(
232            constant.scalar().as_extension().to_storage_scalar(),
233            constant.len(),
234        )
235        .into_array());
236    }
237    let ext = array.clone().execute::<ExtensionArray>(ctx)?;
238    Ok(ext.storage_array().clone())
239}
240
241pub fn scalar_cmp(lhs: &Scalar, rhs: &Scalar, operator: CompareOperator) -> VortexResult<Scalar> {
242    if lhs.is_null() | rhs.is_null() {
243        return Ok(Scalar::null(DType::Bool(Nullability::Nullable)));
244    }
245
246    let nullability = lhs.dtype().nullability() | rhs.dtype().nullability();
247
248    // We use `partial_cmp` to ensure we do not lose a type mismatch error.
249    let ordering = lhs.partial_cmp(rhs).ok_or_else(|| {
250        vortex_err!(
251            "Cannot compare scalars with incompatible types: {} and {}",
252            lhs.dtype(),
253            rhs.dtype()
254        )
255    })?;
256
257    let b = match operator {
258        CompareOperator::Eq => ordering.is_eq(),
259        CompareOperator::NotEq => ordering.is_ne(),
260        CompareOperator::Gt => ordering.is_gt(),
261        CompareOperator::Gte => ordering.is_ge(),
262        CompareOperator::Lt => ordering.is_lt(),
263        CompareOperator::Lte => ordering.is_le(),
264    };
265
266    Ok(Scalar::bool(b, nullability))
267}
268
269/// Returns the predicate `Ordering -> bool` for a comparison operator.
270#[inline]
271pub(super) fn ordering_predicate(op: CompareOperator) -> fn(Ordering) -> bool {
272    match op {
273        CompareOperator::Eq => Ordering::is_eq,
274        CompareOperator::NotEq => Ordering::is_ne,
275        CompareOperator::Gt => Ordering::is_gt,
276        CompareOperator::Gte => Ordering::is_ge,
277        CompareOperator::Lt => Ordering::is_lt,
278        CompareOperator::Lte => Ordering::is_le,
279    }
280}
281
282/// Combine operand validities into the validity of a comparison result with the given result
283/// nullability.
284pub(super) fn compare_validity(
285    lhs: Validity,
286    rhs: Validity,
287    nullability: Nullability,
288) -> VortexResult<Validity> {
289    Ok(lhs.and(rhs)?.union_nullability(nullability))
290}
vortex_array/scalar_fn/fns/binary/compare/mod.rs

vortex_array/scalar_fn/fns/binary/compare/
mod.rs
