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