arrow-ord 54.3.1

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Provides `rank` function to assign a rank to each value in an array

use arrow_array::cast::AsArray;
use arrow_array::types::*;
use arrow_array::{
    downcast_primitive_array, Array, ArrowNativeTypeOp, BooleanArray, GenericByteArray,
};
use arrow_buffer::NullBuffer;
use arrow_schema::{ArrowError, DataType, SortOptions};
use std::cmp::Ordering;

/// Whether `arrow_ord::rank` can rank an array of given data type.
pub(crate) fn can_rank(data_type: &DataType) -> bool {
    data_type.is_primitive()
        || matches!(
            data_type,
            DataType::Boolean
                | DataType::Utf8
                | DataType::LargeUtf8
                | DataType::Binary
                | DataType::LargeBinary
        )
}

/// Assigns a rank to each value in `array` based on its position in the sorted order
///
/// Where values are equal, they will be assigned the highest of their ranks,
/// leaving gaps in the overall rank assignment
///
/// ```
/// # use arrow_array::StringArray;
/// # use arrow_ord::rank::rank;
/// let array = StringArray::from(vec![Some("foo"), None, Some("foo"), None, Some("bar")]);
/// let ranks = rank(&array, None).unwrap();
/// assert_eq!(ranks, &[5, 2, 5, 2, 3]);
/// ```
pub fn rank(array: &dyn Array, options: Option<SortOptions>) -> Result<Vec<u32>, ArrowError> {
    let options = options.unwrap_or_default();
    let ranks = downcast_primitive_array! {
        array => primitive_rank(array.values(), array.nulls(), options),
        DataType::Boolean => boolean_rank(array.as_boolean(), options),
        DataType::Utf8 => bytes_rank(array.as_bytes::<Utf8Type>(), options),
        DataType::LargeUtf8 => bytes_rank(array.as_bytes::<LargeUtf8Type>(), options),
        DataType::Binary => bytes_rank(array.as_bytes::<BinaryType>(), options),
        DataType::LargeBinary => bytes_rank(array.as_bytes::<LargeBinaryType>(), options),
        d => return Err(ArrowError::ComputeError(format!("{d:?} not supported in rank")))
    };
    Ok(ranks)
}

#[inline(never)]
fn primitive_rank<T: ArrowNativeTypeOp>(
    values: &[T],
    nulls: Option<&NullBuffer>,
    options: SortOptions,
) -> Vec<u32> {
    let len: u32 = values.len().try_into().unwrap();
    let to_sort = match nulls.filter(|n| n.null_count() > 0) {
        Some(n) => n
            .valid_indices()
            .map(|idx| (values[idx], idx as u32))
            .collect(),
        None => values.iter().copied().zip(0..len).collect(),
    };
    rank_impl(values.len(), to_sort, options, T::compare, T::is_eq)
}

#[inline(never)]
fn bytes_rank<T: ByteArrayType>(array: &GenericByteArray<T>, options: SortOptions) -> Vec<u32> {
    let to_sort: Vec<(&[u8], u32)> = match array.nulls().filter(|n| n.null_count() > 0) {
        Some(n) => n
            .valid_indices()
            .map(|idx| (array.value(idx).as_ref(), idx as u32))
            .collect(),
        None => (0..array.len())
            .map(|idx| (array.value(idx).as_ref(), idx as u32))
            .collect(),
    };
    rank_impl(array.len(), to_sort, options, Ord::cmp, PartialEq::eq)
}

