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    let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n| n.null_count() > 0) {
110        Some(n) => n
111            .valid_indices()
112            .map(|idx| (array.value(idx).as_ref(), idx as u32))
113            .collect(),
114        None => (0..array.len())
115            .map(|idx| (array.value(idx).as_ref(), idx as u32))
116            .collect(),
117    };
118    rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq)
119}
120
121fn rank_impl<T, C, E>(
122    len: usize,
123    mut valid: Vec<(T, u32)>,
124    options: SortOptions,
125    compare: C,
126    eq: E,
127) -> Vec<u32>
128where
129    T: Copy,
130    C: Fn(T, T) -> Ordering,
131    E: Fn(T, T) -> bool,
132{
133    // We can use an unstable sort as we combine equal values later
134    valid.sort_unstable_by(|a, b| compare(a.0, b.0));
135    if options.descending {
136        valid.reverse();
137    }
138
139    let (mut valid_rank, null_rank) = match options.nulls_first {
140        true => (len as u32, (len - valid.len()) as u32),
141        false => (valid.len() as u32, len as u32),
142    };
143
144    let mut out: Vec<_> = vec![null_rank; len];
145    if let Some(v) = valid.last() {
146        out[v.1 as usize] = valid_rank;
147    }
148
149    let mut count = 1; // Number of values in rank
150    for w in valid.windows(2).rev() {
151        match eq(w[0].0, w[1].0) {
152            true => {
153                count += 1;
154                out[w[0].1 as usize] = valid_rank;
155            }
156            false => {
157                valid_rank -= count;
158                count = 1;
159                out[w[0].1 as usize] = valid_rank
160            }
161        }
162    }
163
164    out
165}
166
167/// Return the index for the rank when ranking boolean array
168///
169/// The index is calculated as follows:
170/// if is_null is true, the index is 2
171/// if is_null is false and the value is true, the index is 1
172/// otherwise, the index is 0
173///
174/// false is 0 and true is 1 because these are the value when cast to number
175#[inline]
176fn get_boolean_rank_index(value: bool, is_null: bool) -> usize {
177    let is_null_num = is_null as usize;
178    (is_null_num << 1) | (value as usize & !is_null_num)
179}
180
181#[inline(never)]
182fn boolean_rank(array: &BooleanArray, options: SortOptions) -> Vec<u32> {
183    let null_count = array.null_count() as u32;
184    let true_count = array.true_count() as u32;
185    let false_count = array.len() as u32 - null_count - true_count;
186
187    // Rank values for [false, true, null] in that order
188    //
189    // The value for a rank is last value rank + own value count
190    // this means that if we have the following order: `false`, `true` and then `null`
191    // the ranks will be:
192    // - false: false_count
193    // - true: false_count + true_count
194    // - null: false_count + true_count + null_count
195    //
196    // If we have the following order: `null`, `false` and then `true`
197    // the ranks will be:
198    // - false: null_count + false_count
199    // - true: null_count + false_count + true_count
200    // - null: null_count
201    //
202    // 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
203    let ranks_index: [u32; 3] = match (options.descending, options.nulls_first) {
204        // The order is null, true, false
205        (true, true) => [
206            null_count + true_count + false_count,
207            null_count + true_count,
208            null_count,
209        ],
210        // The order is true, false, null
211        (true, false) => [
212            true_count + false_count,
213            true_count,
214            true_count + false_count + null_count,
215        ],
216        // The order is null, false, true
217        (false, true) => [
218            null_count + false_count,
219            null_count + false_count + true_count,
220            null_count,
221        ],
222        // The order is false, true, null
223        (false, false) => [
224            false_count,
225            false_count + true_count,
226            false_count + true_count + null_count,
227        ],
228    };
229
230    match array.nulls().filter(|n| n.null_count() > 0) {
231        Some(n) => array
232            .values()
233            .iter()
234            .zip(n.iter())
235            .map(|(value, is_valid)| ranks_index[get_boolean_rank_index(value, !is_valid)])
236            .collect::<Vec<u32>>(),
237        None => array
238            .values()
239            .iter()
240            .map(|value| ranks_index[value as usize])
241            .collect::<Vec<u32>>(),
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248    use arrow_array::*;
249
250    #[test]
251    fn test_primitive() {
252        let descending = SortOptions {
253            descending: true,
254            nulls_first: true,
255        };
256
257        let nulls_last = SortOptions {
258            descending: false,
259            nulls_first: false,
260        };
261
262        let nulls_last_descending = SortOptions {
263            descending: true,
264            nulls_first: false,
265        };
266
267        let a = Int32Array::from(vec![Some(1), Some(1), None, Some(3), Some(3), Some(4)]);
268        let res = rank(&a, None).unwrap();
269        assert_eq!(res, &[3, 3, 1, 5, 5, 6]);
270
271        let res = rank(&a, Some(descending)).unwrap();
272        assert_eq!(res, &[6, 6, 1, 4, 4, 2]);
273
274        let res = rank(&a, Some(nulls_last)).unwrap();
275        assert_eq!(res, &[2, 2, 6, 4, 4, 5]);
276
277        let res = rank(&a, Some(nulls_last_descending)).unwrap();
278        assert_eq!(res, &[5, 5, 6, 3, 3, 1]);
279
280        // Test with non-zero null values
281        let nulls = NullBuffer::from(vec![true, true, false, true, false, false]);
282        let a = Int32Array::new(vec![1, 4, 3, 4, 5, 5].into(), Some(nulls));
283        let res = rank(&a, None).unwrap();
284        assert_eq!(res, &[4, 6, 3, 6, 3, 3]);
285    }
286
287    #[test]
288    fn test_get_boolean_rank_index() {
289        assert_eq!(get_boolean_rank_index(true, true), 2);
290        assert_eq!(get_boolean_rank_index(false, true), 2);
291        assert_eq!(get_boolean_rank_index(true, false), 1);
292        assert_eq!(get_boolean_rank_index(false, false), 0);
293    }
294
295    #[test]
296    fn test_nullable_booleans() {
297        let descending = SortOptions {
298            descending: true,
299            nulls_first: true,
300        };
301
302        let nulls_last = SortOptions {
303            descending: false,
304            nulls_first: false,
305        };
306
307        let nulls_last_descending = SortOptions {
308            descending: true,
309            nulls_first: false,
310        };
311
312        let a = BooleanArray::from(vec![Some(true), Some(true), None, Some(false), Some(false)]);
313        let res = rank(&a, None).unwrap();
314        assert_eq!(res, &[5, 5, 1, 3, 3]);
315
316        let res = rank(&a, Some(descending)).unwrap();
317        assert_eq!(res, &[3, 3, 1, 5, 5]);
318
319        let res = rank(&a, Some(nulls_last)).unwrap();
320        assert_eq!(res, &[4, 4, 5, 2, 2]);
321
322        let res = rank(&a, Some(nulls_last_descending)).unwrap();
323        assert_eq!(res, &[2, 2, 5, 4, 4]);
324
325        // Test with non-zero null values
326        let nulls = NullBuffer::from(vec![true, true, false, true, true]);
327        let a = BooleanArray::new(vec![true, true, true, false, false].into(), Some(nulls));
328        let res = rank(&a, None).unwrap();
329        assert_eq!(res, &[5, 5, 1, 3, 3]);
330    }
331
332    #[test]
333    fn test_booleans() {
334        let descending = SortOptions {
335            descending: true,
336            nulls_first: true,
337        };
338
339        let nulls_last = SortOptions {
340            descending: false,
341            nulls_first: false,
342        };
343
344        let nulls_last_descending = SortOptions {
345            descending: true,
346            nulls_first: false,
347        };
348
349        let a = BooleanArray::from(vec![true, false, false, false, true]);
350        let res = rank(&a, None).unwrap();
351        assert_eq!(res, &[5, 3, 3, 3, 5]);
352
353        let res = rank(&a, Some(descending)).unwrap();
354        assert_eq!(res, &[2, 5, 5, 5, 2]);
355
356        let res = rank(&a, Some(nulls_last)).unwrap();
357        assert_eq!(res, &[5, 3, 3, 3, 5]);
358
359        let res = rank(&a, Some(nulls_last_descending)).unwrap();
360        assert_eq!(res, &[2, 5, 5, 5, 2]);
361    }
362
363    #[test]
364    fn test_bytes() {
365        let v = vec!["foo", "fo", "bar", "bar"];
366        let values = StringArray::from(v.clone());
367        let res = rank(&values, None).unwrap();
368        assert_eq!(res, &[4, 3, 2, 2]);
369
370        let values = LargeStringArray::from(v.clone());
371        let res = rank(&values, None).unwrap();
372        assert_eq!(res, &[4, 3, 2, 2]);
373
374        let values = StringViewArray::from(v);
375        let res = rank(&values, None).unwrap();
376        assert_eq!(res, &[4, 3, 2, 2]);
377
378        let v: Vec<&[u8]> = vec![&[1, 2], &[0], &[1, 2, 3], &[1, 2]];
379        let values = LargeBinaryArray::from(v.clone());
380        let res = rank(&values, None).unwrap();
381        assert_eq!(res, &[3, 1, 4, 3]);
382
383        let values = BinaryArray::from(v.clone());
384        let res = rank(&values, None).unwrap();
385        assert_eq!(res, &[3, 1, 4, 3]);
386
387        let values = BinaryViewArray::from_iter_values(v);
388        let res = rank(&values, None).unwrap();
389        assert_eq!(res, &[3, 1, 4, 3]);
390    }
391
392    #[test]
393    fn test_string_view_with_nulls() {
394        let values = StringViewArray::from(vec![
395            Some("a string longer than twelve bytes"),
396            Some("bar"),
397            None,
398            Some("a string longer than twelve bytes"),
399        ]);
400        let res = rank(&values, None).unwrap();
401        assert_eq!(res, &[3, 4, 1, 3]);
402    }
403
404    #[test]
405    fn test_binary_view_with_nulls() {
406        let long_value = b"a binary value longer than twelve bytes".as_ref();
407        let values = BinaryViewArray::from_iter([
408            Some(long_value),
409            Some(b"bar".as_ref()),
410            None,
411            Some(long_value),
412        ]);
413        let res = rank(&values, None).unwrap();
414        assert_eq!(res, &[3, 4, 1, 3]);
415    }
416}
arrow_ord/rank.rs

arrow_ord/
rank.rs
