1use 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
66pub fn eq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
80 compare_op(Op::Equal, lhs, rhs)
81}
82
83pub fn neq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
97 compare_op(Op::NotEqual, lhs, rhs)
98}
99
100pub fn lt(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
114 compare_op(Op::Less, lhs, rhs)
115}
116
117pub fn lt_eq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
131 compare_op(Op::LessEqual, lhs, rhs)
132}
133
134pub fn gt(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
148 compare_op(Op::Greater, lhs, rhs)
149}
150
151pub fn gt_eq(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
165 compare_op(Op::GreaterEqual, lhs, rhs)
166}
167
168pub fn distinct(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
183 compare_op(Op::Distinct, lhs, rhs)
184}
185
186pub fn not_distinct(lhs: &dyn Datum, rhs: &dyn Datum) -> Result<BooleanArray, ArrowError> {
201 compare_op(Op::NotDistinct, lhs, rhs)
202}
203
204pub(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#[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 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 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 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 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 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 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 match is_scalar {
368 true => match op {
369 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 (None, _, None, _) => BooleanArray::new(values(), None),
390 })
391}
392
393struct 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
406const MAX_LOW_HALF_LEN: u32 = 4;
408
409fn 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 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
446fn 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; }
457
458 if !l_info.is_scalar
459 && !r_info.is_scalar
460 && (l_info.has_indirection() || r_info.has_indirection())
461 {
462 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()); 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 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
513fn 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
554fn 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 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
585fn 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
622fn 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
659fn 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
676fn 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
689fn 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 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 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 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 return l_view == r_view;
824 }
825
826 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 if l_len != r_len {
835 return false;
836 }
837
838 if l_len == 0 {
840 return true;
841 }
842
843 if (*l_view >> 32) as u32 != (*r_view >> 32) as u32 {
845 return false;
846 }
847
848 if l_len <= 12 {
850 return false;
851 }
852
853 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 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 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#[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
958struct ReeInfo<'a> {
961 array: &'a dyn Array,
962 offset: usize,
963 start_physical: usize,
964 len: usize,
965}
966
967impl ReeInfo<'_> {
968 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
978fn 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 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 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 assert!(supports_distinct(&Dictionary(
1138 Box::new(Int16),
1139 Box::new(Utf8),
1140 )));
1141
1142 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 assert_eq!(
1230 lt(&a, &b).unwrap(),
1231 BooleanArray::from(vec![true, false, false])
1232 );
1233 }
1234
1235 #[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 #[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 #[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 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 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 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 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 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 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 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 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 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 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 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 let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]);
1493 let b = ree_str(&[(Some("a"), 2), (Some("b"), 3)]);
1494 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 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 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 let a = ree_str(&[(Some("a"), 3), (Some("b"), 2)]).slice(1, 4);
1536 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 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 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 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 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 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 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}