fn rank_impl<T, C, E>(
    len: usize,
    mut valid: Vec<(T, u32)>,
    options: SortOptions,
    compare: C,
    eq: E,
) -> Vec<u32>
where
    T: Copy,
    C: Fn(T, T) -> Ordering,
    E: Fn(T, T) -> bool,
{
    // We can use an unstable sort as we combine equal values later
    valid.sort_unstable_by(|a, b| compare(a.0, b.0));
    if options.descending {
        valid.reverse();
    }

    let (mut valid_rank, null_rank) = match options.nulls_first {
        true => (len as u32, (len - valid.len()) as u32),
        false => (valid.len() as u32, len as u32),
    };

    let mut out: Vec<_> = vec![null_rank; len];
    if let Some(v) = valid.last() {
        out[v.1 as usize] = valid_rank;
    }

    let mut count = 1; // Number of values in rank
    for w in valid.windows(2).rev() {
        match eq(w[0].0, w[1].0) {
            true => {
                count += 1;
                out[w[0].1 as usize] = valid_rank;
            }
            false => {
                valid_rank -= count;
                count = 1;
                out[w[0].1 as usize] = valid_rank
            }
        }
    }

    out
}

/// Return the index for the rank when ranking boolean array
///
/// The index is calculated as follows:
/// if is_null is true, the index is 2
/// if is_null is false and the value is true, the index is 1
/// otherwise, the index is 0
///
/// false is 0 and true is 1 because these are the value when cast to number
#[inline]
fn get_boolean_rank_index(value: bool, is_null: bool) -> usize {
    let is_null_num = is_null as usize;
    (is_null_num << 1) | (value as usize & !is_null_num)
}

