Skip to main content

arrow_ord/
rank.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Provides `rank` function to assign a rank to each value in an array
19
20use 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
30/// Whether `arrow_ord::rank` can rank an array of given data type.
31pub(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
45/// Assigns a rank to each value in `array` based on its position in the sorted order
46///
47/// Where values are equal, they will be assigned the highest of their ranks,
48/// leaving gaps in the overall rank assignment
49///
50/// ```
51/// # use arrow_array::StringArray;
52/// # use arrow_ord::rank::rank;
53/// let array = StringArray::from(vec![Some("foo"), None, Some("foo"), None, Some("bar")]);
54/// let ranks = rank(&array, None).unwrap();
55/// assert_eq!(ranks, &[5, 2, 5, 2, 3]);
56/// ```
57pub 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    // An inline view already contains the complete value. Convert it once to
110    // a key whose integer ordering matches the byte ordering, as is done by
111    // `sort_byte_view`.
112    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                    // SAFETY: `valid_indices` only yields indices in the array.
118                    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    // Cache a wider prefix than the 4 bytes stored in a non-inline view. This
152    // pays for the backing-buffer access once per value instead of once per
153    // comparison, and only resolves the complete value when two keys collide.
154    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                // SAFETY: `valid_indices` only yields indices in the array.
159                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                // SAFETY: `idx` is in `0..array.len()`.
166                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
180// A 16-byte prefix fits in one `u128` and is wider than the 4-byte view prefix.
181// Shorter keys collide more often; longer keys need another representation.
182const BYTE_VIEW_KEY_LEN: usize = 16;
183
184// Four valid values per window keeps sampling cheap while allowing an early
185// all-collision decision. More samples improve confidence but add buffer reads.
186const BYTE_VIEW_KEY_SAMPLES_PER_WINDOW: usize = 4;
187
188// Capacity for the two windows; it must be at least
189// `2 * BYTE_VIEW_KEY_SAMPLES_PER_WINDOW`. Increasing it alone has no effect;
190// decreasing it without changing the sampling count can overflow the array.
191const BYTE_VIEW_KEY_SAMPLE_SIZE: usize = 8;
192
193// Bound entries inspected when nulls are present. A higher limit finds valid
194// samples more reliably but costs reads; a lower limit is cheaper but less
195// informative for null-heavy arrays.
196const BYTE_VIEW_KEY_MAX_PROBES_PER_WINDOW: usize = 32;
197
198// `colliding_keys * 3 >= sample_len` means roughly one third of sampled keys
199// are duplicates. Lower values fall back earlier; higher values risk keeping
200// the key path for collision-heavy inputs.
201const BYTE_VIEW_KEY_FALLBACK_COLLISION_RATIO: usize = 3;
202
203/// Estimates whether cached byte-view keys collide often enough to make the
204/// key-based ranking path unattractive.
205/// Caching a wider key usually avoids repeated backing-buffer reads for long
206/// views. If sampled keys collide frequently, resolving full values plus the
207/// extra key comparison can be slower than comparing slices directly, so the
208/// caller falls back to that path. The bounded two-window sample keeps this
209/// check inexpensive.
210fn 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    // Probe small local windows in both halves. Keeping each probe local avoids
220    // turning collision detection itself into scattered backing-buffer reads.
221    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            // SAFETY: `idx` is within a window bounded by `array.len()`.
232            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        // Four equal keys already contribute three collisions. Even if every
242        // sample in the other window is distinct, that satisfies the final
243        // one-third threshold, so avoid touching the second backing-buffer
244        // window in the common all-collision case.
245        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    // If at least roughly one third of sampled keys collide, comparing the
263    // wider key before every full-value comparison is likely more expensive
264    // than sorting slices directly.
265    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    // Big-endian conversion makes integer comparison equivalent to comparing
276    // these bytes lexicographically. Equal keys fall back to the full values,
277    // covering values that differ after 16 bytes and prefixes containing zero.
278    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            // SAFETY: both indices were produced from this array above.
290            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    // SAFETY: both indices were produced from this array above.
309    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    // Same ranking and null handling as `rank_impl`, but callbacks receive
326    // tuple references because key collisions use the index to read the full
327    // value. `rank_impl` passes copied values directly through `Fn(T, T)`.
328    // We can use an unstable sort as we combine equal values later
329    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; // Number of values in rank
345    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    // We can use an unstable sort as we combine equal values later
375    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; // Number of values in rank
391    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/// Return the index for the rank when ranking boolean array
409///
410/// The index is calculated as follows:
411/// if is_null is true, the index is 2
412/// if is_null is false and the value is true, the index is 1
413/// otherwise, the index is 0
414///
415/// false is 0 and true is 1 because these are the value when cast to number
416#[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    // Rank values for [false, true, null] in that order
429    //
430    // The value for a rank is last value rank + own value count
431    // this means that if we have the following order: `false`, `true` and then `null`
432    // the ranks will be:
433    // - false: false_count
434    // - true: false_count + true_count
435    // - null: false_count + true_count + null_count
436    //
437    // If we have the following order: `null`, `false` and then `true`
438    // the ranks will be:
439    // - false: null_count + false_count
440    // - true: null_count + false_count + true_count
441    // - null: null_count
442    //
443    // 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
444    let ranks_index: [u32; 3] = match (options.descending, options.nulls_first) {
445        // The order is null, true, false
446        (true, true) => [
447            null_count + true_count + false_count,
448            null_count + true_count,
449            null_count,
450        ],
451        // The order is true, false, null
452        (true, false) => [
453            true_count + false_count,
454            true_count,
455            true_count + false_count + null_count,
456        ],
457        // The order is null, false, true
458        (false, true) => [
459            null_count + false_count,
460            null_count + false_count + true_count,
461            null_count,
462        ],
463        // The order is false, true, null
464        (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        // Test with non-zero null values
537        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        // Test with non-zero null values
582        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"), // exactly the 12-byte inline limit
655            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}