// 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.

//! Defines sort kernel for `ArrayRef`

use crate::ord::{make_comparator, DynComparator};

use arrow_array::builder::BufferBuilder;

use arrow_array::cast::*;

use arrow_array::types::*;

use arrow_array::*;

use arrow_buffer::ArrowNativeType;

use arrow_buffer::BooleanBufferBuilder;

use arrow_data::{ArrayDataBuilder, ByteView, MAX_INLINE_VIEW_LEN};

use arrow_schema::{ArrowError, DataType};

use arrow_select::take::take;

use std::cmp::Ordering;

use std::sync::Arc;

use crate::rank::{can_rank, rank};

pub use arrow_schema::SortOptions;

/// Sort the `ArrayRef` using `SortOptions`.

///

/// Performs a sort on values and indices. Nulls are ordered according

/// to the `nulls_first` flag in `options`.  Floats are sorted using

/// IEEE 754 totalOrder

///

/// Returns an `ArrowError::ComputeError(String)` if the array type is

/// either unsupported by `sort_to_indices` or `take`.

///

/// Note: this is an unstable_sort, meaning it may not preserve the

/// order of equal elements.

///

/// # Example

/// ```rust

/// # use std::sync::Arc;

/// # use arrow_array::Int32Array;

/// # use arrow_ord::sort::sort;

/// let array = Int32Array::from(vec![5, 4, 3, 2, 1]);

/// let sorted_array = sort(&array, None).unwrap();

/// assert_eq!(sorted_array.as_ref(), &Int32Array::from(vec![1, 2, 3, 4, 5]));

/// ```

pub fn sort(values: &dyn Array, options: Option<SortOptions>) -> Result<ArrayRef, ArrowError> {

    downcast_primitive_array!(

        values => sort_native_type(values, options),

        DataType::RunEndEncoded(_, _) => sort_run(values, options, None),

        _ => {

            let indices = sort_to_indices(values, options, None)?;

            take(values, &indices, None)

        }

    )

}

fn sort_native_type<T>(

    primitive_values: &PrimitiveArray<T>,

    options: Option<SortOptions>,

) -> Result<ArrayRef, ArrowError>

where

    T: ArrowPrimitiveType,

{

    let sort_options = options.unwrap_or_default();

    let mut mutable_buffer = vec![T::default_value(); primitive_values.len()];

    let mutable_slice = &mut mutable_buffer;

    let input_values = primitive_values.values().as_ref();

    let nulls_count = primitive_values.null_count();

    let valid_count = primitive_values.len() - nulls_count;

    let null_bit_buffer = match nulls_count > 0 {

        true => {

            let mut validity_buffer = BooleanBufferBuilder::new(primitive_values.len());

            if sort_options.nulls_first {

                validity_buffer.append_n(nulls_count, false);

                validity_buffer.append_n(valid_count, true);

            } else {

                validity_buffer.append_n(valid_count, true);

                validity_buffer.append_n(nulls_count, false);

            }

            Some(validity_buffer.finish().into())

        }

        false => None,

    };

    if let Some(nulls) = primitive_values.nulls().filter(|n| n.null_count() > 0) {

        let values_slice = match sort_options.nulls_first {

            true => &mut mutable_slice[nulls_count..],

            false => &mut mutable_slice[..valid_count],

        };

        for (write_index, index) in nulls.valid_indices().enumerate() {

            values_slice[write_index] = primitive_values.value(index);

        }

        values_slice.sort_unstable_by(|a, b| a.compare(*b));

        if sort_options.descending {

            values_slice.reverse();

        }

    } else {

        mutable_slice.copy_from_slice(input_values);

        mutable_slice.sort_unstable_by(|a, b| a.compare(*b));

        if sort_options.descending {

            mutable_slice.reverse();

        }

    }

    Ok(Arc::new(

        PrimitiveArray::<T>::new(mutable_buffer.into(), null_bit_buffer)

            .with_data_type(primitive_values.data_type().clone()),

    ))

}