#[inline(never)]
fn boolean_rank(array: &BooleanArray, options: SortOptions) -> Vec<u32> {
    let null_count = array.null_count() as u32;
    let true_count = array.true_count() as u32;
    let false_count = array.len() as u32 - null_count - true_count;

    // Rank values for [false, true, null] in that order
    //
    // The value for a rank is last value rank + own value count
    // this means that if we have the following order: `false`, `true` and then `null`
    // the ranks will be:
    // - false: false_count
    // - true: false_count + true_count
    // - null: false_count + true_count + null_count
    //
    // If we have the following order: `null`, `false` and then `true`
    // the ranks will be:
    // - false: null_count + false_count
    // - true: null_count + false_count + true_count
    // - null: null_count
    //
    // 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
    let ranks_index: [u32; 3] = match (options.descending, options.nulls_first) {
        // The order is null, true, false
        (true, true) => [
            null_count + true_count + false_count,
            null_count + true_count,
            null_count,
        ],
        // The order is true, false, null
        (true, false) => [
            true_count + false_count,
            true_count,
            true_count + false_count + null_count,
        ],
        // The order is null, false, true
        (false, true) => [
            null_count + false_count,
            null_count + false_count + true_count,
            null_count,
        ],
        // The order is false, true, null
        (false, false) => [
            false_count,
            false_count + true_count,
            false_count + true_count + null_count,
        ],
    };

    match array.nulls().filter(|n| n.null_count() > 0) {
        Some(n) => array
            .values()
            .iter()
            .zip(n.iter())
            .map(|(value, is_valid)| ranks_index[get_boolean_rank_index(value, !is_valid)])
            .collect::<Vec<u32>>(),
        None => array
            .values()
            .iter()
            .map(|value| ranks_index[value as usize])
            .collect::<Vec<u32>>(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::*;

    #[test]
    fn test_primitive() {
        let descending = SortOptions {
            descending: true,
            nulls_first: true,
        };

        let nulls_last = SortOptions {
            descending: false,
            nulls_first: false,
        };

        let nulls_last_descending = SortOptions {
            descending: true,
            nulls_first: false,
        };

        let a = Int32Array::from(vec![Some(1), Some(1), None, Some(3), Some(3), Some(4)]);
        let res = rank(&a, None).unwrap();
        assert_eq!(res, &[3, 3, 1, 5, 5, 6]);

        let res = rank(&a, Some(descending)).unwrap();
        assert_eq!(res, &[6, 6, 1, 4, 4, 2]);

        let res = rank(&a, Some(nulls_last)).unwrap();
        assert_eq!(res, &[2, 2, 6, 4, 4, 5]);

        let res = rank(&a, Some(nulls_last_descending)).unwrap();
        assert_eq!(res, &[5, 5, 6, 3, 3, 1]);

        // Test with non-zero null values
        let nulls = NullBuffer::from(vec![true, true, false, true, false, false]);
        let a = Int32Array::new(vec![1, 4, 3, 4, 5, 5].into(), Some(nulls));
        let res = rank(&a, None).unwrap();
        assert_eq!(res, &[4, 6, 3, 6, 3, 3]);
    }

    #[test]
    fn test_get_boolean_rank_index() {
        assert_eq!(get_boolean_rank_index(true, true), 2);
        assert_eq!(get_boolean_rank_index(false, true), 2);
        assert_eq!(get_boolean_rank_index(true, false), 1);
        assert_eq!(get_boolean_rank_index(false, false), 0);
    }

    #[test]
    fn test_nullable_booleans() {
        let descending = SortOptions {
            descending: true,
            nulls_first: true,
        };

        let nulls_last = SortOptions {
            descending: false,
            nulls_first: false,
        };

        let nulls_last_descending = SortOptions {
            descending: true,
            nulls_first: false,
        };

        let a = BooleanArray::from(vec![Some(true), Some(true), None, Some(false), Some(false)]);
        let res = rank(&a, None).unwrap();
        assert_eq!(res, &[5, 5, 1, 3, 3]);

        let res = rank(&a, Some(descending)).unwrap();
        assert_eq!(res, &[3, 3, 1, 5, 5]);

        let res = rank(&a, Some(nulls_last)).unwrap();
        assert_eq!(res, &[4, 4, 5, 2, 2]);

        let res = rank(&a, Some(nulls_last_descending)).unwrap();
        assert_eq!(res, &[2, 2, 5, 4, 4]);

        // Test with non-zero null values
        let nulls = NullBuffer::from(vec![true, true, false, true, true]);
        let a = BooleanArray::new(vec![true, true, true, false, false].into(), Some(nulls));
        let res = rank(&a, None).unwrap();
        assert_eq!(res, &[5, 5, 1, 3, 3]);
    }

    #[test]
    fn test_booleans() {
        let descending = SortOptions {
            descending: true,
            nulls_first: true,
        };

        let nulls_last = SortOptions {
            descending: false,
            nulls_first: false,
        };

        let nulls_last_descending = SortOptions {
            descending: true,
            nulls_first: false,
        };

        let a = BooleanArray::from(vec![true, false, false, false, true]);
        let res = rank(&a, None).unwrap();
        assert_eq!(res, &[5, 3, 3, 3, 5]);

        let res = rank(&a, Some(descending)).unwrap();
        assert_eq!(res, &[2, 5, 5, 5, 2]);

        let res = rank(&a, Some(nulls_last)).unwrap();
        assert_eq!(res, &[5, 3, 3, 3, 5]);

        let res = rank(&a, Some(nulls_last_descending)).unwrap();
        assert_eq!(res, &[2, 5, 5, 5, 2]);
    }

    #[test]
    fn test_bytes() {
        let v = vec!["foo", "fo", "bar", "bar"];
        let values = StringArray::from(v.clone());
        let res = rank(&values, None).unwrap();
        assert_eq!(res, &[4, 3, 2, 2]);

        let values = LargeStringArray::from(v.clone());
        let res = rank(&values, None).unwrap();
        assert_eq!(res, &[4, 3, 2, 2]);

        let v: Vec<&[u8]> = vec![&[1, 2], &[0], &[1, 2, 3], &[1, 2]];
        let values = LargeBinaryArray::from(v.clone());
        let res = rank(&values, None).unwrap();
        assert_eq!(res, &[3, 1, 4, 3]);

        let values = BinaryArray::from(v);
        let res = rank(&values, None).unwrap();
        assert_eq!(res, &[3, 1, 4, 3]);
    }
}