1#![deny(rustdoc::broken_intra_doc_links)]
33#![warn(missing_docs)]
34
35use arrow_array::cast::AsArray;
36use arrow_array::types::*;
37use arrow_array::*;
38use arrow_buffer::{ArrowNativeType, NullBuffer};
39use arrow_schema::{ArrowError, DataType, SortOptions};
40use std::{cmp::Ordering, collections::HashMap};
41
42fn compare_run_end_encoded<R: RunEndIndexType>(
43 left: &dyn Array,
44 right: &dyn Array,
45 opts: SortOptions,
46) -> Result<DynComparator, ArrowError> {
47 let left = left.as_run::<R>();
48 let right = right.as_run::<R>();
49
50 let c_opts = child_opts(opts);
51 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
52
53 let l_run_ends = left.run_ends().clone();
54 let r_run_ends = right.run_ends().clone();
55
56 let f = compare(left, right, opts, move |i, j| {
57 let l_physical = l_run_ends.get_physical_index(i);
58 let r_physical = r_run_ends.get_physical_index(j);
59 cmp(l_physical, r_physical)
60 });
61 Ok(f)
62}
63
64pub type DynComparator = Box<dyn Fn(usize, usize) -> Ordering + Send + Sync>;
66
67fn child_opts(opts: SortOptions) -> SortOptions {
70 SortOptions {
71 descending: false,
72 nulls_first: opts.nulls_first != opts.descending,
73 }
74}
75
76fn compare<A, F>(l: &A, r: &A, opts: SortOptions, cmp: F) -> DynComparator
77where
78 A: Array + Clone,
79 F: Fn(usize, usize) -> Ordering + Send + Sync + 'static,
80{
81 let l = l.logical_nulls().filter(|x| x.null_count() > 0);
82 let r = r.logical_nulls().filter(|x| x.null_count() > 0);
83 match (opts.nulls_first, opts.descending) {
84 (true, true) => compare_impl::<true, true, _>(l, r, cmp),
85 (true, false) => compare_impl::<true, false, _>(l, r, cmp),
86 (false, true) => compare_impl::<false, true, _>(l, r, cmp),
87 (false, false) => compare_impl::<false, false, _>(l, r, cmp),
88 }
89}
90
91fn compare_impl<const NULLS_FIRST: bool, const DESCENDING: bool, F>(
92 l: Option<NullBuffer>,
93 r: Option<NullBuffer>,
94 cmp: F,
95) -> DynComparator
96where
97 F: Fn(usize, usize) -> Ordering + Send + Sync + 'static,
98{
99 let cmp = move |i, j| match DESCENDING {
100 true => cmp(i, j).reverse(),
101 false => cmp(i, j),
102 };
103
104 let (left_null, right_null) = match NULLS_FIRST {
105 true => (Ordering::Less, Ordering::Greater),
106 false => (Ordering::Greater, Ordering::Less),
107 };
108
109 match (l, r) {
110 (None, None) => Box::new(cmp),
111 (Some(l), None) => Box::new(move |i, j| match l.is_null(i) {
112 true => left_null,
113 false => cmp(i, j),
114 }),
115 (None, Some(r)) => Box::new(move |i, j| match r.is_null(j) {
116 true => right_null,
117 false => cmp(i, j),
118 }),
119 (Some(l), Some(r)) => Box::new(move |i, j| match (l.is_null(i), r.is_null(j)) {
120 (true, true) => Ordering::Equal,
121 (true, false) => left_null,
122 (false, true) => right_null,
123 (false, false) => cmp(i, j),
124 }),
125 }
126}
127
128fn compare_primitive<T: ArrowPrimitiveType>(
129 left: &dyn Array,
130 right: &dyn Array,
131 opts: SortOptions,
132) -> DynComparator
133where
134 T::Native: ArrowNativeTypeOp,
135{
136 let left = left.as_primitive::<T>();
137 let right = right.as_primitive::<T>();
138 let l_values = left.values().clone();
139 let r_values = right.values().clone();
140
141 compare(&left, &right, opts, move |i, j| {
142 l_values[i].compare(r_values[j])
143 })
144}
145
146fn compare_boolean(left: &dyn Array, right: &dyn Array, opts: SortOptions) -> DynComparator {
147 let left = left.as_boolean();
148 let right = right.as_boolean();
149
150 let l_values = left.values().clone();
151 let r_values = right.values().clone();
152
153 compare(left, right, opts, move |i, j| {
154 l_values.value(i).cmp(&r_values.value(j))
155 })
156}
157
158fn compare_bytes<T: ByteArrayType>(
159 left: &dyn Array,
160 right: &dyn Array,
161 opts: SortOptions,
162) -> DynComparator {
163 let left = left.as_bytes::<T>();
164 let right = right.as_bytes::<T>();
165
166 let l = left.clone();
167 let r = right.clone();
168 compare(left, right, opts, move |i, j| {
169 let l: &[u8] = l.value(i).as_ref();
170 let r: &[u8] = r.value(j).as_ref();
171 l.cmp(r)
172 })
173}
174
175fn compare_byte_view<T: ByteViewType>(
176 left: &dyn Array,
177 right: &dyn Array,
178 opts: SortOptions,
179) -> DynComparator {
180 let left = left.as_byte_view::<T>();
181 let right = right.as_byte_view::<T>();
182
183 let l = left.clone();
184 let r = right.clone();
185 compare(left, right, opts, move |i, j| {
186 compare_byte_view_values(&l, i, &r, j)
187 })
188}
189
190fn compare_byte_view_values<T: ByteViewType>(
191 left: &GenericByteViewArray<T>,
192 left_idx: usize,
193 right: &GenericByteViewArray<T>,
194 right_idx: usize,
195) -> Ordering {
196 assert!(left_idx < left.len());
197 assert!(right_idx < right.len());
198
199 if left.data_buffers().is_empty() && right.data_buffers().is_empty() {
200 let l_view = unsafe { left.views().get_unchecked(left_idx) };
201 let r_view = unsafe { right.views().get_unchecked(right_idx) };
202 return GenericByteViewArray::<T>::inline_key_fast(*l_view)
203 .cmp(&GenericByteViewArray::<T>::inline_key_fast(*r_view));
204 }
205
206 unsafe { GenericByteViewArray::compare_unchecked(left, left_idx, right, right_idx) }
207}
208
209fn compare_dict<K: ArrowDictionaryKeyType>(
210 left: &dyn Array,
211 right: &dyn Array,
212 opts: SortOptions,
213) -> Result<DynComparator, ArrowError> {
214 let left = left.as_dictionary::<K>();
215 let right = right.as_dictionary::<K>();
216
217 let c_opts = child_opts(opts);
218 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
219 let left_keys = left.keys().values().clone();
220 let right_keys = right.keys().values().clone();
221
222 let f = compare(left, right, opts, move |i, j| {
223 let l = left_keys[i].as_usize();
224 let r = right_keys[j].as_usize();
225 cmp(l, r)
226 });
227 Ok(f)
228}
229
230fn compare_list<O: OffsetSizeTrait>(
231 left: &dyn Array,
232 right: &dyn Array,
233 opts: SortOptions,
234) -> Result<DynComparator, ArrowError> {
235 let left = left.as_list::<O>();
236 let right = right.as_list::<O>();
237
238 let c_opts = child_opts(opts);
239 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
240
241 let l_o = left.offsets().clone();
242 let r_o = right.offsets().clone();
243 let f = compare(left, right, opts, move |i, j| {
244 let l_end = l_o[i + 1].as_usize();
245 let l_start = l_o[i].as_usize();
246
247 let r_end = r_o[j + 1].as_usize();
248 let r_start = r_o[j].as_usize();
249
250 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
251 match cmp(i, j) {
252 Ordering::Equal => {}
253 r => return r,
254 }
255 }
256 (l_end - l_start).cmp(&(r_end - r_start))
257 });
258 Ok(f)
259}
260
261fn compare_fixed_list(
262 left: &dyn Array,
263 right: &dyn Array,
264 opts: SortOptions,
265) -> Result<DynComparator, ArrowError> {
266 let left = left.as_fixed_size_list();
267 let right = right.as_fixed_size_list();
268
269 let c_opts = child_opts(opts);
270 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
271
272 let l_size = left.value_length().to_usize().unwrap();
273 let r_size = right.value_length().to_usize().unwrap();
274 let size_cmp = l_size.cmp(&r_size);
275
276 let f = compare(left, right, opts, move |i, j| {
277 let l_start = i * l_size;
278 let l_end = l_start + l_size;
279 let r_start = j * r_size;
280 let r_end = r_start + r_size;
281 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
282 match cmp(i, j) {
283 Ordering::Equal => {}
284 r => return r,
285 }
286 }
287 size_cmp
288 });
289 Ok(f)
290}
291
292fn compare_list_view<O: OffsetSizeTrait>(
293 left: &dyn Array,
294 right: &dyn Array,
295 opts: SortOptions,
296) -> Result<DynComparator, ArrowError> {
297 let left = left.as_list_view::<O>();
298 let right = right.as_list_view::<O>();
299
300 let c_opts = child_opts(opts);
301 let cmp = make_comparator(left.values().as_ref(), right.values().as_ref(), c_opts)?;
302
303 let l_offsets = left.offsets().clone();
304 let l_sizes = left.sizes().clone();
305 let r_offsets = right.offsets().clone();
306 let r_sizes = right.sizes().clone();
307
308 let f = compare(left, right, opts, move |i, j| {
309 let l_start = l_offsets[i].as_usize();
310 let l_len = l_sizes[i].as_usize();
311 let l_end = l_start + l_len;
312
313 let r_start = r_offsets[j].as_usize();
314 let r_len = r_sizes[j].as_usize();
315 let r_end = r_start + r_len;
316
317 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
318 match cmp(i, j) {
319 Ordering::Equal => {}
320 r => return r,
321 }
322 }
323 l_len.cmp(&r_len)
324 });
325 Ok(f)
326}
327
328fn compare_map(
329 left: &dyn Array,
330 right: &dyn Array,
331 opts: SortOptions,
332) -> Result<DynComparator, ArrowError> {
333 let left = left.as_map();
334 let right = right.as_map();
335
336 let c_opts = child_opts(opts);
337 let cmp = make_comparator(left.entries(), right.entries(), c_opts)?;
338
339 let l_o = left.offsets().clone();
340 let r_o = right.offsets().clone();
341 let f = compare(left, right, opts, move |i, j| {
342 let l_end = l_o[i + 1].as_usize();
343 let l_start = l_o[i].as_usize();
344
345 let r_end = r_o[j + 1].as_usize();
346 let r_start = r_o[j].as_usize();
347
348 for (i, j) in (l_start..l_end).zip(r_start..r_end) {
349 match cmp(i, j) {
350 Ordering::Equal => {}
351 r => return r,
352 }
353 }
354 (l_end - l_start).cmp(&(r_end - r_start))
355 });
356 Ok(f)
357}
358
359fn compare_struct(
360 left: &dyn Array,
361 right: &dyn Array,
362 opts: SortOptions,
363) -> Result<DynComparator, ArrowError> {
364 let left = left.as_struct();
365 let right = right.as_struct();
366
367 if left.columns().len() != right.columns().len() {
368 return Err(ArrowError::InvalidArgumentError(
369 "Cannot compare StructArray with different number of columns".to_string(),
370 ));
371 }
372
373 let c_opts = child_opts(opts);
374 let columns = left.columns().iter().zip(right.columns());
375 let comparators = columns
376 .map(|(l, r)| make_comparator(l, r, c_opts))
377 .collect::<Result<Vec<_>, _>>()?;
378
379 let f = compare(left, right, opts, move |i, j| {
380 for cmp in &comparators {
381 match cmp(i, j) {
382 Ordering::Equal => {}
383 r => return r,
384 }
385 }
386 Ordering::Equal
387 });
388 Ok(f)
389}
390
391fn compare_union(
392 left: &dyn Array,
393 right: &dyn Array,
394 opts: SortOptions,
395) -> Result<DynComparator, ArrowError> {
396 let left = left.as_union();
397 let right = right.as_union();
398
399 let DataType::Union(left_fields, left_mode) = left.data_type() else {
400 unreachable!()
401 };
402 let DataType::Union(right_fields, right_mode) = right.data_type() else {
403 unreachable!()
404 };
405
406 if left_fields != right_fields {
407 return Err(ArrowError::InvalidArgumentError(format!(
408 "Cannot compare UnionArrays with different fields: left={left_fields:?}, right={right_fields:?}"
409 )));
410 }
411
412 if left_mode != right_mode {
413 return Err(ArrowError::InvalidArgumentError(format!(
414 "Cannot compare UnionArrays with different modes: left={left_mode:?}, right={right_mode:?}"
415 )));
416 }
417
418 let c_opts = child_opts(opts);
419
420 let mut field_comparators = HashMap::with_capacity(left_fields.len());
421
422 for (type_id, _field) in left_fields.iter() {
423 let left_child = left.child(type_id);
424 let right_child = right.child(type_id);
425 let cmp = make_comparator(left_child.as_ref(), right_child.as_ref(), c_opts)?;
426
427 field_comparators.insert(type_id, cmp);
428 }
429
430 let left_type_ids = left.type_ids().clone();
431 let right_type_ids = right.type_ids().clone();
432
433 let left_offsets = left.offsets().cloned();
434 let right_offsets = right.offsets().cloned();
435
436 let f = compare(left, right, opts, move |i, j| {
437 let left_type_id = left_type_ids[i];
438 let right_type_id = right_type_ids[j];
439
440 match left_type_id.cmp(&right_type_id) {
441 Ordering::Equal => {
442 let left_offset = left_offsets.as_ref().map(|o| o[i] as usize).unwrap_or(i);
443 let right_offset = right_offsets.as_ref().map(|o| o[j] as usize).unwrap_or(j);
444
445 let cmp = field_comparators
446 .get(&left_type_id)
447 .expect("type id not found in field_comparators");
448
449 cmp(left_offset, right_offset)
450 }
451 other => other,
452 }
453 });
454 Ok(f)
455}
456
457pub fn make_comparator(
465 left: &dyn Array,
466 right: &dyn Array,
467 opts: SortOptions,
468) -> Result<DynComparator, ArrowError> {
469 use arrow_schema::DataType::*;
470
471 match (left.data_type(), right.data_type()) {
474 (Decimal32(_, s1), Decimal32(_, s2))
475 | (Decimal64(_, s1), Decimal64(_, s2))
476 | (Decimal128(_, s1), Decimal128(_, s2))
477 | (Decimal256(_, s1), Decimal256(_, s2))
478 if s1 != s2 =>
479 {
480 return Err(ArrowError::InvalidArgumentError(format!(
481 "Can't compare decimal arrays with different scales: {:?}, {:?}",
482 left.data_type(),
483 right.data_type()
484 )));
485 }
486 _ => {}
487 }
488
489 macro_rules! primitive_helper {
490 ($t:ty, $left:expr, $right:expr, $nulls_first:expr) => {
491 Ok(compare_primitive::<$t>($left, $right, $nulls_first))
492 };
493 }
494 downcast_primitive! {
495 left.data_type(), right.data_type() => (primitive_helper, left, right, opts),
496 (Boolean, Boolean) => Ok(compare_boolean(left, right, opts)),
497 (Utf8, Utf8) => Ok(compare_bytes::<Utf8Type>(left, right, opts)),
498 (LargeUtf8, LargeUtf8) => Ok(compare_bytes::<LargeUtf8Type>(left, right, opts)),
499 (Utf8View, Utf8View) => Ok(compare_byte_view::<StringViewType>(left, right, opts)),
500 (Binary, Binary) => Ok(compare_bytes::<BinaryType>(left, right, opts)),
501 (LargeBinary, LargeBinary) => Ok(compare_bytes::<LargeBinaryType>(left, right, opts)),
502 (BinaryView, BinaryView) => Ok(compare_byte_view::<BinaryViewType>(left, right, opts)),
503 (FixedSizeBinary(_), FixedSizeBinary(_)) => {
504 let left = left.as_fixed_size_binary();
505 let right = right.as_fixed_size_binary();
506
507 let l = left.clone();
508 let r = right.clone();
509 Ok(compare(left, right, opts, move |i, j| {
510 l.value(i).cmp(r.value(j))
511 }))
512 },
513 (List(_), List(_)) => compare_list::<i32>(left, right, opts),
514 (LargeList(_), LargeList(_)) => compare_list::<i64>(left, right, opts),
515 (ListView(_), ListView(_)) => compare_list_view::<i32>(left, right, opts),
516 (LargeListView(_), LargeListView(_)) => compare_list_view::<i64>(left, right, opts),
517 (FixedSizeList(_, _), FixedSizeList(_, _)) => compare_fixed_list(left, right, opts),
518 (Struct(_), Struct(_)) => compare_struct(left, right, opts),
519 (Dictionary(l_key, _), Dictionary(r_key, _)) => {
520 macro_rules! dict_helper {
521 ($t:ty, $left:expr, $right:expr, $opts: expr) => {
522 compare_dict::<$t>($left, $right, $opts)
523 };
524 }
525 downcast_integer! {
526 l_key.as_ref(), r_key.as_ref() => (dict_helper, left, right, opts),
527 _ => unreachable!()
528 }
529 },
530 (RunEndEncoded(l_run_ends, _), RunEndEncoded(r_run_ends, _)) => {
531 macro_rules! run_end_helper {
532 ($t:ty, $left:expr, $right:expr, $opts:expr) => {
533 compare_run_end_encoded::<$t>($left, $right, $opts)
534 };
535 }
536 downcast_run_end_index! {
537 l_run_ends.data_type(), r_run_ends.data_type() => (run_end_helper, left, right, opts),
538 _ => Err(ArrowError::InvalidArgumentError(format!(
539 "Cannot compare RunEndEncoded arrays with different run ends types: left={:?}, right={:?}",
540 l_run_ends.data_type(),
541 r_run_ends.data_type()
542 )))
543 }
544 },
545 (Map(_, _), Map(_, _)) => compare_map(left, right, opts),
546 (Null, Null) => Ok(Box::new(|_, _| Ordering::Equal)),
547 (Union(_, _), Union(_, _)) => compare_union(left, right, opts),
548 (lhs, rhs) => Err(ArrowError::InvalidArgumentError(match lhs == rhs {
549 true => format!("The data type type {lhs:?} has no natural order"),
550 false => "Can't compare arrays of different types".to_string(),
551 }))
552 }
553}
554
555#[cfg(test)]
556mod tests {
557 use super::*;
558 use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder};
559 use arrow_buffer::{IntervalDayTime, NullBuffer, OffsetBuffer, ScalarBuffer, i256};
560 use arrow_schema::{ArrowError, DataType, Field, Fields, UnionFields};
561 use half::f16;
562 use std::cmp::Ordering;
563 use std::sync::Arc;
564
565 #[test]
566 fn test_fixed_size_binary() {
567 let items = vec![vec![1u8], vec![2u8]];
568 let array = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
569
570 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
571
572 assert_eq!(Ordering::Less, cmp(0, 1));
573 }
574
575 #[test]
576 fn test_fixed_size_binary_fixed_size_binary() {
577 let items = vec![vec![1u8]];
578 let array1 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
579 let items = vec![vec![2u8]];
580 let array2 = FixedSizeBinaryArray::try_from_iter(items.into_iter()).unwrap();
581
582 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
583
584 assert_eq!(Ordering::Less, cmp(0, 0));
585 }
586
587 #[test]
588 fn test_i32() {
589 let array = Int32Array::from(vec![1, 2]);
590
591 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
592
593 assert_eq!(Ordering::Less, (cmp)(0, 1));
594 }
595
596 #[test]
597 fn test_i32_i32() {
598 let array1 = Int32Array::from(vec![1]);
599 let array2 = Int32Array::from(vec![2]);
600
601 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
602
603 assert_eq!(Ordering::Less, cmp(0, 0));
604 }
605
606 #[cfg_attr(miri, ignore)] #[test]
608 fn test_f16() {
609 let array = Float16Array::from(vec![f16::from_f32(1.0), f16::from_f32(2.0)]);
610
611 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
612
613 assert_eq!(Ordering::Less, cmp(0, 1));
614 }
615
616 #[test]
617 fn test_f64() {
618 let array = Float64Array::from(vec![1.0, 2.0]);
619
620 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
621
622 assert_eq!(Ordering::Less, cmp(0, 1));
623 }
624
625 #[test]
626 fn test_f64_nan() {
627 let array = Float64Array::from(vec![1.0, f64::NAN]);
628
629 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
630
631 assert_eq!(Ordering::Less, cmp(0, 1));
632 assert_eq!(Ordering::Equal, cmp(1, 1));
633 }
634
635 #[test]
636 fn test_f64_zeros() {
637 let array = Float64Array::from(vec![-0.0, 0.0]);
638
639 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
640
641 assert_eq!(Ordering::Less, cmp(0, 1));
642 assert_eq!(Ordering::Greater, cmp(1, 0));
643 }
644
645 #[test]
646 fn test_interval_day_time() {
647 let array = IntervalDayTimeArray::from(vec![
648 IntervalDayTimeType::make_value(0, 1000),
650 IntervalDayTimeType::make_value(1, 2),
652 IntervalDayTimeType::make_value(0, 90_000_000),
654 ]);
655
656 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
657
658 assert_eq!(Ordering::Less, cmp(0, 1));
659 assert_eq!(Ordering::Greater, cmp(1, 0));
660
661 assert_eq!(Ordering::Greater, cmp(1, 2));
665 assert_eq!(Ordering::Less, cmp(2, 1));
666 }
667
668 #[test]
669 fn test_interval_year_month() {
670 let array = IntervalYearMonthArray::from(vec![
671 IntervalYearMonthType::make_value(1, 0),
673 IntervalYearMonthType::make_value(0, 13),
675 IntervalYearMonthType::make_value(1, 1),
677 ]);
678
679 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
680
681 assert_eq!(Ordering::Less, cmp(0, 1));
682 assert_eq!(Ordering::Greater, cmp(1, 0));
683
684 assert_eq!(Ordering::Equal, cmp(1, 2));
686 assert_eq!(Ordering::Equal, cmp(2, 1));
687 }
688
689 #[test]
690 fn test_interval_month_day_nano() {
691 let array = IntervalMonthDayNanoArray::from(vec![
692 IntervalMonthDayNanoType::make_value(0, 100, 0),
694 IntervalMonthDayNanoType::make_value(1, 0, 0),
696 IntervalMonthDayNanoType::make_value(0, 100, 2),
698 ]);
699
700 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
701
702 assert_eq!(Ordering::Less, cmp(0, 1));
703 assert_eq!(Ordering::Greater, cmp(1, 0));
704
705 assert_eq!(Ordering::Greater, cmp(1, 2));
709 assert_eq!(Ordering::Less, cmp(2, 1));
710 }
711
712 #[test]
713 fn test_decimali32() {
714 let array = vec![Some(5_i32), Some(2_i32), Some(3_i32)]
715 .into_iter()
716 .collect::<Decimal32Array>()
717 .with_precision_and_scale(8, 6)
718 .unwrap();
719
720 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
721 assert_eq!(Ordering::Less, cmp(1, 0));
722 assert_eq!(Ordering::Greater, cmp(0, 2));
723 }
724
725 #[test]
726 fn test_decimali64() {
727 let array = vec![Some(5_i64), Some(2_i64), Some(3_i64)]
728 .into_iter()
729 .collect::<Decimal64Array>()
730 .with_precision_and_scale(16, 6)
731 .unwrap();
732
733 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
734 assert_eq!(Ordering::Less, cmp(1, 0));
735 assert_eq!(Ordering::Greater, cmp(0, 2));
736 }
737
738 #[test]
739 fn test_decimali128() {
740 let array = vec![Some(5_i128), Some(2_i128), Some(3_i128)]
741 .into_iter()
742 .collect::<Decimal128Array>()
743 .with_precision_and_scale(23, 6)
744 .unwrap();
745
746 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
747 assert_eq!(Ordering::Less, cmp(1, 0));
748 assert_eq!(Ordering::Greater, cmp(0, 2));
749 }
750
751 #[test]
752 fn test_decimali256() {
753 let array = vec![
754 Some(i256::from_i128(5_i128)),
755 Some(i256::from_i128(2_i128)),
756 Some(i256::from_i128(3_i128)),
757 ]
758 .into_iter()
759 .collect::<Decimal256Array>()
760 .with_precision_and_scale(53, 6)
761 .unwrap();
762
763 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
764 assert_eq!(Ordering::Less, cmp(1, 0));
765 assert_eq!(Ordering::Greater, cmp(0, 2));
766 }
767
768 #[test]
769 fn test_dict() {
770 let data = vec!["a", "b", "c", "a", "a", "c", "c"];
771 let array = data.into_iter().collect::<DictionaryArray<Int16Type>>();
772
773 let cmp = make_comparator(&array, &array, SortOptions::default()).unwrap();
774
775 assert_eq!(Ordering::Less, cmp(0, 1));
776 assert_eq!(Ordering::Equal, cmp(3, 4));
777 assert_eq!(Ordering::Greater, cmp(2, 3));
778 }
779
780 #[test]
781 fn test_multiple_dict() {
782 let d1 = vec!["a", "b", "c", "d"];
783 let a1 = d1.into_iter().collect::<DictionaryArray<Int16Type>>();
784 let d2 = vec!["e", "f", "g", "a"];
785 let a2 = d2.into_iter().collect::<DictionaryArray<Int16Type>>();
786
787 let cmp = make_comparator(&a1, &a2, SortOptions::default()).unwrap();
788
789 assert_eq!(Ordering::Less, cmp(0, 0));
790 assert_eq!(Ordering::Equal, cmp(0, 3));
791 assert_eq!(Ordering::Greater, cmp(1, 3));
792 }
793
794 #[test]
795 fn test_primitive_dict() {
796 let values = Int32Array::from(vec![1_i32, 0, 2, 5]);
797 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
798 let array1 = DictionaryArray::new(keys, Arc::new(values));
799
800 let values = Int32Array::from(vec![2_i32, 3, 4, 5]);
801 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
802 let array2 = DictionaryArray::new(keys, Arc::new(values));
803
804 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
805
806 assert_eq!(Ordering::Less, cmp(0, 0));
807 assert_eq!(Ordering::Less, cmp(0, 3));
808 assert_eq!(Ordering::Equal, cmp(3, 3));
809 assert_eq!(Ordering::Greater, cmp(3, 1));
810 assert_eq!(Ordering::Greater, cmp(3, 2));
811 }
812
813 #[test]
814 fn test_float_dict() {
815 let values = Float32Array::from(vec![1.0, 0.5, 2.1, 5.5]);
816 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
817 let array1 = DictionaryArray::try_new(keys, Arc::new(values)).unwrap();
818
819 let values = Float32Array::from(vec![1.2, 3.2, 4.0, 5.5]);
820 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
821 let array2 = DictionaryArray::new(keys, Arc::new(values));
822
823 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
824
825 assert_eq!(Ordering::Less, cmp(0, 0));
826 assert_eq!(Ordering::Less, cmp(0, 3));
827 assert_eq!(Ordering::Equal, cmp(3, 3));
828 assert_eq!(Ordering::Greater, cmp(3, 1));
829 assert_eq!(Ordering::Greater, cmp(3, 2));
830 }
831
832 #[test]
833 fn test_timestamp_dict() {
834 let values = TimestampSecondArray::from(vec![1, 0, 2, 5]);
835 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
836 let array1 = DictionaryArray::new(keys, Arc::new(values));
837
838 let values = TimestampSecondArray::from(vec![2, 3, 4, 5]);
839 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
840 let array2 = DictionaryArray::new(keys, Arc::new(values));
841
842 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
843
844 assert_eq!(Ordering::Less, cmp(0, 0));
845 assert_eq!(Ordering::Less, cmp(0, 3));
846 assert_eq!(Ordering::Equal, cmp(3, 3));
847 assert_eq!(Ordering::Greater, cmp(3, 1));
848 assert_eq!(Ordering::Greater, cmp(3, 2));
849 }
850
851 #[test]
852 fn test_interval_dict() {
853 let v1 = IntervalDayTime::new(0, 1);
854 let v2 = IntervalDayTime::new(0, 2);
855 let v3 = IntervalDayTime::new(12, 2);
856
857 let values = IntervalDayTimeArray::from(vec![Some(v1), Some(v2), None, Some(v3)]);
858 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
859 let array1 = DictionaryArray::new(keys, Arc::new(values));
860
861 let values = IntervalDayTimeArray::from(vec![Some(v3), Some(v2), None, Some(v1)]);
862 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
863 let array2 = DictionaryArray::new(keys, Arc::new(values));
864
865 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
866
867 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)); }
873
874 #[test]
875 fn test_duration_dict() {
876 let values = DurationSecondArray::from(vec![1, 0, 2, 5]);
877 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
878 let array1 = DictionaryArray::new(keys, Arc::new(values));
879
880 let values = DurationSecondArray::from(vec![2, 3, 4, 5]);
881 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
882 let array2 = DictionaryArray::new(keys, Arc::new(values));
883
884 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
885
886 assert_eq!(Ordering::Less, cmp(0, 0));
887 assert_eq!(Ordering::Less, cmp(0, 3));
888 assert_eq!(Ordering::Equal, cmp(3, 3));
889 assert_eq!(Ordering::Greater, cmp(3, 1));
890 assert_eq!(Ordering::Greater, cmp(3, 2));
891 }
892
893 #[test]
894 fn test_decimal_dict() {
895 let values = Decimal128Array::from(vec![1, 0, 2, 5]);
896 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
897 let array1 = DictionaryArray::new(keys, Arc::new(values));
898
899 let values = Decimal128Array::from(vec![2, 3, 4, 5]);
900 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
901 let array2 = DictionaryArray::new(keys, Arc::new(values));
902
903 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
904
905 assert_eq!(Ordering::Less, cmp(0, 0));
906 assert_eq!(Ordering::Less, cmp(0, 3));
907 assert_eq!(Ordering::Equal, cmp(3, 3));
908 assert_eq!(Ordering::Greater, cmp(3, 1));
909 assert_eq!(Ordering::Greater, cmp(3, 2));
910 }
911
912 #[test]
913 fn test_decimal256_dict() {
914 let values = Decimal256Array::from(vec![
915 i256::from_i128(1),
916 i256::from_i128(0),
917 i256::from_i128(2),
918 i256::from_i128(5),
919 ]);
920 let keys = Int8Array::from_iter_values([0, 0, 1, 3]);
921 let array1 = DictionaryArray::new(keys, Arc::new(values));
922
923 let values = Decimal256Array::from(vec![
924 i256::from_i128(2),
925 i256::from_i128(3),
926 i256::from_i128(4),
927 i256::from_i128(5),
928 ]);
929 let keys = Int8Array::from_iter_values([0, 1, 1, 3]);
930 let array2 = DictionaryArray::new(keys, Arc::new(values));
931
932 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
933
934 assert_eq!(Ordering::Less, cmp(0, 0));
935 assert_eq!(Ordering::Less, cmp(0, 3));
936 assert_eq!(Ordering::Equal, cmp(3, 3));
937 assert_eq!(Ordering::Greater, cmp(3, 1));
938 assert_eq!(Ordering::Greater, cmp(3, 2));
939 }
940
941 #[test]
942 fn test_decimal_different_scale_errors() {
943 let a = Decimal128Array::from(vec![100])
945 .with_precision_and_scale(10, 2)
946 .unwrap();
947 let b = Decimal128Array::from(vec![500])
948 .with_precision_and_scale(10, 3)
949 .unwrap();
950 let err = make_comparator(&a, &b, SortOptions::default())
951 .err()
952 .unwrap();
953 assert_eq!(
954 err.to_string(),
955 "Invalid argument error: Can't compare decimal arrays with different scales: \
956 Decimal128(10, 2), Decimal128(10, 3)"
957 );
958
959 let keys = Int8Array::from_iter_values([0]);
961 let dict_a = DictionaryArray::new(keys.clone(), Arc::new(a.clone()));
962 let dict_b = DictionaryArray::new(keys, Arc::new(b));
963 assert!(make_comparator(&dict_a, &dict_b, SortOptions::default()).is_err());
964
965 let c = Decimal128Array::from(vec![100])
967 .with_precision_and_scale(12, 2)
968 .unwrap();
969 let cmp = make_comparator(&a, &c, SortOptions::default()).unwrap();
970 assert_eq!(Ordering::Equal, cmp(0, 0));
971
972 let a = Decimal32Array::from(vec![1])
974 .with_precision_and_scale(5, 1)
975 .unwrap();
976 let b = Decimal32Array::from(vec![1])
977 .with_precision_and_scale(5, 2)
978 .unwrap();
979 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
980
981 let a = Decimal64Array::from(vec![1])
982 .with_precision_and_scale(10, 1)
983 .unwrap();
984 let b = Decimal64Array::from(vec![1])
985 .with_precision_and_scale(10, 2)
986 .unwrap();
987 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
988
989 let a = Decimal256Array::from(vec![i256::from_i128(1)])
990 .with_precision_and_scale(40, 1)
991 .unwrap();
992 let b = Decimal256Array::from(vec![i256::from_i128(1)])
993 .with_precision_and_scale(40, 2)
994 .unwrap();
995 assert!(make_comparator(&a, &b, SortOptions::default()).is_err());
996 }
997
998 fn test_bytes_impl<T: ByteArrayType>() {
999 let offsets = OffsetBuffer::from_lengths([3, 3, 1]);
1000 let a = GenericByteArray::<T>::new(offsets, b"abcdefa".into(), None);
1001 let cmp = make_comparator(&a, &a, SortOptions::default()).unwrap();
1002
1003 assert_eq!(Ordering::Less, cmp(0, 1));
1004 assert_eq!(Ordering::Greater, cmp(0, 2));
1005 assert_eq!(Ordering::Equal, cmp(1, 1));
1006 }
1007
1008 #[test]
1009 fn test_bytes() {
1010 test_bytes_impl::<Utf8Type>();
1011 test_bytes_impl::<LargeUtf8Type>();
1012 test_bytes_impl::<BinaryType>();
1013 test_bytes_impl::<LargeBinaryType>();
1014 }
1015
1016 fn assert_cmp_cases<A: Array>(
1017 array1: &A,
1018 array2: &A,
1019 opts: SortOptions,
1020 cases: &[(usize, usize, Ordering)],
1021 ) {
1022 let cmp = make_comparator(array1, array2, opts).unwrap();
1023 for (left, right, expected) in cases {
1024 assert_eq!(cmp(*left, *right), *expected);
1025 }
1026 }
1027
1028 #[test]
1029 fn test_lists() {
1030 let mut a = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
1031 a.extend([
1032 Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]),
1033 Some(vec![
1034 Some(vec![Some(1), Some(2), Some(3)]),
1035 Some(vec![Some(1)]),
1036 ]),
1037 Some(vec![]),
1038 ]);
1039 let a = a.finish();
1040 let mut b = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
1041 b.extend([
1042 Some(vec![Some(vec![Some(1), Some(2), None]), Some(vec![None])]),
1043 Some(vec![
1044 Some(vec![Some(1), Some(2), None]),
1045 Some(vec![Some(1)]),
1046 ]),
1047 Some(vec![
1048 Some(vec![Some(1), Some(2), Some(3), Some(4)]),
1049 Some(vec![Some(1)]),
1050 ]),
1051 None,
1052 ]);
1053 let b = b.finish();
1054
1055 assert_cmp_cases(
1057 &a,
1058 &b,
1059 SortOptions {
1060 descending: false,
1061 nulls_first: true,
1062 },
1063 &[
1064 (0, 0, Ordering::Equal),
1065 (0, 1, Ordering::Less),
1066 (0, 2, Ordering::Less),
1067 (1, 2, Ordering::Less),
1068 (1, 3, Ordering::Greater),
1069 (2, 0, Ordering::Less),
1070 ],
1071 );
1072
1073 assert_cmp_cases(
1075 &a,
1076 &b,
1077 SortOptions {
1078 descending: true,
1079 nulls_first: true,
1080 },
1081 &[
1082 (0, 0, Ordering::Equal),
1083 (0, 1, Ordering::Less),
1084 (0, 2, Ordering::Less),
1085 (1, 2, Ordering::Greater),
1086 (1, 3, Ordering::Greater),
1087 (2, 0, Ordering::Greater),
1088 ],
1089 );
1090
1091 assert_cmp_cases(
1093 &a,
1094 &b,
1095 SortOptions {
1096 descending: true,
1097 nulls_first: false,
1098 },
1099 &[
1100 (0, 0, Ordering::Equal),
1101 (0, 1, Ordering::Greater),
1102 (0, 2, Ordering::Greater),
1103 (1, 2, Ordering::Greater),
1104 (1, 3, Ordering::Less),
1105 (2, 0, Ordering::Greater),
1106 ],
1107 );
1108
1109 assert_cmp_cases(
1111 &a,
1112 &b,
1113 SortOptions {
1114 descending: false,
1115 nulls_first: false,
1116 },
1117 &[
1118 (0, 0, Ordering::Equal),
1119 (0, 1, Ordering::Greater),
1120 (0, 2, Ordering::Greater),
1121 (1, 2, Ordering::Less),
1122 (1, 3, Ordering::Less),
1123 (2, 0, Ordering::Less),
1124 ],
1125 );
1126 }
1127
1128 fn list_view_array<O: OffsetSizeTrait>(
1129 values: Vec<i32>,
1130 offsets: &[usize],
1131 sizes: &[usize],
1132 valid: Option<&[bool]>,
1133 ) -> GenericListViewArray<O> {
1134 let offsets = offsets
1135 .iter()
1136 .map(|v| O::from_usize(*v).unwrap())
1137 .collect::<ScalarBuffer<O>>();
1138 let sizes = sizes
1139 .iter()
1140 .map(|v| O::from_usize(*v).unwrap())
1141 .collect::<ScalarBuffer<O>>();
1142 let field = Arc::new(Field::new_list_field(DataType::Int32, true));
1143 let values = Int32Array::from(values);
1144 let nulls = valid.map(NullBuffer::from);
1145 GenericListViewArray::new(field, offsets, sizes, Arc::new(values), nulls)
1146 }
1147
1148 fn test_list_view_comparisons<O: OffsetSizeTrait>() {
1149 let array = list_view_array::<O>(
1150 vec![1, 2, 3, 4, 5],
1151 &[0, 2, 1, 0, 3],
1152 &[2, 2, 2, 0, 2],
1153 Some(&[true, true, true, true, false]),
1154 );
1155
1156 assert_cmp_cases(
1158 &array,
1159 &array,
1160 SortOptions {
1161 descending: false,
1162 nulls_first: true,
1163 },
1164 &[
1165 (0, 2, Ordering::Less), (1, 2, Ordering::Greater), (3, 0, Ordering::Less), (4, 0, Ordering::Less), ],
1170 );
1171
1172 assert_cmp_cases(
1174 &array,
1175 &array,
1176 SortOptions {
1177 descending: false,
1178 nulls_first: false,
1179 },
1180 &[
1181 (0, 2, Ordering::Less),
1182 (1, 2, Ordering::Greater),
1183 (3, 0, Ordering::Less),
1184 (4, 0, Ordering::Greater), ],
1186 );
1187
1188 assert_cmp_cases(
1190 &array,
1191 &array,
1192 SortOptions {
1193 descending: true,
1194 nulls_first: true,
1195 },
1196 &[
1197 (0, 2, Ordering::Greater),
1198 (1, 2, Ordering::Less),
1199 (3, 0, Ordering::Greater),
1200 (4, 0, Ordering::Less),
1201 ],
1202 );
1203
1204 assert_cmp_cases(
1206 &array,
1207 &array,
1208 SortOptions {
1209 descending: true,
1210 nulls_first: false,
1211 },
1212 &[
1213 (0, 2, Ordering::Greater),
1214 (1, 2, Ordering::Less),
1215 (3, 0, Ordering::Greater),
1216 (4, 0, Ordering::Greater),
1217 ],
1218 );
1219 }
1220
1221 #[test]
1222 fn test_list_view() {
1223 test_list_view_comparisons::<i32>();
1224 }
1225
1226 #[test]
1227 fn test_large_list_view() {
1228 test_list_view_comparisons::<i64>();
1229 }
1230
1231 #[test]
1232 fn test_struct() {
1233 let fields = Fields::from(vec![
1234 Field::new("a", DataType::Int32, true),
1235 Field::new_list("b", Field::new_list_field(DataType::Int32, true), true),
1236 ]);
1237
1238 let a = Int32Array::from(vec![Some(1), Some(2), None, None]);
1239 let mut b = ListBuilder::new(Int32Builder::new());
1240 b.extend([Some(vec![Some(1), Some(2)]), Some(vec![None]), None, None]);
1241 let b = b.finish();
1242
1243 let nulls = Some(NullBuffer::from_iter([true, true, true, false]));
1244 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
1245 let s1 = StructArray::new(fields.clone(), values, nulls);
1246
1247 let a = Int32Array::from(vec![None, Some(2), None]);
1248 let mut b = ListBuilder::new(Int32Builder::new());
1249 b.extend([None, None, Some(vec![])]);
1250 let b = b.finish();
1251
1252 let values = vec![Arc::new(a) as _, Arc::new(b) as _];
1253 let s2 = StructArray::new(fields.clone(), values, None);
1254
1255 let opts = SortOptions {
1256 descending: false,
1257 nulls_first: true,
1258 };
1259 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1260 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 {
1269 descending: true,
1270 nulls_first: true,
1271 };
1272 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1273 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 {
1282 descending: true,
1283 nulls_first: false,
1284 };
1285 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1286 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 {
1295 descending: false,
1296 nulls_first: false,
1297 };
1298 let cmp = make_comparator(&s1, &s2, opts).unwrap();
1299 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); }
1307
1308 #[test]
1309 fn test_map() {
1310 let string_builder = StringBuilder::new();
1313 let int_builder = Int32Builder::new();
1314 let mut map1_builder = MapBuilder::new(None, string_builder, int_builder);
1315
1316 map1_builder.keys().append_value("a");
1318 map1_builder.values().append_value(100);
1319 map1_builder.keys().append_value("b");
1320 map1_builder.values().append_value(1);
1321 map1_builder.append(true).unwrap();
1322
1323 map1_builder.keys().append_value("b");
1325 map1_builder.values().append_value(999);
1326 map1_builder.keys().append_value("c");
1327 map1_builder.values().append_value(1);
1328 map1_builder.append(true).unwrap();
1329
1330 map1_builder.append(true).unwrap();
1332
1333 map1_builder.keys().append_value("x");
1335 map1_builder.values().append_value(1);
1336 map1_builder.append(true).unwrap();
1337
1338 let map1 = map1_builder.finish();
1339
1340 let string_builder = StringBuilder::new();
1343 let int_builder = Int32Builder::new();
1344 let mut map2_builder = MapBuilder::new(None, string_builder, int_builder);
1345
1346 map2_builder.keys().append_value("a");
1348 map2_builder.values().append_value(1);
1349 map2_builder.keys().append_value("c");
1350 map2_builder.values().append_value(999);
1351 map2_builder.append(true).unwrap();
1352
1353 map2_builder.keys().append_value("b");
1355 map2_builder.values().append_value(1);
1356 map2_builder.keys().append_value("d");
1357 map2_builder.values().append_value(999);
1358 map2_builder.append(true).unwrap();
1359
1360 map2_builder.keys().append_value("a");
1362 map2_builder.values().append_value(1);
1363 map2_builder.append(true).unwrap();
1364
1365 map2_builder.append(false).unwrap();
1367
1368 let map2 = map2_builder.finish();
1369
1370 let opts = SortOptions {
1371 descending: false,
1372 nulls_first: true,
1373 };
1374 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1375
1376 assert_eq!(cmp(0, 0), Ordering::Greater);
1380
1381 assert_eq!(cmp(1, 1), Ordering::Greater);
1384
1385 assert_eq!(cmp(0, 1), Ordering::Less);
1387
1388 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 {
1401 descending: true,
1402 nulls_first: true,
1403 };
1404 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1405
1406 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 {
1414 descending: false,
1415 nulls_first: false,
1416 };
1417 let cmp = make_comparator(&map1, &map2, opts).unwrap();
1418
1419 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); }
1425
1426 #[test]
1427 fn test_map_vs_list_consistency() {
1428 let string_builder = StringBuilder::new();
1431 let int_builder = Int32Builder::new();
1432 let mut map1_builder = MapBuilder::new(None, string_builder, int_builder);
1433
1434 map1_builder.keys().append_value("a");
1436 map1_builder.values().append_value(1);
1437 map1_builder.keys().append_value("b");
1438 map1_builder.values().append_value(2);
1439 map1_builder.append(true).unwrap();
1440
1441 map1_builder.keys().append_value("x");
1443 map1_builder.values().append_value(10);
1444 map1_builder.append(true).unwrap();
1445
1446 map1_builder.append(true).unwrap();
1448
1449 map1_builder.keys().append_value("c");
1451 map1_builder.values().append_value(3);
1452 map1_builder.append(true).unwrap();
1453
1454 let map1 = map1_builder.finish();
1455
1456 let string_builder = StringBuilder::new();
1458 let int_builder = Int32Builder::new();
1459 let mut map2_builder = MapBuilder::new(None, string_builder, int_builder);
1460
1461 map2_builder.keys().append_value("a");
1463 map2_builder.values().append_value(1);
1464 map2_builder.keys().append_value("b");
1465 map2_builder.values().append_value(2);
1466 map2_builder.append(true).unwrap();
1467
1468 map2_builder.keys().append_value("y");
1470 map2_builder.values().append_value(20);
1471 map2_builder.append(true).unwrap();
1472
1473 map2_builder.keys().append_value("d");
1475 map2_builder.values().append_value(4);
1476 map2_builder.append(true).unwrap();
1477
1478 map2_builder.append(false).unwrap();
1480
1481 let map2 = map2_builder.finish();
1482
1483 let list1: ListArray = map1.clone().into();
1485 let list2: ListArray = map2.clone().into();
1486
1487 let test_cases = [
1488 SortOptions {
1489 descending: false,
1490 nulls_first: true,
1491 },
1492 SortOptions {
1493 descending: true,
1494 nulls_first: true,
1495 },
1496 SortOptions {
1497 descending: false,
1498 nulls_first: false,
1499 },
1500 SortOptions {
1501 descending: true,
1502 nulls_first: false,
1503 },
1504 ];
1505
1506 for opts in test_cases {
1507 let map_cmp = make_comparator(&map1, &map2, opts).unwrap();
1508 let list_cmp = make_comparator(&list1, &list2, opts).unwrap();
1509
1510 for i in 0..map1.len() {
1512 for j in 0..map2.len() {
1513 let map_result = map_cmp(i, j);
1514 let list_result = list_cmp(i, j);
1515 assert_eq!(
1516 map_result, list_result,
1517 "Map comparison and List comparison should be equal for indices ({i}, {j}) with opts {opts:?}. Map: {map_result:?}, List: {list_result:?}"
1518 );
1519 }
1520 }
1521 }
1522 }
1523
1524 #[test]
1525 fn test_dense_union() {
1526 let int_array = Int32Array::from(vec![1, 2, 3]);
1531 let str_array = StringArray::from(vec!["b", "a"]);
1532
1533 let type_ids = [0, 1, 0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1534 let offsets = [0, 0, 1, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
1535
1536 let union_fields = [
1537 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1538 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1539 ]
1540 .into_iter()
1541 .collect::<UnionFields>();
1542
1543 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1544
1545 let array1 =
1546 UnionArray::try_new(union_fields.clone(), type_ids, Some(offsets), children).unwrap();
1547
1548 let int_array2 = Int32Array::from(vec![2, 1]);
1552 let str_array2 = StringArray::from(vec!["a", "c"]);
1553 let type_ids2 = [0, 1, 0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1554 let offsets2 = [0, 0, 1, 1].into_iter().collect::<ScalarBuffer<i32>>();
1555
1556 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)];
1557
1558 let array2 =
1559 UnionArray::try_new(union_fields, type_ids2, Some(offsets2), children2).unwrap();
1560
1561 let opts = SortOptions {
1562 descending: false,
1563 nulls_first: true,
1564 };
1565
1566 let cmp = make_comparator(&array1, &array2, opts).unwrap();
1570
1571 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 {
1596 descending: true,
1597 nulls_first: true,
1598 };
1599 let cmp_desc = make_comparator(&array1, &array2, opts_desc).unwrap();
1600
1601 assert_eq!(cmp_desc(0, 0), Ordering::Greater); assert_eq!(cmp_desc(0, 1), Ordering::Greater); assert_eq!(cmp_desc(1, 1), Ordering::Less); }
1605
1606 #[test]
1607 fn test_sparse_union() {
1608 let int_array = Int32Array::from(vec![Some(1), None, Some(3)]);
1612 let str_array = StringArray::from(vec![None, Some("b"), None]);
1613 let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1614
1615 let union_fields = [
1616 (0, Arc::new(Field::new("a", DataType::Int32, false))),
1617 (1, Arc::new(Field::new("b", DataType::Utf8, false))),
1618 ]
1619 .into_iter()
1620 .collect::<UnionFields>();
1621
1622 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1623
1624 let array = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
1625
1626 let opts = SortOptions::default();
1627 let cmp = make_comparator(&array, &array, opts).unwrap();
1628
1629 assert_eq!(cmp(0, 2), Ordering::Less); assert_eq!(cmp(0, 1), Ordering::Less); }
1634
1635 #[test]
1636 #[should_panic(expected = "index out of bounds")]
1637 fn test_union_out_of_bounds() {
1638 let int_array = Int32Array::from(vec![1, 2]);
1640 let str_array = StringArray::from(vec!["a"]);
1641
1642 let type_ids = [0, 1, 0].into_iter().collect::<ScalarBuffer<i8>>();
1643 let offsets = [0, 0, 1].into_iter().collect::<ScalarBuffer<i32>>();
1644
1645 let union_fields = [
1646 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1647 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1648 ]
1649 .into_iter()
1650 .collect::<UnionFields>();
1651
1652 let children = vec![Arc::new(int_array) as ArrayRef, Arc::new(str_array)];
1653
1654 let array = UnionArray::try_new(union_fields, type_ids, Some(offsets), children).unwrap();
1655
1656 let opts = SortOptions::default();
1657 let cmp = make_comparator(&array, &array, opts).unwrap();
1658
1659 cmp(0, 3);
1661 }
1662
1663 #[test]
1664 fn test_union_incompatible_fields() {
1665 let int_array1 = Int32Array::from(vec![1, 2]);
1667 let str_array1 = StringArray::from(vec!["a", "b"]);
1668
1669 let type_ids1 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1670 let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
1671
1672 let union_fields1 = [
1673 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1674 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1675 ]
1676 .into_iter()
1677 .collect::<UnionFields>();
1678
1679 let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)];
1680
1681 let array1 =
1682 UnionArray::try_new(union_fields1, type_ids1, Some(offsets1), children1).unwrap();
1683
1684 let int_array2 = Int32Array::from(vec![3, 4]);
1686 let float_array2 = Float64Array::from(vec![1.0, 2.0]);
1687
1688 let type_ids2 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1689 let offsets2 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
1690
1691 let union_fields2 = [
1692 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1693 (1, Arc::new(Field::new("C", DataType::Float64, false))),
1694 ]
1695 .into_iter()
1696 .collect::<UnionFields>();
1697
1698 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(float_array2)];
1699
1700 let array2 =
1701 UnionArray::try_new(union_fields2, type_ids2, Some(offsets2), children2).unwrap();
1702
1703 let opts = SortOptions::default();
1704
1705 let Result::Err(ArrowError::InvalidArgumentError(out)) =
1706 make_comparator(&array1, &array2, opts)
1707 else {
1708 panic!("expected error when making comparator of incompatible union arrays");
1709 };
1710
1711 assert_eq!(
1712 &out,
1713 "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 })]"
1714 );
1715 }
1716
1717 #[test]
1718 fn test_union_incompatible_modes() {
1719 let int_array1 = Int32Array::from(vec![1, 2]);
1721 let str_array1 = StringArray::from(vec!["a", "b"]);
1722
1723 let type_ids1 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1724 let offsets1 = [0, 0].into_iter().collect::<ScalarBuffer<i32>>();
1725
1726 let union_fields1 = [
1727 (0, Arc::new(Field::new("A", DataType::Int32, false))),
1728 (1, Arc::new(Field::new("B", DataType::Utf8, false))),
1729 ]
1730 .into_iter()
1731 .collect::<UnionFields>();
1732
1733 let children1 = vec![Arc::new(int_array1) as ArrayRef, Arc::new(str_array1)];
1734
1735 let array1 =
1736 UnionArray::try_new(union_fields1.clone(), type_ids1, Some(offsets1), children1)
1737 .unwrap();
1738
1739 let int_array2 = Int32Array::from(vec![Some(3), None]);
1741 let str_array2 = StringArray::from(vec![None, Some("c")]);
1742
1743 let type_ids2 = [0, 1].into_iter().collect::<ScalarBuffer<i8>>();
1744
1745 let children2 = vec![Arc::new(int_array2) as ArrayRef, Arc::new(str_array2)];
1746
1747 let array2 = UnionArray::try_new(union_fields1, type_ids2, None, children2).unwrap();
1748
1749 let opts = SortOptions::default();
1750
1751 let Result::Err(ArrowError::InvalidArgumentError(out)) =
1752 make_comparator(&array1, &array2, opts)
1753 else {
1754 panic!("expected error when making comparator of union arrays with different modes");
1755 };
1756
1757 assert_eq!(
1758 &out,
1759 "Cannot compare UnionArrays with different modes: left=Dense, right=Sparse"
1760 );
1761 }
1762
1763 #[test]
1764 fn test_null_array_cmp() {
1765 let a = NullArray::new(3);
1766 let b = NullArray::new(3);
1767 let cmp = make_comparator(&a, &b, SortOptions::default()).unwrap();
1768
1769 assert_eq!(cmp(0, 0), Ordering::Equal);
1770 assert_eq!(cmp(0, 1), Ordering::Equal);
1771 assert_eq!(cmp(2, 0), Ordering::Equal);
1772 }
1773
1774 #[test]
1775 fn test_run_end_encoded_int32() {
1776 let run_ends1 = Int32Array::from(vec![2, 5, 6]);
1780 let values1 = Int32Array::from(vec![1, 2, 3]);
1781 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1782
1783 let run_ends2 = Int32Array::from(vec![1, 3, 6]);
1786 let values2 = Int32Array::from(vec![1, 2, 3]);
1787 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1788
1789 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1790
1791 assert_eq!(cmp(0, 0), Ordering::Equal);
1793 assert_eq!(cmp(0, 1), Ordering::Less);
1795 assert_eq!(cmp(2, 1), Ordering::Equal);
1797 assert_eq!(cmp(5, 5), Ordering::Equal);
1799 assert_eq!(cmp(1, 2), Ordering::Less);
1801 assert_eq!(cmp(4, 4), Ordering::Less);
1803 }
1804
1805 #[test]
1806 fn test_run_end_encoded_with_nulls() {
1807 let run_ends1 = Int32Array::from(vec![2, 4, 5]);
1811 let values1 = Int32Array::from(vec![Some(1), None, Some(2)]);
1812 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1813
1814 let run_ends2 = Int32Array::from(vec![1, 3, 4, 5]);
1817 let values2 = Int32Array::from(vec![None, Some(1), Some(2), None]);
1818 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1819
1820 let opts = SortOptions::default();
1821 let cmp = make_comparator(&array1, &array2, opts).unwrap();
1822
1823 assert_eq!(cmp(0, 1), Ordering::Equal);
1825 assert_eq!(cmp(2, 0), Ordering::Equal);
1827 assert_eq!(cmp(0, 0), Ordering::Greater);
1829 assert_eq!(cmp(2, 1), Ordering::Less);
1831 }
1832
1833 #[test]
1834 fn test_run_end_encoded_int16() {
1835 let run_ends1 = Int16Array::from(vec![3_i16, 5, 6]);
1837 let values1 = StringArray::from(vec!["a", "b", "c"]);
1838 let array1 = RunArray::<Int16Type>::try_new(&run_ends1, &values1).unwrap();
1839
1840 let run_ends2 = Int16Array::from(vec![2_i16, 4, 6]);
1841 let values2 = StringArray::from(vec!["a", "b", "c"]);
1842 let array2 = RunArray::<Int16Type>::try_new(&run_ends2, &values2).unwrap();
1843
1844 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1845
1846 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); }
1853
1854 #[test]
1855 fn test_run_end_encoded_int64() {
1856 let run_ends1 = Int64Array::from(vec![2_i64, 4, 6]);
1858 let values1 = Int64Array::from(vec![10_i64, 20, 30]);
1859 let array1 = RunArray::<Int64Type>::try_new(&run_ends1, &values1).unwrap();
1860
1861 let run_ends2 = Int64Array::from(vec![3_i64, 5, 6]);
1862 let values2 = Int64Array::from(vec![10_i64, 20, 30]);
1863 let array2 = RunArray::<Int64Type>::try_new(&run_ends2, &values2).unwrap();
1864
1865 let cmp = make_comparator(&array1, &array2, SortOptions::default()).unwrap();
1866
1867 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); }
1874
1875 #[test]
1876 fn test_run_end_encoded_sliced() {
1877 let run_ends = Int32Array::from(vec![2, 5, 7, 8]);
1881 let values = Int32Array::from(vec![1, 2, 3, 4]);
1882 let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1883
1884 let slice1 = array.slice(1, 4);
1886 let slice2 = array.slice(3, 4);
1888
1889 let cmp = make_comparator(&slice1, &slice2, SortOptions::default()).unwrap();
1890
1891 assert_eq!(cmp(0, 0), Ordering::Less);
1893 assert_eq!(cmp(1, 0), Ordering::Equal);
1895 assert_eq!(cmp(3, 2), Ordering::Less);
1897 assert_eq!(cmp(1, 3), Ordering::Less);
1899
1900 let run_ends2 = Int32Array::from(vec![2, 4]);
1902 let values2 = Int32Array::from(vec![1, 2]);
1903 let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1904
1905 let cmp = make_comparator(&slice1, &array2, SortOptions::default()).unwrap();
1906
1907 assert_eq!(cmp(0, 0), Ordering::Equal);
1909 assert_eq!(cmp(1, 1), Ordering::Greater);
1911 assert_eq!(cmp(3, 3), Ordering::Equal);
1913 }
1914
1915 #[test]
1916 fn test_run_end_encoded_sliced_with_nulls() {
1917 let run_ends = Int32Array::from(vec![2, 4, 6, 7, 8]);
1921 let values = Int32Array::from(vec![Some(1), None, Some(2), None, Some(3)]);
1922 let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1923
1924 let slice1 = array.slice(1, 5);
1926 let slice2 = array.slice(3, 5);
1928
1929 let opts = SortOptions::default(); let cmp = make_comparator(&slice1, &slice2, opts).unwrap();
1931
1932 assert_eq!(cmp(0, 0), Ordering::Greater);
1934 assert_eq!(cmp(1, 0), Ordering::Equal);
1936 assert_eq!(cmp(1, 1), Ordering::Less);
1938 assert_eq!(cmp(3, 1), Ordering::Equal);
1940 assert_eq!(cmp(4, 4), Ordering::Less);
1942 assert_eq!(cmp(3, 3), Ordering::Greater);
1944 }
1945
1946 #[test]
1947 fn test_run_end_encoded_different_types() {
1948 let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1950 let values1 = Int32Array::from(vec![1, 2, 3]);
1951 let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1952
1953 let run_ends2 = Int64Array::from(vec![2_i64, 4, 6]);
1954 let values2 = Int64Array::from(vec![1_i64, 2, 3]);
1955 let array2 = RunArray::<Int64Type>::try_new(&run_ends2, &values2).unwrap();
1956
1957 let result = make_comparator(&array1, &array2, SortOptions::default());
1958 assert!(result.is_err());
1959 let err = match result {
1960 Err(e) => e.to_string(),
1961 Ok(_) => panic!("Expected error"),
1962 };
1963 assert!(err.contains("Cannot compare RunEndEncoded arrays"));
1964 }
1965}