/// Sort the `ArrayRef` partially.

///

/// If `limit` is specified, the resulting array will contain only

/// first `limit` in the sort order. Any data data after the limit

/// will be discarded.

///

/// Note: this is an unstable_sort, meaning it may not preserve the

/// order of equal elements.

///

/// # Example

/// ```rust

/// # use std::sync::Arc;

/// # use arrow_array::Int32Array;

/// # use arrow_ord::sort::{sort_limit, SortOptions};

/// let array = Int32Array::from(vec![5, 4, 3, 2, 1]);

///

/// // Find the the top 2 items

/// let sorted_array = sort_limit(&array, None, Some(2)).unwrap();

/// assert_eq!(sorted_array.as_ref(), &Int32Array::from(vec![1, 2]));

///

/// // Find the bottom top 2 items

/// let options = Some(SortOptions {

///                  descending: true,

///                  ..Default::default()

///               });

/// let sorted_array = sort_limit(&array, options, Some(2)).unwrap();

/// assert_eq!(sorted_array.as_ref(), &Int32Array::from(vec![5, 4]));

/// ```

pub fn sort_limit(

    values: &dyn Array,

    options: Option<SortOptions>,

    limit: Option<usize>,

) -> Result<ArrayRef, ArrowError> {

    if let DataType::RunEndEncoded(_, _) = values.data_type() {

        return sort_run(values, options, limit);

    }

    let indices = sort_to_indices(values, options, limit)?;

    take(values, &indices, None)

}

/// we can only do this if the T is primitive

#[inline]

fn sort_unstable_by<T, F>(array: &mut [T], limit: usize, cmp: F)

where

    F: FnMut(&T, &T) -> Ordering,

{

    if array.len() == limit {

        array.sort_unstable_by(cmp);

    } else {

        partial_sort(array, limit, cmp);

    }

}

/// Partition indices of an Arrow array into two categories:

/// - `valid`: indices of non-null elements

/// - `nulls`: indices of null elements

///

/// Optimized for performance with fast-path for all-valid arrays

/// and bit-parallel scan for null-containing arrays.

#[inline(always)]

pub fn partition_validity(array: &dyn Array) -> (Vec<u32>, Vec<u32>) {

    let len = array.len();

    let null_count = array.null_count();

    // Fast path: if there are no nulls, all elements are valid

    if null_count == 0 {

        // Simply return a range of indices [0, len)

        let valid = (0..len as u32).collect();

        return (valid, Vec::new());

    }

    // null bitmap exists and some values are null

    partition_validity_scan(array, len, null_count)

}

/// Scans the null bitmap and partitions valid/null indices efficiently.

/// Uses bit-level operations to extract bit positions.

/// This function is only called when nulls exist.

#[inline(always)]

fn partition_validity_scan(

    array: &dyn Array,

    len: usize,

    null_count: usize,

) -> (Vec<u32>, Vec<u32>) {

    // SAFETY: Guaranteed by caller that null_count > 0, so bitmap must exist

    let bitmap = array.nulls().unwrap();

    // Preallocate result vectors with exact capacities (avoids reallocations)

    let mut valid = Vec::with_capacity(len - null_count);

    let mut nulls = Vec::with_capacity(null_count);

    unsafe {

        // 1) Write valid indices (bits == 1)

        let valid_slice = valid.spare_capacity_mut();

        for (i, idx) in bitmap.inner().set_indices_u32().enumerate() {

            valid_slice[i].write(idx);

        }

        // 2) Write null indices by inverting

        let inv_buf = !bitmap.inner();

        let null_slice = nulls.spare_capacity_mut();

        for (i, idx) in inv_buf.set_indices_u32().enumerate() {

            null_slice[i].write(idx);

        }

        // Finalize lengths

        valid.set_len(len - null_count);

        nulls.set_len(null_count);

    }

    assert_eq!(valid.len(), len - null_count);

    assert_eq!(nulls.len(), null_count);

    (valid, nulls)

}

/// Whether `sort_to_indices` can sort an array of given data type.

