Skip to main content

arrow_ord/
rank.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Provides `rank` function to assign a rank to each value in an array
6
7use 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
17/// Whether `arrow_ord::rank` can rank an array of given data type.
18pub(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
32/// Assigns a rank to each value in `array` based on its position in the sorted order
33///
34/// Where values are equal, they will be assigned the highest of their ranks,
35/// leaving gaps in the overall rank assignment
36///
37/// ```
38/// # use arrow_array::StringArray;
39/// # use arrow_ord::rank::rank;
40/// let array = StringArray::from(vec![Some("foo"), None, Some("foo"), None, Some("bar")]);
41/// let ranks = rank(&array, None).unwrap();
42/// assert_eq!(ranks, &[5, 2, 5, 2, 3]);
43/// ```
44pub 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    // An inline view already contains the complete value. Convert it once to
97    // a key whose integer ordering matches the byte ordering, as is done by
98    // `sort_byte_view`.
99    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                    // SAFETY: `valid_indices` only yields indices in the array.
105                    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    // Cache a wider prefix than the 4 bytes stored in a non-inline view. This
139    // pays for the backing-buffer access once per value instead of once per
140    // comparison, and only resolves the complete value when two keys collide.
141    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                // SAFETY: `valid_indices` only yields indices in the array.
146                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                // SAFETY: `idx` is in `0..array.len()`.
153                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
167// A 16-byte prefix fits in one `u128` and is wider than the 4-byte view prefix.
168// Shorter keys collide more often; longer keys need another representation.
169const BYTE_VIEW_KEY_LEN: usize = 16;
170
171// Four valid values per window keeps sampling cheap while allowing an early
172// all-collision decision. More samples improve confidence but add buffer reads.
173const BYTE_VIEW_KEY_SAMPLES_PER_WINDOW: usize = 4;
174
175// Capacity for the two windows; it must be at least
176// `2 * BYTE_VIEW_KEY_SAMPLES_PER_WINDOW`. Increasing it alone has no effect;
177// decreasing it without changing the sampling count can overflow the array.
178const BYTE_VIEW_KEY_SAMPLE_SIZE: usize = 8;
179
180// Bound entries inspected when nulls are present. A higher limit finds valid
181// samples more reliably but costs reads; a lower limit is cheaper but less
182// informative for null-heavy arrays.
183const BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW: usize = 32;
184
185// `colliding_keys * 3 >= sample_len` means roughly one third of sampled keys
186// are duplicates. Lower values fall back earlier; higher values risk keeping
187// the key path for collision-heavy inputs.
188const BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO: usize = 3;
189
190/// Estimates whether cached byte-view keys collide often enough to make the
191/// key-based ranking path unattractive.
192/// Caching a wider key usually avoids repeated backing-buffer reads for long
193/// views. If sampled keys collide frequently, resolving full values plus the
194/// extra key comparison can be slower than comparing slices directly, so the
195/// caller falls back to that path. The bounded two-window sample keeps this
196/// check inexpensive.
197fn 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    // Probe small local windows in both halves. Keeping each probe local avoids
207    // turning collision detection itself into scattered backing-buffer reads.
208    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            // SAFETY: `idx` is within a window bounded by `array.len()`.
219            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        // Four equal keys already contribute three collisions. Even if every
229        // sample in the other window is distinct, that satisfies the final
230        // one-third threshold, so avoid touching the second backing-buffer
231        // window in the common all-collision case.
232        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    // If at least roughly one third of sampled keys collide, comparing the
250    // wider key before every full-value comparison is likely more expensive
251    // than sorting slices directly.
252    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    // Big-endian conversion makes integer comparison equivalent to comparing
263    // these bytes lexicographically. Equal keys fall back to the full values,
264    // covering values that differ after 16 bytes and prefixes containing zero.
265    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            // SAFETY: both indices were produced from this array above.
277            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    // SAFETY: both indices were produced from this array above.
296    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    // Same ranking and null handling as `rank_impl`, but callbacks receive
313    // tuple references because key collisions use the index to read the full
314    // value. `rank_impl` passes copied values directly through `Fn(T, T)`.
315    // We can use an unstable sort as we combine equal values later
316    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; // Number of values in rank
332    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    // We can use an unstable sort as we combine equal values later
362    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; // Number of values in rank
378    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/// Return the index for the rank when ranking boolean array
396///
397/// The index is calculated as follows:
398/// if is_null is true, the index is 2
399/// if is_null is false and the value is true, the index is 1
400/// otherwise, the index is 0
401///
402/// false is 0 and true is 1 because these are the value when cast to number
403#[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    // Rank values for [false, true, null] in that order
416    //
417    // The value for a rank is last value rank + own value count
418    // this means that if we have the following order: `false`, `true` and then `null`
419    // the ranks will be:
420    // - false: false_count
421    // - true: false_count + true_count
422    // - null: false_count + true_count + null_count
423    //
424    // If we have the following order: `null`, `false` and then `true`
425    // the ranks will be:
426    // - false: null_count + false_count
427    // - true: null_count + false_count + true_count
428    // - null: null_count
429    //
430    // You will notice that the last rank is always the total length of the array but we don't use it for readability on how the rank is calculated
431    let ranks_index: [u32; 3] = match (options.descending, options.nulls_first) {
432        // The order is null, true, false
433        (true, true) => [
434            null_count + true_count + false_count,
435            null_count + true_count,
436            null_count,
437        ],
438        // The order is true, false, null
439        (true, false) => [
440            true_count + false_count,
441            true_count,
442            true_count + false_count + null_count,
443        ],
444        // The order is null, false, true
445        (false, true) => [
446            null_count + false_count,
447            null_count + false_count + true_count,
448            null_count,
449        ],
450        // The order is false, true, null
451        (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        // Test with non-zero null values
524        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        // Test with non-zero null values
569        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"), // exactly the 12-byte inline limit
642            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}