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
224 let c_opts = child_opts(opts);
225 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
226 let left_keys = left.keys().values().clone();
227 let right_keys = right.keys().values().clone();
228
229 let f = compare(left, right, opts, move |i, j| {
230 let l = left_keys[i].as_usize();
231 let r = right_keys[j].as_usize();
232 cmp(l, r)
233 });
234 Ok(f)
235}
236
237fn compare_list<O: OffsetSizeTrait>(
238 left: &dyn Array,
239 right: &dyn Array,
240 opts: SortOptions,
241) -> Result<DynComparator, ArrowError> {
242 let left = left.as_list::<O>();
243 let right = right.as_list::<O>();
244
245 let c_opts = child_opts(opts);
246 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
247
248 let l_o = left.offsets().clone();
249 let r_o = right.offsets().clone();
250 let f = compare(left, right, opts, move |i, j| {
251 let l_end = l_o[i + 1].as_usize();
252 let l_start = l_o[i].as_usize();
253
254 let r_end = r_o[j + 1].as_usize();
255 let r_start = r_o[j].as_usize();
256
257 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
258 match cmp(i, j) {
259 Ordering::Equal => {}
260 r => return r,
261 }
262 }
263 (l_end - l_start).cmp(&(r_end - r_start))
264 });
265 Ok(f)
266}
267
268fn compare_fixed_list(
269 left: &dyn Array,
270 right: &dyn Array,
271 opts: SortOptions,
272) -> Result<DynComparator, ArrowError> {
273 let left = left.as_fixed_size_list();
274 let right = right.as_fixed_size_list();
275
276 let c_opts = child_opts(opts);
277 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
278
279 let l_size = left.value_length().to_usize().unwrap();
280 let r_size = right.value_length().to_usize().unwrap();
281 let size_cmp = l_size.cmp(&r_size);
282
283 let f = compare(left, right, opts, move |i, j| {
284 let l_start = i * l_size;
285 let l_end = l_start + l_size;
286 let r_start = j * r_size;
287 let r_end = r_start + r_size;
288 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
289 match cmp(i, j) {
290 Ordering::Equal => {}
291 r => return r,
292 }
293 }
294 size_cmp
295 });
296 Ok(f)
297}
298
299fn compare_list_view<O: OffsetSizeTrait>(
300 left: &dyn Array,
301 right: &dyn Array,
302 opts: SortOptions,
303) -> Result<DynComparator, ArrowError> {
304 let left = left.as_list_view::<O>();
305 let right = right.as_list_view::<O>();
306
307 let c_opts = child_opts(opts);
308 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
309
310 let l_offsets = left.offsets().clone();
311 let l_sizes = left.sizes().clone();
312 let r_offsets = right.offsets().clone();
313 let r_sizes = right.sizes().clone();
314
315 let f = compare(left, right, opts, move |i, j| {
316 let l_start = l_offsets[i].as_usize();
317 let l_len = l_sizes[i].as_usize();
318 let l_end = l_start + l_len;
319
320 let r_start = r_offsets[j].as_usize();
321 let r_len = r_sizes[j].as_usize();
322 let r_end = r_start + r_len;
323
324 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
325 match cmp(i, j) {
326 Ordering::Equal => {}
327 r => return r,
328 }
329 }
330 l_len.cmp(&r_len)
331 });
332 Ok(f)
333}
334
335fn compare_map(
336 left: &dyn Array,
337 right: &dyn Array,
338 opts: SortOptions,
339) -> Result<DynComparator, ArrowError> {
340 let left = left.as_map();
341 let right = right.as_map();
342
343 let c_opts = child_opts(opts);
344 let cmp = make_comparator(left.entries(), right.entries(), c_opts)?;
345
346 let l_o = left.offsets().clone();
347 let r_o = right.offsets().clone();
348 let f = compare(left, right, opts, move |i, j| {
349 let l_end = l_o[i + 1].as_usize();
350 let l_start = l_o[i].as_usize();
351
352 let r_end = r_o[j + 1].as_usize();
353 let r_start = r_o[j].as_usize();
354
355 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
356 match cmp(i, j) {
357 Ordering::Equal => {}
358 r => return r,
359 }
360 }
361 (l_end - l_start).cmp(&(r_end - r_start))
362 });
363 Ok(f)
364}
365
366fn compare_struct(
367 left: &dyn Array,
368 right: &dyn Array,
369 opts: SortOptions,
370) -> Result<DynComparator, ArrowError> {
371 let left = left.as_struct();
372 let right = right.as_struct();
373
374 if left.columns().len() != right.columns().len() {
375 return Err(ArrowError::InvalidArgumentError(
376 "Cannot compare StructArray with different number of columns".to_string(),
377 ));
378 }
379
380 let c_opts = child_opts(opts);
381 let columns = left.columns().iter().zip(right.columns());
382 let comparators = columns
383 .map(|(l, r)| make_comparator(l, r, c_opts))
384 .collect::<Result<Vec<_>, _>>()?;
385
386 let f = compare(left, right, opts, move |i, j| {
387 for cmp in &comparators {
388 match cmp(i, j) {
389 Ordering::Equal => {}
390 r => return r,
391 }
392 }
393 Ordering::Equal
394 });
395 Ok(f)
396}
397
398fn compare_union(
399 left: &dyn Array,
400 right: &dyn Array,
401 opts: SortOptions,
402) -> Result<DynComparator, ArrowError> {
403 let left = left.as_union();
404 let right = right.as_union();
405
406 let DataType::Union(left_fields, left_mode) = left.data_type() else {
407 unreachable!()
408 };
409 let DataType::Union(right_fields, right_mode) = right.data_type() else {
410 unreachable!()
411 };
412
413 if left_fields != right_fields {
414 return Err(ArrowError::InvalidArgumentError(format!(
415 "Cannot compare UnionArrays with different fields: left={left_fields:?}, right={right_fields:?}"
416 )));
417 }
418
419 if left_mode != right_mode {
420 return Err(ArrowError::InvalidArgumentError(format!(
421 "Cannot compare UnionArrays with different modes: left={left_mode:?}, right={right_mode:?}"
422 )));
423 }
424
425 let c_opts = child_opts(opts);
426
427 let mut field_comparators = HashMap::with_capacity(left_fields.len());
428
429 for (type_id, _field) in left_fields.iter() {
430 let left_child = left.child(type_id);
431 let right_child = right.child(type_id);
432 let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?;
433
434 field_comparators.insert(type_id, cmp);
435 }
436
437 let left_type_ids = left.type_ids().clone();
438 let right_type_ids = right.type_ids().clone();
439
440 let left_offsets = left.offsets().cloned();
441 let right_offsets = right.offsets().cloned();
442
443 let f = compare(left, right, opts, move |i, j| {
444 let left_type_id = left_type_ids[i];
445 let right_type_id = right_type_ids[j];
446
447 match left_type_id.cmp(&right_type_id) {
448 Ordering::Equal => {
449 let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i);
450 let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j);
451
452 let cmp = field_comparators
453 .get(&left_type_id)
454 .expect("type id not found in field_comparators");
455
456 cmp(left_offset, right_offset)
457 }
458 other => other,
459 }
460 });
461 Ok(f)
462}
463
464pub fn make_comparator(
472 left: &dyn Array,
473 right: &dyn Array,
474 opts: SortOptions,
475) -> Result<DynComparator, ArrowError> {
476 use arrow_schema::DataType::*;
477
478 match (left.data_type(), right.data_type()) {
481 (Decimal32(_, s1), Decimal32(_, s2))
482 | (Decimal64(_, s1), Decimal64(_, s2))
483 | (Decimal128(_, s1), Decimal128(_, s2))
484 | (Decimal256(_, s1), Decimal256(_, s2))
485 if s1 != s2 =>
486 {
487 return Err(ArrowError::InvalidArgumentError(format!(
488 "Can't compare decimal arrays with different scales: {:?}, {:?}",
489 left.data_type(),
490 right.data_type()
491 )));
492 }
493 _ => {}
494 }
495
496 macro_rules! primitive_helper {
497 ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => {
498 Ok(compare_primitive::<$t>($left, $right, $nulls_first))
499 };
500 }
501 downcast_primitive! {
502 left.data_type(), right.data_type() => (primitive_helper, left, right, opts),
503 (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)),
504 (Utf8, Utf8) => Ok(compare_bytes::<Utf8Type>(left, right, opts)),
505 (LargeUtf8, LargeUtf8) => Ok(compare_bytes::<LargeUtf8Type>(left, right, opts)),
506 (Utf8View, Utf8View) => Ok(compare_byte_view::<StringViewType>(left, right, opts)),
507 (Binary, Binary) => Ok(compare_bytes::<BinaryType>(left, right, opts)),
508 (LargeBinary, LargeBinary) => Ok(compare_bytes::<LargeBinaryType>(left, right, opts)),
509 (BinaryView, BinaryView) => Ok(compare_byte_view::<BinaryViewType>(left, right, opts)),
510 (FixedSizeBinary(_), FixedSizeBinary(_)) => {
511 let left = left.as_fixed_size_binary();
512 let right = right.as_fixed_size_binary();
513
514 let l = left.clone();
515 let r = right.clone();
516 Ok(compare(left, right, opts, move |i, j| {
517 l.value(i).cmp(r.value(j))
518 }))
519 },
520 (List(_), List(_)) => compare_list::<i32>(left, right, opts),
521 (LargeList(_), LargeList(_)) => compare_list::<i64>(left, right, opts),
522 (ListView(_), ListView(_)) => compare_list_view::<i32>(left, right, opts),
523 (LargeListView(_), LargeListView(_)) => compare_list_view::<i64>(left, right, opts),
524 (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts),
525 (Struct(_), Struct(_)) => compare_struct(left, right, opts),
526 (Dictionary(l_key, _), Dictionary(r_key, _)) => {
527 macro_rules! dict_helper {
528 ($t:ty, $left:expr, $right:expr, $opts: expr) => {
529 compare_dict::<$t>($left, $right, $opts)
530 };
531 }
532 downcast_integer! {
533 l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts),
534 _ => unreachable!()
535 }
536 },
537 (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => {
538 macro_rules! run_end_helper {
539 ($t:ty, $left:expr, $right:expr, $opts:expr) => {
540 compare_run_end_encoded::<$t>($left, $right, $opts)
541 };
542 }
543 downcast_run_end_index! {
544 l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts),
545 _ => Err(ArrowError::InvalidArgumentError(format!(
546 "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}",
547 l_run_ends.data_type(),
548 r_run_ends.data_type()
549 )))
550 }
551 },
552 (Map(_, _), Map(_, _)) => compare_map(left, right, opts),
553 (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)),
554 (Union(_, _), Union(_, _)) => compare_union(left, right, opts),
555 (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs {
556 true => format!("The data type type {lhs:?} has no natural order"),
557 false => "Can't compare arrays of different types".to_string(),
558 }))
559 }
560}
561
562#[cfg(test)]
563mod tests {
564 use super::*;
565 use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder};
566 use arrow_buffer::{IntervalDayTime, NullBuffer, OffsetBuffer, ScalarBuffer, i256};
567 use arrow_schema::{ArrowError, DataType, Field, Fields, UnionFields};
568 use half::f16;
569 use std::cmp::Ordering;
570 use std::sync::Arc;
571
572 #[test]
573 fn test_fixed_size_binary() {
574 let items = vec![vec![1u8], vec![2u8]];
575 let array = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
576
577 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
578
579 assert_eq!(Ordering::Less, cmp(0, 1));
580 }
581
582 #[test]
583 fn test_fixed_size_binary_fixed_size_binary() {
584 let items = vec![vec![1u8]];
585 let array1 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
586 let items = vec![vec![2u8]];
587 let array2 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
588
589 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
590
591 assert_eq!(Ordering::Less, cmp(0, 0));
592 }
593
594 #[test]
595 fn test_i32() {
596 let array = Int32Array::from(vec![1, 2]);
597
598 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
599
600 assert_eq!(Ordering::Less, (cmp)(0, 1));
601 }
602
603 #[test]
604 fn test_i32_i32() {
605 let array1 = Int32Array::from(vec![1]);
606 let array2 = Int32Array::from(vec![2]);
607
608 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
609
610 assert_eq!(Ordering::Less, cmp(0, 0));
611 }
612
613 #[cfg_attr(miri, ignore)] #[test]
615 fn test_f16() {
616 let array = Float16Array::from(vec![f16::from_f32(1.0), f16::from_f32(2.0)]);
617
618 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
619
620 assert_eq!(Ordering::Less, cmp(0, 1));
621 }
622
623 #[test]
624 fn test_f64() {
625 let array = Float64Array::from(vec![1.0, 2.0]);
626
627 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
628
629 assert_eq!(Ordering::Less, cmp(0, 1));
630 }
631
632 #[test]
633 fn test_f64_nan() {
634 let array = Float64Array::from(vec![1.0, f64::NAN]);
635
636 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
637
638 assert_eq!(Ordering::Less, cmp(0, 1));
639 assert_eq!(Ordering::Equal, cmp(1, 1));
640 }
641
642 #[test]
643 fn test_f64_zeros() {
644 let array = Float64Array::from(vec![-0.0, 0.0]);
645
646 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
647
648 assert_eq!(Ordering::Less, cmp(0, 1));
649 assert_eq!(Ordering::Greater, cmp(1, 0));
650 }
651
652 #[test]
653 fn test_interval_day_time() {
654 let array = IntervalDayTimeArray::from(vec![
655 IntervalDayTimeType::make_value(0, 1000),
657 IntervalDayTimeType::make_value(1, 2),
659 IntervalDayTimeType::make_value(0, 90_000_000),
661 ]);
662
663 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
664
665 assert_eq!(Ordering::Less, cmp(0, 1));
666 assert_eq!(Ordering::Greater, cmp(1, 0));
667
668 assert_eq!(Ordering::Greater, cmp(1, 2));
672 assert_eq!(Ordering::Less, cmp(2, 1));
673 }
674
675 #[test]
676 fn test_interval_year_month() {
677 let array = IntervalYearMonthArray::from(vec![
678 IntervalYearMonthType::make_value(1, 0),
680 IntervalYearMonthType::make_value(0, 13),
682 IntervalYearMonthType::make_value(1, 1),
684 ]);
685
686 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
687
688 assert_eq!(Ordering::Less, cmp(0, 1));
689 assert_eq!(Ordering::Greater, cmp(1, 0));
690
691 assert_eq!(Ordering::Equal, cmp(1, 2));
693 assert_eq!(Ordering::Equal, cmp(2, 1));
694 }
695
696 #[test]
697 fn test_interval_month_day_nano() {
698 let array = IntervalMonthDayNanoArray::from(vec![
699 IntervalMonthDayNanoType::make_value(0, 100, 0),
701 IntervalMonthDayNanoType::make_value(1, 0, 0),
703 IntervalMonthDayNanoType::make_value(0, 100, 2),
705 ]);
706
707 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
708
709 assert_eq!(Ordering::Less, cmp(0, 1));
710 assert_eq!(Ordering::Greater, cmp(1, 0));
711
712 assert_eq!(Ordering::Greater, cmp(1, 2));
716 assert_eq!(Ordering::Less, cmp(2, 1));
717 }
718
719 #[test]
720 fn test_decimali32() {
721 let array = vec![Some(5_i32), Some(2_i32), Some(3_i32)]
722 .into_iter()
723 .collect::<Decimal32Array>()
724 .with_precision_and_scale(8, 6)
725 .unwrap();
726
727 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
728 assert_eq!(Ordering::Less, cmp(1, 0));
729 assert_eq!(Ordering::Greater, cmp(0, 2));
730 }
731
732 #[test]
733 fn test_decimali64() {
734 let array = vec![Some(5_i64), Some(2_i64), Some(3_i64)]
735 .into_iter()
736 .collect::<Decimal64Array>()
737 .with_precision_and_scale(16, 6)
738 .unwrap();
739
740 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
741 assert_eq!(Ordering::Less, cmp(1, 0));
742 assert_eq!(Ordering::Greater, cmp(0, 2));
743 }
744
745 #[test]
746 fn test_decimali128() {
747 let array = vec![Some(5_i128), Some(2_i128), Some(3_i128)]
748 .into_iter()
749 .collect::<Decimal128Array>()
750 .with_precision_and_scale(23, 6)
751 .unwrap();
752
753 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
754 assert_eq!(Ordering::Less, cmp(1, 0));
755 assert_eq!(Ordering::Greater, cmp(0, 2));
756 }
757
758 #[test]
759 fn test_decimali256() {
760 let array = vec![
761 Some(i256::from_i128(5_i128)),
762 Some(i256::from_i128(2_i128)),
763 Some(i256::from_i128(3_i128)),
764 ]
765 .into_iter()
766 .collect::<Decimal256Array>()
767 .with_precision_and_scale(53, 6)
768 .unwrap();
769
770 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
771 assert_eq!(Ordering::Less, cmp(1, 0));
772 assert_eq!(Ordering::Greater, cmp(0, 2));
773 }
774
775 #[test]
776 fn test_dict() {
777 let data = vec!["a", "b", "c", "a", "a", "c", "c"];
778 let array = data.into_iter().collect::<DictionaryArray<Int16Type>>();
779
780 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
781
782 assert_eq!(Ordering::Less, cmp(0, 1));
783 assert_eq!(Ordering::Equal, cmp(3, 4));
784 assert_eq!(Ordering::Greater, cmp(2, 3));
785 }
786
787 #[test]
788 fn test_multiple_dict() {
789 let d1 = vec!["a", "b", "c", "d"];
790 let a1 = d1.into_iter().collect::<DictionaryArray<Int16Type>>();
791 let d2 = vec!["e", "f", "g", "a"];
792 let a2 = d2.into_iter().collect::<DictionaryArray<Int16Type>>();
793
794 let cmp = make_comparator(&a1, &a2, SortOptions::default()).unwrap();
795
796 assert_eq!(Ordering::Less, cmp(0, 0));
797 assert_eq!(Ordering::Equal, cmp(0, 3));
798 assert_eq!(Ordering::Greater, cmp(1, 3));
799 }
800
801 #[test]
802 fn test_primitive_dict() {
803 let values = Int32Array::from(vec![1_i32, 0, 2, 5]);
804 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
805 let array1 = DictionaryArray::new(keys, Arc::new(values));
806
807 let values = Int32Array::from(vec![2_i32, 3, 4, 5]);
808 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
809 let array2 = DictionaryArray::new(keys, Arc::new(values));
810
811 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
812
813 assert_eq!(Ordering::Less, cmp(0, 0));
814 assert_eq!(Ordering::Less, cmp(0, 3));
815 assert_eq!(Ordering::Equal, cmp(3, 3));
816 assert_eq!(Ordering::Greater, cmp(3, 1));
817 assert_eq!(Ordering::Greater, cmp(3, 2));
818 }
819
820 #[test]
821 fn test_float_dict() {
822 let values = Float32Array::from(vec![1.0, 0.5, 2.1, 5.5]);
823 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
824 let array1 = DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
825
826 let values = Float32Array::from(vec![1.2, 3.2, 4.0, 5.5]);
827 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
828 let array2 = DictionaryArray::new(keys, Arc::new(values));
829
830 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
831
832 assert_eq!(Ordering::Less, cmp(0, 0));
833 assert_eq!(Ordering::Less, cmp(0, 3));
834 assert_eq!(Ordering::Equal, cmp(3, 3));
835 assert_eq!(Ordering::Greater, cmp(3, 1));
836 assert_eq!(Ordering::Greater, cmp(3, 2));
837 }
838
839 #[test]
840 fn test_timestamp_dict() {
841 let values = TimestampSecondArray::from(vec![1, 0, 2, 5]);
842 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
843 let array1 = DictionaryArray::new(keys, Arc::new(values));
844
845 let values = TimestampSecondArray::from(vec![2, 3, 4, 5]);
846 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
847 let array2 = DictionaryArray::new(keys, Arc::new(values));
848
849 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
850
851 assert_eq!(Ordering::Less, cmp(0, 0));
852 assert_eq!(Ordering::Less, cmp(0, 3));
853 assert_eq!(Ordering::Equal, cmp(3, 3));
854 assert_eq!(Ordering::Greater, cmp(3, 1));
855 assert_eq!(Ordering::Greater, cmp(3, 2));
856 }
857
858 #[test]
859 fn test_interval_dict() {
860 let v1 = IntervalDayTime::new(0, 1);
861 let v2 = IntervalDayTime::new(0, 2);
862 let v3 = IntervalDayTime::new(12, 2);
863
864 let values = IntervalDayTimeArray::from(vec![Some(v1), Some(v2), None, Some(v3)]);
865 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
866 let array1 = DictionaryArray::new(keys, Arc::new(values));
867
868 let values = IntervalDayTimeArray::from(vec![Some(v3), Some(v2), None, Some(v1)]);
869 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
870 let array2 = DictionaryArray::new(keys, Arc::new(values));
871
872 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
873
874 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)); }
880
881 #[test]
882 fn test_duration_dict() {
883 let values = DurationSecondArray::from(vec![1, 0, 2, 5]);
884 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
885 let array1 = DictionaryArray::new(keys, Arc::new(values));
886
887 let values = DurationSecondArray::from(vec![2, 3, 4, 5]);
888 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
889 let array2 = DictionaryArray::new(keys, Arc::new(values));
890
891 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
892
893 assert_eq!(Ordering::Less, cmp(0, 0));
894 assert_eq!(Ordering::Less, cmp(0, 3));
895 assert_eq!(Ordering::Equal, cmp(3, 3));
896 assert_eq!(Ordering::Greater, cmp(3, 1));
897 assert_eq!(Ordering::Greater, cmp(3, 2));
898 }
899
900 #[test]
901 fn test_decimal_dict() {
902 let values = Decimal128Array::from(vec![1, 0, 2, 5]);
903 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
904 let array1 = DictionaryArray::new(keys, Arc::new(values));
905
906 let values = Decimal128Array::from(vec![2, 3, 4, 5]);
907 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
908 let array2 = DictionaryArray::new(keys, Arc::new(values));
909
910 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
911
912 assert_eq!(Ordering::Less, cmp(0, 0));
913 assert_eq!(Ordering::Less, cmp(0, 3));
914 assert_eq!(Ordering::Equal, cmp(3, 3));
915 assert_eq!(Ordering::Greater, cmp(3, 1));
916 assert_eq!(Ordering::Greater, cmp(3, 2));
917 }
918
919 #[test]
920 fn test_decimal256_dict() {
921 let values = Decimal256Array::from(vec![
922 i256::from_i128(1),
923 i256::from_i128(0),
924 i256::from_i128(2),
925 i256::from_i128(5),
926 ]);
927 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
928 let array1 = DictionaryArray::new(keys, Arc::new(values));
929
930 let values = Decimal256Array::from(vec![
931 i256::from_i128(2),
932 i256::from_i128(3),
933 i256::from_i128(4),
934 i256::from_i128(5),
935 ]);
936 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
937 let array2 = DictionaryArray::new(keys, Arc::new(values));
938
939 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
940
941 assert_eq!(Ordering::Less, cmp(0, 0));
942 assert_eq!(Ordering::Less, cmp(0, 3));
943 assert_eq!(Ordering::Equal, cmp(3, 3));
944 assert_eq!(Ordering::Greater, cmp(3, 1));
945 assert_eq!(Ordering::Greater, cmp(3, 2));
946 }
947
948 #[test]
949 fn test_decimal_different_scale_errors() {
950 let a = Decimal128Array::from(vec![100])
952 .with_precision_and_scale(10, 2)
953 .unwrap();
954 let b = Decimal128Array::from(vec![500])
955 .with_precision_and_scale(10, 3)
956 .unwrap();
957 let err = make_comparator(&a, &b, SortOptions::default())
958 .err()
959 .unwrap();
960 assert_eq!(
961 err.to_string(),
962 "Invalid argument error: Can't compare decimal arrays with different scales: \
963 Decimal128(10, 2), Decimal128(10, 3)"
964 );
965
966 let keys = Int8Array::from_iter_values([0]);
968 let dict_a = DictionaryArray::new(keys.clone(), Arc::new(a.clone()));
969 let dict_b = DictionaryArray::new(keys, Arc::new(b));
970 assert!(make_comparator(&dict_a, &dict_b, SortOptions::default()).is_err());
971
972 let c = Decimal128Array::from(vec![100])
974 .with_precision_and_scale(12, 2)
975 .unwrap();
976 let cmp = make_comparator(&a, &c, SortOptions::default()).unwrap();
977 assert_eq!(Ordering::Equal, cmp(0, 0));
978
979 let a = Decimal32Array::from(vec![1])
981 .with_precision_and_scale(5, 1)
982 .unwrap();
983 let b = Decimal32Array::from(vec![1])
984 .with_precision_and_scale(5, 2)
985 .unwrap();
986 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
987
988 let a = Decimal64Array::from(vec![1])
989 .with_precision_and_scale(10, 1)
990 .unwrap();
991 let b = Decimal64Array::from(vec![1])
992 .with_precision_and_scale(10, 2)
993 .unwrap();
994 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
995
996 let a = Decimal256Array::from(vec![i256::from_i128(1)])
997 .with_precision_and_scale(40, 1)
998 .unwrap();
999 let b = Decimal256Array::from(vec![i256::from_i128(1)])
1000 .with_precision_and_scale(40, 2)
1001 .unwrap();
1002 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
1003 }
1004
1005 fn test_bytes_impl<T: ByteArrayType>() {
1006 let offsets = OffsetBuffer::from_lengths([3, 3, 1]);
1007 let a = GenericByteArray::<T>::new(offsets, b"abcdefa".into(), None);
1008 let cmp = make_comparator(&a, &a, SortOptions::default()).unwrap();
1009
1010 assert_eq!(Ordering::Less, cmp(0, 1));
1011 assert_eq!(Ordering::Greater, cmp(0, 2));
1012 assert_eq!(Ordering::Equal, cmp(1, 1));
1013 }
1014
1015 #[test]
1016 fn test_bytes() {
1017 test_bytes_impl::<Utf8Type>();
1018 test_bytes_impl::<LargeUtf8Type>();
1019 test_bytes_impl::<BinaryType>();
1020 test_bytes_impl::<LargeBinaryType>();
1021 }
1022
1023 fn assert_cmp_cases<A: Array>(
1024 array1: &A,
1025 array2: &A,
1026 opts: SortOptions,
1027 cases: &[(usize, usize, Ordering)],
1028 ) {
1029 let cmp = make_comparator(array1, array2, opts).unwrap();
1030 for (left, right, expected) in cases {
1031 assert_eq!(cmp(*left, *right), *expected);
1032 }
1033 }
1034
1035 #[test]
1036 fn test_lists() {
1037 let mut a = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
1038 a.extend([
1039 Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]),
1040 Some(vec![
1041 Some(vec![Some(1), Some(2), Some(3)]),
1042 Some(vec![Some(1)]),
1043 ]),
1044 Some(vec![]),
1045 ]);
1046 let a = a.finish();
1047 let mut b = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
1048 b.extend([
1049 Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]),
1050 Some(vec![
1051 Some(vec![Some(1), Some(2), None]),
1052 Some(vec![Some(1)]),
1053 ]),
1054 Some(vec![
1055 Some(vec![Some(1), Some(2), Some(3), Some(4)]),
1056 Some(vec![Some(1)]),
1057 ]),
1058 None,
1059 ]);
1060 let b = b.finish();
1061
1062 assert_cmp_cases(
1064 &a,
1065 &b,
1066 SortOptions {
1067 descending: false,
1068 nulls_first: true,
1069 },
1070 &[
1071 (0, 0, Ordering::Equal),
1072 (0, 1, Ordering::Less),
1073 (0, 2, Ordering::Less),
1074 (1, 2, Ordering::Less),
1075 (1, 3, Ordering::Greater),
1076 (2, 0, Ordering::Less),
1077 ],
1078 );
1079
1080 assert_cmp_cases(
1082 &a,
1083 &b,
1084 SortOptions {
1085 descending: true,
1086 nulls_first: true,
1087 },
1088 &[
1089 (0, 0, Ordering::Equal),
1090 (0, 1, Ordering::Less),
1091 (0, 2, Ordering::Less),
1092 (1, 2, Ordering::Greater),
1093 (1, 3, Ordering::Greater),
1094 (2, 0, Ordering::Greater),
1095 ],
1096 );
1097
1098 assert_cmp_cases(
1100 &a,
1101 &b,
1102 SortOptions {
1103 descending: true,
1104 nulls_first: false,
1105 },
1106 &[
1107 (0, 0, Ordering::Equal),
1108 (0, 1, Ordering::Greater),
1109 (0, 2, Ordering::Greater),
1110 (1, 2, Ordering::Greater),
1111 (1, 3, Ordering::Less),
1112 (2, 0, Ordering::Greater),
1113 ],
1114 );
1115
1116 assert_cmp_cases(
1118 &a,
1119 &b,
1120 SortOptions {
1121 descending: false,
1122 nulls_first: false,
1123 },
1124 &[
1125 (0, 0, Ordering::Equal),
1126 (0, 1, Ordering::Greater),
1127 (0, 2, Ordering::Greater),
1128 (1, 2, Ordering::Less),
1129 (1, 3, Ordering::Less),
1130 (2, 0, Ordering::Less),
1131 ],
1132 );
1133 }
1134
1135 fn list_view_array<O: OffsetSizeTrait>(
1136 values: Vec<i32>,
1137 offsets: &[usize],
1138 sizes: &[usize],
1139 valid: Option<&[bool]>,
1140 ) -> GenericListViewArray<O> {
1141 let offsets = offsets
1142 .iter()
1143 .map(|v| O::from_usize(*v).unwrap())
1144 .collect::<ScalarBuffer<O>>();
1145 let sizes = sizes
1146 .iter()
1147 .map(|v| O::from_usize(*v).unwrap())
1148 .collect::<ScalarBuffer<O>>();
1149 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
1150 let values = Int32Array::from(values);
1151 let nulls = valid.map(NullBuffer::from);
1152 GenericListViewArray::new(field, offsets, sizes, Arc::new(values), nulls)
1153 }
1154
1155 fn test_list_view_comparisons<O: OffsetSizeTrait>() {
1156 let array = list_view_array::<O>(
1157 vec![1, 2, 3, 4, 5],
1158 &[0, 2, 1, 0, 3],
1159 &[2, 2, 2, 0, 2],
1160 Some(&[true, true, true, true, false]),
1161 );
1162
1163 assert_cmp_cases(
1165 &array,
1166 &array,
1167 SortOptions {
1168 descending: false,
1169 nulls_first: true,
1170 },
1171 &[
1172 (0, 2, Ordering::Less), (1, 2, Ordering::Greater), (3, 0, Ordering::Less), (4, 0, Ordering::Less), ],
1177 );
1178
1179 assert_cmp_cases(
1181 &array,
1182 &array,
1183 SortOptions {
1184 descending: false,
1185 nulls_first: false,
1186 },
1187 &[
1188 (0, 2, Ordering::Less),
1189 (1, 2, Ordering::Greater),
1190 (3, 0, Ordering::Less),
1191 (4, 0, Ordering::Greater), ],
1193 );
1194
1195 assert_cmp_cases(
1197 &array,
1198 &array,
1199 SortOptions {
1200 descending: true,
1201 nulls_first: true,
1202 },
1203 &[
1204 (0, 2, Ordering::Greater),
1205 (1, 2, Ordering::Less),
1206 (3, 0, Ordering::Greater),
1207 (4, 0, Ordering::Less),
1208 ],
1209 );
1210
1211 assert_cmp_cases(
1213 &array,
1214 &array,
1215 SortOptions {
1216 descending: true,
1217 nulls_first: false,
1218 },
1219 &[
1220 (0, 2, Ordering::Greater),
1221 (1, 2, Ordering::Less),
1222 (3, 0, Ordering::Greater),
1223 (4, 0, Ordering::Greater),
1224 ],
1225 );
1226 }
1227
1228 #[test]
1229 fn test_list_view() {
1230 test_list_view_comparisons::<i32>();
1231 }
1232
1233 #[test]
1234 fn test_large_list_view() {
1235 test_list_view_comparisons::<i64>();
1236 }
1237
1238 #[test]
1239 fn test_struct() {
1240 let fields = Fields::from(vec![
1241 Field::new("a", DataType::Int32, true),
1242 Field::new_list("b", Field::new_list_field(DataType::Int32, true), true),
1243 ]);
1244
1245 let a = Int32Array::from(vec![Some(1), Some(2), None, None]);
1246 let mut b = ListBuilder::new(Int32Builder::new());
1247 b.extend([Some(vec![Some(1), Some(2)]), Some(vec![None]), None, None]);
1248 let b = b.finish();
1249
1250 let nulls = Some(NullBuffer::from_iter([true, true, true, false]));
1251 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
1252 let s1 = StructArray::new(fields.clone(), values, nulls);
1253
1254 let a = Int32Array::from(vec![None, Some(2), None]);
1255 let mut b = ListBuilder::new(Int32Builder::new());
1256 b.extend([None, None, Some(vec![])]);
1257 let b = b.finish();
1258
1259 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
1260 let s2 = StructArray::new(fields.clone(), values, None);
1261
1262 let opts = SortOptions {
1263 descending: false,
1264 nulls_first: true,
1265 };
1266 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1267 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 {
1276 descending: true,
1277 nulls_first: true,
1278 };
1279 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1280 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 {
1289 descending: true,
1290 nulls_first: false,
1291 };
1292 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1293 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 {
1302 descending: false,
1303 nulls_first: false,
1304 };
1305 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1306 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); }
1314
1315 #[test]
1316 fn test_map() {
1317 let string_builder = StringBuilder::new();
1320 let int_builder = Int32Builder::new();
1321 let mut map1_builder = MapBuilder::new(None, string_builder, int_builder);
1322
1323 map1_builder.keys().append_value("a");
1325 map1_builder.values().append_value(100);
1326 map1_builder.keys().append_value("b");
1327 map1_builder.values().append_value(1);
1328 map1_builder.append(true).unwrap();
1329
1330 map1_builder.keys().append_value("b");
1332 map1_builder.values().append_value(999);
1333 map1_builder.keys().append_value("c");
1334 map1_builder.values().append_value(1);
1335 map1_builder.append(true).unwrap();
1336
1337 map1_builder.append(true).unwrap();
1339
1340 map1_builder.keys().append_value("x");
1342 map1_builder.values().append_value(1);
1343 map1_builder.append(true).unwrap();
1344
1345 let map1 = map1_builder.finish();
1346
1347 let string_builder = StringBuilder::new();
1350 let int_builder = Int32Builder::new();
1351 let mut map2_builder = MapBuilder::new(None, string_builder, int_builder);
1352
1353 map2_builder.keys().append_value("a");
1355 map2_builder.values().append_value(1);
1356 map2_builder.keys().append_value("c");
1357 map2_builder.values().append_value(999);
1358 map2_builder.append(true).unwrap();
1359
1360 map2_builder.keys().append_value("b");
1362 map2_builder.values().append_value(1);
1363 map2_builder.keys().append_value("d");
1364 map2_builder.values().append_value(999);
1365 map2_builder.append(true).unwrap();
1366
1367 map2_builder.keys().append_value("a");
1369 map2_builder.values().append_value(1);
1370 map2_builder.append(true).unwrap();
1371
1372 map2_builder.append(false).unwrap();
1374
1375 let map2 = map2_builder.finish();
1376
1377 let opts = SortOptions {
1378 descending: false,
1379 nulls_first: true,
1380 };
1381 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1382
1383 assert_eq!(cmp(0, 0), Ordering::Greater);
1387
1388 assert_eq!(cmp(1, 1), Ordering::Greater);
1391
1392 assert_eq!(cmp(0, 1), Ordering::Less);
1394
1395 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 {
1408 descending: true,
1409 nulls_first: true,
1410 };
1411 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1412
1413 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 {
1421 descending: false,
1422 nulls_first: false,
1423 };
1424 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1425
1426 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); }
1432
1433 #[test]
1434 fn test_map_vs_list_consistency() {
1435 let string_builder = StringBuilder::new();
1438 let int_builder = Int32Builder::new();
1439 let mut map1_builder = MapBuilder::new(None, string_builder, int_builder);
1440
1441 map1_builder.keys().append_value("a");
1443 map1_builder.values().append_value(1);
1444 map1_builder.keys().append_value("b");
1445 map1_builder.values().append_value(2);
1446 map1_builder.append(true).unwrap();
1447
1448 map1_builder.keys().append_value("x");
1450 map1_builder.values().append_value(10);
1451 map1_builder.append(true).unwrap();
1452
1453 map1_builder.append(true).unwrap();
1455
1456 map1_builder.keys().append_value("c");
1458 map1_builder.values().append_value(3);
1459 map1_builder.append(true).unwrap();
1460
1461 let map1 = map1_builder.finish();
1462
1463 let string_builder = StringBuilder::new();
1465 let int_builder = Int32Builder::new();
1466 let mut map2_builder = MapBuilder::new(None, string_builder, int_builder);
1467
1468 map2_builder.keys().append_value("a");
1470 map2_builder.values().append_value(1);
1471 map2_builder.keys().append_value("b");
1472 map2_builder.values().append_value(2);
1473 map2_builder.append(true).unwrap();
1474
1475 map2_builder.keys().append_value("y");
1477 map2_builder.values().append_value(20);
1478 map2_builder.append(true).unwrap();
1479
1480 map2_builder.keys().append_value("d");
1482 map2_builder.values().append_value(4);
1483 map2_builder.append(true).unwrap();
1484
1485 map2_builder.append(false).unwrap();
1487
1488 let map2 = map2_builder.finish();
1489
1490 let list1: ListArray = map1.clone().into();
1492 let list2: ListArray = map2.clone().into();
1493
1494 let test_cases = [
1495 SortOptions {
1496 descending: false,
1497 nulls_first: true,
1498 },
1499 SortOptions {
1500 descending: true,
1501 nulls_first: true,
1502 },
1503 SortOptions {
1504 descending: false,
1505 nulls_first: false,
1506 },
1507 SortOptions {
1508 descending: true,
1509 nulls_first: false,
1510 },
1511 ];
1512
1513 for opts in test_cases {
1514 let map_cmp = make_comparator(&map1, &map2, opts).unwrap();
1515 let list_cmp = make_comparator(&list1, &list2, opts).unwrap();
1516
1517 for i in 0..map1.len() {
1519 for j in 0..map2.len() {
1520 let map_result = map_cmp(i, j);
1521 let list_result = list_cmp(i, j);
1522 assert_eq!(
1523 map_result, list_result,
1524 "Map comparison and List comparison should be equal for indices ({i}, {j}) with opts {opts:?}. Map: {map_result:?}, List: {list_result:?}"
1525 );
1526 }
1527 }
1528 }
1529 }
1530
1531 #[test]
1532 fn test_dense_union() {
1533 let int_array = Int32Array::from(vec![1, 2, 3]);
1538 let str_array = StringArray::from(vec!["b", "a"]);
1539
1540 let type_ids = [0, 1, 0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1541 let offsets = [0, 0, 1, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
1542
1543 let union_fields = [
1544 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1545 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1546 ]
1547 .into_iter()
1548 .collect::<UnionFields>();
1549
1550 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1551
1552 let array1 =
1553 UnionArray::try_new(union_fields.clone(), type_ids, Some(offsets), children).unwrap();
1554
1555 let int_array2 = Int32Array::from(vec![2, 1]);
1559 let str_array2 = StringArray::from(vec!["a", "c"]);
1560 let type_ids2 = [0, 1, 0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1561 let offsets2 = [0, 0, 1, 1].into_iter().collect::<ScalarBuffer<i32>>();
1562
1563 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)];
1564
1565 let array2 =
1566 UnionArray::try_new(union_fields, type_ids2, Some(offsets2), children2).unwrap();
1567
1568 let opts = SortOptions {
1569 descending: false,
1570 nulls_first: true,
1571 };
1572
1573 let cmp = make_comparator(&array1, &array2, opts).unwrap();
1577
1578 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 {
1603 descending: true,
1604 nulls_first: true,
1605 };
1606 let cmp_desc = make_comparator(&array1, &array2, opts_desc).unwrap();
1607
1608 assert_eq!(cmp_desc(0, 0), Ordering::Greater); assert_eq!(cmp_desc(0, 1), Ordering::Greater); assert_eq!(cmp_desc(1, 1), Ordering::Less); }
1612
1613 #[test]
1614 fn test_sparse_union() {
1615 let int_array = Int32Array::from(vec![Some(1), None, Some(3)]);
1619 let str_array = StringArray::from(vec![None, Some("b"), None]);
1620 let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1621
1622 let union_fields = [
1623 (0, Arc::new(Field::new("a", DataType::Int32, false))),
1624 (1, Arc::new(Field::new("b", DataType::Utf8, false))),
1625 ]
1626 .into_iter()
1627 .collect::<UnionFields>();
1628
1629 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1630
1631 let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
1632
1633 let opts = SortOptions::default();
1634 let cmp = make_comparator(&array, &array, opts).unwrap();
1635
1636 assert_eq!(cmp(0, 2), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Less); }
1641
1642 #[test]
1643 #[should_panic(expected = "index out of bounds")]
1644 fn test_union_out_of_bounds() {
1645 let int_array = Int32Array::from(vec![1, 2]);
1647 let str_array = StringArray::from(vec!["a"]);
1648
1649 let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1650 let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
1651
1652 let union_fields = [
1653 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1654 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1655 ]
1656 .into_iter()
1657 .collect::<UnionFields>();
1658
1659 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1660
1661 let array = UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap();
1662
1663 let opts = SortOptions::default();
1664 let cmp = make_comparator(&array, &array, opts).unwrap();
1665
1666 cmp(0, 3);
1668 }
1669
1670 #[test]
1671 fn test_union_incompatible_fields() {
1672 let int_array1 = Int32Array::from(vec![1, 2]);
1674 let str_array1 = StringArray::from(vec!["a", "b"]);
1675
1676 let type_ids1 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1677 let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
1678
1679 let union_fields1 = [
1680 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1681 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1682 ]
1683 .into_iter()
1684 .collect::<UnionFields>();
1685
1686 let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)];
1687
1688 let array1 =
1689 UnionArray::try_new(union_fields1, type_ids1, Some(offsets1), children1).unwrap();
1690
1691 let int_array2 = Int32Array::from(vec![3, 4]);
1693 let float_array2 = Float64Array::from(vec![1.0, 2.0]);
1694
1695 let type_ids2 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1696 let offsets2 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
1697
1698 let union_fields2 = [
1699 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1700 (1, Arc::new(Field::new("C", DataType::Float64, false))),
1701 ]
1702 .into_iter()
1703 .collect::<UnionFields>();
1704
1705 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(float_array2)];
1706
1707 let array2 =
1708 UnionArray::try_new(union_fields2, type_ids2, Some(offsets2), children2).unwrap();
1709
1710 let opts = SortOptions::default();
1711
1712 let Result::Err(ArrowError::InvalidArgumentError(out)) =
1713 make_comparator(&array1, &array2, opts)
1714 else {
1715 panic!("expected error when making comparator of incompatible union arrays");
1716 };
1717
1718 assert_eq!(
1719 &out,
1720 "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 })]"
1721 );
1722 }
1723
1724 #[test]
1725 fn test_union_incompatible_modes() {
1726 let int_array1 = Int32Array::from(vec![1, 2]);
1728 let str_array1 = StringArray::from(vec!["a", "b"]);
1729
1730 let type_ids1 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1731 let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
1732
1733 let union_fields1 = [
1734 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1735 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1736 ]
1737 .into_iter()
1738 .collect::<UnionFields>();
1739
1740 let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)];
1741
1742 let array1 =
1743 UnionArray::try_new(union_fields1.clone(), type_ids1, Some(offsets1), children1)
1744 .unwrap();
1745
1746 let int_array2 = Int32Array::from(vec![Some(3), None]);
1748 let str_array2 = StringArray::from(vec![None, Some("c")]);
1749
1750 let type_ids2 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1751
1752 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)];
1753
1754 let array2 = UnionArray::try_new(union_fields1, type_ids2, None, children2).unwrap();
1755
1756 let opts = SortOptions::default();
1757
1758 let Result::Err(ArrowError::InvalidArgumentError(out)) =
1759 make_comparator(&array1, &array2, opts)
1760 else {
1761 panic!("expected error when making comparator of union arrays with different modes");
1762 };
1763
1764 assert_eq!(
1765 &out,
1766 "Cannot compare UnionArrays with different modes: left=Dense, right=Sparse"
1767 );
1768 }
1769
1770 #[test]
1771 fn test_null_array_cmp() {
1772 let a = NullArray::new(3);
1773 let b = NullArray::new(3);
1774 let cmp = make_comparator(&a, &b, SortOptions::default()).unwrap();
1775
1776 assert_eq!(cmp(0, 0), Ordering::Equal);
1777 assert_eq!(cmp(0, 1), Ordering::Equal);
1778 assert_eq!(cmp(2, 0), Ordering::Equal);
1779 }
1780
1781 #[test]
1782 fn test_run_end_encoded_int32() {
1783 let run_ends1 = Int32Array::from(vec![2, 5, 6]);
1787 let values1 = Int32Array::from(vec![1, 2, 3]);
1788 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1789
1790 let run_ends2 = Int32Array::from(vec![1, 3, 6]);
1793 let values2 = Int32Array::from(vec![1, 2, 3]);
1794 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1795
1796 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1797
1798 assert_eq!(cmp(0, 0), Ordering::Equal);
1800 assert_eq!(cmp(0, 1), Ordering::Less);
1802 assert_eq!(cmp(2, 1), Ordering::Equal);
1804 assert_eq!(cmp(5, 5), Ordering::Equal);
1806 assert_eq!(cmp(1, 2), Ordering::Less);
1808 assert_eq!(cmp(4, 4), Ordering::Less);
1810 }
1811
1812 #[test]
1813 fn test_run_end_encoded_with_nulls() {
1814 let run_ends1 = Int32Array::from(vec![2, 4, 5]);
1818 let values1 = Int32Array::from(vec![Some(1), None, Some(2)]);
1819 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1820
1821 let run_ends2 = Int32Array::from(vec![1, 3, 4, 5]);
1824 let values2 = Int32Array::from(vec![None, Some(1), Some(2), None]);
1825 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1826
1827 let opts = SortOptions::default();
1828 let cmp = make_comparator(&array1, &array2, opts).unwrap();
1829
1830 assert_eq!(cmp(0, 1), Ordering::Equal);
1832 assert_eq!(cmp(2, 0), Ordering::Equal);
1834 assert_eq!(cmp(0, 0), Ordering::Greater);
1836 assert_eq!(cmp(2, 1), Ordering::Less);
1838 }
1839
1840 #[test]
1841 fn test_run_end_encoded_int16() {
1842 let run_ends1 = Int16Array::from(vec![3_i16, 5, 6]);
1844 let values1 = StringArray::from(vec!["a", "b", "c"]);
1845 let array1 = RunArray::<Int16Type>::try_new(&run_ends1, &values1).unwrap();
1846
1847 let run_ends2 = Int16Array::from(vec![2_i16, 4, 6]);
1848 let values2 = StringArray::from(vec!["a", "b", "c"]);
1849 let array2 = RunArray::<Int16Type>::try_new(&run_ends2, &values2).unwrap();
1850
1851 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1852
1853 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); }
1860
1861 #[test]
1862 fn test_run_end_encoded_int64() {
1863 let run_ends1 = Int64Array::from(vec![2_i64, 4, 6]);
1865 let values1 = Int64Array::from(vec![10_i64, 20, 30]);
1866 let array1 = RunArray::<Int64Type>::try_new(&run_ends1, &values1).unwrap();
1867
1868 let run_ends2 = Int64Array::from(vec![3_i64, 5, 6]);
1869 let values2 = Int64Array::from(vec![10_i64, 20, 30]);
1870 let array2 = RunArray::<Int64Type>::try_new(&run_ends2, &values2).unwrap();
1871
1872 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1873
1874 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); }
1881
1882 #[test]
1883 fn test_run_end_encoded_sliced() {
1884 let run_ends = Int32Array::from(vec![2, 5, 7, 8]);
1888 let values = Int32Array::from(vec![1, 2, 3, 4]);
1889 let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1890
1891 let slice1 = array.slice(1, 4);
1893 let slice2 = array.slice(3, 4);
1895
1896 let cmp = make_comparator(&slice1, &slice2, SortOptions::default()).unwrap();
1897
1898 assert_eq!(cmp(0, 0), Ordering::Less);
1900 assert_eq!(cmp(1, 0), Ordering::Equal);
1902 assert_eq!(cmp(3, 2), Ordering::Less);
1904 assert_eq!(cmp(1, 3), Ordering::Less);
1906
1907 let run_ends2 = Int32Array::from(vec![2, 4]);
1909 let values2 = Int32Array::from(vec![1, 2]);
1910 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1911
1912 let cmp = make_comparator(&slice1, &array2, SortOptions::default()).unwrap();
1913
1914 assert_eq!(cmp(0, 0), Ordering::Equal);
1916 assert_eq!(cmp(1, 1), Ordering::Greater);
1918 assert_eq!(cmp(3, 3), Ordering::Equal);
1920 }
1921
1922 #[test]
1923 fn test_run_end_encoded_sliced_with_nulls() {
1924 let run_ends = Int32Array::from(vec![2, 4, 6, 7, 8]);
1928 let values = Int32Array::from(vec![Some(1), None, Some(2), None, Some(3)]);
1929 let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1930
1931 let slice1 = array.slice(1, 5);
1933 let slice2 = array.slice(3, 5);
1935
1936 let opts = SortOptions::default(); let cmp = make_comparator(&slice1, &slice2, opts).unwrap();
1938
1939 assert_eq!(cmp(0, 0), Ordering::Greater);
1941 assert_eq!(cmp(1, 0), Ordering::Equal);
1943 assert_eq!(cmp(1, 1), Ordering::Less);
1945 assert_eq!(cmp(3, 1), Ordering::Equal);
1947 assert_eq!(cmp(4, 4), Ordering::Less);
1949 assert_eq!(cmp(3, 3), Ordering::Greater);
1951 }
1952
1953 #[test]
1954 fn test_run_end_encoded_different_types() {
1955 let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1957 let values1 = Int32Array::from(vec![1, 2, 3]);
1958 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1959
1960 let run_ends2 = Int64Array::from(vec![2_i64, 4, 6]);
1961 let values2 = Int64Array::from(vec![1_i64, 2, 3]);
1962 let array2 = RunArray::<Int64Type>::try_new(&run_ends2, &values2).unwrap();
1963
1964 let result = make_comparator(&array1, &array2, SortOptions::default());
1965 assert!(result.is_err());
1966 let err = match result {
1967 Err(e) => e.to_string(),
1968 Ok(_) => panic!("Expected error"),
1969 };
1970 assert!(err.contains("Cannot compare RunEndEncoded arrays"));
1971 }
1972}