fn can_sort_to_indices(data_type: &DataType) -> bool {

    data_type.is_primitive()

        || matches!(

            data_type,

            DataType::Boolean

                | DataType::Utf8

                | DataType::LargeUtf8

                | DataType::Utf8View

                | DataType::Binary

                | DataType::LargeBinary

                | DataType::BinaryView

                | DataType::FixedSizeBinary(_)

        )

        || match data_type {

            DataType::List(f) if can_rank(f.data_type()) => true,

            DataType::LargeList(f) if can_rank(f.data_type()) => true,

            DataType::FixedSizeList(f, _) if can_rank(f.data_type()) => true,

            DataType::Dictionary(_, values) if can_rank(values.as_ref()) => true,

            DataType::RunEndEncoded(_, f) if can_sort_to_indices(f.data_type()) => true,

            _ => false,

        }

}

/// Sort elements from `ArrayRef` into an unsigned integer (`UInt32Array`) of indices.

/// Floats are sorted using IEEE 754 totalOrder.  `limit` is an option for [partial_sort].

pub fn sort_to_indices(

    array: &dyn Array,

    options: Option<SortOptions>,

    limit: Option<usize>,

) -> Result<UInt32Array, ArrowError> {

    let options = options.unwrap_or_default();

    let (v, n) = partition_validity(array);

    Ok(downcast_primitive_array! {

        array => sort_primitive(array, v, n, options, limit),

        DataType::Boolean => sort_boolean(array.as_boolean(), v, n, options, limit),

        DataType::Utf8 => sort_bytes(array.as_string::<i32>(), v, n, options, limit),

        DataType::LargeUtf8 => sort_bytes(array.as_string::<i64>(), v, n, options, limit),

        DataType::Utf8View => sort_byte_view(array.as_string_view(), v, n, options, limit),

        DataType::Binary => sort_bytes(array.as_binary::<i32>(), v, n, options, limit),

        DataType::LargeBinary => sort_bytes(array.as_binary::<i64>(), v, n, options, limit),

        DataType::BinaryView => sort_byte_view(array.as_binary_view(), v, n, options, limit),

        DataType::FixedSizeBinary(_) => sort_fixed_size_binary(array.as_fixed_size_binary(), v, n, options, limit),

        DataType::List(_) => sort_list(array.as_list::<i32>(), v, n, options, limit)?,

        DataType::LargeList(_) => sort_list(array.as_list::<i64>(), v, n, options, limit)?,

        DataType::FixedSizeList(_, _) => sort_fixed_size_list(array.as_fixed_size_list(), v, n, options, limit)?,

        DataType::Dictionary(_, _) => downcast_dictionary_array!{

            array => sort_dictionary(array, v, n, options, limit)?,

            _ => unreachable!()

        }

        DataType::RunEndEncoded(run_ends_field, _) => match run_ends_field.data_type() {

            DataType::Int16 => sort_run_to_indices::<Int16Type>(array, options, limit),

            DataType::Int32 => sort_run_to_indices::<Int32Type>(array, options, limit),

            DataType::Int64 => sort_run_to_indices::<Int64Type>(array, options, limit),

            dt => {

                return Err(ArrowError::ComputeError(format!(

                    "Invalid run end data type: {dt}"

                )))

            }

        },

        t => {

            return Err(ArrowError::ComputeError(format!(

                "Sort not supported for data type {t:?}"

            )));

        }

    })

}

