Skip to main content

arrow_ord/
cmp.rs

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