1#![deny(rustdoc::broken_intra_doc_links)]
40#![warn(missing_docs)]
41
42use arrow_array::cast::AsArray;
43use arrow_array::types::*;
44use arrow_array::*;
45use arrow_buffer::{ArrowNativeType, NullBuffer};
46use arrow_schema::{ArrowError, DataType, SortOptions};
47use std::{cmp::Ordering, collections::HashMap};
48
49fn compare_run_end_encoded<R: RunEndIndexType>(
50 left: &dyn Array,
51 right: &dyn Array,
52 opts: SortOptions,
53) -> Result<DynComparator, ArrowError> {
54 let left = left.as_run::<R>();
55 let right = right.as_run::<R>();
56
57 let c_opts = child_opts(opts);
58 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
59
60 let l_run_ends = left.run_ends().clone();
61 let r_run_ends = right.run_ends().clone();
62
63 let f = compare(left, right, opts, move |i, j| {
64 let l_physical = l_run_ends.get_physical_index(i);
65 let r_physical = r_run_ends.get_physical_index(j);
66 cmp(l_physical, r_physical)
67 });
68 Ok(f)
69}
70
71pub type DynComparator = Box<dyn Fn(usize, usize) -> Ordering + Send + Sync>;
73
74fn child_opts(opts: SortOptions) -> SortOptions {
77 SortOptions {
78 descending: false,
79 nulls_first: opts.nulls_first != opts.descending,
80 }
81}
82
83fn compare<A, F>(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator
84where
85 A: Array + Clone,
86 F: Fn(usize, usize) -> Ordering + Send + Sync + 'static,
87{
88 let l = l.logical_nulls().filter(|x| x.null_count() > 0);
89 let r = r.logical_nulls().filter(|x| x.null_count() > 0);
90 match (opts.nulls_first, opts.descending) {
91 (true, true) => compare_impl::<true, true, _>(l, r, cmp),
92 (true, false) => compare_impl::<true, false, _>(l, r, cmp),
93 (false, true) => compare_impl::<false, true, _>(l, r, cmp),
94 (false, false) => compare_impl::<false, false, _>(l, r, cmp),
95 }
96}
97
98fn compare_impl<const NULLS_FIRST: bool, const DESCENDING: bool, F>(
99 l: Option<NullBuffer>,
100 r: Option<NullBuffer>,
101 cmp: F,
102) -> DynComparator
103where
104 F: Fn(usize, usize) -> Ordering + Send + Sync + 'static,
105{
106 let cmp = move |i, j| match DESCENDING {
107 true => cmp(i, j).reverse(),
108 false => cmp(i, j),
109 };
110
111 let (left_null, right_null) = match NULLS_FIRST {
112 true => (Ordering::Less, Ordering::Greater),
113 false => (Ordering::Greater, Ordering::Less),
114 };
115
116 match (l, r) {
117 (None, None) => Box::new(cmp),
118 (Some(l), None) => Box::new(move |i, j| match l.is_null(i) {
119 true => left_null,
120 false => cmp(i, j),
121 }),
122 (None, Some(r)) => Box::new(move |i, j| match r.is_null(j) {
123 true => right_null,
124 false => cmp(i, j),
125 }),
126 (Some(l), Some(r)) => Box::new(move |i, j| match (l.is_null(i), r.is_null(j)) {
127 (true, true) => Ordering::Equal,
128 (true, false) => left_null,
129 (false, true) => right_null,
130 (false, false) => cmp(i, j),
131 }),
132 }
133}
134
135fn compare_primitive<T: ArrowPrimitiveType>(
136 left: &dyn Array,
137 right: &dyn Array,
138 opts: SortOptions,
139) -> DynComparator
140where
141 T::Native: ArrowNativeTypeOp,
142{
143 let left = left.as_primitive::<T>();
144 let right = right.as_primitive::<T>();
145 let l_values = left.values().clone();
146 let r_values = right.values().clone();
147
148 compare(&left, &right, opts, move |i, j| {
149 l_values[i].compare(r_values[j])
150 })
151}
152
153fn compare_boolean(left: &dyn Array, right: &dyn Array, opts: SortOptions) -> DynComparator {
154 let left = left.as_boolean();
155 let right = right.as_boolean();
156
157 let l_values = left.values().clone();
158 let r_values = right.values().clone();
159
160 compare(left, right, opts, move |i, j| {
161 l_values.value(i).cmp(&r_values.value(j))
162 })
163}
164
165fn compare_bytes<T: ByteArrayType>(
166 left: &dyn Array,
167 right: &dyn Array,
168 opts: SortOptions,
169) -> DynComparator {
170 let left = left.as_bytes::<T>();
171 let right = right.as_bytes::<T>();
172
173 let l = left.clone();
174 let r = right.clone();
175 compare(left, right, opts, move |i, j| {
176 let l: &[u8] = l.value(i).as_ref();
177 let r: &[u8] = r.value(j).as_ref();
178 l.cmp(r)
179 })
180}
181
182fn compare_byte_view<T: ByteViewType>(
183 left: &dyn Array,
184 right: &dyn Array,
185 opts: SortOptions,
186) -> DynComparator {
187 let left = left.as_byte_view::<T>();
188 let right = right.as_byte_view::<T>();
189
190 let l = left.clone();
191 let r = right.clone();
192 compare(left, right, opts, move |i, j| {
193 compare_byte_view_values(&l, i, &r, j)
194 })
195}
196
197fn compare_byte_view_values<T: ByteViewType>(
198 left: &GenericByteViewArray<T>,
199 left_idx: usize,
200 right: &GenericByteViewArray<T>,
201 right_idx: usize,
202) -> Ordering {
203 assert!(left_idx < left.len());
204 assert!(right_idx < right.len());
205
206 if left.data_buffers().is_empty() && right.data_buffers().is_empty() {
207 let l_view = unsafe { left.views().get_unchecked(left_idx) };
208 let r_view = unsafe { right.views().get_unchecked(right_idx) };
209 return GenericByteViewArray::<T>::inline_key_fast(*l_view)
210 .cmp(&GenericByteViewArray::<T>::inline_key_fast(*r_view));
211 }
212
213 unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) }
214}
215
216fn compare_dict<K: ArrowDictionaryKeyType>(
217 left: &dyn Array,
218 right: &dyn Array,
219 opts: SortOptions,
220) -> Result<DynComparator, ArrowError> {
221 let left = left.as_dictionary::<K>();
222 let right = right.as_dictionary::<K>();
223 let value_dtype = left.values().data_type();
224 if value_dtype == right.values().data_type() {
225 match value_dtype {
226 DataType::Utf8 => return Ok(compare_dict_bytes::<K, Utf8Type>(left, right, opts)),
227 DataType::LargeUtf8 => {
228 return Ok(compare_dict_bytes::<K, LargeUtf8Type>(left, right, opts));
229 }
230 DataType::Binary => {
231 return Ok(compare_dict_bytes::<K, BinaryType>(left, right, opts));
232 }
233 DataType::LargeBinary => {
234 return Ok(compare_dict_bytes::<K, LargeBinaryType>(left, right, opts));
235 }
236 DataType::FixedSizeBinary(_) => {
237 return Ok(compare_dict_fixed_size_binary::<K>(left, right, opts));
238 }
239 _ => {}
240 }
241 }
242
243 let c_opts = child_opts(opts);
244 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
245 let left_keys = left.keys().values().clone();
246 let right_keys = right.keys().values().clone();
247
248 let f = compare(left, right, opts, move |i, j| {
249 let l = left_keys[i].as_usize();
250 let r = right_keys[j].as_usize();
251 cmp(l, r)
252 });
253 Ok(f)
254}
255
256fn compare_dict_bytes<K: ArrowDictionaryKeyType, T: ByteArrayType>(
257 left: &DictionaryArray<K>,
258 right: &DictionaryArray<K>,
259 opts: SortOptions,
260) -> DynComparator {
261 let left_values = left.values().as_bytes::<T>().clone();
262 let right_values = right.values().as_bytes::<T>().clone();
263 let left_keys = left.keys().values().clone();
264 let right_keys = right.keys().values().clone();
265 compare(left, right, opts, move |left_row, right_row| {
266 let left_key = left_keys[left_row].as_usize();
267 let right_key = right_keys[right_row].as_usize();
268 debug_assert!(left_key < left_values.len() && right_key < right_values.len());
269 let (left_bytes, right_bytes): (&[u8], &[u8]) = unsafe {
272 (
273 left_values.value_unchecked(left_key).as_ref(),
274 right_values.value_unchecked(right_key).as_ref(),
275 )
276 };
277 left_bytes.cmp(right_bytes)
278 })
279}
280
281fn compare_dict_fixed_size_binary<K: ArrowDictionaryKeyType>(
282 left: &DictionaryArray<K>,
283 right: &DictionaryArray<K>,
284 opts: SortOptions,
285) -> DynComparator {
286 let left_values = left.values().as_fixed_size_binary().clone();
287 let right_values = right.values().as_fixed_size_binary().clone();
288 let left_keys = left.keys().values().clone();
289 let right_keys = right.keys().values().clone();
290 compare(left, right, opts, move |left_row, right_row| {
291 let left_key = left_keys[left_row].as_usize();
292 let right_key = right_keys[right_row].as_usize();
293 debug_assert!(left_key < left_values.len() && right_key < right_values.len());
294 let (left_bytes, right_bytes) = unsafe {
297 (
298 left_values.value_unchecked(left_key),
299 right_values.value_unchecked(right_key),
300 )
301 };
302 left_bytes.cmp(right_bytes)
303 })
304}
305
306fn compare_list<O: OffsetSizeTrait>(
307 left: &dyn Array,
308 right: &dyn Array,
309 opts: SortOptions,
310) -> Result<DynComparator, ArrowError> {
311 let left = left.as_list::<O>();
312 let right = right.as_list::<O>();
313
314 let c_opts = child_opts(opts);
315 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
316
317 let l_o = left.offsets().clone();
318 let r_o = right.offsets().clone();
319 let f = compare(left, right, opts, move |i, j| {
320 let l_end = l_o[i + 1].as_usize();
321 let l_start = l_o[i].as_usize();
322
323 let r_end = r_o[j + 1].as_usize();
324 let r_start = r_o[j].as_usize();
325
326 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
327 match cmp(i, j) {
328 Ordering::Equal => {}
329 r => return r,
330 }
331 }
332 (l_end - l_start).cmp(&(r_end - r_start))
333 });
334 Ok(f)
335}
336
337fn compare_fixed_list(
338 left: &dyn Array,
339 right: &dyn Array,
340 opts: SortOptions,
341) -> Result<DynComparator, ArrowError> {
342 let left = left.as_fixed_size_list();
343 let right = right.as_fixed_size_list();
344
345 let c_opts = child_opts(opts);
346 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
347
348 let l_size = left.value_length().to_usize().unwrap();
349 let r_size = right.value_length().to_usize().unwrap();
350 let size_cmp = l_size.cmp(&r_size);
351
352 let f = compare(left, right, opts, move |i, j| {
353 let l_start = i * l_size;
354 let l_end = l_start + l_size;
355 let r_start = j * r_size;
356 let r_end = r_start + r_size;
357 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
358 match cmp(i, j) {
359 Ordering::Equal => {}
360 r => return r,
361 }
362 }
363 size_cmp
364 });
365 Ok(f)
366}
367
368fn compare_list_view<O: OffsetSizeTrait>(
369 left: &dyn Array,
370 right: &dyn Array,
371 opts: SortOptions,
372) -> Result<DynComparator, ArrowError> {
373 let left = left.as_list_view::<O>();
374 let right = right.as_list_view::<O>();
375
376 let c_opts = child_opts(opts);
377 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
378
379 let l_offsets = left.offsets().clone();
380 let l_sizes = left.sizes().clone();
381 let r_offsets = right.offsets().clone();
382 let r_sizes = right.sizes().clone();
383
384 let f = compare(left, right, opts, move |i, j| {
385 let l_start = l_offsets[i].as_usize();
386 let l_len = l_sizes[i].as_usize();
387 let l_end = l_start + l_len;
388
389 let r_start = r_offsets[j].as_usize();
390 let r_len = r_sizes[j].as_usize();
391 let r_end = r_start + r_len;
392
393 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
394 match cmp(i, j) {
395 Ordering::Equal => {}
396 r => return r,
397 }
398 }
399 l_len.cmp(&r_len)
400 });
401 Ok(f)
402}
403
404fn compare_map(
405 left: &dyn Array,
406 right: &dyn Array,
407 opts: SortOptions,
408) -> Result<DynComparator, ArrowError> {
409 let left = left.as_map();
410 let right = right.as_map();
411
412 let c_opts = child_opts(opts);
413 let cmp = make_comparator(left.entries(), right.entries(), c_opts)?;
414
415 let l_o = left.offsets().clone();
416 let r_o = right.offsets().clone();
417 let f = compare(left, right, opts, move |i, j| {
418 let l_end = l_o[i + 1].as_usize();
419 let l_start = l_o[i].as_usize();
420
421 let r_end = r_o[j + 1].as_usize();
422 let r_start = r_o[j].as_usize();
423
424 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
425 match cmp(i, j) {
426 Ordering::Equal => {}
427 r => return r,
428 }
429 }
430 (l_end - l_start).cmp(&(r_end - r_start))
431 });
432 Ok(f)
433}
434
435fn compare_struct(
436 left: &dyn Array,
437 right: &dyn Array,
438 opts: SortOptions,
439) -> Result<DynComparator, ArrowError> {
440 let left = left.as_struct();
441 let right = right.as_struct();
442
443 if left.columns().len() != right.columns().len() {
444 return Err(ArrowError::InvalidArgumentError(
445 "Cannot compare StructArray with different number of columns".to_string(),
446 ));
447 }
448
449 let c_opts = child_opts(opts);
450 let columns = left.columns().iter().zip(right.columns());
451 let comparators = columns
452 .map(|(l, r)| make_comparator(l, r, c_opts))
453 .collect::<Result<Vec<_>, _>>()?;
454
455 let f = compare(left, right, opts, move |i, j| {
456 for cmp in &comparators {
457 match cmp(i, j) {
458 Ordering::Equal => {}
459 r => return r,
460 }
461 }
462 Ordering::Equal
463 });
464 Ok(f)
465}
466
467fn compare_union(
468 left: &dyn Array,
469 right: &dyn Array,
470 opts: SortOptions,
471) -> Result<DynComparator, ArrowError> {
472 let left = left.as_union();
473 let right = right.as_union();
474
475 let DataType::Union(left_fields, left_mode) = left.data_type() else {
476 unreachable!()
477 };
478 let DataType::Union(right_fields, right_mode) = right.data_type() else {
479 unreachable!()
480 };
481
482 if left_fields != right_fields {
483 return Err(ArrowError::InvalidArgumentError(format!(
484 "Cannot compare UnionArrays with different fields: left={left_fields:?}, right={right_fields:?}"
485 )));
486 }
487
488 if left_mode != right_mode {
489 return Err(ArrowError::InvalidArgumentError(format!(
490 "Cannot compare UnionArrays with different modes: left={left_mode:?}, right={right_mode:?}"
491 )));
492 }
493
494 let c_opts = child_opts(opts);
495
496 let mut field_comparators = HashMap::with_capacity(left_fields.len());
497
498 for (type_id, _field) in left_fields.iter() {
499 let left_child = left.child(type_id);
500 let right_child = right.child(type_id);
501 let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?;
502
503 field_comparators.insert(type_id, cmp);
504 }
505
506 let left_type_ids = left.type_ids().clone();
507 let right_type_ids = right.type_ids().clone();
508
509 let left_offsets = left.offsets().cloned();
510 let right_offsets = right.offsets().cloned();
511
512 let f = compare(left, right, opts, move |i, j| {
513 let left_type_id = left_type_ids[i];
514 let right_type_id = right_type_ids[j];
515
516 match left_type_id.cmp(&right_type_id) {
517 Ordering::Equal => {
518 let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i);
519 let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j);
520
521 let cmp = field_comparators
522 .get(&left_type_id)
523 .expect("type id not found in field_comparators");
524
525 cmp(left_offset, right_offset)
526 }
527 other => other,
528 }
529 });
530 Ok(f)
531}
532
533pub fn make_comparator(
541 left: &dyn Array,
542 right: &dyn Array,
543 opts: SortOptions,
544) -> Result<DynComparator, ArrowError> {
545 use arrow_schema::DataType::*;
546
547 match (left.data_type(), right.data_type()) {
550 (Decimal32(_, s1), Decimal32(_, s2))
551 | (Decimal64(_, s1), Decimal64(_, s2))
552 | (Decimal128(_, s1), Decimal128(_, s2))
553 | (Decimal256(_, s1), Decimal256(_, s2))
554 if s1 != s2 =>
555 {
556 return Err(ArrowError::InvalidArgumentError(format!(
557 "Can't compare decimal arrays with different scales: {:?}, {:?}",
558 left.data_type(),
559 right.data_type()
560 )));
561 }
562 _ => {}
563 }
564
565 macro_rules! primitive_helper {
566 ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => {
567 Ok(compare_primitive::<$t>($left, $right, $nulls_first))
568 };
569 }
570 downcast_primitive! {
571 left.data_type(), right.data_type() => (primitive_helper, left, right, opts),
572 (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)),
573 (Utf8, Utf8) => Ok(compare_bytes::<Utf8Type>(left, right, opts)),
574 (LargeUtf8, LargeUtf8) => Ok(compare_bytes::<LargeUtf8Type>(left, right, opts)),
575 (Utf8View, Utf8View) => Ok(compare_byte_view::<StringViewType>(left, right, opts)),
576 (Binary, Binary) => Ok(compare_bytes::<BinaryType>(left, right, opts)),
577 (LargeBinary, LargeBinary) => Ok(compare_bytes::<LargeBinaryType>(left, right, opts)),
578 (BinaryView, BinaryView) => Ok(compare_byte_view::<BinaryViewType>(left, right, opts)),
579 (FixedSizeBinary(_), FixedSizeBinary(_)) => {
580 let left = left.as_fixed_size_binary();
581 let right = right.as_fixed_size_binary();
582
583 let l = left.clone();
584 let r = right.clone();
585 Ok(compare(left, right, opts, move |i, j| {
586 l.value(i).cmp(r.value(j))
587 }))
588 },
589 (List(_), List(_)) => compare_list::<i32>(left, right, opts),
590 (LargeList(_), LargeList(_)) => compare_list::<i64>(left, right, opts),
591 (ListView(_), ListView(_)) => compare_list_view::<i32>(left, right, opts),
592 (LargeListView(_), LargeListView(_)) => compare_list_view::<i64>(left, right, opts),
593 (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts),
594 (Struct(_), Struct(_)) => compare_struct(left, right, opts),
595 (Dictionary(l_key, _), Dictionary(r_key, _)) => {
596 macro_rules! dict_helper {
597 ($t:ty, $left:expr, $right:expr, $opts: expr) => {
598 compare_dict::<$t>($left, $right, $opts)
599 };
600 }
601 downcast_integer! {
602 l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts),
603 _ => unreachable!()
604 }
605 },
606 (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => {
607 macro_rules! run_end_helper {
608 ($t:ty, $left:expr, $right:expr, $opts:expr) => {
609 compare_run_end_encoded::<$t>($left, $right, $opts)
610 };
611 }
612 downcast_run_end_index! {
613 l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts),
614 _ => Err(ArrowError::InvalidArgumentError(format!(
615 "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}",
616 l_run_ends.data_type(),
617 r_run_ends.data_type()
618 )))
619 }
620 },
621 (Map(_, _), Map(_, _)) => compare_map(left, right, opts),
622 (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)),
623 (Union(_, _), Union(_, _)) => compare_union(left, right, opts),
624 (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs {
625 true => format!("The data type type {lhs:?} has no natural order"),
626 false => "Can't compare arrays of different types".to_string(),
627 }))
628 }
629}
630
631#[cfg(test)]
632mod tests {
633 use super::*;
634 use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder};
635 use arrow_buffer::{IntervalDayTime, NullBuffer, OffsetBuffer, ScalarBuffer, i256};
636 use arrow_schema::{ArrowError, DataType, Field, Fields, UnionFields};
637 use half::f16;
638 use std::cmp::Ordering;
639 use std::sync::Arc;
640
641 #[test]
642 fn test_fixed_size_binary() {
643 let items = vec![vec![1u8], vec![2u8]];
644 let array = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
645
646 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
647
648 assert_eq!(Ordering::Less, cmp(0, 1));
649 }
650
651 #[test]
652 fn test_fixed_size_binary_fixed_size_binary() {
653 let items = vec![vec![1u8]];
654 let array1 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
655 let items = vec![vec![2u8]];
656 let array2 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
657
658 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
659
660 assert_eq!(Ordering::Less, cmp(0, 0));
661 }
662
663 #[test]
664 fn test_i32() {
665 let array = Int32Array::from(vec![1, 2]);
666
667 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
668
669 assert_eq!(Ordering::Less, (cmp)(0, 1));
670 }
671
672 #[test]
673 fn test_i32_i32() {
674 let array1 = Int32Array::from(vec![1]);
675 let array2 = Int32Array::from(vec![2]);
676
677 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
678
679 assert_eq!(Ordering::Less, cmp(0, 0));
680 }
681
682 #[cfg_attr(miri, ignore)] #[test]
684 fn test_f16() {
685 let array = Float16Array::from(vec![f16::from_f32(1.0), f16::from_f32(2.0)]);
686
687 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
688
689 assert_eq!(Ordering::Less, cmp(0, 1));
690 }
691
692 #[test]
693 fn test_f64() {
694 let array = Float64Array::from(vec![1.0, 2.0]);
695
696 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
697
698 assert_eq!(Ordering::Less, cmp(0, 1));
699 }
700
701 #[test]
702 fn test_f64_nan() {
703 let array = Float64Array::from(vec![1.0, f64::NAN]);
704
705 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
706
707 assert_eq!(Ordering::Less, cmp(0, 1));
708 assert_eq!(Ordering::Equal, cmp(1, 1));
709 }
710
711 #[test]
712 fn test_f64_zeros() {
713 let array = Float64Array::from(vec![-0.0, 0.0]);
714
715 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
716
717 assert_eq!(Ordering::Less, cmp(0, 1));
718 assert_eq!(Ordering::Greater, cmp(1, 0));
719 }
720
721 #[test]
722 fn test_interval_day_time() {
723 let array = IntervalDayTimeArray::from(vec![
724 IntervalDayTimeType::make_value(0, 1000),
726 IntervalDayTimeType::make_value(1, 2),
728 IntervalDayTimeType::make_value(0, 90_000_000),
730 ]);
731
732 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
733
734 assert_eq!(Ordering::Less, cmp(0, 1));
735 assert_eq!(Ordering::Greater, cmp(1, 0));
736
737 assert_eq!(Ordering::Greater, cmp(1, 2));
741 assert_eq!(Ordering::Less, cmp(2, 1));
742 }
743
744 #[test]
745 fn test_interval_year_month() {
746 let array = IntervalYearMonthArray::from(vec![
747 IntervalYearMonthType::make_value(1, 0),
749 IntervalYearMonthType::make_value(0, 13),
751 IntervalYearMonthType::make_value(1, 1),
753 ]);
754
755 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
756
757 assert_eq!(Ordering::Less, cmp(0, 1));
758 assert_eq!(Ordering::Greater, cmp(1, 0));
759
760 assert_eq!(Ordering::Equal, cmp(1, 2));
762 assert_eq!(Ordering::Equal, cmp(2, 1));
763 }
764
765 #[test]
766 fn test_interval_month_day_nano() {
767 let array = IntervalMonthDayNanoArray::from(vec![
768 IntervalMonthDayNanoType::make_value(0, 100, 0),
770 IntervalMonthDayNanoType::make_value(1, 0, 0),
772 IntervalMonthDayNanoType::make_value(0, 100, 2),
774 ]);
775
776 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
777
778 assert_eq!(Ordering::Less, cmp(0, 1));
779 assert_eq!(Ordering::Greater, cmp(1, 0));
780
781 assert_eq!(Ordering::Greater, cmp(1, 2));
785 assert_eq!(Ordering::Less, cmp(2, 1));
786 }
787
788 #[test]
789 fn test_decimali32() {
790 let array = vec![Some(5_i32), Some(2_i32), Some(3_i32)]
791 .into_iter()
792 .collect::<Decimal32Array>()
793 .with_precision_and_scale(8, 6)
794 .unwrap();
795
796 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
797 assert_eq!(Ordering::Less, cmp(1, 0));
798 assert_eq!(Ordering::Greater, cmp(0, 2));
799 }
800
801 #[test]
802 fn test_decimali64() {
803 let array = vec![Some(5_i64), Some(2_i64), Some(3_i64)]
804 .into_iter()
805 .collect::<Decimal64Array>()
806 .with_precision_and_scale(16, 6)
807 .unwrap();
808
809 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
810 assert_eq!(Ordering::Less, cmp(1, 0));
811 assert_eq!(Ordering::Greater, cmp(0, 2));
812 }
813
814 #[test]
815 fn test_decimali128() {
816 let array = vec![Some(5_i128), Some(2_i128), Some(3_i128)]
817 .into_iter()
818 .collect::<Decimal128Array>()
819 .with_precision_and_scale(23, 6)
820 .unwrap();
821
822 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
823 assert_eq!(Ordering::Less, cmp(1, 0));
824 assert_eq!(Ordering::Greater, cmp(0, 2));
825 }
826
827 #[test]
828 fn test_decimali256() {
829 let array = vec![
830 Some(i256::from_i128(5_i128)),
831 Some(i256::from_i128(2_i128)),
832 Some(i256::from_i128(3_i128)),
833 ]
834 .into_iter()
835 .collect::<Decimal256Array>()
836 .with_precision_and_scale(53, 6)
837 .unwrap();
838
839 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
840 assert_eq!(Ordering::Less, cmp(1, 0));
841 assert_eq!(Ordering::Greater, cmp(0, 2));
842 }
843
844 #[test]
845 fn test_dict() {
846 let data = vec!["a", "b", "c", "a", "a", "c", "c"];
847 let array = data.into_iter().collect::<DictionaryArray<Int16Type>>();
848
849 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
850
851 assert_eq!(Ordering::Less, cmp(0, 1));
852 assert_eq!(Ordering::Equal, cmp(3, 4));
853 assert_eq!(Ordering::Greater, cmp(2, 3));
854 }
855
856 #[test]
857 fn test_multiple_dict() {
858 let d1 = vec!["a", "b", "c", "d"];
859 let a1 = d1.into_iter().collect::<DictionaryArray<Int16Type>>();
860 let d2 = vec!["e", "f", "g", "a"];
861 let a2 = d2.into_iter().collect::<DictionaryArray<Int16Type>>();
862
863 let cmp = make_comparator(&a1, &a2, SortOptions::default()).unwrap();
864
865 assert_eq!(Ordering::Less, cmp(0, 0));
866 assert_eq!(Ordering::Equal, cmp(0, 3));
867 assert_eq!(Ordering::Greater, cmp(1, 3));
868 }
869
870 #[test]
871 fn dict_fixed_size_binary_fast_path_derefs_values_both_directions() {
872 let shared_keys = || Int32Array::from(vec![0_i32, 1]);
873 let left_values =
874 FixedSizeBinaryArray::try_from_iter([b"zz".as_ref(), b"aa".as_ref()].into_iter())
875 .unwrap();
876 let right_values =
877 FixedSizeBinaryArray::try_from_iter([b"aa".as_ref(), b"zz".as_ref()].into_iter())
878 .unwrap();
879 let left = DictionaryArray::<Int32Type>::new(shared_keys(), Arc::new(left_values));
880 let right = DictionaryArray::<Int32Type>::new(shared_keys(), Arc::new(right_values));
881
882 let ascending = make_comparator(&left, &right, SortOptions::default()).unwrap();
883 assert_eq!(ascending(0, 0), Ordering::Greater);
884 assert_eq!(ascending(1, 1), Ordering::Less);
885
886 let descending = make_comparator(
887 &left,
888 &right,
889 SortOptions {
890 descending: true,
891 nulls_first: true,
892 },
893 )
894 .unwrap();
895 assert_eq!(descending(0, 0), Ordering::Less);
896 assert_eq!(descending(1, 1), Ordering::Greater);
897 }
898
899 fn dict_bytes_fast_path_derefs_values_both_directions_impl<T: ByteArrayType>() {
900 let shared_keys = || Int32Array::from(vec![0_i32, 1]);
901 let left_values =
902 GenericByteArray::<T>::new(OffsetBuffer::from_lengths([1, 1]), b"za".into(), None);
903 let right_values =
904 GenericByteArray::<T>::new(OffsetBuffer::from_lengths([1, 1]), b"az".into(), None);
905 let left = DictionaryArray::<Int32Type>::new(shared_keys(), Arc::new(left_values));
906 let right = DictionaryArray::<Int32Type>::new(shared_keys(), Arc::new(right_values));
907
908 let ascending = make_comparator(&left, &right, SortOptions::default()).unwrap();
909 assert_eq!(ascending(0, 0), Ordering::Greater);
910 assert_eq!(ascending(1, 1), Ordering::Less);
911
912 let descending = make_comparator(
913 &left,
914 &right,
915 SortOptions {
916 descending: true,
917 nulls_first: true,
918 },
919 )
920 .unwrap();
921 assert_eq!(descending(0, 0), Ordering::Less);
922 assert_eq!(descending(1, 1), Ordering::Greater);
923 }
924
925 #[test]
926 fn dict_bytes_fast_path_derefs_values_both_directions() {
927 dict_bytes_fast_path_derefs_values_both_directions_impl::<Utf8Type>();
928 dict_bytes_fast_path_derefs_values_both_directions_impl::<LargeUtf8Type>();
929 dict_bytes_fast_path_derefs_values_both_directions_impl::<BinaryType>();
930 dict_bytes_fast_path_derefs_values_both_directions_impl::<LargeBinaryType>();
931 }
932
933 fn dict_bytes_handles_key_and_value_nulls_impl<T: ByteArrayType>() {
934 let left_keys = Int32Array::from(vec![Some(0_i32), None, Some(1)]);
936 let left_values = GenericByteArray::<T>::new(
937 OffsetBuffer::from_lengths([5, 6]),
938 b"applebanana".into(),
939 None,
940 );
941 let left = DictionaryArray::<Int32Type>::new(left_keys, Arc::new(left_values));
942
943 let right_keys = Int32Array::from(vec![Some(0_i32), Some(1), Some(0)]);
945 let right_values = GenericByteArray::<T>::new(
946 OffsetBuffer::from_lengths([5, 0]),
947 b"apple".into(),
948 Some(NullBuffer::from(vec![true, false])),
949 );
950 let right = DictionaryArray::<Int32Type>::new(right_keys, Arc::new(right_values));
951
952 let cmp = make_comparator(&left, &right, SortOptions::default()).unwrap();
953 assert_eq!(cmp(0, 0), Ordering::Equal); assert_eq!(cmp(1, 0), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Greater); assert_eq!(cmp(1, 1), Ordering::Equal); }
958
959 #[test]
960 fn dict_bytes_handles_key_and_value_nulls() {
961 dict_bytes_handles_key_and_value_nulls_impl::<Utf8Type>();
962 dict_bytes_handles_key_and_value_nulls_impl::<LargeUtf8Type>();
963 dict_bytes_handles_key_and_value_nulls_impl::<BinaryType>();
964 dict_bytes_handles_key_and_value_nulls_impl::<LargeBinaryType>();
965 }
966
967 fn dict_bytes_sort_options_with_nulls_impl<T: ByteArrayType>() {
968 let null_row_on_left = Int32Array::from(vec![Some(0_i32), None, Some(1)]);
970 let left_values = GenericByteArray::<T>::new(
971 OffsetBuffer::from_lengths([5, 6]),
972 b"applebanana".into(),
973 None,
974 );
975 let left = DictionaryArray::<Int32Type>::new(null_row_on_left, Arc::new(left_values));
976
977 let keys_pointing_at_null = Int32Array::from(vec![Some(0_i32), Some(1), Some(0)]);
978 let right_values = GenericByteArray::<T>::new(
979 OffsetBuffer::from_lengths([6, 0]),
980 b"cherry".into(),
981 Some(NullBuffer::from(vec![true, false])),
982 );
983 let right =
984 DictionaryArray::<Int32Type>::new(keys_pointing_at_null, Arc::new(right_values));
985
986 let asc_nulls_last = SortOptions {
987 descending: false,
988 nulls_first: false,
989 };
990 let desc_nulls_first = SortOptions {
991 descending: true,
992 nulls_first: true,
993 };
994 let desc_nulls_last = SortOptions {
995 descending: true,
996 nulls_first: false,
997 };
998
999 let cmp = make_comparator(&left, &right, asc_nulls_last).unwrap();
1000 assert_eq!(cmp(0, 0), Ordering::Less); assert_eq!(cmp(1, 0), Ordering::Greater); assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Equal);
1004
1005 let cmp = make_comparator(&left, &right, desc_nulls_first).unwrap();
1006 assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 0), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Greater); assert_eq!(cmp(1, 1), Ordering::Equal);
1010
1011 let cmp = make_comparator(&left, &right, desc_nulls_last).unwrap();
1012 assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 0), Ordering::Greater); assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Equal);
1016 }
1017
1018 #[test]
1019 fn dict_bytes_sort_options_with_nulls() {
1020 dict_bytes_sort_options_with_nulls_impl::<Utf8Type>();
1021 dict_bytes_sort_options_with_nulls_impl::<LargeUtf8Type>();
1022 dict_bytes_sort_options_with_nulls_impl::<BinaryType>();
1023 dict_bytes_sort_options_with_nulls_impl::<LargeBinaryType>();
1024 }
1025
1026 fn dict_bytes_fast_path_panics_on_out_of_bounds_row_impl<T: ByteArrayType>() {
1027 let two_row_keys = Int32Array::from(vec![0_i32, 1]);
1028 let values =
1029 GenericByteArray::<T>::new(OffsetBuffer::from_lengths([1, 1]), b"ab".into(), None);
1030 let dict = DictionaryArray::<Int32Type>::new(two_row_keys, Arc::new(values));
1031
1032 let cmp = make_comparator(&dict, &dict, SortOptions::default()).unwrap();
1033 cmp(2, 0);
1034 }
1035
1036 #[test]
1037 #[should_panic(expected = "index out of bounds")]
1038 fn dict_utf8_fast_path_panics_on_out_of_bounds_row() {
1039 dict_bytes_fast_path_panics_on_out_of_bounds_row_impl::<Utf8Type>();
1040 }
1041
1042 #[test]
1043 #[should_panic(expected = "index out of bounds")]
1044 fn dict_large_utf8_fast_path_panics_on_out_of_bounds_row() {
1045 dict_bytes_fast_path_panics_on_out_of_bounds_row_impl::<LargeUtf8Type>();
1046 }
1047
1048 #[test]
1049 #[should_panic(expected = "index out of bounds")]
1050 fn dict_binary_fast_path_panics_on_out_of_bounds_row() {
1051 dict_bytes_fast_path_panics_on_out_of_bounds_row_impl::<BinaryType>();
1052 }
1053
1054 #[test]
1055 #[should_panic(expected = "index out of bounds")]
1056 fn dict_large_binary_fast_path_panics_on_out_of_bounds_row() {
1057 dict_bytes_fast_path_panics_on_out_of_bounds_row_impl::<LargeBinaryType>();
1058 }
1059
1060 #[test]
1061 #[should_panic(expected = "index out of bounds")]
1062 fn dict_fixed_size_binary_fast_path_panics_on_out_of_bounds_row() {
1063 let two_row_keys = Int32Array::from(vec![0_i32, 1]);
1064 let values = Arc::new(
1065 FixedSizeBinaryArray::try_from_iter([b"aa".as_ref(), b"bb".as_ref()].into_iter())
1066 .unwrap(),
1067 );
1068 let dict = DictionaryArray::<Int32Type>::new(two_row_keys, values);
1069
1070 let cmp = make_comparator(&dict, &dict, SortOptions::default()).unwrap();
1071 cmp(2, 0);
1072 }
1073
1074 #[test]
1075 fn dict_fixed_size_binary_sort_options_with_nulls() {
1076 let null_row_on_left = Int32Array::from(vec![Some(0_i32), None, Some(1)]);
1077 let left_values =
1078 FixedSizeBinaryArray::try_from_iter([b"aa".as_ref(), b"bb".as_ref()].into_iter())
1079 .unwrap();
1080 let left = DictionaryArray::<Int32Type>::new(null_row_on_left, Arc::new(left_values));
1081
1082 let keys_pointing_at_null = Int32Array::from(vec![Some(0_i32), Some(1), Some(0)]);
1083 let right_values = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1084 [Some(b"cc".as_ref()), None].into_iter(),
1085 2,
1086 )
1087 .unwrap();
1088 let right =
1089 DictionaryArray::<Int32Type>::new(keys_pointing_at_null, Arc::new(right_values));
1090
1091 let asc_nulls_last = SortOptions {
1092 descending: false,
1093 nulls_first: false,
1094 };
1095 let desc_nulls_first = SortOptions {
1096 descending: true,
1097 nulls_first: true,
1098 };
1099 let desc_nulls_last = SortOptions {
1100 descending: true,
1101 nulls_first: false,
1102 };
1103
1104 let cmp = make_comparator(&left, &right, asc_nulls_last).unwrap();
1105 assert_eq!(cmp(0, 0), Ordering::Less); assert_eq!(cmp(1, 0), Ordering::Greater); assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Equal);
1109
1110 let cmp = make_comparator(&left, &right, desc_nulls_first).unwrap();
1111 assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 0), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Greater); assert_eq!(cmp(1, 1), Ordering::Equal);
1115
1116 let cmp = make_comparator(&left, &right, desc_nulls_last).unwrap();
1117 assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 0), Ordering::Greater); assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Equal);
1121 }
1122
1123 #[test]
1124 fn test_primitive_dict() {
1125 let values = Int32Array::from(vec![1_i32, 0, 2, 5]);
1126 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1127 let array1 = DictionaryArray::new(keys, Arc::new(values));
1128
1129 let values = Int32Array::from(vec![2_i32, 3, 4, 5]);
1130 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1131 let array2 = DictionaryArray::new(keys, Arc::new(values));
1132
1133 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1134
1135 assert_eq!(Ordering::Less, cmp(0, 0));
1136 assert_eq!(Ordering::Less, cmp(0, 3));
1137 assert_eq!(Ordering::Equal, cmp(3, 3));
1138 assert_eq!(Ordering::Greater, cmp(3, 1));
1139 assert_eq!(Ordering::Greater, cmp(3, 2));
1140 }
1141
1142 #[test]
1143 fn test_float_dict() {
1144 let values = Float32Array::from(vec![1.0, 0.5, 2.1, 5.5]);
1145 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1146 let array1 = DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
1147
1148 let values = Float32Array::from(vec![1.2, 3.2, 4.0, 5.5]);
1149 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1150 let array2 = DictionaryArray::new(keys, Arc::new(values));
1151
1152 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1153
1154 assert_eq!(Ordering::Less, cmp(0, 0));
1155 assert_eq!(Ordering::Less, cmp(0, 3));
1156 assert_eq!(Ordering::Equal, cmp(3, 3));
1157 assert_eq!(Ordering::Greater, cmp(3, 1));
1158 assert_eq!(Ordering::Greater, cmp(3, 2));
1159 }
1160
1161 #[test]
1162 fn test_timestamp_dict() {
1163 let values = TimestampSecondArray::from(vec![1, 0, 2, 5]);
1164 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1165 let array1 = DictionaryArray::new(keys, Arc::new(values));
1166
1167 let values = TimestampSecondArray::from(vec![2, 3, 4, 5]);
1168 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1169 let array2 = DictionaryArray::new(keys, Arc::new(values));
1170
1171 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1172
1173 assert_eq!(Ordering::Less, cmp(0, 0));
1174 assert_eq!(Ordering::Less, cmp(0, 3));
1175 assert_eq!(Ordering::Equal, cmp(3, 3));
1176 assert_eq!(Ordering::Greater, cmp(3, 1));
1177 assert_eq!(Ordering::Greater, cmp(3, 2));
1178 }
1179
1180 #[test]
1181 fn test_interval_dict() {
1182 let v1 = IntervalDayTime::new(0, 1);
1183 let v2 = IntervalDayTime::new(0, 2);
1184 let v3 = IntervalDayTime::new(12, 2);
1185
1186 let values = IntervalDayTimeArray::from(vec![Some(v1), Some(v2), None, Some(v3)]);
1187 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1188 let array1 = DictionaryArray::new(keys, Arc::new(values));
1189
1190 let values = IntervalDayTimeArray::from(vec![Some(v3), Some(v2), None, Some(v1)]);
1191 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1192 let array2 = DictionaryArray::new(keys, Arc::new(values));
1193
1194 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1195
1196 assert_eq!(Ordering::Less, cmp(0, 0)); assert_eq!(Ordering::Equal, cmp(0, 3)); assert_eq!(Ordering::Greater, cmp(3, 3)); assert_eq!(Ordering::Greater, cmp(3, 1)); assert_eq!(Ordering::Greater, cmp(3, 2)); }
1202
1203 #[test]
1204 fn test_duration_dict() {
1205 let values = DurationSecondArray::from(vec![1, 0, 2, 5]);
1206 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1207 let array1 = DictionaryArray::new(keys, Arc::new(values));
1208
1209 let values = DurationSecondArray::from(vec![2, 3, 4, 5]);
1210 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1211 let array2 = DictionaryArray::new(keys, Arc::new(values));
1212
1213 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1214
1215 assert_eq!(Ordering::Less, cmp(0, 0));
1216 assert_eq!(Ordering::Less, cmp(0, 3));
1217 assert_eq!(Ordering::Equal, cmp(3, 3));
1218 assert_eq!(Ordering::Greater, cmp(3, 1));
1219 assert_eq!(Ordering::Greater, cmp(3, 2));
1220 }
1221
1222 #[test]
1223 fn test_decimal_dict() {
1224 let values = Decimal128Array::from(vec![1, 0, 2, 5]);
1225 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1226 let array1 = DictionaryArray::new(keys, Arc::new(values));
1227
1228 let values = Decimal128Array::from(vec![2, 3, 4, 5]);
1229 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1230 let array2 = DictionaryArray::new(keys, Arc::new(values));
1231
1232 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1233
1234 assert_eq!(Ordering::Less, cmp(0, 0));
1235 assert_eq!(Ordering::Less, cmp(0, 3));
1236 assert_eq!(Ordering::Equal, cmp(3, 3));
1237 assert_eq!(Ordering::Greater, cmp(3, 1));
1238 assert_eq!(Ordering::Greater, cmp(3, 2));
1239 }
1240
1241 #[test]
1242 fn test_decimal256_dict() {
1243 let values = Decimal256Array::from(vec![
1244 i256::from_i128(1),
1245 i256::from_i128(0),
1246 i256::from_i128(2),
1247 i256::from_i128(5),
1248 ]);
1249 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
1250 let array1 = DictionaryArray::new(keys, Arc::new(values));
1251
1252 let values = Decimal256Array::from(vec![
1253 i256::from_i128(2),
1254 i256::from_i128(3),
1255 i256::from_i128(4),
1256 i256::from_i128(5),
1257 ]);
1258 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
1259 let array2 = DictionaryArray::new(keys, Arc::new(values));
1260
1261 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1262
1263 assert_eq!(Ordering::Less, cmp(0, 0));
1264 assert_eq!(Ordering::Less, cmp(0, 3));
1265 assert_eq!(Ordering::Equal, cmp(3, 3));
1266 assert_eq!(Ordering::Greater, cmp(3, 1));
1267 assert_eq!(Ordering::Greater, cmp(3, 2));
1268 }
1269
1270 #[test]
1271 fn test_decimal_different_scale_errors() {
1272 let a = Decimal128Array::from(vec![100])
1274 .with_precision_and_scale(10, 2)
1275 .unwrap();
1276 let b = Decimal128Array::from(vec![500])
1277 .with_precision_and_scale(10, 3)
1278 .unwrap();
1279 let err = make_comparator(&a, &b, SortOptions::default())
1280 .err()
1281 .unwrap();
1282 assert_eq!(
1283 err.to_string(),
1284 "Invalid argument error: Can't compare decimal arrays with different scales: \
1285 Decimal128(10, 2), Decimal128(10, 3)"
1286 );
1287
1288 let keys = Int8Array::from_iter_values([0]);
1290 let dict_a = DictionaryArray::new(keys.clone(), Arc::new(a.clone()));
1291 let dict_b = DictionaryArray::new(keys, Arc::new(b));
1292 assert!(make_comparator(&dict_a, &dict_b, SortOptions::default()).is_err());
1293
1294 let c = Decimal128Array::from(vec![100])
1296 .with_precision_and_scale(12, 2)
1297 .unwrap();
1298 let cmp = make_comparator(&a, &c, SortOptions::default()).unwrap();
1299 assert_eq!(Ordering::Equal, cmp(0, 0));
1300
1301 let a = Decimal32Array::from(vec![1])
1303 .with_precision_and_scale(5, 1)
1304 .unwrap();
1305 let b = Decimal32Array::from(vec![1])
1306 .with_precision_and_scale(5, 2)
1307 .unwrap();
1308 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
1309
1310 let a = Decimal64Array::from(vec![1])
1311 .with_precision_and_scale(10, 1)
1312 .unwrap();
1313 let b = Decimal64Array::from(vec![1])
1314 .with_precision_and_scale(10, 2)
1315 .unwrap();
1316 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
1317
1318 let a = Decimal256Array::from(vec![i256::from_i128(1)])
1319 .with_precision_and_scale(40, 1)
1320 .unwrap();
1321 let b = Decimal256Array::from(vec![i256::from_i128(1)])
1322 .with_precision_and_scale(40, 2)
1323 .unwrap();
1324 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
1325 }
1326
1327 fn test_bytes_impl<T: ByteArrayType>() {
1328 let offsets = OffsetBuffer::from_lengths([3, 3, 1]);
1329 let a = GenericByteArray::<T>::new(offsets, b"abcdefa".into(), None);
1330 let cmp = make_comparator(&a, &a, SortOptions::default()).unwrap();
1331
1332 assert_eq!(Ordering::Less, cmp(0, 1));
1333 assert_eq!(Ordering::Greater, cmp(0, 2));
1334 assert_eq!(Ordering::Equal, cmp(1, 1));
1335 }
1336
1337 #[test]
1338 fn test_bytes() {
1339 test_bytes_impl::<Utf8Type>();
1340 test_bytes_impl::<LargeUtf8Type>();
1341 test_bytes_impl::<BinaryType>();
1342 test_bytes_impl::<LargeBinaryType>();
1343 }
1344
1345 fn assert_cmp_cases<A: Array>(
1346 array1: &A,
1347 array2: &A,
1348 opts: SortOptions,
1349 cases: &[(usize, usize, Ordering)],
1350 ) {
1351 let cmp = make_comparator(array1, array2, opts).unwrap();
1352 for (left, right, expected) in cases {
1353 assert_eq!(cmp(*left, *right), *expected);
1354 }
1355 }
1356
1357 #[test]
1358 fn test_lists() {
1359 let mut a = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
1360 a.extend([
1361 Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]),
1362 Some(vec![
1363 Some(vec![Some(1), Some(2), Some(3)]),
1364 Some(vec![Some(1)]),
1365 ]),
1366 Some(vec![]),
1367 ]);
1368 let a = a.finish();
1369 let mut b = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
1370 b.extend([
1371 Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]),
1372 Some(vec![
1373 Some(vec![Some(1), Some(2), None]),
1374 Some(vec![Some(1)]),
1375 ]),
1376 Some(vec![
1377 Some(vec![Some(1), Some(2), Some(3), Some(4)]),
1378 Some(vec![Some(1)]),
1379 ]),
1380 None,
1381 ]);
1382 let b = b.finish();
1383
1384 assert_cmp_cases(
1386 &a,
1387 &b,
1388 SortOptions {
1389 descending: false,
1390 nulls_first: true,
1391 },
1392 &[
1393 (0, 0, Ordering::Equal),
1394 (0, 1, Ordering::Less),
1395 (0, 2, Ordering::Less),
1396 (1, 2, Ordering::Less),
1397 (1, 3, Ordering::Greater),
1398 (2, 0, Ordering::Less),
1399 ],
1400 );
1401
1402 assert_cmp_cases(
1404 &a,
1405 &b,
1406 SortOptions {
1407 descending: true,
1408 nulls_first: true,
1409 },
1410 &[
1411 (0, 0, Ordering::Equal),
1412 (0, 1, Ordering::Less),
1413 (0, 2, Ordering::Less),
1414 (1, 2, Ordering::Greater),
1415 (1, 3, Ordering::Greater),
1416 (2, 0, Ordering::Greater),
1417 ],
1418 );
1419
1420 assert_cmp_cases(
1422 &a,
1423 &b,
1424 SortOptions {
1425 descending: true,
1426 nulls_first: false,
1427 },
1428 &[
1429 (0, 0, Ordering::Equal),
1430 (0, 1, Ordering::Greater),
1431 (0, 2, Ordering::Greater),
1432 (1, 2, Ordering::Greater),
1433 (1, 3, Ordering::Less),
1434 (2, 0, Ordering::Greater),
1435 ],
1436 );
1437
1438 assert_cmp_cases(
1440 &a,
1441 &b,
1442 SortOptions {
1443 descending: false,
1444 nulls_first: false,
1445 },
1446 &[
1447 (0, 0, Ordering::Equal),
1448 (0, 1, Ordering::Greater),
1449 (0, 2, Ordering::Greater),
1450 (1, 2, Ordering::Less),
1451 (1, 3, Ordering::Less),
1452 (2, 0, Ordering::Less),
1453 ],
1454 );
1455 }
1456
1457 fn list_view_array<O: OffsetSizeTrait>(
1458 values: Vec<i32>,
1459 offsets: &[usize],
1460 sizes: &[usize],
1461 valid: Option<&[bool]>,
1462 ) -> GenericListViewArray<O> {
1463 let offsets = offsets
1464 .iter()
1465 .map(|v| O::from_usize(*v).unwrap())
1466 .collect::<ScalarBuffer<O>>();
1467 let sizes = sizes
1468 .iter()
1469 .map(|v| O::from_usize(*v).unwrap())
1470 .collect::<ScalarBuffer<O>>();
1471 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
1472 let values = Int32Array::from(values);
1473 let nulls = valid.map(NullBuffer::from);
1474 GenericListViewArray::new(field, offsets, sizes, Arc::new(values), nulls)
1475 }
1476
1477 fn test_list_view_comparisons<O: OffsetSizeTrait>() {
1478 let array = list_view_array::<O>(
1479 vec![1, 2, 3, 4, 5],
1480 &[0, 2, 1, 0, 3],
1481 &[2, 2, 2, 0, 2],
1482 Some(&[true, true, true, true, false]),
1483 );
1484
1485 assert_cmp_cases(
1487 &array,
1488 &array,
1489 SortOptions {
1490 descending: false,
1491 nulls_first: true,
1492 },
1493 &[
1494 (0, 2, Ordering::Less), (1, 2, Ordering::Greater), (3, 0, Ordering::Less), (4, 0, Ordering::Less), ],
1499 );
1500
1501 assert_cmp_cases(
1503 &array,
1504 &array,
1505 SortOptions {
1506 descending: false,
1507 nulls_first: false,
1508 },
1509 &[
1510 (0, 2, Ordering::Less),
1511 (1, 2, Ordering::Greater),
1512 (3, 0, Ordering::Less),
1513 (4, 0, Ordering::Greater), ],
1515 );
1516
1517 assert_cmp_cases(
1519 &array,
1520 &array,
1521 SortOptions {
1522 descending: true,
1523 nulls_first: true,
1524 },
1525 &[
1526 (0, 2, Ordering::Greater),
1527 (1, 2, Ordering::Less),
1528 (3, 0, Ordering::Greater),
1529 (4, 0, Ordering::Less),
1530 ],
1531 );
1532
1533 assert_cmp_cases(
1535 &array,
1536 &array,
1537 SortOptions {
1538 descending: true,
1539 nulls_first: false,
1540 },
1541 &[
1542 (0, 2, Ordering::Greater),
1543 (1, 2, Ordering::Less),
1544 (3, 0, Ordering::Greater),
1545 (4, 0, Ordering::Greater),
1546 ],
1547 );
1548 }
1549
1550 #[test]
1551 fn test_list_view() {
1552 test_list_view_comparisons::<i32>();
1553 }
1554
1555 #[test]
1556 fn test_large_list_view() {
1557 test_list_view_comparisons::<i64>();
1558 }
1559
1560 #[test]
1561 fn test_struct() {
1562 let fields = Fields::from(vec![
1563 Field::new("a", DataType::Int32, true),
1564 Field::new_list("b", Field::new_list_field(DataType::Int32, true), true),
1565 ]);
1566
1567 let a = Int32Array::from(vec![Some(1), Some(2), None, None]);
1568 let mut b = ListBuilder::new(Int32Builder::new());
1569 b.extend([Some(vec![Some(1), Some(2)]), Some(vec![None]), None, None]);
1570 let b = b.finish();
1571
1572 let nulls = Some(NullBuffer::from_iter([true, true, true, false]));
1573 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
1574 let s1 = StructArray::new(fields.clone(), values, nulls);
1575
1576 let a = Int32Array::from(vec![None, Some(2), None]);
1577 let mut b = ListBuilder::new(Int32Builder::new());
1578 b.extend([None, None, Some(vec![])]);
1579 let b = b.finish();
1580
1581 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
1582 let s2 = StructArray::new(fields.clone(), values, None);
1583
1584 let opts = SortOptions {
1585 descending: false,
1586 nulls_first: true,
1587 };
1588 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1589 assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 1), Ordering::Greater); assert_eq!(cmp(2, 2), Ordering::Less); assert_eq!(cmp(3, 0), Ordering::Less); assert_eq!(cmp(2, 0), Ordering::Equal); assert_eq!(cmp(3, 0), Ordering::Less); let opts = SortOptions {
1598 descending: true,
1599 nulls_first: true,
1600 };
1601 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1602 assert_eq!(cmp(0, 1), Ordering::Greater); assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 1), Ordering::Greater); assert_eq!(cmp(2, 2), Ordering::Less); assert_eq!(cmp(3, 0), Ordering::Less); assert_eq!(cmp(2, 0), Ordering::Equal); assert_eq!(cmp(3, 0), Ordering::Less); let opts = SortOptions {
1611 descending: true,
1612 nulls_first: false,
1613 };
1614 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1615 assert_eq!(cmp(0, 1), Ordering::Greater); assert_eq!(cmp(0, 0), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Less); assert_eq!(cmp(2, 2), Ordering::Greater); assert_eq!(cmp(3, 0), Ordering::Greater); assert_eq!(cmp(2, 0), Ordering::Equal); assert_eq!(cmp(3, 0), Ordering::Greater); let opts = SortOptions {
1624 descending: false,
1625 nulls_first: false,
1626 };
1627 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1628 assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(0, 0), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Less); assert_eq!(cmp(2, 2), Ordering::Greater); assert_eq!(cmp(3, 0), Ordering::Greater); assert_eq!(cmp(2, 0), Ordering::Equal); assert_eq!(cmp(3, 0), Ordering::Greater); }
1636
1637 #[test]
1638 fn test_map() {
1639 let string_builder = StringBuilder::new();
1642 let int_builder = Int32Builder::new();
1643 let mut map1_builder = MapBuilder::new(None, string_builder, int_builder);
1644
1645 map1_builder.keys().append_value("a");
1647 map1_builder.values().append_value(100);
1648 map1_builder.keys().append_value("b");
1649 map1_builder.values().append_value(1);
1650 map1_builder.append(true).unwrap();
1651
1652 map1_builder.keys().append_value("b");
1654 map1_builder.values().append_value(999);
1655 map1_builder.keys().append_value("c");
1656 map1_builder.values().append_value(1);
1657 map1_builder.append(true).unwrap();
1658
1659 map1_builder.append(true).unwrap();
1661
1662 map1_builder.keys().append_value("x");
1664 map1_builder.values().append_value(1);
1665 map1_builder.append(true).unwrap();
1666
1667 let map1 = map1_builder.finish();
1668
1669 let string_builder = StringBuilder::new();
1672 let int_builder = Int32Builder::new();
1673 let mut map2_builder = MapBuilder::new(None, string_builder, int_builder);
1674
1675 map2_builder.keys().append_value("a");
1677 map2_builder.values().append_value(1);
1678 map2_builder.keys().append_value("c");
1679 map2_builder.values().append_value(999);
1680 map2_builder.append(true).unwrap();
1681
1682 map2_builder.keys().append_value("b");
1684 map2_builder.values().append_value(1);
1685 map2_builder.keys().append_value("d");
1686 map2_builder.values().append_value(999);
1687 map2_builder.append(true).unwrap();
1688
1689 map2_builder.keys().append_value("a");
1691 map2_builder.values().append_value(1);
1692 map2_builder.append(true).unwrap();
1693
1694 map2_builder.append(false).unwrap();
1696
1697 let map2 = map2_builder.finish();
1698
1699 let opts = SortOptions {
1700 descending: false,
1701 nulls_first: true,
1702 };
1703 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1704
1705 assert_eq!(cmp(0, 0), Ordering::Greater);
1709
1710 assert_eq!(cmp(1, 1), Ordering::Greater);
1713
1714 assert_eq!(cmp(0, 1), Ordering::Less);
1716
1717 assert_eq!(cmp(2, 2), Ordering::Less); assert_eq!(cmp(3, 3), Ordering::Greater); assert_eq!(cmp(3, 0), Ordering::Greater); assert_eq!(cmp(2, 0), Ordering::Less); let opts = SortOptions {
1730 descending: true,
1731 nulls_first: true,
1732 };
1733 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1734
1735 assert_eq!(cmp(0, 0), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Greater); assert_eq!(cmp(3, 3), Ordering::Greater); assert_eq!(cmp(2, 2), Ordering::Greater); let opts = SortOptions {
1743 descending: false,
1744 nulls_first: false,
1745 };
1746 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1747
1748 assert_eq!(cmp(0, 0), Ordering::Greater); assert_eq!(cmp(1, 1), Ordering::Greater); assert_eq!(cmp(3, 3), Ordering::Less); assert_eq!(cmp(2, 2), Ordering::Less); }
1754
1755 #[test]
1756 fn test_map_vs_list_consistency() {
1757 let string_builder = StringBuilder::new();
1760 let int_builder = Int32Builder::new();
1761 let mut map1_builder = MapBuilder::new(None, string_builder, int_builder);
1762
1763 map1_builder.keys().append_value("a");
1765 map1_builder.values().append_value(1);
1766 map1_builder.keys().append_value("b");
1767 map1_builder.values().append_value(2);
1768 map1_builder.append(true).unwrap();
1769
1770 map1_builder.keys().append_value("x");
1772 map1_builder.values().append_value(10);
1773 map1_builder.append(true).unwrap();
1774
1775 map1_builder.append(true).unwrap();
1777
1778 map1_builder.keys().append_value("c");
1780 map1_builder.values().append_value(3);
1781 map1_builder.append(true).unwrap();
1782
1783 let map1 = map1_builder.finish();
1784
1785 let string_builder = StringBuilder::new();
1787 let int_builder = Int32Builder::new();
1788 let mut map2_builder = MapBuilder::new(None, string_builder, int_builder);
1789
1790 map2_builder.keys().append_value("a");
1792 map2_builder.values().append_value(1);
1793 map2_builder.keys().append_value("b");
1794 map2_builder.values().append_value(2);
1795 map2_builder.append(true).unwrap();
1796
1797 map2_builder.keys().append_value("y");
1799 map2_builder.values().append_value(20);
1800 map2_builder.append(true).unwrap();
1801
1802 map2_builder.keys().append_value("d");
1804 map2_builder.values().append_value(4);
1805 map2_builder.append(true).unwrap();
1806
1807 map2_builder.append(false).unwrap();
1809
1810 let map2 = map2_builder.finish();
1811
1812 let list1: ListArray = map1.clone().into();
1814 let list2: ListArray = map2.clone().into();
1815
1816 let test_cases = [
1817 SortOptions {
1818 descending: false,
1819 nulls_first: true,
1820 },
1821 SortOptions {
1822 descending: true,
1823 nulls_first: true,
1824 },
1825 SortOptions {
1826 descending: false,
1827 nulls_first: false,
1828 },
1829 SortOptions {
1830 descending: true,
1831 nulls_first: false,
1832 },
1833 ];
1834
1835 for opts in test_cases {
1836 let map_cmp = make_comparator(&map1, &map2, opts).unwrap();
1837 let list_cmp = make_comparator(&list1, &list2, opts).unwrap();
1838
1839 for i in 0..map1.len() {
1841 for j in 0..map2.len() {
1842 let map_result = map_cmp(i, j);
1843 let list_result = list_cmp(i, j);
1844 assert_eq!(
1845 map_result, list_result,
1846 "Map comparison and List comparison should be equal for indices ({i}, {j}) with opts {opts:?}. Map: {map_result:?}, List: {list_result:?}"
1847 );
1848 }
1849 }
1850 }
1851 }
1852
1853 #[test]
1854 fn test_dense_union() {
1855 let int_array = Int32Array::from(vec![1, 2, 3]);
1860 let str_array = StringArray::from(vec!["b", "a"]);
1861
1862 let type_ids = [0, 1, 0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1863 let offsets = [0, 0, 1, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
1864
1865 let union_fields = [
1866 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1867 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1868 ]
1869 .into_iter()
1870 .collect::<UnionFields>();
1871
1872 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1873
1874 let array1 =
1875 UnionArray::try_new(union_fields.clone(), type_ids, Some(offsets), children).unwrap();
1876
1877 let int_array2 = Int32Array::from(vec![2, 1]);
1881 let str_array2 = StringArray::from(vec!["a", "c"]);
1882 let type_ids2 = [0, 1, 0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1883 let offsets2 = [0, 0, 1, 1].into_iter().collect::<ScalarBuffer<i32>>();
1884
1885 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)];
1886
1887 let array2 =
1888 UnionArray::try_new(union_fields, type_ids2, Some(offsets2), children2).unwrap();
1889
1890 let opts = SortOptions {
1891 descending: false,
1892 nulls_first: true,
1893 };
1894
1895 let cmp = make_comparator(&array1, &array2, opts).unwrap();
1899
1900 assert_eq!(cmp(0, 0), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Less); assert_eq!(cmp(1, 1), Ordering::Greater); assert_eq!(cmp(2, 0), Ordering::Equal); assert_eq!(cmp(3, 1), Ordering::Equal); assert_eq!(cmp(1, 3), Ordering::Less); let opts_desc = SortOptions {
1925 descending: true,
1926 nulls_first: true,
1927 };
1928 let cmp_desc = make_comparator(&array1, &array2, opts_desc).unwrap();
1929
1930 assert_eq!(cmp_desc(0, 0), Ordering::Greater); assert_eq!(cmp_desc(0, 1), Ordering::Greater); assert_eq!(cmp_desc(1, 1), Ordering::Less); }
1934
1935 #[test]
1936 fn test_sparse_union() {
1937 let int_array = Int32Array::from(vec![Some(1), None, Some(3)]);
1941 let str_array = StringArray::from(vec![None, Some("b"), None]);
1942 let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1943
1944 let union_fields = [
1945 (0, Arc::new(Field::new("a", DataType::Int32, false))),
1946 (1, Arc::new(Field::new("b", DataType::Utf8, false))),
1947 ]
1948 .into_iter()
1949 .collect::<UnionFields>();
1950
1951 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1952
1953 let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
1954
1955 let opts = SortOptions::default();
1956 let cmp = make_comparator(&array, &array, opts).unwrap();
1957
1958 assert_eq!(cmp(0, 2), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Less); }
1963
1964 #[test]
1965 #[should_panic(expected = "index out of bounds")]
1966 fn test_union_out_of_bounds() {
1967 let int_array = Int32Array::from(vec![1, 2]);
1969 let str_array = StringArray::from(vec!["a"]);
1970
1971 let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1972 let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
1973
1974 let union_fields = [
1975 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1976 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1977 ]
1978 .into_iter()
1979 .collect::<UnionFields>();
1980
1981 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1982
1983 let array = UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap();
1984
1985 let opts = SortOptions::default();
1986 let cmp = make_comparator(&array, &array, opts).unwrap();
1987
1988 cmp(0, 3);
1990 }
1991
1992 #[test]
1993 fn test_union_incompatible_fields() {
1994 let int_array1 = Int32Array::from(vec![1, 2]);
1996 let str_array1 = StringArray::from(vec!["a", "b"]);
1997
1998 let type_ids1 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1999 let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
2000
2001 let union_fields1 = [
2002 (0, Arc::new(Field::new("A", DataType::Int32, false))),
2003 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
2004 ]
2005 .into_iter()
2006 .collect::<UnionFields>();
2007
2008 let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)];
2009
2010 let array1 =
2011 UnionArray::try_new(union_fields1, type_ids1, Some(offsets1), children1).unwrap();
2012
2013 let int_array2 = Int32Array::from(vec![3, 4]);
2015 let float_array2 = Float64Array::from(vec![1.0, 2.0]);
2016
2017 let type_ids2 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
2018 let offsets2 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
2019
2020 let union_fields2 = [
2021 (0, Arc::new(Field::new("A", DataType::Int32, false))),
2022 (1, Arc::new(Field::new("C", DataType::Float64, false))),
2023 ]
2024 .into_iter()
2025 .collect::<UnionFields>();
2026
2027 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(float_array2)];
2028
2029 let array2 =
2030 UnionArray::try_new(union_fields2, type_ids2, Some(offsets2), children2).unwrap();
2031
2032 let opts = SortOptions::default();
2033
2034 let Result::Err(ArrowError::InvalidArgumentError(out)) =
2035 make_comparator(&array1, &array2, opts)
2036 else {
2037 panic!("expected error when making comparator of incompatible union arrays");
2038 };
2039
2040 assert_eq!(
2041 &out,
2042 "Cannot compare UnionArrays with different fields: left=[(0, Field { name: \"A\", data_type: Int32 }), (1, Field { name: \"B\", data_type: Utf8 })], right=[(0, Field { name: \"A\", data_type: Int32 }), (1, Field { name: \"C\", data_type: Float64 })]"
2043 );
2044 }
2045
2046 #[test]
2047 fn test_union_incompatible_modes() {
2048 let int_array1 = Int32Array::from(vec![1, 2]);
2050 let str_array1 = StringArray::from(vec!["a", "b"]);
2051
2052 let type_ids1 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
2053 let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
2054
2055 let union_fields1 = [
2056 (0, Arc::new(Field::new("A", DataType::Int32, false))),
2057 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
2058 ]
2059 .into_iter()
2060 .collect::<UnionFields>();
2061
2062 let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)];
2063
2064 let array1 =
2065 UnionArray::try_new(union_fields1.clone(), type_ids1, Some(offsets1), children1)
2066 .unwrap();
2067
2068 let int_array2 = Int32Array::from(vec![Some(3), None]);
2070 let str_array2 = StringArray::from(vec![None, Some("c")]);
2071
2072 let type_ids2 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
2073
2074 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)];
2075
2076 let array2 = UnionArray::try_new(union_fields1, type_ids2, None, children2).unwrap();
2077
2078 let opts = SortOptions::default();
2079
2080 let Result::Err(ArrowError::InvalidArgumentError(out)) =
2081 make_comparator(&array1, &array2, opts)
2082 else {
2083 panic!("expected error when making comparator of union arrays with different modes");
2084 };
2085
2086 assert_eq!(
2087 &out,
2088 "Cannot compare UnionArrays with different modes: left=Dense, right=Sparse"
2089 );
2090 }
2091
2092 #[test]
2093 fn test_null_array_cmp() {
2094 let a = NullArray::new(3);
2095 let b = NullArray::new(3);
2096 let cmp = make_comparator(&a, &b, SortOptions::default()).unwrap();
2097
2098 assert_eq!(cmp(0, 0), Ordering::Equal);
2099 assert_eq!(cmp(0, 1), Ordering::Equal);
2100 assert_eq!(cmp(2, 0), Ordering::Equal);
2101 }
2102
2103 #[test]
2104 fn test_run_end_encoded_int32() {
2105 let run_ends1 = Int32Array::from(vec![2, 5, 6]);
2109 let values1 = Int32Array::from(vec![1, 2, 3]);
2110 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
2111
2112 let run_ends2 = Int32Array::from(vec![1, 3, 6]);
2115 let values2 = Int32Array::from(vec![1, 2, 3]);
2116 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
2117
2118 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
2119
2120 assert_eq!(cmp(0, 0), Ordering::Equal);
2122 assert_eq!(cmp(0, 1), Ordering::Less);
2124 assert_eq!(cmp(2, 1), Ordering::Equal);
2126 assert_eq!(cmp(5, 5), Ordering::Equal);
2128 assert_eq!(cmp(1, 2), Ordering::Less);
2130 assert_eq!(cmp(4, 4), Ordering::Less);
2132 }
2133
2134 #[test]
2135 fn test_run_end_encoded_with_nulls() {
2136 let run_ends1 = Int32Array::from(vec![2, 4, 5]);
2140 let values1 = Int32Array::from(vec![Some(1), None, Some(2)]);
2141 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
2142
2143 let run_ends2 = Int32Array::from(vec![1, 3, 4, 5]);
2146 let values2 = Int32Array::from(vec![None, Some(1), Some(2), None]);
2147 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
2148
2149 let opts = SortOptions::default();
2150 let cmp = make_comparator(&array1, &array2, opts).unwrap();
2151
2152 assert_eq!(cmp(0, 1), Ordering::Equal);
2154 assert_eq!(cmp(2, 0), Ordering::Equal);
2156 assert_eq!(cmp(0, 0), Ordering::Greater);
2158 assert_eq!(cmp(2, 1), Ordering::Less);
2160 }
2161
2162 #[test]
2163 fn test_run_end_encoded_int16() {
2164 let run_ends1 = Int16Array::from(vec![3_i16, 5, 6]);
2166 let values1 = StringArray::from(vec!["a", "b", "c"]);
2167 let array1 = RunArray::<Int16Type>::try_new(&run_ends1, &values1).unwrap();
2168
2169 let run_ends2 = Int16Array::from(vec![2_i16, 4, 6]);
2170 let values2 = StringArray::from(vec!["a", "b", "c"]);
2171 let array2 = RunArray::<Int16Type>::try_new(&run_ends2, &values2).unwrap();
2172
2173 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
2174
2175 assert_eq!(cmp(0, 0), Ordering::Equal); assert_eq!(cmp(2, 2), Ordering::Less); assert_eq!(cmp(3, 2), Ordering::Equal); assert_eq!(cmp(5, 4), Ordering::Equal); }
2182
2183 #[test]
2184 fn test_run_end_encoded_int64() {
2185 let run_ends1 = Int64Array::from(vec![2_i64, 4, 6]);
2187 let values1 = Int64Array::from(vec![10_i64, 20, 30]);
2188 let array1 = RunArray::<Int64Type>::try_new(&run_ends1, &values1).unwrap();
2189
2190 let run_ends2 = Int64Array::from(vec![3_i64, 5, 6]);
2191 let values2 = Int64Array::from(vec![10_i64, 20, 30]);
2192 let array2 = RunArray::<Int64Type>::try_new(&run_ends2, &values2).unwrap();
2193
2194 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
2195
2196 assert_eq!(cmp(0, 0), Ordering::Equal); assert_eq!(cmp(1, 2), Ordering::Equal); assert_eq!(cmp(2, 3), Ordering::Equal); assert_eq!(cmp(4, 4), Ordering::Greater); }
2203
2204 #[test]
2205 fn test_run_end_encoded_sliced() {
2206 let run_ends = Int32Array::from(vec![2, 5, 7, 8]);
2210 let values = Int32Array::from(vec![1, 2, 3, 4]);
2211 let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
2212
2213 let slice1 = array.slice(1, 4);
2215 let slice2 = array.slice(3, 4);
2217
2218 let cmp = make_comparator(&slice1, &slice2, SortOptions::default()).unwrap();
2219
2220 assert_eq!(cmp(0, 0), Ordering::Less);
2222 assert_eq!(cmp(1, 0), Ordering::Equal);
2224 assert_eq!(cmp(3, 2), Ordering::Less);
2226 assert_eq!(cmp(1, 3), Ordering::Less);
2228
2229 let run_ends2 = Int32Array::from(vec![2, 4]);
2231 let values2 = Int32Array::from(vec![1, 2]);
2232 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
2233
2234 let cmp = make_comparator(&slice1, &array2, SortOptions::default()).unwrap();
2235
2236 assert_eq!(cmp(0, 0), Ordering::Equal);
2238 assert_eq!(cmp(1, 1), Ordering::Greater);
2240 assert_eq!(cmp(3, 3), Ordering::Equal);
2242 }
2243
2244 #[test]
2245 fn test_run_end_encoded_sliced_with_nulls() {
2246 let run_ends = Int32Array::from(vec![2, 4, 6, 7, 8]);
2250 let values = Int32Array::from(vec![Some(1), None, Some(2), None, Some(3)]);
2251 let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
2252
2253 let slice1 = array.slice(1, 5);
2255 let slice2 = array.slice(3, 5);
2257
2258 let opts = SortOptions::default(); let cmp = make_comparator(&slice1, &slice2, opts).unwrap();
2260
2261 assert_eq!(cmp(0, 0), Ordering::Greater);
2263 assert_eq!(cmp(1, 0), Ordering::Equal);
2265 assert_eq!(cmp(1, 1), Ordering::Less);
2267 assert_eq!(cmp(3, 1), Ordering::Equal);
2269 assert_eq!(cmp(4, 4), Ordering::Less);
2271 assert_eq!(cmp(3, 3), Ordering::Greater);
2273 }
2274
2275 #[test]
2276 fn test_run_end_encoded_different_types() {
2277 let run_ends1 = Int32Array::from(vec![2, 4, 6]);
2279 let values1 = Int32Array::from(vec![1, 2, 3]);
2280 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
2281
2282 let run_ends2 = Int64Array::from(vec![2_i64, 4, 6]);
2283 let values2 = Int64Array::from(vec![1_i64, 2, 3]);
2284 let array2 = RunArray::<Int64Type>::try_new(&run_ends2, &values2).unwrap();
2285
2286 let result = make_comparator(&array1, &array2, SortOptions::default());
2287 assert!(result.is_err());
2288 let err = match result {
2289 Err(e) => e.to_string(),
2290 Ok(_) => panic!("Expected error"),
2291 };
2292 assert!(err.contains("Cannot compare RunEndEncoded arrays"));
2293 }
2294}