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