1use std::cmp::Ordering;
14
15use minarrow::{Bitmask, BooleanArray, CategoricalArray, Integer, MaskedArray, Vec64};
16use num_traits::{Float, NumCast, Zero};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
20pub enum SortAlgorithm {
21 #[default]
23 Auto,
24 Comparison,
26 Radix,
28 #[cfg(feature = "simd_sort")]
30 Simd,
31}
32
33#[derive(Debug, Clone, Default)]
35pub struct ArgsortConfig {
36 pub descending: bool,
37 pub algorithm: SortAlgorithm,
38 pub parallel: bool,
39}
40
41impl ArgsortConfig {
42 pub fn new() -> Self {
44 Self::default()
45 }
46
47 pub fn descending(mut self, descending: bool) -> Self {
49 self.descending = descending;
50 self
51 }
52
53 pub fn algorithm(mut self, algorithm: SortAlgorithm) -> Self {
55 self.algorithm = algorithm;
56 self
57 }
58
59 pub fn parallel(mut self, parallel: bool) -> Self {
61 self.parallel = parallel;
62 self
63 }
64}
65
66#[inline]
73pub fn sort_float<T: Float>(slice: &mut [T]) {
74 slice.sort_unstable_by(total_cmp_f);
75}
76#[inline(always)]
81pub fn total_cmp_f<T: Float>(a: &T, b: &T) -> Ordering {
82 if std::mem::size_of::<T>() == 4 {
84 let ua = unsafe { *(a as *const T as *const u32) };
86 let ub = unsafe { *(b as *const T as *const u32) };
87 let ka = if ua & 0x8000_0000 != 0 {
89 !ua
90 } else {
91 ua ^ 0x8000_0000
92 };
93 let kb = if ub & 0x8000_0000 != 0 {
94 !ub
95 } else {
96 ub ^ 0x8000_0000
97 };
98 ka.cmp(&kb)
99 } else if std::mem::size_of::<T>() == 8 {
100 let ua = unsafe { *(a as *const T as *const u64) };
102 let ub = unsafe { *(b as *const T as *const u64) };
103 let ka = if ua & 0x8000_0000_0000_0000 != 0 {
104 !ua
105 } else {
106 ua ^ 0x8000_0000_0000_0000
107 };
108 let kb = if ub & 0x8000_0000_0000_0000 != 0 {
109 !ub
110 } else {
111 ub ^ 0x8000_0000_0000_0000
112 };
113 ka.cmp(&kb)
114 } else {
115 unreachable!("Only f32 and f64 are valid Float types.")
116 }
117}
118
119pub fn sorted_float<T: Float>(data: &[T]) -> Vec64<T> {
121 let mut v = Vec64::from_slice(data);
122 sort_float(&mut v);
123 v
124}
125
126#[inline]
146pub fn sort_int<T: Ord + Copy>(slice: &mut [T]) {
147 slice.sort_unstable();
148}
149
150pub fn sorted_int<T: Ord + Copy>(data: &[T]) -> Vec64<T> {
175 let mut v = Vec64::from_slice(data);
176 sort_int(&mut v);
177 v
178}
179
180#[inline]
197pub fn sort_str(slice: &mut [&str]) {
198 slice.sort_unstable();
199}
200
201pub fn sorted_str(data: &[&str]) -> Vec64<String> {
238 let mut v = Vec64::from_slice(data);
239 sort_str(&mut v);
240 v.iter().map(|s| s.to_string()).collect()
241}
242
243pub fn sort_string_array(offsets: &[usize], values: &str) -> (Vec64<usize>, String) {
246 let n = offsets.len() - 1;
247 let mut indices: Vec<usize> = (0..n).collect();
248
249 indices.sort_unstable_by(|&i, &j| {
250 let a = &values[offsets[i]..offsets[i + 1]];
251 let b = &values[offsets[j]..offsets[j + 1]];
252 a.cmp(b)
253 });
254
255 let total_bytes: usize = indices
257 .iter()
258 .map(|&idx| offsets[idx + 1] - offsets[idx])
259 .sum();
260
261 let mut new_values = String::with_capacity(total_bytes);
263 let mut new_offsets = Vec64::<usize>::with_capacity(n + 1);
264
265 unsafe {
267 new_offsets.set_len(n + 1);
268 }
269 unsafe {
270 new_values.as_mut_vec().set_len(total_bytes);
271 }
272
273 let values_bytes = values.as_bytes();
274 let out_bytes = unsafe { new_values.as_mut_vec() };
275
276 unsafe {
278 *new_offsets.get_unchecked_mut(0) = 0;
279 }
280 let mut current_offset = 0;
281 let mut out_ptr = 0;
282 for (i, &idx) in indices.iter().enumerate() {
283 let start = offsets[idx];
284 let end = offsets[idx + 1];
285 let len = end - start;
286 unsafe {
288 std::ptr::copy_nonoverlapping(
289 values_bytes.as_ptr().add(start),
290 out_bytes.as_mut_ptr().add(out_ptr),
291 len,
292 );
293 current_offset += len;
294 *new_offsets.get_unchecked_mut(i + 1) = current_offset;
295 }
296 out_ptr += len;
297 }
298 unsafe {
300 new_values.as_mut_vec().set_len(current_offset);
301 }
302
303 (new_offsets, new_values)
304}
305
306pub fn sort_categorical_lexical<T: Ord + Copy + Zero + NumCast + Integer>(
309 cat: &CategoricalArray<T>,
310) -> (Vec64<T>, Option<Bitmask>) {
311 let len = cat.data.len();
312 let mut indices: Vec<usize> = (0..len).collect();
313
314 indices.sort_by(|&i, &j| {
316 let a_is_null = cat.is_null(i);
317 let b_is_null = cat.is_null(j);
318 match (a_is_null, b_is_null) {
319 (true, true) => Ordering::Equal,
320 (true, false) => Ordering::Less,
321 (false, true) => Ordering::Greater,
322 (false, false) => {
323 let a_key = &cat.unique_values[cat.data[i].to_usize()];
325 let b_key = &cat.unique_values[cat.data[j].to_usize()];
326 a_key.cmp(b_key)
327 }
328 }
329 });
330
331 let mut sorted_data = Vec64::<T>::with_capacity(len);
333 let mut sorted_mask = cat
334 .null_mask
335 .as_ref()
336 .map(|_| Bitmask::new_set_all(len, false));
337
338 for (out_i, &orig_i) in indices.iter().enumerate() {
339 sorted_data.push(cat.data[orig_i]);
340 if let Some(ref mask) = cat.null_mask {
341 if let Some(ref mut sm) = sorted_mask {
342 if unsafe { mask.get_unchecked(orig_i) } {
343 unsafe { sm.set_unchecked(out_i, true) };
344 }
345 }
346 }
347 }
348
349 (sorted_data, sorted_mask)
350}
351
352#[inline]
354fn unpack_bitmask(mask: &Bitmask) -> Vec64<bool> {
355 let mut out = Vec64::with_capacity(mask.len);
356 for i in 0..mask.len {
357 out.push(unsafe { mask.get_unchecked(i) });
358 }
359 out
360}
361
362#[inline]
364fn pack_bitmask(bits: &[bool]) -> Bitmask {
365 let mut mask = Bitmask::new_set_all(bits.len(), false);
366 for (i, &b) in bits.iter().enumerate() {
367 if b {
368 unsafe { mask.set_unchecked(i, true) };
369 }
370 }
371 mask
372}
373
374pub fn sort_boolean_array(arr: &mut BooleanArray<()>) {
377 let bits: Vec64<bool> = unpack_bitmask(&arr.data);
378 let nulls: Option<Vec64<bool>> = arr.null_mask.as_ref().map(unpack_bitmask);
379
380 let mut indices: Vec<usize> = (0..arr.len).collect();
381
382 indices.sort_unstable_by(|&i, &j| {
384 let a_null = !nulls.as_ref().map_or(true, |n| n[i]);
385 let b_null = !nulls.as_ref().map_or(true, |n| n[j]);
386
387 match (a_null, b_null) {
388 (true, true) => Ordering::Equal, (true, false) => Ordering::Less, (false, true) => Ordering::Greater, (false, false) => {
392 match (bits[i], bits[j]) {
394 (false, false) => Ordering::Equal,
395 (false, true) => Ordering::Less,
396 (true, false) => Ordering::Greater,
397 (true, true) => Ordering::Equal,
398 }
399 }
400 }
401 });
402
403 let sorted_bits: Vec<bool> = indices
405 .iter()
406 .map(|&idx| {
407 let is_null = !nulls.as_ref().map_or(true, |n| n[idx]);
408 if is_null { false } else { bits[idx] }
409 })
410 .collect();
411 arr.data = pack_bitmask(&sorted_bits);
412
413 if let Some(null_mask) = arr.null_mask.as_mut() {
415 let sorted_nulls: Vec<bool> = indices
416 .iter()
417 .map(|&idx| nulls.as_ref().unwrap()[idx])
418 .collect();
419 *null_mask = pack_bitmask(&sorted_nulls);
420 }
421}
422
423pub fn sort_slice_with_mask<T: Ord + Copy>(
425 data: &[T],
426 mask: Option<&Bitmask>,
427) -> (Vec64<T>, Option<Bitmask>) {
428 let n = data.len();
429 let mut indices: Vec<usize> = (0..n).collect();
430 indices.sort_unstable_by(|&i, &j| data[i].cmp(&data[j]));
431
432 let mut sorted_data = Vec64::<T>::with_capacity(n);
433 unsafe {
434 sorted_data.set_len(n);
435 }
436 for (i, &idx) in indices.iter().enumerate() {
437 unsafe {
438 *sorted_data.get_unchecked_mut(i) = data[idx];
439 }
440 }
441
442 let sorted_mask = mask.map(|m| {
443 let mut out = Bitmask::new_set_all(n, false);
444 for (i, &idx) in indices.iter().enumerate() {
445 if unsafe { m.get_unchecked(idx) } {
446 unsafe { out.set_unchecked(i, true) };
447 }
448 }
449 out
450 });
451
452 (sorted_data, sorted_mask)
453}
454
455#[inline]
459pub fn argsort<T: Ord>(data: &[T], descending: bool) -> Vec<usize> {
460 let n = data.len();
461 if n == 0 {
462 return vec![];
463 }
464
465 let mut indices: Vec<usize> = (0..n).collect();
466 if descending {
467 indices.sort_unstable_by(|&i, &j| data[j].cmp(&data[i]));
468 } else {
469 indices.sort_unstable_by(|&i, &j| data[i].cmp(&data[j]));
470 }
471 indices
472}
473
474#[inline]
476pub fn argsort_float<T: Float>(data: &[T], descending: bool) -> Vec<usize> {
477 let n = data.len();
478 if n == 0 {
479 return vec![];
480 }
481
482 let mut indices: Vec<usize> = (0..n).collect();
483 if descending {
484 indices.sort_unstable_by(|&i, &j| total_cmp_f(&data[j], &data[i]));
485 } else {
486 indices.sort_unstable_by(|&i, &j| total_cmp_f(&data[i], &data[j]));
487 }
488 indices
489}
490
491#[inline]
493pub fn argsort_str(data: &[&str], descending: bool) -> Vec<usize> {
494 let n = data.len();
495 if n == 0 {
496 return vec![];
497 }
498
499 let mut indices: Vec<usize> = (0..n).collect();
500 if descending {
501 indices.sort_unstable_by(|&i, &j| data[j].cmp(data[i]));
502 } else {
503 indices.sort_unstable_by(|&i, &j| data[i].cmp(data[j]));
504 }
505 indices
506}
507
508pub fn argsort_string_array(offsets: &[usize], values: &str, descending: bool) -> Vec<usize> {
512 let n = if offsets.is_empty() { 0 } else { offsets.len() - 1 };
513 if n == 0 {
514 return vec![];
515 }
516
517 let mut indices: Vec<usize> = (0..n).collect();
518 if descending {
519 indices.sort_unstable_by(|&i, &j| {
520 let a = &values[offsets[i]..offsets[i + 1]];
521 let b = &values[offsets[j]..offsets[j + 1]];
522 b.cmp(a)
523 });
524 } else {
525 indices.sort_unstable_by(|&i, &j| {
526 let a = &values[offsets[i]..offsets[i + 1]];
527 let b = &values[offsets[j]..offsets[j + 1]];
528 a.cmp(b)
529 });
530 }
531 indices
532}
533
534pub fn argsort_categorical_lexical<T: Ord + Copy + Zero + NumCast + Integer>(
539 cat: &CategoricalArray<T>,
540 descending: bool,
541) -> Vec<usize> {
542 let len = cat.data.len();
543 if len == 0 {
544 return vec![];
545 }
546
547 let mut indices: Vec<usize> = (0..len).collect();
548
549 indices.sort_by(|&i, &j| {
550 let a_is_null = cat.is_null(i);
551 let b_is_null = cat.is_null(j);
552
553 let ordering = match (a_is_null, b_is_null) {
554 (true, true) => Ordering::Equal,
555 (true, false) => Ordering::Less, (false, true) => Ordering::Greater,
557 (false, false) => {
558 let a_key = &cat.unique_values[cat.data[i].to_usize()];
559 let b_key = &cat.unique_values[cat.data[j].to_usize()];
560 a_key.cmp(b_key)
561 }
562 };
563
564 if descending {
565 ordering.reverse()
566 } else {
567 ordering
568 }
569 });
570
571 indices
572}
573
574pub fn argsort_categorical_custom<T: Ord + Copy + Zero + NumCast + Integer>(
580 cat: &CategoricalArray<T>,
581 category_order: &[&str],
582 descending: bool,
583) -> Vec<usize> {
584 let len = cat.data.len();
585 if len == 0 {
586 return vec![];
587 }
588
589 use std::collections::HashMap;
591 let order_map: HashMap<&str, usize> = category_order
592 .iter()
593 .enumerate()
594 .map(|(i, &s)| (s, i))
595 .collect();
596
597 let mut indices: Vec<usize> = (0..len).collect();
598
599 indices.sort_by(|&i, &j| {
600 let a_is_null = cat.is_null(i);
601 let b_is_null = cat.is_null(j);
602
603 let ordering = match (a_is_null, b_is_null) {
604 (true, true) => Ordering::Equal,
605 (true, false) => Ordering::Less,
606 (false, true) => Ordering::Greater,
607 (false, false) => {
608 let a_key = &cat.unique_values[cat.data[i].to_usize()];
609 let b_key = &cat.unique_values[cat.data[j].to_usize()];
610
611 let a_pos = order_map.get(a_key.as_str());
613 let b_pos = order_map.get(b_key.as_str());
614
615 match (a_pos, b_pos) {
616 (Some(ap), Some(bp)) => ap.cmp(bp),
617 (Some(_), None) => Ordering::Less, (None, Some(_)) => Ordering::Greater,
619 (None, None) => a_key.cmp(b_key), }
621 }
622 };
623
624 if descending {
625 ordering.reverse()
626 } else {
627 ordering
628 }
629 });
630
631 indices
632}
633
634pub fn argsort_boolean(data: &Bitmask, null_mask: Option<&Bitmask>, descending: bool) -> Vec<usize> {
638 let n = data.len;
639 if n == 0 {
640 return vec![];
641 }
642
643 let mut indices: Vec<usize> = (0..n).collect();
644
645 indices.sort_unstable_by(|&i, &j| {
646 let a_null = null_mask.map_or(false, |m| !m.get(i));
647 let b_null = null_mask.map_or(false, |m| !m.get(j));
648
649 let ordering = match (a_null, b_null) {
650 (true, true) => Ordering::Equal,
651 (true, false) => Ordering::Less,
652 (false, true) => Ordering::Greater,
653 (false, false) => {
654 let a_val = data.get(i);
655 let b_val = data.get(j);
656 a_val.cmp(&b_val)
657 }
658 };
659
660 if descending {
661 ordering.reverse()
662 } else {
663 ordering
664 }
665 });
666
667 indices
668}
669
670pub fn argsort_radix_u32(data: &[u32], descending: bool) -> Vec<usize> {
677 let n = data.len();
678 if n == 0 {
679 return vec![];
680 }
681
682 let mut indices: Vec<usize> = (0..n).collect();
683 let mut temp_indices = vec![0usize; n];
684
685 for shift in (0..32).step_by(8) {
687 let mut counts = [0usize; 256];
688 for &idx in &indices {
689 let byte = ((data[idx] >> shift) & 0xFF) as usize;
690 counts[byte] += 1;
691 }
692
693 let mut positions = [0usize; 256];
694 if descending {
695 let mut sum = 0;
696 for i in (0..256).rev() {
697 positions[i] = sum;
698 sum += counts[i];
699 }
700 } else {
701 let mut sum = 0;
702 for i in 0..256 {
703 positions[i] = sum;
704 sum += counts[i];
705 }
706 }
707
708 for &idx in &indices {
709 let byte = ((data[idx] >> shift) & 0xFF) as usize;
710 temp_indices[positions[byte]] = idx;
711 positions[byte] += 1;
712 }
713
714 std::mem::swap(&mut indices, &mut temp_indices);
715 }
716
717 indices
718}
719
720pub fn argsort_radix_u64(data: &[u64], descending: bool) -> Vec<usize> {
722 let n = data.len();
723 if n == 0 {
724 return vec![];
725 }
726
727 let mut indices: Vec<usize> = (0..n).collect();
728 let mut temp_indices = vec![0usize; n];
729
730 for shift in (0..64).step_by(8) {
731 let mut counts = [0usize; 256];
732 for &idx in &indices {
733 let byte = ((data[idx] >> shift) & 0xFF) as usize;
734 counts[byte] += 1;
735 }
736
737 let mut positions = [0usize; 256];
738 if descending {
739 let mut sum = 0;
740 for i in (0..256).rev() {
741 positions[i] = sum;
742 sum += counts[i];
743 }
744 } else {
745 let mut sum = 0;
746 for i in 0..256 {
747 positions[i] = sum;
748 sum += counts[i];
749 }
750 }
751
752 for &idx in &indices {
753 let byte = ((data[idx] >> shift) & 0xFF) as usize;
754 temp_indices[positions[byte]] = idx;
755 positions[byte] += 1;
756 }
757
758 std::mem::swap(&mut indices, &mut temp_indices);
759 }
760
761 indices
762}
763
764pub fn argsort_radix_i32(data: &[i32], descending: bool) -> Vec<usize> {
766 let unsigned: Vec<u32> = data.iter().map(|&x| (x as u32) ^ 0x8000_0000).collect();
768 argsort_radix_u32(&unsigned, descending)
769}
770
771pub fn argsort_radix_i64(data: &[i64], descending: bool) -> Vec<usize> {
773 let unsigned: Vec<u64> = data.iter().map(|&x| (x as u64) ^ 0x8000_0000_0000_0000).collect();
774 argsort_radix_u64(&unsigned, descending)
775}
776
777#[cfg(feature = "simd_sort")]
780pub mod simd_argsort {
781 use std::cmp::Ordering;
782 use voracious_radix_sort::{RadixSort, Radixable};
783
784 #[derive(Copy, Clone, Debug)]
786 struct IndexedU32 {
787 value: u32,
788 index: usize,
789 }
790
791 impl PartialOrd for IndexedU32 {
792 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
793 self.value.partial_cmp(&other.value)
794 }
795 }
796
797 impl PartialEq for IndexedU32 {
798 fn eq(&self, other: &Self) -> bool {
799 self.value == other.value
800 }
801 }
802
803 impl Radixable<u32> for IndexedU32 {
804 type Key = u32;
805 #[inline]
806 fn key(&self) -> Self::Key {
807 self.value
808 }
809 }
810
811 #[derive(Copy, Clone, Debug)]
813 struct IndexedU64 {
814 value: u64,
815 index: usize,
816 }
817
818 impl PartialOrd for IndexedU64 {
819 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
820 self.value.partial_cmp(&other.value)
821 }
822 }
823
824 impl PartialEq for IndexedU64 {
825 fn eq(&self, other: &Self) -> bool {
826 self.value == other.value
827 }
828 }
829
830 impl Radixable<u64> for IndexedU64 {
831 type Key = u64;
832 #[inline]
833 fn key(&self) -> Self::Key {
834 self.value
835 }
836 }
837
838 #[derive(Copy, Clone, Debug)]
840 struct IndexedI32 {
841 value: i32,
842 index: usize,
843 }
844
845 impl PartialOrd for IndexedI32 {
846 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
847 self.value.partial_cmp(&other.value)
848 }
849 }
850
851 impl PartialEq for IndexedI32 {
852 fn eq(&self, other: &Self) -> bool {
853 self.value == other.value
854 }
855 }
856
857 impl Radixable<i32> for IndexedI32 {
858 type Key = i32;
859 #[inline]
860 fn key(&self) -> Self::Key {
861 self.value
862 }
863 }
864
865 #[derive(Copy, Clone, Debug)]
867 struct IndexedI64 {
868 value: i64,
869 index: usize,
870 }
871
872 impl PartialOrd for IndexedI64 {
873 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
874 self.value.partial_cmp(&other.value)
875 }
876 }
877
878 impl PartialEq for IndexedI64 {
879 fn eq(&self, other: &Self) -> bool {
880 self.value == other.value
881 }
882 }
883
884 impl Radixable<i64> for IndexedI64 {
885 type Key = i64;
886 #[inline]
887 fn key(&self) -> Self::Key {
888 self.value
889 }
890 }
891
892 pub fn argsort_simd_u32(data: &[u32], descending: bool) -> Vec<usize> {
894 let n = data.len();
895 if n == 0 {
896 return vec![];
897 }
898
899 let mut indexed: Vec<IndexedU32> = data
900 .iter()
901 .enumerate()
902 .map(|(index, &value)| IndexedU32 { value, index })
903 .collect();
904
905 indexed.voracious_sort();
906
907 if descending {
908 indexed.iter().rev().map(|x| x.index).collect()
909 } else {
910 indexed.iter().map(|x| x.index).collect()
911 }
912 }
913
914 pub fn argsort_simd_u64(data: &[u64], descending: bool) -> Vec<usize> {
916 let n = data.len();
917 if n == 0 {
918 return vec![];
919 }
920
921 let mut indexed: Vec<IndexedU64> = data
922 .iter()
923 .enumerate()
924 .map(|(index, &value)| IndexedU64 { value, index })
925 .collect();
926
927 indexed.voracious_sort();
928
929 if descending {
930 indexed.iter().rev().map(|x| x.index).collect()
931 } else {
932 indexed.iter().map(|x| x.index).collect()
933 }
934 }
935
936 pub fn argsort_simd_i32(data: &[i32], descending: bool) -> Vec<usize> {
938 let n = data.len();
939 if n == 0 {
940 return vec![];
941 }
942
943 let mut indexed: Vec<IndexedI32> = data
944 .iter()
945 .enumerate()
946 .map(|(index, &value)| IndexedI32 { value, index })
947 .collect();
948
949 indexed.voracious_sort();
950
951 if descending {
952 indexed.iter().rev().map(|x| x.index).collect()
953 } else {
954 indexed.iter().map(|x| x.index).collect()
955 }
956 }
957
958 pub fn argsort_simd_i64(data: &[i64], descending: bool) -> Vec<usize> {
960 let n = data.len();
961 if n == 0 {
962 return vec![];
963 }
964
965 let mut indexed: Vec<IndexedI64> = data
966 .iter()
967 .enumerate()
968 .map(|(index, &value)| IndexedI64 { value, index })
969 .collect();
970
971 indexed.voracious_sort();
972
973 if descending {
974 indexed.iter().rev().map(|x| x.index).collect()
975 } else {
976 indexed.iter().map(|x| x.index).collect()
977 }
978 }
979}
980
981#[cfg(feature = "parallel_sort")]
984pub mod parallel_argsort {
985 use rayon::prelude::*;
986
987 pub fn argsort_parallel<T: Ord + Sync>(data: &[T], descending: bool, stable: bool) -> Vec<usize> {
989 let n = data.len();
990 if n == 0 {
991 return vec![];
992 }
993
994 let mut indices: Vec<usize> = (0..n).collect();
995
996 if stable {
997 if descending {
998 indices.par_sort_by(|&i, &j| data[j].cmp(&data[i]));
999 } else {
1000 indices.par_sort_by(|&i, &j| data[i].cmp(&data[j]));
1001 }
1002 } else {
1003 if descending {
1004 indices.par_sort_unstable_by(|&i, &j| data[j].cmp(&data[i]));
1005 } else {
1006 indices.par_sort_unstable_by(|&i, &j| data[i].cmp(&data[j]));
1007 }
1008 }
1009
1010 indices
1011 }
1012}
1013
1014pub fn argsort_auto_i32(data: &[i32], config: &ArgsortConfig) -> Vec<usize> {
1022 match config.algorithm {
1023 SortAlgorithm::Comparison => argsort(data, config.descending),
1024 SortAlgorithm::Radix => argsort_radix_i32(data, config.descending),
1025 #[cfg(feature = "simd_sort")]
1026 SortAlgorithm::Simd => simd_argsort::argsort_simd_i32(data, config.descending),
1027 SortAlgorithm::Auto => {
1028 #[cfg(feature = "simd_sort")]
1029 {
1030 simd_argsort::argsort_simd_i32(data, config.descending)
1031 }
1032 #[cfg(not(feature = "simd_sort"))]
1033 {
1034 argsort_radix_i32(data, config.descending)
1035 }
1036 }
1037 }
1038}
1039
1040pub fn argsort_auto_i64(data: &[i64], config: &ArgsortConfig) -> Vec<usize> {
1042 match config.algorithm {
1043 SortAlgorithm::Comparison => argsort(data, config.descending),
1044 SortAlgorithm::Radix => argsort_radix_i64(data, config.descending),
1045 #[cfg(feature = "simd_sort")]
1046 SortAlgorithm::Simd => simd_argsort::argsort_simd_i64(data, config.descending),
1047 SortAlgorithm::Auto => {
1048 #[cfg(feature = "simd_sort")]
1049 {
1050 simd_argsort::argsort_simd_i64(data, config.descending)
1051 }
1052 #[cfg(not(feature = "simd_sort"))]
1053 {
1054 argsort_radix_i64(data, config.descending)
1055 }
1056 }
1057 }
1058}
1059
1060pub fn argsort_auto_u32(data: &[u32], config: &ArgsortConfig) -> Vec<usize> {
1062 match config.algorithm {
1063 SortAlgorithm::Comparison => argsort(data, config.descending),
1064 SortAlgorithm::Radix => argsort_radix_u32(data, config.descending),
1065 #[cfg(feature = "simd_sort")]
1066 SortAlgorithm::Simd => simd_argsort::argsort_simd_u32(data, config.descending),
1067 SortAlgorithm::Auto => {
1068 #[cfg(feature = "simd_sort")]
1069 {
1070 simd_argsort::argsort_simd_u32(data, config.descending)
1071 }
1072 #[cfg(not(feature = "simd_sort"))]
1073 {
1074 argsort_radix_u32(data, config.descending)
1075 }
1076 }
1077 }
1078}
1079
1080pub fn argsort_auto_u64(data: &[u64], config: &ArgsortConfig) -> Vec<usize> {
1082 match config.algorithm {
1083 SortAlgorithm::Comparison => argsort(data, config.descending),
1084 SortAlgorithm::Radix => argsort_radix_u64(data, config.descending),
1085 #[cfg(feature = "simd_sort")]
1086 SortAlgorithm::Simd => simd_argsort::argsort_simd_u64(data, config.descending),
1087 SortAlgorithm::Auto => {
1088 #[cfg(feature = "simd_sort")]
1089 {
1090 simd_argsort::argsort_simd_u64(data, config.descending)
1091 }
1092 #[cfg(not(feature = "simd_sort"))]
1093 {
1094 argsort_radix_u64(data, config.descending)
1095 }
1096 }
1097 }
1098}
1099
1100#[cfg(test)]
1101mod tests {
1102 use minarrow::vec64;
1103
1104 use super::*;
1105
1106 #[test]
1107 fn test_total_cmp_f32_ordering_basic() {
1108 let a = 1.0f32;
1109 let b = 2.0f32;
1110 assert_eq!(total_cmp_f(&a, &b), std::cmp::Ordering::Less);
1111 assert_eq!(total_cmp_f(&b, &a), std::cmp::Ordering::Greater);
1112 assert_eq!(total_cmp_f(&a, &a), std::cmp::Ordering::Equal);
1113 }
1114
1115 #[test]
1116 fn test_total_cmp_f64_ordering_basic() {
1117 let a = 1.0f64;
1118 let b = 2.0f64;
1119 assert_eq!(total_cmp_f(&a, &b), std::cmp::Ordering::Less);
1120 assert_eq!(total_cmp_f(&b, &a), std::cmp::Ordering::Greater);
1121 assert_eq!(total_cmp_f(&a, &a), std::cmp::Ordering::Equal);
1122 }
1123
1124 #[test]
1125 fn test_total_cmp_zero_and_negzero() {
1126 let pz = 0.0f32;
1127 let nz = -0.0f32;
1128 assert_ne!(f32::to_bits(pz), f32::to_bits(nz));
1130 assert_ne!(total_cmp_f(&pz, &nz), std::cmp::Ordering::Equal);
1131 }
1132
1133 #[test]
1134 fn test_total_cmp_nan_ordering_f32() {
1135 let nan = f32::NAN;
1136 let pos_inf = f32::INFINITY;
1137 let neg_inf = f32::NEG_INFINITY;
1138 let one = 1.0f32;
1139
1140 assert_eq!(total_cmp_f(&nan, &pos_inf), std::cmp::Ordering::Greater);
1142 assert_eq!(total_cmp_f(&nan, &neg_inf), std::cmp::Ordering::Greater);
1143 assert_eq!(total_cmp_f(&nan, &one), std::cmp::Ordering::Greater);
1144 assert_eq!(total_cmp_f(&nan, &nan), std::cmp::Ordering::Equal);
1146 }
1147
1148 #[test]
1149 fn test_total_cmp_nan_ordering_f64() {
1150 let nan = f64::NAN;
1151 let pos_inf = f64::INFINITY;
1152 let neg_inf = f64::NEG_INFINITY;
1153 let one = 1.0f64;
1154
1155 assert_eq!(total_cmp_f(&nan, &pos_inf), std::cmp::Ordering::Greater);
1156 assert_eq!(total_cmp_f(&nan, &neg_inf), std::cmp::Ordering::Greater);
1157 assert_eq!(total_cmp_f(&nan, &one), std::cmp::Ordering::Greater);
1158 assert_eq!(total_cmp_f(&nan, &nan), std::cmp::Ordering::Equal);
1159 }
1160
1161 #[test]
1162 fn test_sort_float_f32_all_edge_cases() {
1163 let mut v = vec64![
1164 3.0f32,
1165 -0.0,
1166 0.0,
1167 f32::INFINITY,
1168 f32::NEG_INFINITY,
1169 1.0,
1170 -1.0,
1171 f32::NAN,
1172 2.0,
1173 -2.0,
1174 ];
1175 sort_float(&mut v);
1176 assert_eq!(v[0], f32::NEG_INFINITY);
1179 assert_eq!(v[1], -2.0);
1180 assert_eq!(v[2], -1.0);
1181 assert_eq!(v[3], -0.0);
1182 assert_eq!(v[4], 0.0);
1183 assert_eq!(v[5], 1.0);
1184 assert_eq!(v[6], 2.0);
1185 assert_eq!(v[7], 3.0);
1186 assert_eq!(v[8], f32::INFINITY);
1187 assert!(v[9].is_nan());
1188 }
1189
1190 #[test]
1191 fn test_sort_float_f64_all_edge_cases() {
1192 let mut v = vec64![
1193 3.0f64,
1194 -0.0,
1195 0.0,
1196 f64::INFINITY,
1197 f64::NEG_INFINITY,
1198 1.0,
1199 -1.0,
1200 f64::NAN,
1201 2.0,
1202 -2.0,
1203 ];
1204 sort_float(&mut v);
1205 assert_eq!(v[0], f64::NEG_INFINITY);
1206 assert_eq!(v[1], -2.0);
1207 assert_eq!(v[2], -1.0);
1208 assert_eq!(v[3], -0.0);
1209 assert_eq!(v[4], 0.0);
1210 assert_eq!(v[5], 1.0);
1211 assert_eq!(v[6], 2.0);
1212 assert_eq!(v[7], 3.0);
1213 assert_eq!(v[8], f64::INFINITY);
1214 assert!(v[9].is_nan());
1215 }
1216
1217 #[test]
1218 fn test_sorted_float_immutability_and_return_type() {
1219 let v = vec64![1.0f32, 2.0, 3.0];
1220 let out = sorted_float(&v);
1221 assert_eq!(out.as_slice(), &[1.0, 2.0, 3.0]);
1222 assert_eq!(*v, [1.0, 2.0, 3.0]);
1224 }
1225
1226 #[test]
1227 fn test_sorted_float_correct_for_f64() {
1228 let v = vec64![3.0f64, 2.0, 1.0];
1229 let out = sorted_float(&v);
1230 assert_eq!(out.as_slice(), &[1.0, 2.0, 3.0]);
1231 }
1232
1233 #[test]
1234 fn test_sort_float_empty_and_single() {
1235 let mut v: [f32; 0] = [];
1236 sort_float(&mut v);
1237 let mut v2 = [42.0f32];
1238 sort_float(&mut v2);
1239 assert_eq!(v2, [42.0]);
1240 }
1241
1242 #[cfg(test)]
1243 mod tests {
1244 use super::*;
1245 use minarrow::{Vec64, vec64};
1246
1247 #[test]
1248 fn test_sort_int_empty_and_single() {
1249 let mut arr: [i32; 0] = [];
1250 sort_int(&mut arr);
1251 assert_eq!(arr, [] as [i32; 0]);
1252 let mut arr2 = vec64![42];
1253 sort_int(&mut arr2);
1254 assert_eq!(*arr2, [42]);
1255 }
1256
1257 #[test]
1258 fn test_sort_int_order() {
1259 let mut arr = vec64![4, 2, 7, 1, 1, 6, 0, -5];
1260 sort_int(&mut arr);
1261 assert_eq!(*arr, [-5, 0, 1, 1, 2, 4, 6, 7]);
1262 }
1263
1264 #[test]
1265 fn test_sorted_int_immutability_and_output() {
1266 let arr = vec64![5, 3, 7, 1, 2];
1267 let sorted = sorted_int(&arr);
1268 assert_eq!(sorted.as_slice(), &[1, 2, 3, 5, 7]);
1269 assert_eq!(*arr, [5, 3, 7, 1, 2]);
1271 }
1272
1273 #[test]
1274 fn test_sort_str_basic() {
1275 let mut arr = vec64!["z", "b", "a", "d"];
1276 sort_str(&mut arr);
1277 assert_eq!(*arr, ["a", "b", "d", "z"]);
1278 }
1279
1280 #[test]
1281 fn test_sorted_str_and_non_ascii() {
1282 let arr = vec64!["z", "ä", "b", "a", "d"];
1283 let sorted = sorted_str(&arr);
1284 assert_eq!(sorted.as_slice(), &["a", "b", "d", "z", "ä"]); assert_eq!(*arr, ["z", "ä", "b", "a", "d"]);
1287 }
1288
1289 #[test]
1290 fn test_sort_string_array_basic() {
1291 let offsets = vec![0, 1, 3, 5, 5, 6]; let values = "abcde".to_string() + "f";
1293 let (new_offsets, new_values) = sort_string_array(&offsets, &values);
1294 let slices: Vec<_> = new_offsets
1296 .windows(2)
1297 .map(|w| &new_values[w[0]..w[1]])
1298 .collect();
1299 assert_eq!(slices, &["", "a", "bc", "de", "f"]);
1300 }
1301
1302 #[test]
1303 fn test_sort_categorical_lexical_basic_and_nulls() {
1304 let unique = Vec64::from_slice_clone(&[
1306 "apple".to_string(),
1307 "banana".to_string(),
1308 "pear".to_string(),
1309 ]);
1310 let data = Vec64::from_slice(&[2u8, 0, 1, 1, 2, 0, 1]);
1311 let mask = Bitmask::from_bools(&[true, false, true, true, true, false, true]);
1312 let cat = CategoricalArray {
1313 data: data.into(),
1314 unique_values: unique,
1315 null_mask: Some(mask.clone()),
1316 };
1317 let (sorted_data, sorted_mask) = sort_categorical_lexical(&cat);
1318
1319 let mask_out = sorted_mask.unwrap();
1321 let null_pos: Vec<_> = (0..mask_out.len()).filter(|&i| !mask_out.get(i)).collect();
1322 assert_eq!(null_pos, &[0, 1]);
1323
1324 let sorted_as_str: Vec<_> = sorted_data
1326 .iter()
1327 .map(|&k| &cat.unique_values[k.to_usize()][..])
1328 .collect();
1329 let vals = &sorted_as_str[null_pos.len()..];
1330 assert_eq!(vals, &["banana", "banana", "banana", "pear", "pear"]);
1331 }
1332
1333 #[test]
1334 fn test_sort_categorical_all_nulls_and_no_nulls() {
1335 let unique = Vec64::from_slice_clone(&["x".to_string()]);
1337 let data = Vec64::from_slice(&[0u8, 0, 0]);
1338 let mask = Bitmask::from_bools(&[false, false, false]);
1339 let cat = CategoricalArray {
1340 data: data.clone().into(),
1341 unique_values: unique.clone(),
1342 null_mask: Some(mask.clone()),
1343 };
1344 let (_, sorted_mask) = sort_categorical_lexical(&cat);
1345 assert_eq!(
1346 unpack_bitmask(&sorted_mask.unwrap()),
1347 vec64![false, false, false]
1348 );
1349 let mask2 = Bitmask::from_bools(&[true, true, true]);
1351 let cat2 = CategoricalArray {
1352 data: data.clone().into(),
1353 unique_values: unique,
1354 null_mask: Some(mask2.clone()),
1355 };
1356 let (_, sorted_mask2) = sort_categorical_lexical(&cat2);
1357 assert_eq!(
1358 unpack_bitmask(&sorted_mask2.unwrap()),
1359 vec64![true, true, true]
1360 );
1361 }
1362 #[test]
1363 fn test_sort_boolean_array_with_nulls() {
1364 let mut arr = BooleanArray {
1365 data: pack_bitmask(&[false, true, false, true, true, false]),
1366 null_mask: Some(pack_bitmask(&[true, false, true, true, false, true])),
1367 len: 6,
1368 _phantom: std::marker::PhantomData,
1369 };
1370 sort_boolean_array(&mut arr);
1371 let bits = unpack_bitmask(&arr.data);
1373 let nulls = unpack_bitmask(arr.null_mask.as_ref().unwrap());
1374 assert_eq!(nulls, vec64![false, false, true, true, true, true]);
1375 assert_eq!(&bits[2..], [false, false, false, true]);
1377 }
1378
1379 #[test]
1380 fn test_sort_boolean_array_all_true_and_all_false() {
1381 let mut arr = BooleanArray {
1382 data: pack_bitmask(&[true, true, true]),
1383 null_mask: None,
1384 len: 3,
1385 _phantom: std::marker::PhantomData,
1386 };
1387 sort_boolean_array(&mut arr);
1388 assert_eq!(unpack_bitmask(&arr.data), vec64![true, true, true]);
1389
1390 let mut arr2 = BooleanArray {
1391 data: pack_bitmask(&[false, false, false]),
1392 null_mask: None,
1393 len: 3,
1394 _phantom: std::marker::PhantomData,
1395 };
1396 sort_boolean_array(&mut arr2);
1397 assert_eq!(unpack_bitmask(&arr2.data), vec64![false, false, false]);
1398 }
1399
1400 #[test]
1401 fn test_sort_boolean_array_all_nulls_and_none() {
1402 let mut arr = BooleanArray {
1403 data: pack_bitmask(&[true, false, true]),
1404 null_mask: Some(pack_bitmask(&[false, false, false])),
1405 len: 3,
1406 _phantom: std::marker::PhantomData,
1407 };
1408 sort_boolean_array(&mut arr);
1409 assert_eq!(
1410 unpack_bitmask(arr.null_mask.as_ref().unwrap()),
1411 vec64![false, false, false]
1412 );
1413 }
1414
1415 #[test]
1416 fn test_sort_slice_with_mask_basic() {
1417 let data = vec64![3, 1, 2, 1];
1418 let mask = pack_bitmask(&[true, false, true, true]);
1419 let (sorted, mask_out) = sort_slice_with_mask(&data, Some(&mask));
1420 assert_eq!(sorted.as_slice(), &[1, 1, 2, 3]);
1421 assert_eq!(
1422 unpack_bitmask(&mask_out.unwrap()),
1423 vec64![false, true, true, true]
1424 );
1425 }
1426
1427 #[test]
1428 fn test_sort_slice_with_mask_no_mask() {
1429 let data = vec64![3, 2, 1, 1, 0];
1430 let (sorted, mask_out) = sort_slice_with_mask(&data, None);
1431 assert_eq!(sorted.as_slice(), &[0, 1, 1, 2, 3]);
1432 assert!(mask_out.is_none());
1433 }
1434
1435 #[test]
1436 fn test_sort_slice_with_mask_all_true_and_all_false() {
1437 let data = vec64![3, 2, 1, 0];
1438 let mask_true = pack_bitmask(&[true; 4]);
1439 let mask_false = pack_bitmask(&[false; 4]);
1440 let (_, mask_true_out) = sort_slice_with_mask(&data, Some(&mask_true));
1441 let (_, mask_false_out) = sort_slice_with_mask(&data, Some(&mask_false));
1442 assert_eq!(
1443 unpack_bitmask(&mask_true_out.unwrap()),
1444 vec64![true, true, true, true]
1445 );
1446 assert_eq!(
1447 unpack_bitmask(&mask_false_out.unwrap()),
1448 vec64![false, false, false, false]
1449 );
1450 }
1451
1452 #[test]
1453 fn test_sort_int_with_duplicates_and_negatives() {
1454 let mut arr = vec64![-10, -1, 5, 0, 5, -10];
1455 sort_int(&mut arr);
1456 assert_eq!(*arr, [-10, -10, -1, 0, 5, 5]);
1457 }
1458
1459 #[test]
1460 fn test_sort_str_empty_and_duplicate() {
1461 let mut arr = vec64!["", "a", "b", "a", ""];
1462 sort_str(&mut arr);
1463 assert_eq!(*arr, ["", "", "a", "a", "b"]);
1464 }
1465
1466 #[test]
1471 fn test_argsort_basic() {
1472 let data = [5i32, 2, 8, 1, 9];
1473 let indices = argsort(&data, false);
1474 assert_eq!(indices, vec![3, 1, 0, 2, 4]); let sorted: Vec<i32> = indices.iter().map(|&i| data[i]).collect();
1477 assert_eq!(sorted, vec![1, 2, 5, 8, 9]);
1478 }
1479
1480 #[test]
1481 fn test_argsort_descending() {
1482 let data = [5i32, 2, 8, 1, 9];
1483 let indices = argsort(&data, true);
1484 let sorted: Vec<i32> = indices.iter().map(|&i| data[i]).collect();
1485 assert_eq!(sorted, vec![9, 8, 5, 2, 1]);
1486 }
1487
1488 #[test]
1489 fn test_argsort_empty() {
1490 let data: [i32; 0] = [];
1491 let indices = argsort(&data, false);
1492 assert!(indices.is_empty());
1493 }
1494
1495 #[test]
1496 fn test_argsort_float_basic() {
1497 let data = [3.0f64, 1.0, 4.0, 1.5, 2.0];
1498 let indices = argsort_float(&data, false);
1499 let sorted: Vec<f64> = indices.iter().map(|&i| data[i]).collect();
1500 assert_eq!(sorted, vec![1.0, 1.5, 2.0, 3.0, 4.0]);
1501 }
1502
1503 #[test]
1504 fn test_argsort_float_with_nan() {
1505 let data = [3.0f64, f64::NAN, 1.0, f64::NEG_INFINITY];
1506 let indices = argsort_float(&data, false);
1507 let sorted: Vec<f64> = indices.iter().map(|&i| data[i]).collect();
1508 assert_eq!(sorted[0], f64::NEG_INFINITY);
1510 assert_eq!(sorted[1], 1.0);
1511 assert_eq!(sorted[2], 3.0);
1512 assert!(sorted[3].is_nan());
1513 }
1514
1515 #[test]
1516 fn test_argsort_str_basic() {
1517 let data = ["cherry", "apple", "banana"];
1518 let indices = argsort_str(&data, false);
1519 let sorted: Vec<&str> = indices.iter().map(|&i| data[i]).collect();
1520 assert_eq!(sorted, vec!["apple", "banana", "cherry"]);
1521 }
1522
1523 #[test]
1524 fn test_argsort_str_descending() {
1525 let data = ["cherry", "apple", "banana"];
1526 let indices = argsort_str(&data, true);
1527 let sorted: Vec<&str> = indices.iter().map(|&i| data[i]).collect();
1528 assert_eq!(sorted, vec!["cherry", "banana", "apple"]);
1529 }
1530
1531 #[test]
1532 fn test_argsort_string_array_basic() {
1533 let values = "applecherrybanana";
1535 let offsets = [0usize, 5, 11, 17];
1536 let indices = argsort_string_array(&offsets, values, false);
1537 assert_eq!(indices, vec![0, 2, 1]);
1539 }
1540
1541 #[test]
1542 fn test_argsort_radix_u32_basic() {
1543 let data = vec![5u32, 2, 8, 1, 9, 3, 7, 4, 6, 0];
1544 let indices = argsort_radix_u32(&data, false);
1545 let sorted: Vec<u32> = indices.iter().map(|&i| data[i]).collect();
1546 assert_eq!(sorted, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
1547 }
1548
1549 #[test]
1550 fn test_argsort_radix_u32_descending() {
1551 let data = vec![5u32, 2, 8, 1, 9, 3, 7, 4, 6, 0];
1552 let indices = argsort_radix_u32(&data, true);
1553 let sorted: Vec<u32> = indices.iter().map(|&i| data[i]).collect();
1554 assert_eq!(sorted, vec![9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);
1555 }
1556
1557 #[test]
1558 fn test_argsort_radix_i32_with_negatives() {
1559 let data = vec![5i32, -2, 8, -1, 9, 3, -7, 4, 6, 0];
1560 let indices = argsort_radix_i32(&data, false);
1561 let sorted: Vec<i32> = indices.iter().map(|&i| data[i]).collect();
1562 assert_eq!(sorted, vec![-7, -2, -1, 0, 3, 4, 5, 6, 8, 9]);
1563 }
1564
1565 #[test]
1566 fn test_argsort_radix_u64_basic() {
1567 let data = vec![100u64, 50, 200, 25];
1568 let indices = argsort_radix_u64(&data, false);
1569 let sorted: Vec<u64> = indices.iter().map(|&i| data[i]).collect();
1570 assert_eq!(sorted, vec![25, 50, 100, 200]);
1571 }
1572
1573 #[test]
1574 fn test_argsort_radix_i64_with_negatives() {
1575 let data = vec![5i64, -2, 8, -1, 0];
1576 let indices = argsort_radix_i64(&data, false);
1577 let sorted: Vec<i64> = indices.iter().map(|&i| data[i]).collect();
1578 assert_eq!(sorted, vec![-2, -1, 0, 5, 8]);
1579 }
1580
1581 #[test]
1582 fn test_argsort_boolean_basic() {
1583 use minarrow::Bitmask;
1584 let mut data = Bitmask::new_set_all(4, false);
1586 data.set(0, true);
1587 data.set(2, true);
1588
1589 let indices = argsort_boolean(&data, None, false);
1590 let sorted: Vec<bool> = indices.iter().map(|&i| data.get(i)).collect();
1592 assert_eq!(sorted, vec![false, false, true, true]);
1593 }
1594
1595 #[test]
1596 fn test_argsort_boolean_descending() {
1597 use minarrow::Bitmask;
1598 let mut data = Bitmask::new_set_all(4, false);
1599 data.set(0, true);
1600 data.set(2, true);
1601
1602 let indices = argsort_boolean(&data, None, true);
1603 let sorted: Vec<bool> = indices.iter().map(|&i| data.get(i)).collect();
1604 assert_eq!(sorted, vec![true, true, false, false]);
1605 }
1606
1607 #[test]
1608 fn test_argsort_boolean_with_nulls() {
1609 use minarrow::Bitmask;
1610 let mut data = Bitmask::new_set_all(4, false);
1613 data.set(0, true);
1614 data.set(2, true);
1615
1616 let mut null_mask = Bitmask::new_set_all(4, true);
1617 null_mask.set(1, false); let indices = argsort_boolean(&data, Some(&null_mask), false);
1620 assert_eq!(indices[0], 1); }
1623
1624 #[test]
1625 fn test_argsort_auto_i32_basic() {
1626 let data = [5i32, 2, 8, 1, 9];
1627 let config = ArgsortConfig::default();
1628 let indices = argsort_auto_i32(&data, &config);
1629 let sorted: Vec<i32> = indices.iter().map(|&i| data[i]).collect();
1630 assert_eq!(sorted, vec![1, 2, 5, 8, 9]);
1631 }
1632
1633 #[test]
1634 fn test_argsort_auto_with_comparison() {
1635 let data = [5i32, 2, 8, 1, 9];
1636 let config = ArgsortConfig::new().algorithm(SortAlgorithm::Comparison);
1637 let indices = argsort_auto_i32(&data, &config);
1638 let sorted: Vec<i32> = indices.iter().map(|&i| data[i]).collect();
1639 assert_eq!(sorted, vec![1, 2, 5, 8, 9]);
1640 }
1641
1642 #[test]
1643 fn test_argsort_auto_descending() {
1644 let data = [5i32, 2, 8, 1, 9];
1645 let config = ArgsortConfig::new().descending(true);
1646 let indices = argsort_auto_i32(&data, &config);
1647 let sorted: Vec<i32> = indices.iter().map(|&i| data[i]).collect();
1648 assert_eq!(sorted, vec![9, 8, 5, 2, 1]);
1649 }
1650
1651 fn apply_indices<T: Copy>(data: &[T], indices: &[usize]) -> Vec<T> {
1659 indices.iter().map(|&i| data[i]).collect()
1660 }
1661
1662 #[test]
1663 fn test_i32_sorts_produce_same_results() {
1664 let data = vec![
1666 5i32, -2, 8, -1, 9, 3, -7, 4, 6, 0,
1667 i32::MAX, i32::MIN, 0, -100, 100,
1668 1, 1, 1, -1, -1, ];
1670
1671 let comparison_asc = argsort(&data, false);
1673 let radix_asc = argsort_radix_i32(&data, false);
1674 let auto_comparison = argsort_auto_i32(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Comparison));
1675 let auto_radix = argsort_auto_i32(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Radix));
1676 let auto_default = argsort_auto_i32(&data, &ArgsortConfig::default());
1677
1678 let sorted_comparison: Vec<i32> = apply_indices(&data, &comparison_asc);
1680 let sorted_radix: Vec<i32> = apply_indices(&data, &radix_asc);
1681 let sorted_auto_comparison: Vec<i32> = apply_indices(&data, &auto_comparison);
1682 let sorted_auto_radix: Vec<i32> = apply_indices(&data, &auto_radix);
1683 let sorted_auto_default: Vec<i32> = apply_indices(&data, &auto_default);
1684
1685 assert_eq!(sorted_comparison, sorted_radix, "comparison vs radix mismatch");
1687 assert_eq!(sorted_comparison, sorted_auto_comparison, "comparison vs auto_comparison mismatch");
1688 assert_eq!(sorted_comparison, sorted_auto_radix, "comparison vs auto_radix mismatch");
1689 assert_eq!(sorted_comparison, sorted_auto_default, "comparison vs auto_default mismatch");
1690
1691 for i in 1..sorted_comparison.len() {
1693 assert!(sorted_comparison[i-1] <= sorted_comparison[i], "not sorted at index {}", i);
1694 }
1695
1696 let comparison_desc = argsort(&data, true);
1698 let radix_desc = argsort_radix_i32(&data, true);
1699
1700 let sorted_comparison_desc: Vec<i32> = apply_indices(&data, &comparison_desc);
1701 let sorted_radix_desc: Vec<i32> = apply_indices(&data, &radix_desc);
1702
1703 assert_eq!(sorted_comparison_desc, sorted_radix_desc, "descending comparison vs radix mismatch");
1704
1705 for i in 1..sorted_comparison_desc.len() {
1707 assert!(sorted_comparison_desc[i-1] >= sorted_comparison_desc[i], "not sorted descending at index {}", i);
1708 }
1709 }
1710
1711 #[test]
1712 fn test_i64_sorts_produce_same_results() {
1713 let data = vec![
1714 5i64, -2, 8, -1, 9, 3, -7, 4, 6, 0,
1715 i64::MAX, i64::MIN, 0, -100, 100,
1716 i64::MAX - 1, i64::MIN + 1,
1717 ];
1718
1719 let comparison_asc = argsort(&data, false);
1720 let radix_asc = argsort_radix_i64(&data, false);
1721 let auto_comparison = argsort_auto_i64(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Comparison));
1722 let auto_radix = argsort_auto_i64(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Radix));
1723
1724 let sorted_comparison: Vec<i64> = apply_indices(&data, &comparison_asc);
1725 let sorted_radix: Vec<i64> = apply_indices(&data, &radix_asc);
1726 let sorted_auto_comparison: Vec<i64> = apply_indices(&data, &auto_comparison);
1727 let sorted_auto_radix: Vec<i64> = apply_indices(&data, &auto_radix);
1728
1729 assert_eq!(sorted_comparison, sorted_radix, "i64: comparison vs radix mismatch");
1730 assert_eq!(sorted_comparison, sorted_auto_comparison, "i64: comparison vs auto_comparison mismatch");
1731 assert_eq!(sorted_comparison, sorted_auto_radix, "i64: comparison vs auto_radix mismatch");
1732
1733 for i in 1..sorted_comparison.len() {
1735 assert!(sorted_comparison[i-1] <= sorted_comparison[i], "i64: not sorted at index {}", i);
1736 }
1737 }
1738
1739 #[test]
1740 fn test_u32_sorts_produce_same_results() {
1741 let data = vec![
1742 5u32, 2, 8, 1, 9, 3, 7, 4, 6, 0,
1743 u32::MAX, 0, 100, u32::MAX - 1,
1744 1, 1, 1, ];
1746
1747 let comparison_asc = argsort(&data, false);
1748 let radix_asc = argsort_radix_u32(&data, false);
1749 let auto_comparison = argsort_auto_u32(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Comparison));
1750 let auto_radix = argsort_auto_u32(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Radix));
1751
1752 let sorted_comparison: Vec<u32> = apply_indices(&data, &comparison_asc);
1753 let sorted_radix: Vec<u32> = apply_indices(&data, &radix_asc);
1754 let sorted_auto_comparison: Vec<u32> = apply_indices(&data, &auto_comparison);
1755 let sorted_auto_radix: Vec<u32> = apply_indices(&data, &auto_radix);
1756
1757 assert_eq!(sorted_comparison, sorted_radix, "u32: comparison vs radix mismatch");
1758 assert_eq!(sorted_comparison, sorted_auto_comparison, "u32: comparison vs auto_comparison mismatch");
1759 assert_eq!(sorted_comparison, sorted_auto_radix, "u32: comparison vs auto_radix mismatch");
1760
1761 for i in 1..sorted_comparison.len() {
1763 assert!(sorted_comparison[i-1] <= sorted_comparison[i], "u32: not sorted at index {}", i);
1764 }
1765 }
1766
1767 #[test]
1768 fn test_u64_sorts_produce_same_results() {
1769 let data = vec![
1770 5u64, 2, 8, 1, 9, 3, 7, 4, 6, 0,
1771 u64::MAX, 0, 100, u64::MAX - 1,
1772 ];
1773
1774 let comparison_asc = argsort(&data, false);
1775 let radix_asc = argsort_radix_u64(&data, false);
1776 let auto_comparison = argsort_auto_u64(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Comparison));
1777 let auto_radix = argsort_auto_u64(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Radix));
1778
1779 let sorted_comparison: Vec<u64> = apply_indices(&data, &comparison_asc);
1780 let sorted_radix: Vec<u64> = apply_indices(&data, &radix_asc);
1781 let sorted_auto_comparison: Vec<u64> = apply_indices(&data, &auto_comparison);
1782 let sorted_auto_radix: Vec<u64> = apply_indices(&data, &auto_radix);
1783
1784 assert_eq!(sorted_comparison, sorted_radix, "u64: comparison vs radix mismatch");
1785 assert_eq!(sorted_comparison, sorted_auto_comparison, "u64: comparison vs auto_comparison mismatch");
1786 assert_eq!(sorted_comparison, sorted_auto_radix, "u64: comparison vs auto_radix mismatch");
1787
1788 for i in 1..sorted_comparison.len() {
1790 assert!(sorted_comparison[i-1] <= sorted_comparison[i], "u64: not sorted at index {}", i);
1791 }
1792 }
1793
1794 #[cfg(feature = "simd_sort")]
1795 #[test]
1796 fn test_i32_sorts_produce_same_results_with_simd() {
1797 use super::simd_argsort;
1798
1799 let data = vec![
1800 5i32, -2, 8, -1, 9, 3, -7, 4, 6, 0,
1801 i32::MAX, i32::MIN, 0, -100, 100,
1802 1, 1, 1, -1, -1,
1803 ];
1804
1805 let comparison_asc = argsort(&data, false);
1806 let radix_asc = argsort_radix_i32(&data, false);
1807 let simd_asc = simd_argsort::argsort_simd_i32(&data, false);
1808 let auto_simd = argsort_auto_i32(&data, &ArgsortConfig::new().algorithm(SortAlgorithm::Simd));
1809
1810 let sorted_comparison: Vec<i32> = apply_indices(&data, &comparison_asc);
1811 let sorted_radix: Vec<i32> = apply_indices(&data, &radix_asc);
1812 let sorted_simd: Vec<i32> = apply_indices(&data, &simd_asc);
1813 let sorted_auto_simd: Vec<i32> = apply_indices(&data, &auto_simd);
1814
1815 assert_eq!(sorted_comparison, sorted_radix, "simd test: comparison vs radix mismatch");
1816 assert_eq!(sorted_comparison, sorted_simd, "simd test: comparison vs simd mismatch");
1817 assert_eq!(sorted_comparison, sorted_auto_simd, "simd test: comparison vs auto_simd mismatch");
1818 }
1819
1820 #[cfg(feature = "simd_sort")]
1821 #[test]
1822 fn test_i64_sorts_produce_same_results_with_simd() {
1823 use super::simd_argsort;
1824
1825 let data = vec![
1826 5i64, -2, 8, -1, 9, 3, -7, 4, 6, 0,
1827 i64::MAX, i64::MIN, 0, -100, 100,
1828 ];
1829
1830 let comparison_asc = argsort(&data, false);
1831 let radix_asc = argsort_radix_i64(&data, false);
1832 let simd_asc = simd_argsort::argsort_simd_i64(&data, false);
1833
1834 let sorted_comparison: Vec<i64> = apply_indices(&data, &comparison_asc);
1835 let sorted_radix: Vec<i64> = apply_indices(&data, &radix_asc);
1836 let sorted_simd: Vec<i64> = apply_indices(&data, &simd_asc);
1837
1838 assert_eq!(sorted_comparison, sorted_radix, "simd i64: comparison vs radix mismatch");
1839 assert_eq!(sorted_comparison, sorted_simd, "simd i64: comparison vs simd mismatch");
1840 }
1841
1842 #[cfg(feature = "simd_sort")]
1843 #[test]
1844 fn test_u32_sorts_produce_same_results_with_simd() {
1845 use super::simd_argsort;
1846
1847 let data = vec![
1848 5u32, 2, 8, 1, 9, 3, 7, 4, 6, 0,
1849 u32::MAX, 0, 100, u32::MAX - 1,
1850 ];
1851
1852 let comparison_asc = argsort(&data, false);
1853 let radix_asc = argsort_radix_u32(&data, false);
1854 let simd_asc = simd_argsort::argsort_simd_u32(&data, false);
1855
1856 let sorted_comparison: Vec<u32> = apply_indices(&data, &comparison_asc);
1857 let sorted_radix: Vec<u32> = apply_indices(&data, &radix_asc);
1858 let sorted_simd: Vec<u32> = apply_indices(&data, &simd_asc);
1859
1860 assert_eq!(sorted_comparison, sorted_radix, "simd u32: comparison vs radix mismatch");
1861 assert_eq!(sorted_comparison, sorted_simd, "simd u32: comparison vs simd mismatch");
1862 }
1863
1864 #[cfg(feature = "simd_sort")]
1865 #[test]
1866 fn test_u64_sorts_produce_same_results_with_simd() {
1867 use super::simd_argsort;
1868
1869 let data = vec![
1870 5u64, 2, 8, 1, 9, 3, 7, 4, 6, 0,
1871 u64::MAX, 0, 100, u64::MAX - 1,
1872 ];
1873
1874 let comparison_asc = argsort(&data, false);
1875 let radix_asc = argsort_radix_u64(&data, false);
1876 let simd_asc = simd_argsort::argsort_simd_u64(&data, false);
1877
1878 let sorted_comparison: Vec<u64> = apply_indices(&data, &comparison_asc);
1879 let sorted_radix: Vec<u64> = apply_indices(&data, &radix_asc);
1880 let sorted_simd: Vec<u64> = apply_indices(&data, &simd_asc);
1881
1882 assert_eq!(sorted_comparison, sorted_radix, "simd u64: comparison vs radix mismatch");
1883 assert_eq!(sorted_comparison, sorted_simd, "simd u64: comparison vs simd mismatch");
1884 }
1885
1886 #[cfg(feature = "parallel_sort")]
1887 #[test]
1888 fn test_i32_sorts_produce_same_results_with_parallel() {
1889 use super::parallel_argsort;
1890
1891 let data = vec![
1892 5i32, -2, 8, -1, 9, 3, -7, 4, 6, 0,
1893 i32::MAX, i32::MIN, 0, -100, 100,
1894 1, 1, 1, -1, -1,
1895 ];
1896
1897 let comparison_asc = argsort(&data, false);
1898 let parallel_stable = parallel_argsort::argsort_parallel(&data, false, true);
1899 let parallel_unstable = parallel_argsort::argsort_parallel(&data, false, false);
1900
1901 let sorted_comparison: Vec<i32> = apply_indices(&data, &comparison_asc);
1902 let sorted_parallel_stable: Vec<i32> = apply_indices(&data, ¶llel_stable);
1903 let sorted_parallel_unstable: Vec<i32> = apply_indices(&data, ¶llel_unstable);
1904
1905 assert_eq!(sorted_comparison, sorted_parallel_stable, "parallel: comparison vs parallel_stable mismatch");
1906 assert_eq!(sorted_comparison, sorted_parallel_unstable, "parallel: comparison vs parallel_unstable mismatch");
1907 }
1908
1909 #[test]
1911 fn test_numeric_sorts_produce_same_results_large_dataset() {
1912 fn lcg_next(state: &mut u64) -> i32 {
1914 *state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
1915 (*state >> 33) as i32
1916 }
1917
1918 let mut state = 12345u64;
1919 let data: Vec<i32> = (0..10_000).map(|_| lcg_next(&mut state)).collect();
1920
1921 let comparison_asc = argsort(&data, false);
1922 let radix_asc = argsort_radix_i32(&data, false);
1923
1924 let sorted_comparison: Vec<i32> = apply_indices(&data, &comparison_asc);
1925 let sorted_radix: Vec<i32> = apply_indices(&data, &radix_asc);
1926
1927 assert_eq!(sorted_comparison, sorted_radix, "large dataset: comparison vs radix mismatch");
1928
1929 for i in 1..sorted_comparison.len() {
1931 assert!(sorted_comparison[i-1] <= sorted_comparison[i], "large dataset: not sorted at index {}", i);
1932 }
1933 }
1934 }
1935}