Skip to main content

arrow_ord/
cmp.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Comparison kernels for `Array`s.
6//!
7//! These kernels can leverage SIMD if available on your system.  Currently no runtime
8//! detection is provided, you should enable the specific SIMD intrinsics using
9//! `RUSTFLAGS="-C target-feature=+avx2"` for example.  See the documentation
10//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
11//!
12
13use arrow_array::cast::AsArray;
14use arrow_array::types::{ByteArrayType, ByteViewType};
15use arrow_array::{
16    AnyDictionaryArray, Array, ArrowNativeTypeOp, BooleanArray, Datum, FixedSizeBinaryArray,
17    GenericByteArray, GenericByteViewArray, downcast_primitive_array, downcast_run_array,
18};
19use arrow_buffer::bit_util::ceil;
20use arrow_buffer::{ArrowNativeType, BooleanBuffer, BooleanBufferBuilder, NullBuffer};
21use arrow_schema::{ArrowError, DataType};
22use arrow_select::take::take;
23use std::cmp::Ordering;
24use std::ops::Not;
25
26#[derive(Debug, Copy, Clone)]
27enum Op {
28    Equal,
29    NotEqual,
30    Less,
31    LessEqual,
32    Greater,
33    GreaterEqual,
34    Distinct,
35    NotDistinct,
36}
37
38impl std::fmt::Display for Op {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        match self {
41            Op::Equal => write!(f, "=="),
42            Op::NotEqual => write!(f, "!="),
43            Op::Less => write!(f, "<"),
44            Op::LessEqual => write!(f, "<="),
45            Op::Greater => write!(f, ">"),
46            Op::GreaterEqual => write!(f, ">="),
47            Op::Distinct => write!(f, "IS DISTINCT FROM"),
48            Op::NotDistinct => write!(f, "IS NOT DISTINCT FROM"),
49        }
50    }
51}
52
53/// Perform `left == right` operation on two [`Datum`].
54///
55/// Comparing null values on either side will yield a null in the corresponding
56/// slot of the resulting [`BooleanArray`].
57///
58/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
59/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
60/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
61/// to treat them as equal, please normalize zeros before calling this kernel. See
62/// [`f32::total_cmp`] and [`f64::total_cmp`].
63///
64/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
65/// For comparisons involving nested types see [`crate::ord::make_comparator`]
66pub fn eq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
67    compare_op(Op::Equal, lhs, rhs)
68}
69
70/// Perform `left != right` operation on two [`Datum`].
71///
72/// Comparing null values on either side will yield a null in the corresponding
73/// slot of the resulting [`BooleanArray`].
74///
75/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
76/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
77/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
78/// to treat them as equal, please normalize zeros before calling this kernel. See
79/// [`f32::total_cmp`] and [`f64::total_cmp`].
80///
81/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
82/// For comparisons involving nested types see [`crate::ord::make_comparator`]
83pub fn neq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
84    compare_op(Op::NotEqual, lhs, rhs)
85}
86
87/// Perform `left < right` operation on two [`Datum`].
88///
89/// Comparing null values on either side will yield a null in the corresponding
90/// slot of the resulting [`BooleanArray`].
91///
92/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
93/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
94/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
95/// to treat them as equal, please normalize zeros before calling this kernel. See
96/// [`f32::total_cmp`] and [`f64::total_cmp`].
97///
98/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
99/// For comparisons involving nested types see [`crate::ord::make_comparator`]
100pub fn lt(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
101    compare_op(Op::Less, lhs, rhs)
102}
103
104/// Perform `left <= right` operation on two [`Datum`].
105///
106/// Comparing null values on either side will yield a null in the corresponding
107/// slot of the resulting [`BooleanArray`].
108///
109/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
110/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
111/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
112/// to treat them as equal, please normalize zeros before calling this kernel. See
113/// [`f32::total_cmp`] and [`f64::total_cmp`].
114///
115/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
116/// For comparisons involving nested types see [`crate::ord::make_comparator`]
117pub fn lt_eq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
118    compare_op(Op::LessEqual, lhs, rhs)
119}
120
121/// Perform `left > right` operation on two [`Datum`].
122///
123/// Comparing null values on either side will yield a null in the corresponding
124/// slot of the resulting [`BooleanArray`].
125///
126/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
127/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
128/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
129/// to treat them as equal, please normalize zeros before calling this kernel. See
130/// [`f32::total_cmp`] and [`f64::total_cmp`].
131///
132/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
133/// For comparisons involving nested types see [`crate::ord::make_comparator`]
134pub fn gt(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
135    compare_op(Op::Greater, lhs, rhs)
136}
137
138/// Perform `left >= right` operation on two [`Datum`].
139///
140/// Comparing null values on either side will yield a null in the corresponding
141/// slot of the resulting [`BooleanArray`].
142///
143/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
144/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
145/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
146/// to treat them as equal, please normalize zeros before calling this kernel. See
147/// [`f32::total_cmp`] and [`f64::total_cmp`].
148///
149/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
150/// For comparisons involving nested types see [`crate::ord::make_comparator`]
151pub fn gt_eq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
152    compare_op(Op::GreaterEqual, lhs, rhs)
153}
154
155/// Perform `left IS DISTINCT FROM right` operation on two [`Datum`]
156///
157/// [`distinct`] is similar to [`neq`], only differing in null handling. In particular, two
158/// operands are considered DISTINCT if they have a different value or if one of them is NULL
159/// and the other isn't. The result of [`distinct`] is never NULL.
160///
161/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
162/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
163/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
164/// to treat them as equal, please normalize zeros before calling this kernel. See
165/// [`f32::total_cmp`] and [`f64::total_cmp`].
166///
167/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
168/// For comparisons involving nested types see [`crate::ord::make_comparator`]
169pub fn distinct(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
170    compare_op(Op::Distinct, lhs, rhs)
171}
172
173/// Perform `left IS NOT DISTINCT FROM right` operation on two [`Datum`]
174///
175/// [`not_distinct`] is similar to [`eq`], only differing in null handling. In particular, two
176/// operands are considered `NOT DISTINCT` if they have the same value or if both of them
177/// is NULL. The result of [`not_distinct`] is never NULL.
178///
179/// For floating values like f32 and f64, this comparison produces an ordering in accordance to
180/// the totalOrder predicate as defined in the IEEE 754 (2008 revision) floating point standard.
181/// Note that totalOrder treats positive and negative zeros as different. If it is necessary
182/// to treat them as equal, please normalize zeros before calling this kernel. See
183/// [`f32::total_cmp`] and [`f64::total_cmp`].
184///
185/// Nested types, such as lists, are not supported as the null semantics are not well-defined.
186/// For comparisons involving nested types see [`crate::ord::make_comparator`]
187pub fn not_distinct(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
188    compare_op(Op::NotDistinct, lhs, rhs)
189}
190
191/// Returns true if `distinct` (via `compare_op`) can handle this data type.
192///
193/// `compare_op` unwraps at most one level of dictionary, then dispatches on
194/// the leaf type. Anything else (REE, nested dictionary, nested/complex types)
195/// must go through `make_comparator` instead.
196pub(crate) fn supports_distinct(dt: &DataType) -> bool {
197    use arrow_schema::DataType::*;
198    let leaf = match dt {
199        Dictionary(_, v) => v.as_ref(),
200        dt => dt,
201    };
202    !leaf.is_nested() && !matches!(leaf, Dictionary(_, _) | RunEndEncoded(_, _))
203}
204
205/// Perform `op` on the provided `Datum`
206#[inline(never)]
207fn compare_op(op: Op, lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
208    use arrow_schema::DataType::*;
209    let (l, l_s) = lhs.get();
210    let (r, r_s) = rhs.get();
211
212    let l_len = l.len();
213    let r_len = r.len();
214
215    if l_len != r_len && !l_s && !r_s {
216        return Err(ArrowError::InvalidArgumentError(format!(
217            "Cannot compare arrays of different lengths, got {l_len} vs {r_len}"
218        )));
219    }
220
221    let len = match l_s {
222        true => r_len,
223        false => l_len,
224    };
225
226    let l_nulls = l.logical_nulls();
227    let r_nulls = r.logical_nulls();
228
229    let l_ree_info = ree_info_opt(l);
230    let l = l_ree_info.as_ref().map(|(vals, _)| *vals).unwrap_or(l);
231
232    let r_ree_info = ree_info_opt(r);
233    let r = r_ree_info.as_ref().map(|(vals, _)| *vals).unwrap_or(r);
234
235    let l_v = l.as_any_dictionary_opt();
236    let l = l_v.map(|x| x.values().as_ref()).unwrap_or(l);
237    let l_t = l.data_type();
238
239    let r_v = r.as_any_dictionary_opt();
240    let r = r_v.map(|x| x.values().as_ref()).unwrap_or(r);
241    let r_t = r.data_type();
242
243    if r_t.is_nested() || l_t.is_nested() {
244        return Err(ArrowError::InvalidArgumentError(format!(
245            "Nested comparison: {l_t} {op} {r_t} (hint: use make_comparator instead)"
246        )));
247    } else if matches!(l_t, Dictionary(_, _) | RunEndEncoded(_, _))
248        || matches!(r_t, Dictionary(_, _) | RunEndEncoded(_, _))
249    {
250        // One dictionary/REE layer is unwrapped above. A leftover dictionary
251        // (nested dict) or REE is not nested according to `DataType::is_nested`,
252        // so without this check `downcast_primitive_array!` panics.
253        return Err(ArrowError::InvalidArgumentError(format!(
254            "Invalid comparison operation: {l_t} {op} {r_t}"
255        )));
256    } else if l_t != r_t {
257        return Err(ArrowError::InvalidArgumentError(format!(
258            "Invalid comparison operation: {l_t} {op} {r_t}"
259        )));
260    }
261
262    let l_side = SideInfo {
263        is_scalar: l_s,
264        dict: l_v,
265        ree: l_ree_info.as_ref().map(|(_, info)| info),
266    };
267    let r_side = SideInfo {
268        is_scalar: r_s,
269        dict: r_v,
270        ree: r_ree_info.as_ref().map(|(_, info)| info),
271    };
272
273    // Special case: equality against a short scalar that fits the inlined prefix
274    // of a view array, which reduces the comparison to a scan of fixed-width
275    // integers
276    if matches!(op, Op::Equal | Op::NotEqual)
277        && l_side.dict.is_none()
278        && r_side.dict.is_none()
279        && l_side.ree.is_none()
280        && r_side.ree.is_none()
281    {
282        let sides = match (l_s, r_s) {
283            (false, true) => Some((l, r, &l_nulls, &r_nulls)),
284            (true, false) => Some((r, l, &r_nulls, &l_nulls)),
285            _ => None,
286        };
287        // A null constant makes every row null, which this path cannot express
288        if let Some((values, scalar, values_nulls, scalar_nulls)) = sides
289            && scalar_nulls.as_ref().is_none_or(|n| n.null_count() == 0)
290            && let Some(mask) = eq_inline_scalar(values, scalar, l_t, matches!(op, Op::NotEqual))
291        {
292            let nulls = values_nulls.clone().filter(|n| n.null_count() > 0);
293            return Ok(BooleanArray::new(mask, nulls));
294        }
295    }
296
297    // Defer computation as may not be necessary
298    let values = || -> BooleanBuffer {
299        let d = downcast_primitive_array! {
300            (l, r) => apply(op, l.values().as_ref(), &l_side, r.values().as_ref(), &r_side),
301            (Boolean, Boolean) => apply(op, l.as_boolean(), &l_side, r.as_boolean(), &r_side),
302            (Utf8, Utf8) => apply(op, l.as_string::<i32>(), &l_side, r.as_string::<i32>(), &r_side),
303            (Utf8View, Utf8View) => apply(op, l.as_string_view(), &l_side, r.as_string_view(), &r_side),
304            (LargeUtf8, LargeUtf8) => apply(op, l.as_string::<i64>(), &l_side, r.as_string::<i64>(), &r_side),
305            (Binary, Binary) => apply(op, l.as_binary::<i32>(), &l_side, r.as_binary::<i32>(), &r_side),
306            (BinaryView, BinaryView) => apply(op, l.as_binary_view(), &l_side, r.as_binary_view(), &r_side),
307            (LargeBinary, LargeBinary) => apply(op, l.as_binary::<i64>(), &l_side, r.as_binary::<i64>(), &r_side),
308            (FixedSizeBinary(_), FixedSizeBinary(_)) => apply(op, l.as_fixed_size_binary(), &l_side, r.as_fixed_size_binary(), &r_side),
309            (Null, Null) => None,
310            _ => unreachable!(),
311        };
312        d.unwrap_or_else(|| BooleanBuffer::new_unset(len))
313    };
314
315    let l_nulls = l_nulls.filter(|n| n.null_count() > 0);
316    let r_nulls = r_nulls.filter(|n| n.null_count() > 0);
317    Ok(match (l_nulls, l_s, r_nulls, r_s) {
318        (Some(l), true, Some(r), true) | (Some(l), false, Some(r), false) => {
319            // Either both sides are scalar or neither side is scalar
320            match op {
321                Op::Distinct => {
322                    let values = values();
323                    let l = l.inner().bit_chunks().iter_padded();
324                    let r = r.inner().bit_chunks().iter_padded();
325                    let ne = values.bit_chunks().iter_padded();
326
327                    let c = |((l, r), n)| (l ^ r) | (l & r & n);
328                    let buffer = l.zip(r).zip(ne).map(c).collect();
329                    BooleanBuffer::new(buffer, 0, len).into()
330                }
331                Op::NotDistinct => {
332                    let values = values();
333                    let l = l.inner().bit_chunks().iter_padded();
334                    let r = r.inner().bit_chunks().iter_padded();
335                    let e = values.bit_chunks().iter_padded();
336
337                    let c = |((l, r), e)| u64::not(l | r) | (l & r & e);
338                    let buffer = l.zip(r).zip(e).map(c).collect();
339                    BooleanBuffer::new(buffer, 0, len).into()
340                }
341                _ => BooleanArray::new(values(), NullBuffer::union(Some(&l), Some(&r))),
342            }
343        }
344        (Some(_), true, Some(a), false) | (Some(a), false, Some(_), true) => {
345            // Scalar is null, other side is non-scalar and nullable
346            match op {
347                Op::Distinct => a.into_inner().into(),
348                Op::NotDistinct => a.into_inner().not().into(),
349                _ => BooleanArray::new_null(len),
350            }
351        }
352        (Some(nulls), is_scalar, None, _) | (None, _, Some(nulls), is_scalar) => {
353            // Only one side is nullable
354            match is_scalar {
355                true => match op {
356                    // Scalar is null, other side is not nullable
357                    Op::Distinct => BooleanBuffer::new_set(len).into(),
358                    Op::NotDistinct => BooleanBuffer::new_unset(len).into(),
359                    _ => BooleanArray::new_null(len),
360                },
361                false => match op {
362                    Op::Distinct => {
363                        let values = values();
364                        let l = nulls.inner().bit_chunks().iter_padded();
365                        let ne = values.bit_chunks().iter_padded();
366                        let c = |(l, n)| u64::not(l) | n;
367                        let buffer = l.zip(ne).map(c).collect();
368                        BooleanBuffer::new(buffer, 0, len).into()
369                    }
370                    Op::NotDistinct => (nulls.inner() & &values()).into(),
371                    _ => BooleanArray::new(values(), Some(nulls)),
372                },
373            }
374        }
375        // Neither side is nullable
376        (None, _, None, _) => BooleanArray::new(values(), None),
377    })
378}
379
380/// Per-side metadata for a comparison operand.
381struct SideInfo<'a> {
382    is_scalar: bool,
383    dict: Option<&'a dyn AnyDictionaryArray>,
384    ree: Option<&'a ReeInfo<'a>>,
385}
386
387impl SideInfo<'_> {
388    fn has_indirection(&self) -> bool {
389        self.dict.is_some() || self.ree.is_some()
390    }
391}
392
393/// Longest constant whose length and bytes both fit a view's low 64 bits
394const MAX_LOW_HALF_LEN: u32 = 4;
395
396/// Attempts a special case: comparing every value of a byte-view array against a
397/// short constant that fits entirely within a view's inlined prefix
398///
399/// Returns `None` for any other shape, which the caller then handles on the
400/// generic path.
401fn eq_inline_scalar(
402    l: &dyn Array,
403    r: &dyn Array,
404    data_type: &DataType,
405    negate: bool,
406) -> Option<BooleanBuffer> {
407    let (values, needle) = match data_type {
408        DataType::Utf8View => (
409            l.as_string_view().views(),
410            *r.as_string_view().views().first()?,
411        ),
412        DataType::BinaryView => (
413            l.as_binary_view().views(),
414            *r.as_binary_view().views().first()?,
415        ),
416        _ => return None,
417    };
418    // Only a constant whose length and bytes fit the view's low half. Wider
419    // constants need the whole 128-bit view, and comparing that is slower than
420    // the generic path's early exit on a length mismatch.
421    let needle_len = needle as u32;
422    if needle_len > MAX_LOW_HALF_LEN {
423        return None;
424    }
425    let significant = u64::MAX >> (32 - needle_len * 8);
426    let needle = needle as u64 & significant;
427    Some(collect_bool(values.len(), negate, |idx| {
428        let view = unsafe { *values.get_unchecked(idx) };
429        view as u64 & significant == needle
430    }))
431}
432
433/// Perform a potentially vectored `op` on the provided `ArrayOrd`
434fn apply<T: ArrayOrd>(
435    op: Op,
436    l: T,
437    l_info: &SideInfo,
438    r: T,
439    r_info: &SideInfo,
440) -> Option<BooleanBuffer> {
441    if l.len() == 0 || r.len() == 0 {
442        return None; // Handle empty dictionaries
443    }
444
445    if !l_info.is_scalar
446        && !r_info.is_scalar
447        && (l_info.has_indirection() || r_info.has_indirection())
448    {
449        // Both non-scalar with indirection. Pure REE-vs-REE uses segment-based
450        // bulk comparison; other combinations fall back to index vectors.
451        if let (Some(li), None, Some(ri), None) = (l_info.ree, l_info.dict, r_info.ree, r_info.dict)
452        {
453            return Some(apply_ree(op, l, li, r, ri));
454        }
455
456        let l_v = logical_indices(l.len(), l_info.dict, l_info.ree);
457        let r_v = logical_indices(r.len(), r_info.dict, r_info.ree);
458
459        assert_eq!(l_v.len(), r_v.len()); // Sanity check
460
461        Some(match op {
462            Op::Equal | Op::NotDistinct => apply_op_vectored(l, &l_v, r, &r_v, false, T::is_eq),
463            Op::NotEqual | Op::Distinct => apply_op_vectored(l, &l_v, r, &r_v, true, T::is_eq),
464            Op::Less => apply_op_vectored(l, &l_v, r, &r_v, false, T::is_lt),
465            Op::LessEqual => apply_op_vectored(r, &r_v, l, &l_v, true, T::is_lt),
466            Op::Greater => apply_op_vectored(r, &r_v, l, &l_v, false, T::is_lt),
467            Op::GreaterEqual => apply_op_vectored(l, &l_v, r, &r_v, true, T::is_lt),
468        })
469    } else {
470        let l_s = l_info
471            .is_scalar
472            .then(|| scalar_index(l_info.dict, l_info.ree));
473        let r_s = r_info
474            .is_scalar
475            .then(|| scalar_index(r_info.dict, r_info.ree));
476
477        let buffer = match op {
478            Op::Equal | Op::NotDistinct => apply_op(l, l_s, r, r_s, false, T::is_eq),
479            Op::NotEqual | Op::Distinct => apply_op(l, l_s, r, r_s, true, T::is_eq),
480            Op::Less => apply_op(l, l_s, r, r_s, false, T::is_lt),
481            Op::LessEqual => apply_op(r, r_s, l, l_s, true, T::is_lt),
482            Op::Greater => apply_op(r, r_s, l, l_s, false, T::is_lt),
483            Op::GreaterEqual => apply_op(l, l_s, r, r_s, true, T::is_lt),
484        };
485
486        // Expand the physical-length result back to logical length.
487        // Find the non-scalar side that needs expansion (at most one).
488        let side = if l_s.is_none() { l_info } else { r_info };
489        let buffer = match side.dict {
490            Some(d) => take_bits(d, buffer),
491            None => buffer,
492        };
493        Some(match side.ree {
494            Some(info) => expand_from_runs(info, buffer),
495            None => buffer,
496        })
497    }
498}
499
500/// Build a logical→physical index vector for one side of a non-scalar comparison.
501fn logical_indices(
502    len: usize,
503    dict: Option<&dyn AnyDictionaryArray>,
504    ree: Option<&ReeInfo>,
505) -> Vec<usize> {
506    match (dict, ree) {
507        (Some(d), Some(info)) => {
508            let keys = d.normalized_keys();
509            ree_physical_indices(info)
510                .iter()
511                .map(|&i| keys[i])
512                .collect()
513        }
514        (Some(d), None) => d.normalized_keys(),
515        (None, Some(info)) => ree_physical_indices(info),
516        (None, None) => (0..len).collect(),
517    }
518}
519
520fn ree_physical_indices(info: &ReeInfo) -> Vec<usize> {
521    let run_ends = info.run_ends_as_usize();
522    let end = info.offset + info.len;
523    let mut indices = Vec::with_capacity(info.len);
524    let mut pos = info.offset;
525    for (physical, &run_end) in run_ends.iter().enumerate().skip(info.start_physical) {
526        let run_end = run_end.min(end);
527        indices.extend(std::iter::repeat_n(physical, run_end - pos));
528        pos = run_end;
529        if pos >= end {
530            break;
531        }
532    }
533    indices
534}
535
536fn scalar_index(dict: Option<&dyn AnyDictionaryArray>, ree: Option<&ReeInfo>) -> usize {
537    let idx = ree.map(|r| r.start_physical).unwrap_or_default();
538    dict.map(|d| d.normalized_keys()[idx]).unwrap_or(idx)
539}
540
541/// Expand a physical-length BooleanBuffer to logical length by bulk-appending
542/// each run's result. Zero allocation — downcasts internally to access typed
543/// run_ends directly.
544fn expand_from_runs(info: &ReeInfo, buffer: BooleanBuffer) -> BooleanBuffer {
545    let array = info.array;
546    downcast_run_array!(
547        array => {
548            let run_ends = array.run_ends().values();
549            let end = info.offset + info.len;
550            let mut builder = BooleanBufferBuilder::new(info.len);
551            let mut pos = info.offset;
552            for (physical, re) in run_ends.iter().enumerate().skip(info.start_physical) {
553                let run_end = re.as_usize().min(end);
554                // SAFETY: physical < buffer.len() (one entry per run in the values array)
555                builder.append_n(run_end - pos, unsafe { buffer.value_unchecked(physical) });
556                pos = run_end;
557                if pos >= end {
558                    break;
559                }
560            }
561            builder.finish()
562        },
563        _ => unreachable!()
564    )
565}
566
567fn take_bits(v: &dyn AnyDictionaryArray, buffer: BooleanBuffer) -> BooleanBuffer {
568    let array = take(&BooleanArray::new(buffer, None), v.keys(), None).unwrap();
569    array.as_boolean().values().clone()
570}
571
572/// Invokes `f` with values `0..len` collecting the boolean results into a new `BooleanBuffer`
573///
574/// This is similar to [`arrow_buffer::MutableBuffer::collect_bool`] but with
575/// the option to efficiently negate the result
576fn collect_bool(len: usize, neg: bool, f: impl Fn(usize) -> bool) -> BooleanBuffer {
577    let mut buffer = Vec::with_capacity(ceil(len, 64));
578
579    let chunks = len / 64;
580    let remainder = len % 64;
581    buffer.extend((0..chunks).map(|chunk| {
582        let mut packed = 0;
583        for bit_idx in 0..64 {
584            let i = bit_idx + chunk * 64;
585            packed |= (f(i) as u64) << bit_idx;
586        }
587        if neg {
588            packed = !packed
589        }
590
591        packed
592    }));
593
594    if remainder != 0 {
595        let mut packed = 0;
596        for bit_idx in 0..remainder {
597            let i = bit_idx + chunks * 64;
598            packed |= (f(i) as u64) << bit_idx;
599        }
600        if neg {
601            packed = !packed
602        }
603
604        buffer.push(packed);
605    }
606    BooleanBuffer::new(buffer.into(), 0, len)
607}
608
609/// Applies `op` to possibly scalar `ArrayOrd`
610///
611/// If l is scalar `l_s` will be `Some(idx)` where `idx` is the index of the scalar value in `l`
612/// If r is scalar `r_s` will be `Some(idx)` where `idx` is the index of the scalar value in `r`
613///
614/// If `neg` is true the result of `op` will be negated
615fn apply_op<T: ArrayOrd>(
616    l: T,
617    l_s: Option<usize>,
618    r: T,
619    r_s: Option<usize>,
620    neg: bool,
621    op: impl Fn(T::Item, T::Item) -> bool,
622) -> BooleanBuffer {
623    match (l_s, r_s) {
624        (None, None) => {
625            assert_eq!(l.len(), r.len());
626            collect_bool(l.len(), neg, |idx| unsafe {
627                op(l.value_unchecked(idx), r.value_unchecked(idx))
628            })
629        }
630        (Some(l_s), Some(r_s)) => {
631            let a = l.value(l_s);
632            let b = r.value(r_s);
633            std::iter::once(op(a, b) ^ neg).collect()
634        }
635        (Some(l_s), None) => {
636            let v = l.value(l_s);
637            collect_bool(r.len(), neg, |idx| op(v, unsafe { r.value_unchecked(idx) }))
638        }
639        (None, Some(r_s)) => {
640            let v = r.value(r_s);
641            collect_bool(l.len(), neg, |idx| op(unsafe { l.value_unchecked(idx) }, v))
642        }
643    }
644}
645
646/// Applies `op` to possibly scalar `ArrayOrd` with the given indices
647fn apply_op_vectored<T: ArrayOrd>(
648    l: T,
649    l_v: &[usize],
650    r: T,
651    r_v: &[usize],
652    neg: bool,
653    op: impl Fn(T::Item, T::Item) -> bool,
654) -> BooleanBuffer {
655    assert_eq!(l_v.len(), r_v.len());
656    collect_bool(l_v.len(), neg, |idx| unsafe {
657        let l_idx = *l_v.get_unchecked(idx);
658        let r_idx = *r_v.get_unchecked(idx);
659        op(l.value_unchecked(l_idx), r.value_unchecked(r_idx))
660    })
661}
662
663/// Dispatch `op` for a REE-vs-REE comparison (no dictionary on either side)
664/// using segment-based bulk comparison.
665fn apply_ree<T: ArrayOrd>(op: Op, l: T, l_info: &ReeInfo, r: T, r_info: &ReeInfo) -> BooleanBuffer {
666    match op {
667        Op::Equal | Op::NotDistinct => apply_op_ree_segments(l, l_info, r, r_info, false, T::is_eq),
668        Op::NotEqual | Op::Distinct => apply_op_ree_segments(l, l_info, r, r_info, true, T::is_eq),
669        Op::Less => apply_op_ree_segments(l, l_info, r, r_info, false, T::is_lt),
670        Op::LessEqual => apply_op_ree_segments(r, r_info, l, l_info, true, T::is_lt),
671        Op::Greater => apply_op_ree_segments(r, r_info, l, l_info, false, T::is_lt),
672        Op::GreaterEqual => apply_op_ree_segments(l, l_info, r, r_info, true, T::is_lt),
673    }
674}
675
676/// Compare two REE arrays by walking both run_ends simultaneously, comparing
677/// once per aligned segment and bulk-filling the result.
678fn apply_op_ree_segments<T: ArrayOrd>(
679    l: T,
680    l_info: &ReeInfo,
681    r: T,
682    r_info: &ReeInfo,
683    neg: bool,
684    op: fn(T::Item, T::Item) -> bool,
685) -> BooleanBuffer {
686    let l_re = l_info.run_ends_as_usize();
687    let r_re = r_info.run_ends_as_usize();
688    let end = l_info.len;
689    let mut builder = BooleanBufferBuilder::new(l_info.len);
690    let mut l_phys = l_info.start_physical;
691    let mut r_phys = r_info.start_physical;
692    let mut pos = 0usize;
693
694    while pos < end {
695        // SAFETY: l_phys/r_phys are bounded by their respective run counts —
696        // they advance only when pos reaches a run boundary, and pos < end
697        // guarantees we haven't exhausted all runs.
698        // Subtract each side's offset to convert absolute run-ends to logical
699        // coordinates so that arrays with different offsets align correctly.
700        let l_end = (unsafe { *l_re.get_unchecked(l_phys) } - l_info.offset).min(end);
701        let r_end = (unsafe { *r_re.get_unchecked(r_phys) } - r_info.offset).min(end);
702        let seg_end = l_end.min(r_end);
703
704        let result = unsafe { op(l.value_unchecked(l_phys), r.value_unchecked(r_phys)) } ^ neg;
705        builder.append_n(seg_end - pos, result);
706
707        pos = seg_end;
708        if pos >= l_end {
709            l_phys += 1;
710        }
711        if pos >= r_end {
712            r_phys += 1;
713        }
714    }
715
716    builder.finish()
717}
718
719trait ArrayOrd {
720    type Item: Copy;
721
722    fn len(&self) -> usize;
723
724    fn value(&self, idx: usize) -> Self::Item {
725        assert!(idx < self.len());
726        unsafe { self.value_unchecked(idx) }
727    }
728
729    /// # Safety
730    ///
731    /// Safe if `idx < self.len()`
732    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item;
733
734    fn is_eq(l: Self::Item, r: Self::Item) -> bool;
735
736    fn is_lt(l: Self::Item, r: Self::Item) -> bool;
737}
738
739impl ArrayOrd for &BooleanArray {
740    type Item = bool;
741
742    fn len(&self) -> usize {
743        Array::len(self)
744    }
745
746    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
747        unsafe { BooleanArray::value_unchecked(self, idx) }
748    }
749
750    fn is_eq(l: Self::Item, r: Self::Item) -> bool {
751        l == r
752    }
753
754    fn is_lt(l: Self::Item, r: Self::Item) -> bool {
755        !l & r
756    }
757}
758
759impl<T: ArrowNativeTypeOp> ArrayOrd for &[T] {
760    type Item = T;
761
762    fn len(&self) -> usize {
763        (*self).len()
764    }
765
766    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
767        unsafe { *self.get_unchecked(idx) }
768    }
769
770    fn is_eq(l: Self::Item, r: Self::Item) -> bool {
771        l.is_eq(r)
772    }
773
774    fn is_lt(l: Self::Item, r: Self::Item) -> bool {
775        l.is_lt(r)
776    }
777}
778
779impl<'a, T: ByteArrayType> ArrayOrd for &'a GenericByteArray<T> {
780    type Item = &'a [u8];
781
782    fn len(&self) -> usize {
783        Array::len(self)
784    }
785
786    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
787        unsafe { GenericByteArray::value_unchecked(self, idx).as_ref() }
788    }
789
790    fn is_eq(l: Self::Item, r: Self::Item) -> bool {
791        l == r
792    }
793
794    fn is_lt(l: Self::Item, r: Self::Item) -> bool {
795        l < r
796    }
797}
798
799impl<'a, T: ByteViewType> ArrayOrd for &'a GenericByteViewArray<T> {
800    /// This is the item type for the GenericByteViewArray::compare
801    /// Item.0 is the array, Item.1 is the index
802    type Item = (&'a GenericByteViewArray<T>, usize);
803
804    #[inline(always)]
805    fn is_eq(l: Self::Item, r: Self::Item) -> bool {
806        let l_view = unsafe { l.0.views().get_unchecked(l.1) };
807        let r_view = unsafe { r.0.views().get_unchecked(r.1) };
808        if l.0.data_buffers().is_empty() && r.0.data_buffers().is_empty() {
809            // For eq case, we can directly compare the inlined bytes
810            return l_view == r_view;
811        }
812
813        // Fast path for same view (and both inlined)
814        if l_view == r_view && *l_view as u32 <= 12 {
815            return true;
816        }
817
818        let l_len = *l_view as u32;
819        let r_len = *r_view as u32;
820        // Lengths differ
821        if l_len != r_len {
822            return false;
823        }
824
825        // Both are empty
826        if l_len == 0 {
827            return true;
828        }
829
830        // Check prefix
831        if (*l_view >> 32) as u32 != (*r_view >> 32) as u32 {
832            return false;
833        }
834
835        // Both are inlined, and prefixes are equal (so they differ in rest of inlined bytes)
836        if l_len <= 12 {
837            return false;
838        }
839
840        // # Safety
841        // The index is within bounds as it is checked in value()
842        unsafe {
843            let l_buffer_idx = (*l_view >> 64) as u32;
844            let l_offset = (*l_view >> 96) as u32;
845            let r_buffer_idx = (*r_view >> 64) as u32;
846            let r_offset = (*r_view >> 96) as u32;
847
848            let l_data = l.0.data_buffers().get_unchecked(l_buffer_idx as usize);
849            let r_data = r.0.data_buffers().get_unchecked(r_buffer_idx as usize);
850
851            let l_slice = l_data
852                .as_slice()
853                .get_unchecked(l_offset as usize..(l_offset + l_len) as usize);
854            let r_slice = r_data
855                .as_slice()
856                .get_unchecked(r_offset as usize..(r_offset + r_len) as usize);
857            l_slice == r_slice
858        }
859    }
860
861    #[inline(always)]
862    fn is_lt(l: Self::Item, r: Self::Item) -> bool {
863        let l_view = unsafe { l.0.views().get_unchecked(l.1) };
864        let r_view = unsafe { r.0.views().get_unchecked(r.1) };
865
866        if l.0.data_buffers().is_empty() && r.0.data_buffers().is_empty() {
867            // For lt case, we can directly compare the inlined bytes
868            return GenericByteViewArray::<T>::inline_key_fast(*l_view)
869                < GenericByteViewArray::<T>::inline_key_fast(*r_view);
870        }
871
872        if (*l_view as u32) <= 12 && (*r_view as u32) <= 12 {
873            return GenericByteViewArray::<T>::inline_key_fast(*l_view)
874                < GenericByteViewArray::<T>::inline_key_fast(*r_view);
875        }
876
877        let l_prefix = (*l_view >> 32) as u32;
878        let r_prefix = (*r_view >> 32) as u32;
879        if l_prefix != r_prefix {
880            return l_prefix.swap_bytes() < r_prefix.swap_bytes();
881        }
882
883        // Fallback to the generic, unchecked comparison for mixed cases
884        // # Safety
885        // The index is within bounds as it is checked in value()
886        unsafe {
887            let l_data: &[u8] = l.0.value_unchecked(l.1).as_ref();
888            let r_data: &[u8] = r.0.value_unchecked(r.1).as_ref();
889            l_data < r_data
890        }
891    }
892
893    fn len(&self) -> usize {
894        Array::len(self)
895    }
896
897    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
898        (self, idx)
899    }
900}
901
902impl<'a> ArrayOrd for &'a FixedSizeBinaryArray {
903    type Item = &'a [u8];
904
905    fn len(&self) -> usize {
906        Array::len(self)
907    }
908
909    unsafe fn value_unchecked(&self, idx: usize) -> Self::Item {
910        unsafe { FixedSizeBinaryArray::value_unchecked(self, idx) }
911    }
912
913    fn is_eq(l: Self::Item, r: Self::Item) -> bool {
914        l == r
915    }
916
917    fn is_lt(l: Self::Item, r: Self::Item) -> bool {
918        l < r
919    }
920}
921
922/// Compares two [`GenericByteViewArray`] at index `left_idx` and `right_idx`
923///
924/// # Panics
925///
926/// Panics if `left_idx >= left.len()` or `right_idx >= right.len()`
927#[inline(always)]
928pub fn compare_byte_view<T: ByteViewType>(
929    left: &GenericByteViewArray<T>,
930    left_idx: usize,
931    right: &GenericByteViewArray<T>,
932    right_idx: usize,
933) -> Ordering {
934    assert!(left_idx < left.len());
935    assert!(right_idx < right.len());
936    if left.data_buffers().is_empty() && right.data_buffers().is_empty() {
937        let l_view = unsafe { left.views().get_unchecked(left_idx) };
938        let r_view = unsafe { right.views().get_unchecked(right_idx) };
939        return GenericByteViewArray::<T>::inline_key_fast(*l_view)
940            .cmp(&GenericByteViewArray::<T>::inline_key_fast(*r_view));
941    }
942    unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) }
943}
944
945/// Run-end encoding metadata for one side of a comparison. Holds a reference
946/// to the original REE array for deferred typed access to run_ends.
947struct ReeInfo<'a> {
948    array: &'a dyn Array,
949    offset: usize,
950    start_physical: usize,
951    len: usize,
952}
953
954impl ReeInfo<'_> {
955    /// Materialize run_ends as `Vec<usize>`.
956    fn run_ends_as_usize(&self) -> Vec<usize> {
957        let array = self.array;
958        downcast_run_array!(
959            array => array.run_ends().values().iter().map(|v| v.as_usize()).collect(),
960            _ => unreachable!()
961        )
962    }
963}
964
965/// If `array` is RunEndEncoded, return its physical values array and run metadata.
966fn ree_info_opt(array: &dyn Array) -> Option<(&dyn Array, ReeInfo<'_>)> {
967    downcast_run_array!(
968        array => {
969            let run_ends = array.run_ends();
970            let info = ReeInfo {
971                array,
972                offset: run_ends.offset(),
973                start_physical: run_ends.get_start_physical_index(),
974                len: run_ends.len(),
975            };
976            Some((array.values().as_ref(), info))
977        },
978        _ => None
979    )
980}
981
982#[cfg(test)]
983mod tests {
984    use std::sync::Arc;
985
986    use arrow_array::types::Int32Type;
987    use arrow_array::{DictionaryArray, Int32Array, Scalar, StringArray};
988    use arrow_buffer::{Buffer, ScalarBuffer};
989
990    use super::*;
991
992    #[test]
993    fn test_null_dict() {
994        let a = DictionaryArray::new(Int32Array::new_null(10), Arc::new(Int32Array::new_null(0)));
995        let r = eq(&a, &a).unwrap();
996        assert_eq!(r.null_count(), 10);
997
998        let a = DictionaryArray::new(
999            Int32Array::from(vec![1, 2, 3, 4, 5, 6]),
1000            Arc::new(Int32Array::new_null(10)),
1001        );
1002        let r = eq(&a, &a).unwrap();
1003        assert_eq!(r.null_count(), 6);
1004
1005        let scalar =
1006            DictionaryArray::new(Int32Array::new_null(1), Arc::new(Int32Array::new_null(0)));
1007        let r = eq(&a, &Scalar::new(&scalar)).unwrap();
1008        assert_eq!(r.null_count(), 6);
1009
1010        let scalar =
1011            DictionaryArray::new(Int32Array::new_null(1), Arc::new(Int32Array::new_null(0)));
1012        let r = eq(&Scalar::new(&scalar), &Scalar::new(&scalar)).unwrap();
1013        assert_eq!(r.null_count(), 1);
1014
1015        let a = DictionaryArray::new(
1016            Int32Array::from(vec![0, 1, 2]),
1017            Arc::new(Int32Array::from(vec![3, 2, 1])),
1018        );
1019        let r = eq(&a, &Scalar::new(&scalar)).unwrap();
1020        assert_eq!(r.null_count(), 3);
1021    }
1022
1023    #[test]
1024    fn is_distinct_from_non_nulls() {
1025        let left_int_array = Int32Array::from(vec![0, 1, 2, 3, 4]);
1026        let right_int_array = Int32Array::from(vec![4, 3, 2, 1, 0]);
1027
1028        assert_eq!(
1029            BooleanArray::from(vec![true, true, false, true, true,]),
1030            distinct(&left_int_array, &right_int_array).unwrap()
1031        );
1032        assert_eq!(
1033            BooleanArray::from(vec![false, false, true, false, false,]),
1034            not_distinct(&left_int_array, &right_int_array).unwrap()
1035        );
1036    }
1037
1038    #[test]
1039    fn is_distinct_from_nulls() {
1040        // [0, 0, NULL, 0, 0, 0]
1041        let left_int_array = Int32Array::new(
1042            vec![0, 0, 1, 3, 0, 0].into(),
1043            Some(NullBuffer::from(vec![true, true, false, true, true, true])),
1044        );
1045        // [0, NULL, NULL, NULL, 0, NULL]
1046        let right_int_array = Int32Array::new(
1047            vec![0; 6].into(),
1048            Some(NullBuffer::from(vec![
1049                true, false, false, false, true, false,
1050            ])),
1051        );
1052
1053        assert_eq!(
1054            BooleanArray::from(vec![false, true, false, true, false, true,]),
1055            distinct(&left_int_array, &right_int_array).unwrap()
1056        );
1057
1058        assert_eq!(
1059            BooleanArray::from(vec![true, false, true, false, true, false,]),
1060            not_distinct(&left_int_array, &right_int_array).unwrap()
1061        );
1062    }
1063
1064    #[test]
1065    fn test_distinct_scalar() {
1066        let a = Int32Array::new_scalar(12);
1067        let b = Int32Array::new_scalar(12);
1068        assert!(!distinct(&a, &b).unwrap().value(0));
1069        assert!(not_distinct(&a, &b).unwrap().value(0));
1070
1071        let a = Int32Array::new_scalar(12);
1072        let b = Int32Array::new_null(1);
1073        assert!(distinct(&a, &b).unwrap().value(0));
1074        assert!(!not_distinct(&a, &b).unwrap().value(0));
1075        assert!(distinct(&b, &a).unwrap().value(0));
1076        assert!(!not_distinct(&b, &a).unwrap().value(0));
1077
1078        let b = Scalar::new(b);
1079        assert!(distinct(&a, &b).unwrap().value(0));
1080        assert!(!not_distinct(&a, &b).unwrap().value(0));
1081
1082        assert!(!distinct(&b, &b).unwrap().value(0));
1083        assert!(not_distinct(&b, &b).unwrap().value(0));
1084
1085        let a = Int32Array::new(
1086            vec![0, 1, 2, 3].into(),
1087            Some(vec![false, false, true, true].into()),
1088        );
1089        let expected = BooleanArray::from(vec![false, false, true, true]);
1090        assert_eq!(distinct(&a, &b).unwrap(), expected);
1091        assert_eq!(distinct(&b, &a).unwrap(), expected);
1092
1093        let expected = BooleanArray::from(vec![true, true, false, false]);
1094        assert_eq!(not_distinct(&a, &b).unwrap(), expected);
1095        assert_eq!(not_distinct(&b, &a).unwrap(), expected);
1096
1097        let b = Int32Array::new_scalar(1);
1098        let expected = BooleanArray::from(vec![true; 4]);
1099        assert_eq!(distinct(&a, &b).unwrap(), expected);
1100        assert_eq!(distinct(&b, &a).unwrap(), expected);
1101        let expected = BooleanArray::from(vec![false; 4]);
1102        assert_eq!(not_distinct(&a, &b).unwrap(), expected);
1103        assert_eq!(not_distinct(&b, &a).unwrap(), expected);
1104
1105        let b = Int32Array::new_scalar(3);
1106        let expected = BooleanArray::from(vec![true, true, true, false]);
1107        assert_eq!(distinct(&a, &b).unwrap(), expected);
1108        assert_eq!(distinct(&b, &a).unwrap(), expected);
1109        let expected = BooleanArray::from(vec![false, false, false, true]);
1110        assert_eq!(not_distinct(&a, &b).unwrap(), expected);
1111        assert_eq!(not_distinct(&b, &a).unwrap(), expected);
1112    }
1113
1114    #[test]
1115    fn test_supports_distinct() {
1116        use arrow_schema::{DataType::*, Field};
1117
1118        assert!(supports_distinct(&Int32));
1119        assert!(supports_distinct(&Float64));
1120        assert!(supports_distinct(&Utf8));
1121        assert!(supports_distinct(&Boolean));
1122
1123        // One level of dictionary unwrap is supported.
1124        assert!(supports_distinct(&Dictionary(
1125            Box::new(Int16),
1126            Box::new(Utf8),
1127        )));
1128
1129        // REE, nested dictionary, and complex types are not supported.
1130        assert!(!supports_distinct(&RunEndEncoded(
1131            Arc::new(Field::new(
1132                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
1133                Int32,
1134                false
1135            )),
1136            Arc::new(Field::new(
1137                Field::REE_VALUES_FIELD_DEFAULT_NAME,
1138                Int32,
1139                true
1140            )),
1141        )));
1142        assert!(!supports_distinct(&Dictionary(
1143            Box::new(Int16),
1144            Box::new(Dictionary(Box::new(Int8), Box::new(Utf8))),
1145        )));
1146        assert!(!supports_distinct(&List(Arc::new(Field::new(
1147            "item", Int32, true,
1148        )))));
1149        assert!(!supports_distinct(&Struct(
1150            vec![Field::new("a", Int32, true)].into()
1151        )));
1152    }
1153
1154    #[test]
1155    fn test_scalar_negation() {
1156        let a = Int32Array::new_scalar(54);
1157        let b = Int32Array::new_scalar(54);
1158        let r = eq(&a, &b).unwrap();
1159        assert!(r.value(0));
1160
1161        let r = neq(&a, &b).unwrap();
1162        assert!(!r.value(0))
1163    }
1164
1165    #[test]
1166    fn test_scalar_empty() {
1167        let a = Int32Array::new_null(0);
1168        let b = Int32Array::new_scalar(23);
1169        let r = eq(&a, &b).unwrap();
1170        assert_eq!(r.len(), 0);
1171        let r = eq(&b, &a).unwrap();
1172        assert_eq!(r.len(), 0);
1173    }
1174
1175    #[test]
1176    fn test_dictionary_nulls() {
1177        let values = StringArray::from(vec![Some("us-west"), Some("us-east")]);
1178        let nulls = NullBuffer::from(vec![false, true, true]);
1179
1180        let key_values = vec![100i32, 1i32, 0i32].into();
1181        let keys = Int32Array::new(key_values, Some(nulls));
1182        let col = DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
1183
1184        neq(&col.slice(0, col.len() - 1), &col.slice(1, col.len() - 1)).unwrap();
1185    }
1186
1187    #[test]
1188    fn test_nested_dictionary_eq_returns_error() {
1189        let inner = DictionaryArray::new(
1190            Int32Array::from(vec![0, 1]),
1191            Arc::new(StringArray::from(vec!["a", "b"])),
1192        );
1193        let outer = DictionaryArray::new(Int32Array::from(vec![0, 1]), Arc::new(inner));
1194        let err = eq(&outer, &outer).expect_err("nested dictionary must not panic");
1195        assert!(
1196            matches!(err, ArrowError::InvalidArgumentError(_)),
1197            "unexpected error: {err}"
1198        );
1199    }
1200
1201    #[test]
1202    fn test_string_view_mixed_lt() {
1203        let a = arrow_array::StringViewArray::from(vec![
1204            Some("apple"),
1205            Some("apple"),
1206            Some("apple_long_string"),
1207        ]);
1208        let b = arrow_array::StringViewArray::from(vec![
1209            Some("apple_long_string"),
1210            Some("appl"),
1211            Some("apple"),
1212        ]);
1213        // "apple" < "apple_long_string" -> true
1214        // "apple" < "appl" -> false
1215        // "apple_long_string" < "apple" -> false
1216        assert_eq!(
1217            lt(&a, &b).unwrap(),
1218            BooleanArray::from(vec![true, false, false])
1219        );
1220    }
1221
1222    /// A null constant makes every row null, whatever the values are
1223    #[test]
1224    fn test_byte_view_eq_null_scalar() {
1225        let a = arrow_array::StringViewArray::from(vec![Some(""), Some("x"), None]);
1226        let scalar = arrow_array::StringViewArray::new_null(1);
1227
1228        let r = neq(&a, &Scalar::new(&scalar)).unwrap();
1229
1230        assert_eq!(r.null_count(), 3);
1231    }
1232
1233    /// A null row's view is arbitrary, so validity decides, not bytes
1234    #[test]
1235    fn test_byte_view_eq_null_row() {
1236        let a = arrow_array::StringViewArray::from(vec![Some(""), Some("x"), None]);
1237        let scalar = arrow_array::StringViewArray::from(vec![""]);
1238
1239        let r = neq(&a, &Scalar::new(&scalar)).unwrap();
1240
1241        assert!(!r.value(0));
1242        assert!(r.value(1));
1243        assert!(r.is_null(2));
1244    }
1245
1246    /// The constant may sit on either side; equality is symmetric
1247    #[test]
1248    fn test_byte_view_eq_scalar_either_side() {
1249        let a =
1250            arrow_array::StringViewArray::from(vec![Some(""), Some("xxxx"), Some("xxxxxxxxxxxxx")]);
1251        let scalar = Scalar::new(arrow_array::StringViewArray::from(vec!["xxxx"]));
1252
1253        assert_eq!(eq(&a, &scalar).unwrap(), eq(&scalar, &a).unwrap());
1254        assert_eq!(eq(&a, &scalar).unwrap().true_count(), 1);
1255    }
1256
1257    #[test]
1258    fn test_string_view_eq() {
1259        let a = arrow_array::StringViewArray::from(vec![
1260            Some("hello"),
1261            Some("world"),
1262            None,
1263            Some("very long string exceeding 12 bytes"),
1264        ]);
1265        let b = arrow_array::StringViewArray::from(vec![
1266            Some("hello"),
1267            Some("world"),
1268            None,
1269            Some("very long string exceeding 12 bytes"),
1270        ]);
1271        assert_eq!(
1272            eq(&a, &b).unwrap(),
1273            BooleanArray::from(vec![Some(true), Some(true), None, Some(true)])
1274        );
1275
1276        let c = arrow_array::StringViewArray::from(vec![
1277            Some("hello"),
1278            Some("world!"),
1279            None,
1280            Some("very long string exceeding 12 bytes!"),
1281        ]);
1282        assert_eq!(
1283            eq(&a, &c).unwrap(),
1284            BooleanArray::from(vec![Some(true), Some(false), None, Some(false)])
1285        );
1286    }
1287
1288    #[test]
1289    fn test_string_view_lt() {
1290        let a = arrow_array::StringViewArray::from(vec![
1291            Some("apple"),
1292            Some("banana"),
1293            Some("very long apple exceeding 12 bytes"),
1294            Some("very long banana exceeding 12 bytes"),
1295        ]);
1296        let b = arrow_array::StringViewArray::from(vec![
1297            Some("banana"),
1298            Some("apple"),
1299            Some("very long banana exceeding 12 bytes"),
1300            Some("very long apple exceeding 12 bytes"),
1301        ]);
1302        assert_eq!(
1303            lt(&a, &b).unwrap(),
1304            BooleanArray::from(vec![true, false, true, false])
1305        );
1306    }
1307
1308    #[test]
1309    fn test_string_view_eq_prefix_mismatch() {
1310        // Prefix mismatch should short-circuit equality for long values.
1311        let a =
1312            arrow_array::StringViewArray::from(vec![Some("very long apple exceeding 12 bytes")]);
1313        let b =
1314            arrow_array::StringViewArray::from(vec![Some("very long banana exceeding 12 bytes")]);
1315        assert_eq!(eq(&a, &b).unwrap(), BooleanArray::from(vec![Some(false)]));
1316    }
1317
1318    #[test]
1319    fn test_string_view_lt_prefix_mismatch() {
1320        // Prefix mismatch should decide ordering without full compare for long values.
1321        let a =
1322            arrow_array::StringViewArray::from(vec![Some("apple long string exceeding 12 bytes")]);
1323        let b =
1324            arrow_array::StringViewArray::from(vec![Some("banana long string exceeding 12 bytes")]);
1325        assert_eq!(lt(&a, &b).unwrap(), BooleanArray::from(vec![true]));
1326    }
1327
1328    #[test]
1329    fn test_string_view_eq_inline_fast_path() {
1330        // Inline-only arrays should compare by view equality fast path.
1331        let a = arrow_array::StringViewArray::from(vec![Some("ab")]);
1332        let b = arrow_array::StringViewArray::from(vec![Some("ab")]);
1333        assert!(!has_buffers(&a));
1334        assert!(!has_buffers(&b));
1335        assert_eq!(eq(&a, &b).unwrap(), BooleanArray::from(vec![Some(true)]));
1336    }
1337
1338    #[test]
1339    fn test_string_view_eq_inline_prefix_mismatch_with_buffers() {
1340        // Non-empty buffers force the prefix mismatch branch for inline values.
1341        let a = arrow_array::StringViewArray::from(vec![
1342            Some("ab"),
1343            Some("long string to allocate buffers"),
1344        ]);
1345        let b = arrow_array::StringViewArray::from(vec![
1346            Some("ac"),
1347            Some("long string to allocate buffers"),
1348        ]);
1349        assert!(has_buffers(&a));
1350        assert!(has_buffers(&b));
1351        assert_eq!(
1352            eq(&a, &b).unwrap(),
1353            BooleanArray::from(vec![Some(false), Some(true)])
1354        );
1355    }
1356
1357    #[test]
1358    fn test_string_view_eq_empty_len_branch() {
1359        // Reach the zero-length branch by bypassing the inline fast path with a dummy buffer.
1360        let raw_a = 0u128;
1361        let raw_b = 1u128 << 96;
1362        let views_a = ScalarBuffer::from(vec![raw_a]);
1363        let views_b = ScalarBuffer::from(vec![raw_b]);
1364        let buffers: Arc<[Buffer]> = Arc::from([Buffer::from_slice_ref([0u8])]);
1365        let a =
1366            unsafe { arrow_array::StringViewArray::new_unchecked(views_a, buffers.clone(), None) };
1367        let b = unsafe { arrow_array::StringViewArray::new_unchecked(views_b, buffers, None) };
1368        assert!(has_buffers(&a));
1369        assert!(has_buffers(&b));
1370        assert!(<&arrow_array::StringViewArray as ArrayOrd>::is_eq(
1371            (&a, 0),
1372            (&b, 0)
1373        ));
1374    }
1375
1376    #[test]
1377    fn test_string_view_long_prefix_mismatch_array_ord() {
1378        // Long strings with differing prefixes should short-circuit on prefix ordering.
1379        let a =
1380            arrow_array::StringViewArray::from(vec![Some("apple long string exceeding 12 bytes")]);
1381        let b =
1382            arrow_array::StringViewArray::from(vec![Some("banana long string exceeding 12 bytes")]);
1383        assert!(has_buffers(&a));
1384        assert!(has_buffers(&b));
1385        assert!(<&arrow_array::StringViewArray as ArrayOrd>::is_lt(
1386            (&a, 0),
1387            (&b, 0)
1388        ));
1389    }
1390
1391    #[test]
1392    fn test_string_view_inline_mismatch_array_ord() {
1393        // Long strings with differing prefixes should short-circuit on prefix ordering.
1394        let a = arrow_array::StringViewArray::from(vec![Some("ap")]);
1395        let b = arrow_array::StringViewArray::from(vec![Some("ba")]);
1396        assert!(!has_buffers(&a));
1397        assert!(!has_buffers(&b));
1398        assert!(<&arrow_array::StringViewArray as ArrayOrd>::is_lt(
1399            (&a, 0),
1400            (&b, 0)
1401        ));
1402    }
1403    #[test]
1404    fn test_compare_byte_view_inline_fast_path() {
1405        // Inline-only views should compare via inline key in compare_byte_view.
1406        let a = arrow_array::StringViewArray::from(vec![Some("ab")]);
1407        let b = arrow_array::StringViewArray::from(vec![Some("ac")]);
1408        assert!(!has_buffers(&a));
1409        assert!(!has_buffers(&b));
1410        assert_eq!(compare_byte_view(&a, 0, &b, 0), Ordering::Less);
1411    }
1412
1413    fn has_buffers<T: ByteViewType>(array: &GenericByteViewArray<T>) -> bool {
1414        !array.data_buffers().is_empty()
1415    }
1416
1417    fn ree_str(runs: &[(Option<&str>, i32)]) -> arrow_array::RunArray<Int32Type> {
1418        let mut ends = Vec::new();
1419        let mut vals = Vec::new();
1420        let mut end = 0i32;
1421        for &(v, n) in runs {
1422            end += n;
1423            ends.push(end);
1424            vals.push(v);
1425        }
1426        arrow_array::RunArray::try_new(&Int32Array::from(ends), &StringArray::from(vals)).unwrap()
1427    }
1428
1429    #[test]
1430    fn test_ree_scalar() {
1431        let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]);
1432
1433        let s = Scalar::new(StringArray::from(vec!["b"]));
1434        assert_eq!(
1435            eq(&a, &s).unwrap(),
1436            BooleanArray::from(vec![false, false, false, true, true])
1437        );
1438        assert_eq!(
1439            neq(&a, &s).unwrap(),
1440            BooleanArray::from(vec![true, true, true, false, false])
1441        );
1442        assert_eq!(
1443            lt(&a, &s).unwrap(),
1444            BooleanArray::from(vec![true, true, true, false, false])
1445        );
1446        assert_eq!(lt_eq(&a, &s).unwrap(), BooleanArray::from(vec![true; 5]));
1447        assert_eq!(gt(&a, &s).unwrap(), BooleanArray::from(vec![false; 5]));
1448        assert_eq!(
1449            gt_eq(&a, &s).unwrap(),
1450            BooleanArray::from(vec![false, false, false, true, true])
1451        );
1452
1453        // Scalar on left side
1454        let scalar = Scalar::new(ree_str(&[(Some("a"), 1)]));
1455        assert_eq!(
1456            eq(&scalar, &a).unwrap(),
1457            BooleanArray::from(vec![true, true, true, false, false])
1458        );
1459        assert_eq!(
1460            lt_eq(&scalar, &a).unwrap(),
1461            BooleanArray::from(vec![true; 5])
1462        );
1463
1464        // REE-wrapped scalar (DataFusion's ScalarValue::RunEndEncoded)
1465        assert_eq!(
1466            eq(&a, &Scalar::new(ree_str(&[(Some("a"), 1)]))).unwrap(),
1467            BooleanArray::from(vec![true, true, true, false, false]),
1468        );
1469
1470        // Single run
1471        let a = ree_str(&[(Some("x"), 100)]);
1472        let r = eq(&a, &Scalar::new(StringArray::from(vec!["x"]))).unwrap();
1473        assert_eq!(r.true_count(), 100);
1474    }
1475
1476    #[test]
1477    fn test_ree_ree() {
1478        // Different run boundaries, all ops.
1479        let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]);
1480        let b = ree_str(&[(Some("a"), 2), (Some("b"), 3)]);
1481        // a=[a,a,a,b,b] vs b=[a,a,b,b,b]
1482        assert_eq!(
1483            eq(&a, &b).unwrap(),
1484            BooleanArray::from(vec![true, true, false, true, true])
1485        );
1486        assert_eq!(
1487            neq(&a, &b).unwrap(),
1488            BooleanArray::from(vec![false, false, true, false, false])
1489        );
1490        assert_eq!(
1491            lt(&a, &b).unwrap(),
1492            BooleanArray::from(vec![false, false, true, false, false])
1493        );
1494        assert_eq!(
1495            gt_eq(&a, &b).unwrap(),
1496            BooleanArray::from(vec![true, true, false, true, true])
1497        );
1498    }
1499
1500    #[test]
1501    fn test_ree_sliced() {
1502        // Scalar with sliced REE
1503        let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]).slice(2, 3);
1504        let s = Scalar::new(StringArray::from(vec!["b"]));
1505        assert_eq!(
1506            eq(&a, &s).unwrap(),
1507            BooleanArray::from(vec![false, true, true])
1508        );
1509
1510        // Both sides sliced, REE vs REE
1511        let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]).slice(1, 4);
1512        let b = ree_str(&[(Some("a"), 2), (Some("b"), 3)]).slice(1, 4);
1513        assert_eq!(
1514            eq(&a, &b).unwrap(),
1515            BooleanArray::from(vec![true, false, true, true])
1516        );
1517    }
1518
1519    #[test]
1520    fn test_ree_sliced_different_offsets() {
1521        // left expands to ["a", "a", "b", "b"]
1522        let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]).slice(1, 4);
1523        // right expands to ["a", "a", "b", "b"]
1524        let b = ree_str(&[(Some("a"), 2), (Some("b"), 3)]).slice(0, 4);
1525        assert_eq!(
1526            eq(&a, &b).unwrap(),
1527            BooleanArray::from(vec![true, true, true, true])
1528        );
1529    }
1530
1531    #[test]
1532    fn test_ree_nullable() {
1533        let a = ree_str(&[(Some("a"), 2), (None, 1), (Some("b"), 2)]);
1534
1535        // Scalar: null-aware ops
1536        let s = Scalar::new(StringArray::from(vec!["a"]));
1537        assert_eq!(
1538            not_distinct(&a, &s).unwrap(),
1539            BooleanArray::from(vec![true, true, false, false, false])
1540        );
1541        assert_eq!(
1542            distinct(&a, &s).unwrap(),
1543            BooleanArray::from(vec![false, false, true, true, true])
1544        );
1545
1546        // REE vs REE with nulls
1547        let b = ree_str(&[(Some("a"), 3), (None, 2)]);
1548        assert_eq!(
1549            eq(&a, &b).unwrap(),
1550            BooleanArray::from(vec![Some(true), Some(true), None, None, None])
1551        );
1552    }
1553
1554    #[test]
1555    fn test_ree_mixed() {
1556        let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]);
1557
1558        // REE vs plain array
1559        let b = StringArray::from(vec!["a", "a", "b", "b", "b"]);
1560        assert_eq!(
1561            eq(&a, &b).unwrap(),
1562            BooleanArray::from(vec![true, true, false, true, true])
1563        );
1564
1565        // REE wrapping a DictionaryArray
1566        let dict = DictionaryArray::new(
1567            Int32Array::from(vec![1, 0]),
1568            Arc::new(StringArray::from(vec!["x", "y"])),
1569        );
1570        let ree_dict =
1571            arrow_array::RunArray::try_new(&Int32Array::from(vec![3, 5]), &dict).unwrap();
1572        let s = Scalar::new(StringArray::from(vec!["y"]));
1573        assert_eq!(
1574            eq(&ree_dict, &s).unwrap(),
1575            BooleanArray::from(vec![true, true, true, false, false])
1576        );
1577
1578        // Numeric REE (Int32 values)
1579        let ree_int = arrow_array::RunArray::try_new(
1580            &Int32Array::from(vec![3, 5]),
1581            &Int32Array::from(vec![10, 20]),
1582        )
1583        .unwrap();
1584        assert_eq!(
1585            eq(&ree_int, &Scalar::new(Int32Array::from(vec![10]))).unwrap(),
1586            BooleanArray::from(vec![true, true, true, false, false])
1587        );
1588
1589        // Empty REE
1590        let empty_a = ree_str(&[(Some("a"), 1)]).slice(0, 0);
1591        let empty_b = ree_str(&[(Some("b"), 1)]).slice(0, 0);
1592        assert_eq!(eq(&empty_a, &empty_b).unwrap().len(), 0);
1593    }
1594
1595    #[test]
1596    fn test_compare_byte_view() {
1597        let a = arrow_array::StringViewArray::from(vec![
1598            Some("apple"),
1599            Some("banana"),
1600            Some("very long apple exceeding 12 bytes"),
1601            Some("very long banana exceeding 12 bytes"),
1602        ]);
1603        let b = arrow_array::StringViewArray::from(vec![
1604            Some("apple"),
1605            Some("apple"),
1606            Some("very long apple exceeding 12 bytes"),
1607            Some("very long apple exceeding 12 bytes"),
1608        ]);
1609
1610        assert_eq!(compare_byte_view(&a, 0, &b, 0), Ordering::Equal);
1611        assert_eq!(compare_byte_view(&a, 1, &b, 1), Ordering::Greater);
1612        assert_eq!(compare_byte_view(&a, 2, &b, 2), Ordering::Equal);
1613        assert_eq!(compare_byte_view(&a, 3, &b, 3), Ordering::Greater);
1614    }
1615}