1use arrow_array::cast::AsArray;
21use arrow_array::types::*;
22use arrow_array::{
23 Array, ArrowNativeTypeOp, BooleanArray, GenericByteArray, GenericByteViewArray,
24 downcast_primitive_array,
25};
26use arrow_buffer::NullBuffer;
27use arrow_schema::{ArrowError, DataType, SortOptions};
28use std::cmp::Ordering;
29
30pub(crate) fn can_rank(data_type: &DataType) -> bool {
32 data_type.is_primitive()
33 || matches!(
34 data_type,
35 DataType::Boolean
36 | DataType::Utf8
37 | DataType::LargeUtf8
38 | DataType::Binary
39 | DataType::LargeBinary
40 | DataType::Utf8View
41 | DataType::BinaryView
42 )
43}
44
45pub fn rank(array: &dyn Array, options: Option<SortOptions>) -> Result<Vec<u32>, ArrowError> {
58 let options = options.unwrap_or_default();
59 let ranks = downcast_primitive_array! {
60 array => primitive_rank(array.values(), array.nulls(), options),
61 DataType::Boolean => boolean_rank(array.as_boolean(), options),
62 DataType::Utf8 => bytes_rank(array.as_bytes::<Utf8Type>(), options),
63 DataType::LargeUtf8 => bytes_rank(array.as_bytes::<LargeUtf8Type>(), options),
64 DataType::Binary => bytes_rank(array.as_bytes::<BinaryType>(), options),
65 DataType::LargeBinary => bytes_rank(array.as_bytes::<LargeBinaryType>(), options),
66 DataType::Utf8View => byte_view_rank(array.as_string_view(), options),
67 DataType::BinaryView => byte_view_rank(array.as_binary_view(), options),
68 d => return Err(ArrowError::ComputeError(format!("{d:?} not supported in rank")))
69 };
70 Ok(ranks)
71}
72
73#[inline(never)]
74fn primitive_rank<T: ArrowNativeTypeOp>(
75 values: &[T],
76 nulls: Option<&NullBuffer>,
77 options: SortOptions,
78) -> Vec<u32> {
79 let len: u32 = values.len().try_into().unwrap();
80 let to_sort = match nulls.filter(|n| n.null_count() > 0) {
81 Some(n) => n
82 .valid_indices()
83 .map(|idx| (values[idx], idx as u32))
84 .collect(),
85 None => values.iter().copied().zip(0..len).collect(),
86 };
87 rank_impl(values.len(), to_sort, options, T::compare, T::is_eq)
88}
89
90#[inline(never)]
91fn bytes_rank<T: ByteArrayType>(array: &GenericByteArray<T>, options: SortOptions) -> Vec<u32> {
92 let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n| n.null_count() > 0) {
93 Some(n) => n
94 .valid_indices()
95 .map(|idx| (array.value(idx).as_ref(), idx as u32))
96 .collect(),
97 None => (0..array.len())
98 .map(|idx| (array.value(idx).as_ref(), idx as u32))
99 .collect(),
100 };
101 rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq)
102}
103
104#[inline(never)]
105fn byte_view_rank<T: ByteViewType>(
106 array: &GenericByteViewArray<T>,
107 options: SortOptions,
108) -> Vec<u32> {
109 if array.data_buffers().is_empty() {
113 let to_sort: Vec<(u128, u32)> = match array.nulls().filter(|n| n.null_count() > 0) {
114 Some(n) => n
115 .valid_indices()
116 .map(|idx| {
117 let raw = unsafe { *array.views().get_unchecked(idx) };
119 (GenericByteViewArray::<T>::inline_key_fast(raw), idx as u32)
120 })
121 .collect(),
122 None => array
123 .views()
124 .iter()
125 .enumerate()
126 .map(|(idx, raw)| (GenericByteViewArray::<T>::inline_key_fast(*raw), idx as u32))
127 .collect(),
128 };
129 return rank_impl(
130 array.len(),
131 to_sort,
132 options,
133 |a, b| a.cmp(&b),
134 |a, b| a == b,
135 );
136 }
137
138 if has_high_byte_view_key_collision_rate(array) {
139 let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n| n.null_count() > 0) {
140 Some(n) => n
141 .valid_indices()
142 .map(|idx| (array.value(idx).as_ref(), idx as u32))
143 .collect(),
144 None => (0..array.len())
145 .map(|idx| (array.value(idx).as_ref(), idx as u32))
146 .collect(),
147 };
148 return rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq);
149 }
150
151 let to_sort: Vec<(u128, u32)> = match array.nulls().filter(|n| n.null_count() > 0) {
155 Some(n) => n
156 .valid_indices()
157 .map(|idx| {
158 let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() };
160 (byte_view_key(value), idx as u32)
161 })
162 .collect(),
163 None => (0..array.len())
164 .map(|idx| {
165 let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() };
167 (byte_view_key(value), idx as u32)
168 })
169 .collect(),
170 };
171 rank_impl_by(
172 array.len(),
173 to_sort,
174 options,
175 |a, b| compare_view_key(array, a, b),
176 |a, b| equal_view_key(array, a, b),
177 )
178}
179
180const BYTE_VIEW_KEY_LEN: usize = 16;
183
184const BYTE_VIEW_KEY_SAMPLES_PER_WINDOW: usize = 4;
187
188const BYTE_VIEW_KEY_SAMPLE_SIZE: usize = 8;
192
193const BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW: usize = 32;
197
198const BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO: usize = 3;
202
203fn has_high_byte_view_key_collision_rate<T: ByteViewType>(array: &GenericByteViewArray<T>) -> bool {
211 if array.len() < 2 {
212 return false;
213 }
214
215 let mut keys = [0_u128; BYTE_VIEW_KEY_SAMPLE_SIZE];
216 let mut sample_len = 0;
217 let midpoint = array.len() / 2;
218
219 for (start, end) in [(0, midpoint), (midpoint, array.len())] {
222 let probe_end = end.min(start.saturating_add(BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW));
223 let window_start = sample_len;
224 let mut window_samples = 0;
225
226 for idx in start..probe_end {
227 if array.is_null(idx) {
228 continue;
229 }
230
231 let value: &[u8] = unsafe { array.value_unchecked(idx).as_ref() };
233 keys[sample_len] = byte_view_key(value);
234 sample_len += 1;
235 window_samples += 1;
236 if window_samples == BYTE_VIEW_KEY_SAMPLES_PER_WINDOW {
237 break;
238 }
239 }
240
241 if window_samples == BYTE_VIEW_KEY_SAMPLES_PER_WINDOW
246 && keys[window_start..sample_len]
247 .windows(2)
248 .all(|w| w[0] == w[1])
249 {
250 return true;
251 }
252 }
253
254 if sample_len < 2 {
255 return false;
256 }
257
258 let keys = &mut keys[..sample_len];
259 keys.sort_unstable();
260 let unique_keys = 1 + keys.windows(2).filter(|w| w[0] != w[1]).count();
261
262 let colliding_keys = sample_len - unique_keys;
266 colliding_keys * BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO >= sample_len
267}
268
269#[inline(always)]
270fn byte_view_key(value: &[u8]) -> u128 {
271 let mut key = [0_u8; BYTE_VIEW_KEY_LEN];
272 let key_len = value.len().min(key.len());
273 key[..key_len].copy_from_slice(&value[..key_len]);
274
275 u128::from_be_bytes(key)
279}
280
281#[inline(always)]
282fn compare_view_key<T: ByteViewType>(
283 array: &GenericByteViewArray<T>,
284 a: &(u128, u32),
285 b: &(u128, u32),
286) -> Ordering {
287 match a.0.cmp(&b.0) {
288 Ordering::Equal => {
289 let full_a: &[u8] = unsafe { array.value_unchecked(a.1 as usize).as_ref() };
291 let full_b: &[u8] = unsafe { array.value_unchecked(b.1 as usize).as_ref() };
292 full_a.cmp(full_b)
293 }
294 ordering => ordering,
295 }
296}
297
298#[inline(always)]
299fn equal_view_key<T: ByteViewType>(
300 array: &GenericByteViewArray<T>,
301 a: &(u128, u32),
302 b: &(u128, u32),
303) -> bool {
304 if a.0 != b.0 {
305 return false;
306 }
307
308 let full_a: &[u8] = unsafe { array.value_unchecked(a.1 as usize).as_ref() };
310 let full_b: &[u8] = unsafe { array.value_unchecked(b.1 as usize).as_ref() };
311 full_a == full_b
312}
313
314fn rank_impl_by<T, C, E>(
315 len: usize,
316 mut valid: Vec<(T, u32)>,
317 options: SortOptions,
318 compare: C,
319 eq: E,
320) -> Vec<u32>
321where
322 C: Fn(&(T, u32), &(T, u32)) -> Ordering,
323 E: Fn(&(T, u32), &(T, u32)) -> bool,
324{
325 valid.sort_unstable_by(compare);
330 if options.descending {
331 valid.reverse();
332 }
333
334 let (mut valid_rank, null_rank) = match options.nulls_first {
335 true => (len as u32, (len - valid.len()) as u32),
336 false => (valid.len() as u32, len as u32),
337 };
338
339 let mut out: Vec<_> = vec![null_rank; len];
340 if let Some(v) = valid.last() {
341 out[v.1 as usize] = valid_rank;
342 }
343
344 let mut count = 1; for w in valid.windows(2).rev() {
346 match eq(&w[0], &w[1]) {
347 true => {
348 count += 1;
349 out[w[0].1 as usize] = valid_rank;
350 }
351 false => {
352 valid_rank -= count;
353 count = 1;
354 out[w[0].1 as usize] = valid_rank
355 }
356 }
357 }
358
359 out
360}
361
362fn rank_impl<T, C, E>(
363 len: usize,
364 mut valid: Vec<(T, u32)>,
365 options: SortOptions,
366 compare: C,
367 eq: E,
368) -> Vec<u32>
369where
370 T: Copy,
371 C: Fn(T, T) -> Ordering,
372 E: Fn(T, T) -> bool,
373{
374 valid.sort_unstable_by(|a, b| compare(a.0, b.0));
376 if options.descending {
377 valid.reverse();
378 }
379
380 let (mut valid_rank, null_rank) = match options.nulls_first {
381 true => (len as u32, (len - valid.len()) as u32),
382 false => (valid.len() as u32, len as u32),
383 };
384
385 let mut out: Vec<_> = vec![null_rank; len];
386 if let Some(v) = valid.last() {
387 out[v.1 as usize] = valid_rank;
388 }
389
390 let mut count = 1; for w in valid.windows(2).rev() {
392 match eq(w[0].0, w[1].0) {
393 true => {
394 count += 1;
395 out[w[0].1 as usize] = valid_rank;
396 }
397 false => {
398 valid_rank -= count;
399 count = 1;
400 out[w[0].1 as usize] = valid_rank
401 }
402 }
403 }
404
405 out
406}
407
408#[inline]
417fn get_boolean_rank_index(value: bool, is_null: bool) -> usize {
418 let is_null_num = is_null as usize;
419 (is_null_num << 1) | (value as usize & !is_null_num)
420}
421
422#[inline(never)]
423fn boolean_rank(array: &BooleanArray, options: SortOptions) -> Vec<u32> {
424 let null_count = array.null_count() as u32;
425 let true_count = array.true_count() as u32;
426 let false_count = array.len() as u32 - null_count - true_count;
427
428 let ranks_index: [u32; 3] = match (options.descending, options.nulls_first) {
445 (true, true) => [
447 null_count + true_count + false_count,
448 null_count + true_count,
449 null_count,
450 ],
451 (true, false) => [
453 true_count + false_count,
454 true_count,
455 true_count + false_count + null_count,
456 ],
457 (false, true) => [
459 null_count + false_count,
460 null_count + false_count + true_count,
461 null_count,
462 ],
463 (false, false) => [
465 false_count,
466 false_count + true_count,
467 false_count + true_count + null_count,
468 ],
469 };
470
471 match array.nulls().filter(|n| n.null_count() > 0) {
472 Some(n) => array
473 .values()
474 .iter()
475 .zip(n.iter())
476 .map(|(value, is_valid)| ranks_index[get_boolean_rank_index(value, !is_valid)])
477 .collect::<Vec<u32>>(),
478 None => array
479 .values()
480 .iter()
481 .map(|value| ranks_index[value as usize])
482 .collect::<Vec<u32>>(),
483 }
484}
485
486#[cfg(test)]
487mod tests {
488 use super::*;
489 use arrow_array::*;
490
491 fn assert_same_rank(left: &dyn Array, right: &dyn Array) {
492 for descending in [false, true] {
493 for nulls_first in [false, true] {
494 let options = SortOptions {
495 descending,
496 nulls_first,
497 };
498 assert_eq!(
499 rank(left, Some(options)).unwrap(),
500 rank(right, Some(options)).unwrap()
501 );
502 }
503 }
504 }
505
506 #[test]
507 fn test_primitive() {
508 let descending = SortOptions {
509 descending: true,
510 nulls_first: true,
511 };
512
513 let nulls_last = SortOptions {
514 descending: false,
515 nulls_first: false,
516 };
517
518 let nulls_last_descending = SortOptions {
519 descending: true,
520 nulls_first: false,
521 };
522
523 let a = Int32Array::from(vec![Some(1), Some(1), None, Some(3), Some(3), Some(4)]);
524 let res = rank(&a, None).unwrap();
525 assert_eq!(res, &[3, 3, 1, 5, 5, 6]);
526
527 let res = rank(&a, Some(descending)).unwrap();
528 assert_eq!(res, &[6, 6, 1, 4, 4, 2]);
529
530 let res = rank(&a, Some(nulls_last)).unwrap();
531 assert_eq!(res, &[2, 2, 6, 4, 4, 5]);
532
533 let res = rank(&a, Some(nulls_last_descending)).unwrap();
534 assert_eq!(res, &[5, 5, 6, 3, 3, 1]);
535
536 let nulls = NullBuffer::from(vec![true, true, false, true, false, false]);
538 let a = Int32Array::new(vec![1, 4, 3, 4, 5, 5].into(), Some(nulls));
539 let res = rank(&a, None).unwrap();
540 assert_eq!(res, &[4, 6, 3, 6, 3, 3]);
541 }
542
543 #[test]
544 fn test_get_boolean_rank_index() {
545 assert_eq!(get_boolean_rank_index(true, true), 2);
546 assert_eq!(get_boolean_rank_index(false, true), 2);
547 assert_eq!(get_boolean_rank_index(true, false), 1);
548 assert_eq!(get_boolean_rank_index(false, false), 0);
549 }
550
551 #[test]
552 fn test_nullable_booleans() {
553 let descending = SortOptions {
554 descending: true,
555 nulls_first: true,
556 };
557
558 let nulls_last = SortOptions {
559 descending: false,
560 nulls_first: false,
561 };
562
563 let nulls_last_descending = SortOptions {
564 descending: true,
565 nulls_first: false,
566 };
567
568 let a = BooleanArray::from(vec![Some(true), Some(true), None, Some(false), Some(false)]);
569 let res = rank(&a, None).unwrap();
570 assert_eq!(res, &[5, 5, 1, 3, 3]);
571
572 let res = rank(&a, Some(descending)).unwrap();
573 assert_eq!(res, &[3, 3, 1, 5, 5]);
574
575 let res = rank(&a, Some(nulls_last)).unwrap();
576 assert_eq!(res, &[4, 4, 5, 2, 2]);
577
578 let res = rank(&a, Some(nulls_last_descending)).unwrap();
579 assert_eq!(res, &[2, 2, 5, 4, 4]);
580
581 let nulls = NullBuffer::from(vec![true, true, false, true, true]);
583 let a = BooleanArray::new(vec![true, true, true, false, false].into(), Some(nulls));
584 let res = rank(&a, None).unwrap();
585 assert_eq!(res, &[5, 5, 1, 3, 3]);
586 }
587
588 #[test]
589 fn test_booleans() {
590 let descending = SortOptions {
591 descending: true,
592 nulls_first: true,
593 };
594
595 let nulls_last = SortOptions {
596 descending: false,
597 nulls_first: false,
598 };
599
600 let nulls_last_descending = SortOptions {
601 descending: true,
602 nulls_first: false,
603 };
604
605 let a = BooleanArray::from(vec![true, false, false, false, true]);
606 let res = rank(&a, None).unwrap();
607 assert_eq!(res, &[5, 3, 3, 3, 5]);
608
609 let res = rank(&a, Some(descending)).unwrap();
610 assert_eq!(res, &[2, 5, 5, 5, 2]);
611
612 let res = rank(&a, Some(nulls_last)).unwrap();
613 assert_eq!(res, &[5, 3, 3, 3, 5]);
614
615 let res = rank(&a, Some(nulls_last_descending)).unwrap();
616 assert_eq!(res, &[2, 5, 5, 5, 2]);
617 }
618
619 #[test]
620 fn test_bytes() {
621 let v = vec!["foo", "fo", "bar", "bar"];
622 let values = StringArray::from(v.clone());
623 let res = rank(&values, None).unwrap();
624 assert_eq!(res, &[4, 3, 2, 2]);
625
626 let values = LargeStringArray::from(v.clone());
627 let res = rank(&values, None).unwrap();
628 assert_eq!(res, &[4, 3, 2, 2]);
629
630 let values = StringViewArray::from(v);
631 let res = rank(&values, None).unwrap();
632 assert_eq!(res, &[4, 3, 2, 2]);
633
634 let v: Vec<&[u8]> = vec![&[1, 2], &[0], &[1, 2, 3], &[1, 2]];
635 let values = LargeBinaryArray::from(v.clone());
636 let res = rank(&values, None).unwrap();
637 assert_eq!(res, &[3, 1, 4, 3]);
638
639 let values = BinaryArray::from(v.clone());
640 let res = rank(&values, None).unwrap();
641 assert_eq!(res, &[3, 1, 4, 3]);
642
643 let values = BinaryViewArray::from_iter_values(v);
644 let res = rank(&values, None).unwrap();
645 assert_eq!(res, &[3, 1, 4, 3]);
646 }
647
648 #[test]
649 fn test_inline_byte_views() {
650 let string_values = vec![
651 Some(""),
652 Some("short"),
653 None,
654 Some("0123456789qa"), Some("short"),
656 ];
657 let string_view = StringViewArray::from(string_values.clone());
658 let string = StringArray::from(string_values);
659
660 assert!(string_view.data_buffers().is_empty());
661 assert_same_rank(&string_view, &string);
662
663 let binary_values: Vec<Option<&[u8]>> = vec![
664 Some(b""),
665 Some(b"short"),
666 None,
667 Some(b"0123456789qa"),
668 Some(b"short"),
669 ];
670 let binary_view = BinaryViewArray::from_iter(binary_values.clone());
671 let binary = BinaryArray::from_opt_vec(binary_values);
672
673 assert!(binary_view.data_buffers().is_empty());
674 assert_same_rank(&binary_view, &binary);
675 }
676
677 #[test]
678 fn test_string_view_with_nulls() {
679 let values = StringViewArray::from(vec![
680 Some("a string longer than twelve bytes"),
681 Some("bar"),
682 None,
683 Some("a string longer than twelve bytes"),
684 ]);
685 let res = rank(&values, None).unwrap();
686 assert_eq!(res, &[3, 4, 1, 3]);
687 }
688
689 #[test]
690 fn test_binary_view_with_nulls() {
691 let long_value = b"a binary value longer than twelve bytes".as_ref();
692 let values = BinaryViewArray::from_iter([
693 Some(long_value),
694 Some(b"bar".as_ref()),
695 None,
696 Some(long_value),
697 ]);
698 let res = rank(&values, None).unwrap();
699 assert_eq!(res, &[3, 4, 1, 3]);
700 }
701
702 #[test]
703 fn test_string_view_key_collisions() {
704 let values = vec![
705 Some("abcdefghijklmnop"),
706 Some("abcdefghijklmnopA"),
707 Some("abcdefghijklmnopB"),
708 Some("abcdefghijklmnopA"),
709 Some("abcdefghijklmno"),
710 Some("short"),
711 None,
712 ];
713 let expected = StringArray::from(values.clone());
714 let actual = StringViewArray::from(values);
715
716 assert_same_rank(&actual, &expected);
717 }
718
719 #[test]
720 fn test_binary_view_key_collisions() {
721 let zeroes_16 = [0_u8; 16];
722 let zeroes_17 = [0_u8; 17];
723 let mut zeroes_then_one = [0_u8; 17];
724 zeroes_then_one[16] = 1;
725 let mut zeroes_then_two = [0_u8; 17];
726 zeroes_then_two[16] = 2;
727
728 let values: Vec<Option<&[u8]>> = vec![
729 Some(b""),
730 Some(b"\0"),
731 Some(&zeroes_16),
732 Some(&zeroes_17),
733 Some(&zeroes_then_one),
734 Some(&zeroes_then_two),
735 Some(&zeroes_then_one),
736 None,
737 ];
738 let expected = BinaryArray::from_opt_vec(values.clone());
739 let actual = BinaryViewArray::from_iter(values);
740
741 assert_same_rank(&actual, &expected);
742 }
743
744 #[test]
745 fn test_byte_view_high_key_collision_detection() {
746 const SIZE: u32 = 64;
747
748 let same_key: StringViewArray = (0..SIZE)
749 .map(|i| {
750 let suffix = i.wrapping_mul(2_654_435_761);
751 Some(format!("abcdefghijklmnop{suffix:08x}"))
752 })
753 .collect();
754
755 assert_eq!(
756 byte_view_key(same_key.value(0).as_bytes()),
757 byte_view_key(same_key.value(1).as_bytes())
758 );
759 assert!(has_high_byte_view_key_collision_rate(&same_key));
760
761 let clustered: StringViewArray = (0..SIZE)
762 .map(|i| {
763 let suffix = i.wrapping_mul(2_654_435_761);
764 let value = if i < SIZE / 2 {
765 format!("{suffix:016x}abcdefgh")
766 } else {
767 format!("abcdefghijklmnop{suffix:08x}")
768 };
769 Some(value)
770 })
771 .collect();
772 assert!(has_high_byte_view_key_collision_rate(&clustered));
773
774 let with_nulls: StringViewArray = (0..SIZE)
775 .map(|i| {
776 (i % 2 == 0).then(|| {
777 let suffix = i.wrapping_mul(2_654_435_761);
778 format!("abcdefghijklmnop{suffix:08x}")
779 })
780 })
781 .collect();
782 assert!(has_high_byte_view_key_collision_rate(&with_nulls));
783
784 let distinct: StringViewArray = (0..SIZE)
785 .map(|i| {
786 let suffix = i.wrapping_mul(2_654_435_761);
787 Some(format!("{suffix:016x}abcdefgh"))
788 })
789 .collect();
790 assert!(!has_high_byte_view_key_collision_rate(&distinct));
791 }
792}