fn sort_boolean(

    values: &BooleanArray,

    value_indices: Vec<u32>,

    null_indices: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> UInt32Array {

    let mut valids = value_indices

        .into_iter()

        .map(|index| (index, values.value(index as usize)))

        .collect::<Vec<(u32, bool)>>();

    sort_impl(options, &mut valids, &null_indices, limit, |a, b| a.cmp(&b)).into()

}

fn sort_primitive<T: ArrowPrimitiveType>(

    values: &PrimitiveArray<T>,

    value_indices: Vec<u32>,

    nulls: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> UInt32Array {

    let mut valids = value_indices

        .into_iter()

        .map(|index| (index, values.value(index as usize)))

        .collect::<Vec<(u32, T::Native)>>();

    sort_impl(options, &mut valids, &nulls, limit, T::Native::compare).into()

}

fn sort_bytes<T: ByteArrayType>(

    values: &GenericByteArray<T>,

    value_indices: Vec<u32>,

    nulls: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> UInt32Array {

    // Note: Why do we use 4‑byte prefix?

    // Compute the 4‑byte prefix in BE order, or left‑pad if shorter.

    // Most byte‐sequences differ in their first few bytes, so by

    // comparing up to 4 bytes as a single u32 we avoid the overhead

    // of a full lexicographical compare for the vast majority of cases.

    // 1. Build a vector of (index, prefix, length) tuples

    let mut valids: Vec<(u32, u32, u64)> = value_indices

        .into_iter()

        .map(|idx| unsafe {

            let slice: &[u8] = values.value_unchecked(idx as usize).as_ref();

            let len = slice.len() as u64;

            // Compute the 4‑byte prefix in BE order, or left‑pad if shorter

            let prefix = if slice.len() >= 4 {

                let raw = std::ptr::read_unaligned(slice.as_ptr() as *const u32);

                u32::from_be(raw)

            } else if slice.is_empty() {

                // Handle empty slice case to avoid shift overflow

                0u32

            } else {

                let mut v = 0u32;

                for &b in slice {

                    v = (v << 8) | (b as u32);

                }

                // Safe shift: slice.len() is in range [1, 3], so shift is in range [8, 24]

                v << (8 * (4 - slice.len()))

            };

            (idx, prefix, len)

        })

        .collect();

    // 2. compute the number of non-null entries to partially sort

    let vlimit = match (limit, options.nulls_first) {

        (Some(l), true) => l.saturating_sub(nulls.len()).min(valids.len()),

        _ => valids.len(),

    };

    // 3. Comparator: compare prefix, then (when both slices shorter than 4) length, otherwise full slice

    let cmp_bytes = |a: &(u32, u32, u64), b: &(u32, u32, u64)| unsafe {

        let (ia, pa, la) = *a;

        let (ib, pb, lb) = *b;

        // 3.1 prefix (first 4 bytes)

        let ord = pa.cmp(&pb);

        if ord != Ordering::Equal {

            return ord;

        }

        // 3.2 only if both slices had length < 4 (so prefix was padded)

        if la < 4 || lb < 4 {

            let ord = la.cmp(&lb);

            if ord != Ordering::Equal {

                return ord;

            }

        }

        // 3.3 full lexicographical compare

        let a_bytes: &[u8] = values.value_unchecked(ia as usize).as_ref();

        let b_bytes: &[u8] = values.value_unchecked(ib as usize).as_ref();

        a_bytes.cmp(b_bytes)

    };

    // 4. Partially sort according to ascending/descending

    if !options.descending {

        sort_unstable_by(&mut valids, vlimit, cmp_bytes);

    } else {

        sort_unstable_by(&mut valids, vlimit, |x, y| cmp_bytes(x, y).reverse());

    }

    // 5. Assemble nulls and sorted indices into final output

    let total = valids.len() + nulls.len();

    let out_limit = limit.unwrap_or(total).min(total);

    let mut out = Vec::with_capacity(out_limit);

    if options.nulls_first {

        out.extend_from_slice(&nulls[..nulls.len().min(out_limit)]);

        let rem = out_limit - out.len();

        out.extend(valids.iter().map(|&(i, _, _)| i).take(rem));

    } else {

        out.extend(valids.iter().map(|&(i, _, _)| i).take(out_limit));

        let rem = out_limit - out.len();

        out.extend_from_slice(&nulls[..rem]);

    }

    out.into()

}

fn sort_byte_view<T: ByteViewType>(

    values: &GenericByteViewArray<T>,

    value_indices: Vec<u32>,

    nulls: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> UInt32Array {

    // 1. Build a list of (index, raw_view, length)

    let mut valids: Vec<_>;

    // 2. Compute the number of non-null entries to partially sort

    let vlimit: usize = match (limit, options.nulls_first) {

        (Some(l), true) => l.saturating_sub(nulls.len()).min(value_indices.len()),

        _ => value_indices.len(),

    };

    // 3.a Check if all views are inline (no data buffers)

    if values.data_buffers().is_empty() {

        valids = value_indices

            .into_iter()

            .map(|idx| {

                // SAFETY: we know idx < values.len()

                let raw = unsafe { *values.views().get_unchecked(idx as usize) };

                let inline_key = GenericByteViewArray::<T>::inline_key_fast(raw);

                (idx, inline_key)

            })

            .collect();

        let cmp_inline = |a: &(u32, u128), b: &(u32, u128)| a.1.cmp(&b.1);

        // Partially sort according to ascending/descending

        if !options.descending {

            sort_unstable_by(&mut valids, vlimit, cmp_inline);

        } else {

            sort_unstable_by(&mut valids, vlimit, |x, y| cmp_inline(x, y).reverse());

        }

    } else {

        valids = value_indices

            .into_iter()

            .map(|idx| {

                // SAFETY: we know idx < values.len()

                let raw = unsafe { *values.views().get_unchecked(idx as usize) };

                (idx, raw)

            })

            .collect();

        // 3.b Mixed comparator: first prefix, then inline vs full comparison

        let cmp_mixed = |a: &(u32, u128), b: &(u32, u128)| {

            let (_, raw_a) = *a;

            let (_, raw_b) = *b;

            let len_a = raw_a as u32;

            let len_b = raw_b as u32;

            // 3.b.1 Both inline (≤12 bytes): compare full 128-bit key including length

            if len_a <= MAX_INLINE_VIEW_LEN && len_b <= MAX_INLINE_VIEW_LEN {

                return GenericByteViewArray::<T>::inline_key_fast(raw_a)

                    .cmp(&GenericByteViewArray::<T>::inline_key_fast(raw_b));

            }

            // 3.b.2 Compare 4-byte prefix in big-endian order

            let pref_a = ByteView::from(raw_a).prefix.swap_bytes();

            let pref_b = ByteView::from(raw_b).prefix.swap_bytes();

            if pref_a != pref_b {

                return pref_a.cmp(&pref_b);

            }

            // 3.b.3 Fallback to full byte-slice comparison

            let full_a: &[u8] = unsafe { values.value_unchecked(a.0 as usize).as_ref() };

            let full_b: &[u8] = unsafe { values.value_unchecked(b.0 as usize).as_ref() };

            full_a.cmp(full_b)

        };

        // 3.b.4 Partially sort according to ascending/descending

        if !options.descending {

            sort_unstable_by(&mut valids, vlimit, cmp_mixed);

        } else {

            sort_unstable_by(&mut valids, vlimit, |x, y| cmp_mixed(x, y).reverse());

        }

    }

    // 5. Assemble nulls and sorted indices into final output

    let total = valids.len() + nulls.len();

    let out_limit = limit.unwrap_or(total).min(total);

    let mut out = Vec::with_capacity(total);

    if options.nulls_first {

        // Place null indices first

        out.extend_from_slice(&nulls[..nulls.len().min(out_limit)]);

        let rem = out_limit - out.len();

        out.extend(valids.iter().map(|&(i, _)| i).take(rem));

    } else {

        // Place non-null indices first

        out.extend(valids.iter().map(|&(i, _)| i).take(out_limit));

        let rem = out_limit - out.len();

        out.extend_from_slice(&nulls[..rem]);

    }

    out.into()

}

fn sort_fixed_size_binary(

    values: &FixedSizeBinaryArray,

    value_indices: Vec<u32>,

    nulls: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> UInt32Array {

    let mut valids = value_indices

        .iter()

        .copied()

        .map(|index| (index, values.value(index as usize)))

        .collect::<Vec<(u32, &[u8])>>();

    sort_impl(options, &mut valids, &nulls, limit, Ord::cmp).into()

}

fn sort_dictionary<K: ArrowDictionaryKeyType>(

    dict: &DictionaryArray<K>,

    value_indices: Vec<u32>,

    null_indices: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> Result<UInt32Array, ArrowError> {

    let keys: &PrimitiveArray<K> = dict.keys();

    let rank = child_rank(dict.values().as_ref(), options)?;

    // create tuples that are used for sorting

    let mut valids = value_indices

        .into_iter()

        .map(|index| {

            let key: K::Native = keys.value(index as usize);

            (index, rank[key.as_usize()])

        })

        .collect::<Vec<(u32, u32)>>();

    Ok(sort_impl(options, &mut valids, &null_indices, limit, |a, b| a.cmp(&b)).into())

}

fn sort_list<O: OffsetSizeTrait>(

    array: &GenericListArray<O>,

    value_indices: Vec<u32>,

    null_indices: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> Result<UInt32Array, ArrowError> {

    let rank = child_rank(array.values().as_ref(), options)?;

    let offsets = array.value_offsets();

    let mut valids = value_indices

        .into_iter()

        .map(|index| {

            let end = offsets[index as usize + 1].as_usize();

            let start = offsets[index as usize].as_usize();

            (index, &rank[start..end])

        })

        .collect::<Vec<(u32, &[u32])>>();

    Ok(sort_impl(options, &mut valids, &null_indices, limit, Ord::cmp).into())

}

fn sort_fixed_size_list(

    array: &FixedSizeListArray,

    value_indices: Vec<u32>,

    null_indices: Vec<u32>,

    options: SortOptions,

    limit: Option<usize>,

) -> Result<UInt32Array, ArrowError> {

    let rank = child_rank(array.values().as_ref(), options)?;

    let size = array.value_length() as usize;

    let mut valids = value_indices

        .into_iter()

        .map(|index| {

            let start = index as usize * size;

            (index, &rank[start..start + size])

        })

        .collect::<Vec<(u32, &[u32])>>();

    Ok(sort_impl(options, &mut valids, &null_indices, limit, Ord::cmp).into())

}

#[inline(never)]

fn sort_impl<T: Copy>(

    options: SortOptions,

    valids: &mut [(u32, T)],

    nulls: &[u32],

    limit: Option<usize>,

    mut cmp: impl FnMut(T, T) -> Ordering,

) -> Vec<u32> {

    let v_limit = match (limit, options.nulls_first) {

        (Some(l), true) => l.saturating_sub(nulls.len()).min(valids.len()),

        _ => valids.len(),

    };

    match options.descending {

        false => sort_unstable_by(valids, v_limit, |a, b| cmp(a.1, b.1)),

        true => sort_unstable_by(valids, v_limit, |a, b| cmp(a.1, b.1).reverse()),

    }

    let len = valids.len() + nulls.len();

    let limit = limit.unwrap_or(len).min(len);

    let mut out = Vec::with_capacity(len);

    match options.nulls_first {

        true => {

            out.extend_from_slice(&nulls[..nulls.len().min(limit)]);

            let remaining = limit - out.len();

            out.extend(valids.iter().map(|x| x.0).take(remaining));

        }

        false => {

            out.extend(valids.iter().map(|x| x.0).take(limit));

            let remaining = limit - out.len();

            out.extend_from_slice(&nulls[..remaining])

        }

    }

    out

}

/// Computes the rank for a set of child values

fn child_rank(values: &dyn Array, options: SortOptions) -> Result<Vec<u32>, ArrowError> {

    // If parent sort order is descending we need to invert the value of nulls_first so that

    // when the parent is sorted based on the produced ranks, nulls are still ordered correctly

    let value_options = Some(SortOptions {

        descending: false,

        nulls_first: options.nulls_first != options.descending,

    });

    rank(values, value_options)

}

// Sort run array and return sorted run array.

// The output RunArray will be encoded at the same level as input run array.

// For e.g. an input RunArray { run_ends = [2,4,6,8], values = [1,2,1,2] }

// will result in output RunArray { run_ends = [2,4,6,8], values = [1,1,2,2] }

// and not RunArray { run_ends = [4,8], values = [1,2] }

fn sort_run(

    values: &dyn Array,

    options: Option<SortOptions>,

    limit: Option<usize>,

) -> Result<ArrayRef, ArrowError> {

    match values.data_type() {

        DataType::RunEndEncoded(run_ends_field, _) => match run_ends_field.data_type() {

            DataType::Int16 => sort_run_downcasted::<Int16Type>(values, options, limit),

            DataType::Int32 => sort_run_downcasted::<Int32Type>(values, options, limit),

            DataType::Int64 => sort_run_downcasted::<Int64Type>(values, options, limit),

            dt => unreachable!("Not valid run ends data type {dt}"),

        },

        dt => Err(ArrowError::InvalidArgumentError(format!(

            "Input is not a run encoded array. Input data type {dt}"

        ))),

    }

}

fn sort_run_downcasted<R: RunEndIndexType>(

    values: &dyn Array,

    options: Option<SortOptions>,

    limit: Option<usize>,

) -> Result<ArrayRef, ArrowError> {

    let run_array = values.as_any().downcast_ref::<RunArray<R>>().unwrap();

    // Determine the length of output run array.

    let output_len = if let Some(limit) = limit {

        limit.min(run_array.len())

    } else {

        run_array.len()

    };

    let run_ends = run_array.run_ends();

    let mut new_run_ends_builder = BufferBuilder::<R::Native>::new(run_ends.len());

    let mut new_run_end: usize = 0;

    let mut new_physical_len: usize = 0;

    let consume_runs = |run_length, _| {

        new_run_end += run_length;

        new_physical_len += 1;

        new_run_ends_builder.append(R::Native::from_usize(new_run_end).unwrap());

    };

    let (values_indices, run_values) = sort_run_inner(run_array, options, output_len, consume_runs);

    let new_run_ends = unsafe {

        // Safety:

        // The function builds a valid run_ends array and hence need not be validated.

        ArrayDataBuilder::new(R::DATA_TYPE)

            .len(new_physical_len)

            .add_buffer(new_run_ends_builder.finish())

            .build_unchecked()

    };

    // slice the sorted value indices based on limit.

    let new_values_indices: PrimitiveArray<UInt32Type> = values_indices

        .slice(0, new_run_ends.len())

        .into_data()

        .into();

    let new_values = take(&run_values, &new_values_indices, None)?;

    // Build sorted run array

    let builder = ArrayDataBuilder::new(run_array.data_type().clone())

        .len(new_run_end)

        .add_child_data(new_run_ends)

        .add_child_data(new_values.into_data());

    let array_data: RunArray<R> = unsafe {

        // Safety:

        //  This function builds a valid run array and hence can skip validation.

        builder.build_unchecked().into()

    };

    Ok(Arc::new(array_data))

}

// Sort to indices for run encoded array.

// This function will be slow for run array as it decodes the physical indices to

// logical indices and to get the run array back, the logical indices has to be

// encoded back to run array.

fn sort_run_to_indices<R: RunEndIndexType>(

    values: &dyn Array,

    options: SortOptions,

    limit: Option<usize>,

) -> UInt32Array {

    let run_array = values.as_any().downcast_ref::<RunArray<R>>().unwrap();

    let output_len = if let Some(limit) = limit {

        limit.min(run_array.len())

    } else {

        run_array.len()

    };

    let mut result: Vec<u32> = Vec::with_capacity(output_len);

    //Add all logical indices belonging to a physical index to the output

    let consume_runs = |run_length, logical_start| {

        result.extend(logical_start as u32..(logical_start + run_length) as u32);

    };

    sort_run_inner(run_array, Some(options), output_len, consume_runs);

    UInt32Array::from(result)

}

fn sort_run_inner<R: RunEndIndexType, F>(

    run_array: &RunArray<R>,

    options: Option<SortOptions>,

    output_len: usize,

    mut consume_runs: F,

) -> (PrimitiveArray<UInt32Type>, ArrayRef)

where

    F: FnMut(usize, usize),

{

    // slice the run_array.values based on offset and length.

    let start_physical_index = run_array.get_start_physical_index();

    let end_physical_index = run_array.get_end_physical_index();

    let physical_len = end_physical_index - start_physical_index + 1;

    let run_values = run_array.values().slice(start_physical_index, physical_len);

    // All the values have to be sorted irrespective of input limit.

    let values_indices = sort_to_indices(&run_values, options, None).unwrap();

    let mut remaining_len = output_len;

    let run_ends = run_array.run_ends().values();

    assert_eq!(

        0,

        values_indices.null_count(),

        "The output of sort_to_indices should not have null values. Its values is {}",

        values_indices.null_count()

    );

    // Calculate `run length` of sorted value indices.

    // Find the `logical index` at which the run starts.

    // Call the consumer using the run length and starting logical index.

    for physical_index in values_indices.values() {

        // As the values were sliced with offset = start_physical_index, it has to be added back

        // before accessing `RunArray::run_ends`

        let physical_index = *physical_index as usize + start_physical_index;

        // calculate the run length and logical index of sorted values

        let (run_length, logical_index_start) = unsafe {

            // Safety:

            // The index will be within bounds as its in bounds of start_physical_index

            // and len, both of which are within bounds of run_array

            if physical_index == start_physical_index {

                (

                    run_ends.get_unchecked(physical_index).as_usize() - run_array.offset(),

                    0,

                )

            } else if physical_index == end_physical_index {

                let prev_run_end = run_ends.get_unchecked(physical_index - 1).as_usize();

                (

                    run_array.offset() + run_array.len() - prev_run_end,

                    prev_run_end - run_array.offset(),

                )

            } else {

                let prev_run_end = run_ends.get_unchecked(physical_index - 1).as_usize();

                (

                    run_ends.get_unchecked(physical_index).as_usize() - prev_run_end,

                    prev_run_end - run_array.offset(),

                )

            }

        };

        let new_run_length = run_length.min(remaining_len);

        consume_runs(new_run_length, logical_index_start);

        remaining_len -= new_run_length;

        if remaining_len == 0 {

            break;

        }

    }

    if remaining_len > 0 {

        panic!("Remaining length should be zero its values is {remaining_len}")

    }

    (values_indices, run_values)

}

/// One column to be used in lexicographical sort

#[derive(Clone, Debug)]

pub struct SortColumn {

    /// The column to sort

    pub values: ArrayRef,

    /// Sort options for this column

    pub options: Option<SortOptions>,

}

/// Sort a list of `ArrayRef` using `SortOptions` provided for each array.

///

/// Performs a stable lexicographical sort on values and indices.

///

/// Returns an `ArrowError::ComputeError(String)` if any of the array type is either unsupported by

/// `lexsort_to_indices` or `take`.

///

/// Example:

///

/// ```

/// # use std::convert::From;

/// # use std::sync::Arc;

/// # use arrow_array::{ArrayRef, StringArray, PrimitiveArray};

/// # use arrow_array::types::Int64Type;

/// # use arrow_array::cast::AsArray;

/// # use arrow_ord::sort::{SortColumn, SortOptions, lexsort};

///

/// let sorted_columns = lexsort(&vec![

///     SortColumn {

///         values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![

///             None,

///             Some(-2),

///             Some(89),

///             Some(-64),

///             Some(101),

///         ])) as ArrayRef,

///         options: None,

///     },

///     SortColumn {

///         values: Arc::new(StringArray::from(vec![

///             Some("hello"),

///             Some("world"),

///             Some(","),

///             Some("foobar"),

///             Some("!"),

///         ])) as ArrayRef,

///         options: Some(SortOptions {

///             descending: true,

///             nulls_first: false,

///         }),

///     },

/// ], None).unwrap();

///