arrow_ord/
sort.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//! Defines sort kernel for `ArrayRef`
6
7use crate::ord::{DynComparator, make_comparator};
8use arrow_array::cast::*;
9use arrow_array::types::*;
10use arrow_array::*;
11use arrow_buffer::ArrowNativeType;
12use arrow_buffer::BooleanBufferBuilder;
13use arrow_data::{ArrayDataBuilder, ByteView, MAX_INLINE_VIEW_LEN};
14use arrow_schema::{ArrowError, DataType};
15use arrow_select::take::take;
16use std::cmp::Ordering;
17use std::sync::Arc;
18
19use crate::rank::{can_rank, rank};
20pub use arrow_schema::SortOptions;
21
22/// Sort the `ArrayRef` using `SortOptions`.
23///
24/// Performs a sort on values and indices. Nulls are ordered according
25/// to the `nulls_first` flag in `options`.  Floats are sorted using
26/// IEEE 754 totalOrder
27///
28/// Returns an `ArrowError::ComputeError(String)` if the array type is
29/// either unsupported by `sort_to_indices` or `take`.
30///
31/// Note: this is an unstable_sort, meaning it may not preserve the
32/// order of equal elements.
33///
34/// # Example
35/// ```rust
36/// # use std::sync::Arc;
37/// # use arrow_array::Int32Array;
38/// # use arrow_ord::sort::sort;
39/// let array = Int32Array::from(vec![5, 4, 3, 2, 1]);
40/// let sorted_array = sort(&array, None).unwrap();
41/// assert_eq!(sorted_array.as_ref(), &Int32Array::from(vec![1, 2, 3, 4, 5]));
42/// ```
43pub fn sort(values: &dyn Array, options: Option<SortOptions>) -> Result<ArrayRef, ArrowError> {
44    if values.is_empty() {
45        return Ok(new_empty_array(values.data_type()));
46    }
47    downcast_primitive_array!(
48        values => sort_native_type(values, options),
49        DataType::RunEndEncoded(_, _) => sort_run(values, options, None),
50        _ => {
51            let indices = sort_to_indices(values, options, None)?;
52            take(values, &indices, None)
53        }
54    )
55}
56
57fn sort_native_type<T>(
58    primitive_values: &PrimitiveArray<T>,
59    options: Option<SortOptions>,
60) -> Result<ArrayRef, ArrowError>
61where
62    T: ArrowPrimitiveType,
63{
64    let sort_options = options.unwrap_or_default();
65
66    let mut mutable_buffer = vec![T::default_value(); primitive_values.len()];
67    let mutable_slice = &mut mutable_buffer;
68
69    let input_values = primitive_values.values().as_ref();
70
71    let nulls_count = primitive_values.null_count();
72    let valid_count = primitive_values.len() - nulls_count;
73
74    let null_bit_buffer = match nulls_count > 0 {
75        true => {
76            let mut validity_buffer = BooleanBufferBuilder::new(primitive_values.len());
77            if sort_options.nulls_first {
78                validity_buffer.append_n(nulls_count, false);
79                validity_buffer.append_n(valid_count, true);
80            } else {
81                validity_buffer.append_n(valid_count, true);
82                validity_buffer.append_n(nulls_count, false);
83            }
84            Some(validity_buffer.finish().into())
85        }
86        false => None,
87    };
88
89    if let Some(nulls) = primitive_values.nulls().filter(|n| n.null_count() > 0) {
90        let values_slice = match sort_options.nulls_first {
91            true => &mut mutable_slice[nulls_count..],
92            false => &mut mutable_slice[..valid_count],
93        };
94
95        for (write_index, index) in nulls.valid_indices().enumerate() {
96            values_slice[write_index] = primitive_values.value(index);
97        }
98
99        values_slice.sort_unstable_by(|a, b| a.compare(*b));
100        if sort_options.descending {
101            values_slice.reverse();
102        }
103    } else {
104        mutable_slice.copy_from_slice(input_values);
105        mutable_slice.sort_unstable_by(|a, b| a.compare(*b));
106        if sort_options.descending {
107            mutable_slice.reverse();
108        }
109    }
110
111    Ok(Arc::new(
112        PrimitiveArray::<T>::try_new(mutable_buffer.into(), null_bit_buffer)?
113            .with_data_type(primitive_values.data_type().clone()),
114    ))
115}
116
117/// Sort the `ArrayRef` partially.
118///
119/// If `limit` is specified, the resulting array will contain only
120/// first `limit` in the sort order. Any data data after the limit
121/// will be discarded.
122///
123/// Note: this is an unstable_sort, meaning it may not preserve the
124/// order of equal elements.
125///
126/// # Example
127/// ```rust
128/// # use std::sync::Arc;
129/// # use arrow_array::Int32Array;
130/// # use arrow_ord::sort::{sort_limit, SortOptions};
131/// let array = Int32Array::from(vec![5, 4, 3, 2, 1]);
132///
133/// // Find the the top 2 items
134/// let sorted_array = sort_limit(&array, None, Some(2)).unwrap();
135/// assert_eq!(sorted_array.as_ref(), &Int32Array::from(vec![1, 2]));
136///
137/// // Find the bottom top 2 items
138/// let options = Some(SortOptions {
139///                  descending: true,
140///                  ..Default::default()
141///               });
142/// let sorted_array = sort_limit(&array, options, Some(2)).unwrap();
143/// assert_eq!(sorted_array.as_ref(), &Int32Array::from(vec![5, 4]));
144/// ```
145pub fn sort_limit(
146    values: &dyn Array,
147    options: Option<SortOptions>,
148    limit: Option<usize>,
149) -> Result<ArrayRef, ArrowError> {
150    if values.is_empty() || limit == Some(0) {
151        return Ok(new_empty_array(values.data_type()));
152    }
153    if let DataType::RunEndEncoded(_, _) = values.data_type() {
154        return sort_run(values, options, limit);
155    }
156    let indices = sort_to_indices(values, options, limit)?;
157    take(values, &indices, None)
158}
159
160/// we can only do this if the T is primitive
161#[inline]
162fn sort_unstable_by<T, F>(array: &mut [T], limit: usize, cmp: F)
163where
164    F: FnMut(&T, &T) -> Ordering,
165{
166    if array.len() == limit {
167        array.sort_unstable_by(cmp);
168    } else {
169        partial_sort(array, limit, cmp);
170    }
171}
172
173/// Partition indices of an Arrow array into two categories:
174/// - `valid`: indices of non-null elements
175/// - `nulls`: indices of null elements
176///
177/// Optimized for performance with fast-path for all-valid arrays
178/// and bit-parallel scan for null-containing arrays.
179#[inline(always)]
180pub fn partition_validity(array: &dyn Array) -> (Vec<u32>, Vec<u32>) {
181    let len = array.len();
182    let null_count = array.null_count();
183
184    // Fast path: if there are no nulls, all elements are valid
185    if null_count == 0 {
186        // Simply return a range of indices [0, len)
187        let valid = (0..len as u32).collect();
188        return (valid, Vec::new());
189    }
190
191    // null bitmap exists and some values are null
192    partition_validity_scan(array, len, null_count)
193}
194
195/// Scans the null bitmap and partitions valid/null indices efficiently.
196/// Uses bit-level operations to extract bit positions.
197/// This function is only called when nulls exist.
198#[inline(always)]
199fn partition_validity_scan(
200    array: &dyn Array,
201    len: usize,
202    null_count: usize,
203) -> (Vec<u32>, Vec<u32>) {
204    // SAFETY: Guaranteed by caller that null_count > 0, so bitmap must exist
205    let bitmap = array.nulls().unwrap();
206
207    // Preallocate result vectors with exact capacities (avoids reallocations)
208    let mut valid = Vec::with_capacity(len - null_count);
209    let mut nulls = Vec::with_capacity(null_count);
210
211    unsafe {
212        // 1) Write valid indices (bits == 1)
213        let valid_slice = valid.spare_capacity_mut();
214        for (i, idx) in bitmap.inner().set_indices_u32().enumerate() {
215            valid_slice[i].write(idx);
216        }
217
218        // 2) Write null indices by inverting
219        let inv_buf = !bitmap.inner();
220        let null_slice = nulls.spare_capacity_mut();
221        for (i, idx) in inv_buf.set_indices_u32().enumerate() {
222            null_slice[i].write(idx);
223        }
224
225        // Finalize lengths
226        valid.set_len(len - null_count);
227        nulls.set_len(null_count);
228    }
229
230    assert_eq!(valid.len(), len - null_count);
231    assert_eq!(nulls.len(), null_count);
232    (valid, nulls)
233}
234
235/// Whether `sort_to_indices` can sort an array of given data type.
236fn can_sort_to_indices(data_type: &DataType) -> bool {
237    data_type.is_primitive()
238        || matches!(
239            data_type,
240            DataType::Boolean
241                | DataType::Utf8
242                | DataType::LargeUtf8
243                | DataType::Utf8View
244                | DataType::Binary
245                | DataType::LargeBinary
246                | DataType::BinaryView
247                | DataType::FixedSizeBinary(_)
248        )
249        || match data_type {
250            DataType::List(f) if can_rank(f.data_type()) => true,
251            DataType::ListView(f) if can_rank(f.data_type()) => true,
252            DataType::LargeList(f) if can_rank(f.data_type()) => true,
253            DataType::LargeListView(f) if can_rank(f.data_type()) => true,
254            DataType::FixedSizeList(f, _) if can_rank(f.data_type()) => true,
255            DataType::Dictionary(_, values) if can_rank(values.as_ref()) => true,
256            DataType::RunEndEncoded(_, f) if can_sort_to_indices(f.data_type()) => true,
257            _ => false,
258        }
259}
260
261/// Sort elements from `ArrayRef` into an unsigned integer (`UInt32Array`) of indices.
262/// Floats are sorted using IEEE 754 totalOrder.  `limit` is an option for [partial_sort].
263pub fn sort_to_indices(
264    array: &dyn Array,
265    options: Option<SortOptions>,
266    limit: Option<usize>,
267) -> Result<UInt32Array, ArrowError> {
268    if array.is_empty() || limit == Some(0) {
269        return Ok(UInt32Array::from(Vec::<u32>::new()));
270    }
271
272    let options = options.unwrap_or_default();
273
274    let (v, n) = partition_validity(array);
275
276    Ok(downcast_primitive_array! {
277        array => sort_primitive(array, v, n, options, limit),
278        DataType::Boolean => sort_boolean(array.as_boolean(), v, n, options, limit),
279        DataType::Utf8 => sort_bytes(array.as_string::<i32>(), v, n, options, limit),
280        DataType::LargeUtf8 => sort_bytes(array.as_string::<i64>(), v, n, options, limit),
281        DataType::Utf8View => sort_byte_view(array.as_string_view(), v, n, options, limit),
282        DataType::Binary => sort_bytes(array.as_binary::<i32>(), v, n, options, limit),
283        DataType::LargeBinary => sort_bytes(array.as_binary::<i64>(), v, n, options, limit),
284        DataType::BinaryView => sort_byte_view(array.as_binary_view(), v, n, options, limit),
285        DataType::FixedSizeBinary(_) => sort_fixed_size_binary(array.as_fixed_size_binary(), v, n, options, limit),
286        DataType::List(_) => sort_list(array.as_list::<i32>(), v, n, options, limit)?,
287        DataType::ListView(_) => sort_list_view(array.as_list_view::<i32>(), v, n, options, limit)?,
288        DataType::LargeList(_) => sort_list(array.as_list::<i64>(), v, n, options, limit)?,
289        DataType::LargeListView(_) => sort_list_view(array.as_list_view::<i64>(), v, n, options, limit)?,
290        DataType::FixedSizeList(_, _) => sort_fixed_size_list(array.as_fixed_size_list(), v, n, options, limit)?,
291        DataType::Dictionary(_, _) => downcast_dictionary_array!{
292            array => sort_dictionary(array, v, n, options, limit)?,
293            _ => unreachable!()
294        }
295        DataType::RunEndEncoded(run_ends_field, _) => match run_ends_field.data_type() {
296            DataType::Int16 => sort_run_to_indices::<Int16Type>(array, options, limit),
297            DataType::Int32 => sort_run_to_indices::<Int32Type>(array, options, limit),
298            DataType::Int64 => sort_run_to_indices::<Int64Type>(array, options, limit),
299            dt => {
300                return Err(ArrowError::ComputeError(format!(
301                    "Invalid run end data type: {dt}"
302                )))
303            }
304        },
305        t => {
306            return Err(ArrowError::ComputeError(format!(
307                "Sort not supported for data type {t}"
308            )));
309        }
310    })
311}
312
313fn sort_boolean(
314    values: &BooleanArray,
315    value_indices: Vec<u32>,
316    null_indices: Vec<u32>,
317    options: SortOptions,
318    limit: Option<usize>,
319) -> UInt32Array {
320    let mut valids = value_indices
321        .into_iter()
322        .map(|index| (index, values.value(index as usize)))
323        .collect::<Vec<(u32, bool)>>();
324    sort_impl(options, &mut valids, &null_indices, limit, |a, b| a.cmp(&b)).into()
325}
326
327fn sort_primitive<T: ArrowPrimitiveType>(
328    values: &PrimitiveArray<T>,
329    value_indices: Vec<u32>,
330    nulls: Vec<u32>,
331    options: SortOptions,
332    limit: Option<usize>,
333) -> UInt32Array {
334    let mut valids = value_indices
335        .into_iter()
336        .map(|index| (index, values.value(index as usize)))
337        .collect::<Vec<(u32, T::Native)>>();
338    sort_impl(options, &mut valids, &nulls, limit, T::Native::compare).into()
339}
340
341fn sort_bytes<T: ByteArrayType>(
342    values: &GenericByteArray<T>,
343    value_indices: Vec<u32>,
344    nulls: Vec<u32>,
345    options: SortOptions,
346    limit: Option<usize>,
347) -> UInt32Array {
348    // Note: Why do we use 4‑byte prefix?
349    // Compute the 4‑byte prefix in BE order, or left‑pad if shorter.
350    // Most byte‐sequences differ in their first few bytes, so by
351    // comparing up to 4 bytes as a single u32 we avoid the overhead
352    // of a full lexicographical compare for the vast majority of cases.
353
354    // 1. Build a vector of (index, prefix, length) tuples
355    let mut valids: Vec<(u32, u32, u64)> = value_indices
356        .into_iter()
357        .map(|idx| unsafe {
358            let slice: &[u8] = values.value_unchecked(idx as usize).as_ref();
359            let len = slice.len() as u64;
360            // Compute the 4‑byte prefix in BE order, or left‑pad if shorter
361            let prefix = if slice.len() >= 4 {
362                let raw = std::ptr::read_unaligned(slice.as_ptr().cast::<u32>());
363                u32::from_be(raw)
364            } else if slice.is_empty() {
365                // Handle empty slice case to avoid shift overflow
366                0u32
367            } else {
368                let mut v = 0u32;
369                for &b in slice {
370                    v = (v << 8) | (b as u32);
371                }
372                // Safe shift: slice.len() is in range [1, 3], so shift is in range [8, 24]
373                v << (8 * (4 - slice.len()))
374            };
375            (idx, prefix, len)
376        })
377        .collect();
378
379    // 2. compute the number of non-null entries to partially sort
380    let vlimit = match (limit, options.nulls_first) {
381        (Some(l), true) => l.saturating_sub(nulls.len()).min(valids.len()),
382        _ => valids.len(),
383    };
384
385    // 3. Comparator: compare prefix, then (when both slices shorter than 4) length, otherwise full slice
386    let cmp_bytes = |a: &(u32, u32, u64), b: &(u32, u32, u64)| unsafe {
387        let (ia, pa, la) = *a;
388        let (ib, pb, lb) = *b;
389        // 3.1 prefix (first 4 bytes)
390        let ord = pa.cmp(&pb);
391        if ord != Ordering::Equal {
392            return ord;
393        }
394        // 3.2 only if both slices had length < 4 (so prefix was padded)
395        if la < 4 || lb < 4 {
396            let ord = la.cmp(&lb);
397            if ord != Ordering::Equal {
398                return ord;
399            }
400        }
401        // 3.3 full lexicographical compare
402        let a_bytes: &[u8] = values.value_unchecked(ia as usize).as_ref();
403        let b_bytes: &[u8] = values.value_unchecked(ib as usize).as_ref();
404        a_bytes.cmp(b_bytes)
405    };
406
407    // 4. Partially sort according to ascending/descending
408    if !options.descending {
409        sort_unstable_by(&mut valids, vlimit, cmp_bytes);
410    } else {
411        sort_unstable_by(&mut valids, vlimit, |x, y| cmp_bytes(x, y).reverse());
412    }
413
414    // 5. Assemble nulls and sorted indices into final output
415    let total = valids.len() + nulls.len();
416    let out_limit = limit.unwrap_or(total).min(total);
417    let mut out = Vec::with_capacity(out_limit);
418
419    if options.nulls_first {
420        out.extend_from_slice(&nulls[..nulls.len().min(out_limit)]);
421        let rem = out_limit - out.len();
422        out.extend(valids.iter().map(|&(i, _, _)| i).take(rem));
423    } else {
424        out.extend(valids.iter().map(|&(i, _, _)| i).take(out_limit));
425        let rem = out_limit - out.len();
426        out.extend_from_slice(&nulls[..rem]);
427    }
428
429    out.into()
430}
431
432fn sort_byte_view<T: ByteViewType>(
433    values: &GenericByteViewArray<T>,
434    value_indices: Vec<u32>,
435    nulls: Vec<u32>,
436    options: SortOptions,
437    limit: Option<usize>,
438) -> UInt32Array {
439    // 1. Build a list of (index, raw_view, length)
440    let mut valids: Vec<_>;
441    // 2. Compute the number of non-null entries to partially sort
442    let vlimit: usize = match (limit, options.nulls_first) {
443        (Some(l), true) => l.saturating_sub(nulls.len()).min(value_indices.len()),
444        _ => value_indices.len(),
445    };
446    // 3.a Check if all views are inline (no data buffers)
447    if values.data_buffers().is_empty() {
448        valids = value_indices
449            .into_iter()
450            .map(|idx| {
451                // SAFETY: we know idx < values.len()
452                let raw = unsafe { *values.views().get_unchecked(idx as usize) };
453                let inline_key = GenericByteViewArray::<T>::inline_key_fast(raw);
454                (idx, inline_key)
455            })
456            .collect();
457        let cmp_inline = |a: &(u32, u128), b: &(u32, u128)| a.1.cmp(&b.1);
458
459        // Partially sort according to ascending/descending
460        if !options.descending {
461            sort_unstable_by(&mut valids, vlimit, cmp_inline);
462        } else {
463            sort_unstable_by(&mut valids, vlimit, |x, y| cmp_inline(x, y).reverse());
464        }
465    } else {
466        valids = value_indices
467            .into_iter()
468            .map(|idx| {
469                // SAFETY: we know idx < values.len()
470                let raw = unsafe { *values.views().get_unchecked(idx as usize) };
471                (idx, raw)
472            })
473            .collect();
474        // 3.b Mixed comparator: first prefix, then inline vs full comparison
475        let cmp_mixed = |a: &(u32, u128), b: &(u32, u128)| {
476            let (_, raw_a) = *a;
477            let (_, raw_b) = *b;
478            let len_a = raw_a as u32;
479            let len_b = raw_b as u32;
480            // 3.b.1 Both inline (≤12 bytes): compare full 128-bit key including length
481            if len_a <= MAX_INLINE_VIEW_LEN && len_b <= MAX_INLINE_VIEW_LEN {
482                return GenericByteViewArray::<T>::inline_key_fast(raw_a)
483                    .cmp(&GenericByteViewArray::<T>::inline_key_fast(raw_b));
484            }
485
486            // 3.b.2 Compare 4-byte prefix in big-endian order
487            let pref_a = ByteView::from(raw_a).prefix.swap_bytes();
488            let pref_b = ByteView::from(raw_b).prefix.swap_bytes();
489            if pref_a != pref_b {
490                return pref_a.cmp(&pref_b);
491            }
492
493            // 3.b.3 Fallback to full byte-slice comparison
494            let full_a: &[u8] = unsafe { values.value_unchecked(a.0 as usize).as_ref() };
495            let full_b: &[u8] = unsafe { values.value_unchecked(b.0 as usize).as_ref() };
496            full_a.cmp(full_b)
497        };
498
499        // 3.b.4 Partially sort according to ascending/descending
500        if !options.descending {
501            sort_unstable_by(&mut valids, vlimit, cmp_mixed);
502        } else {
503            sort_unstable_by(&mut valids, vlimit, |x, y| cmp_mixed(x, y).reverse());
504        }
505    }
506
507    // 5. Assemble nulls and sorted indices into final output
508    let total = valids.len() + nulls.len();
509    let out_limit = limit.unwrap_or(total).min(total);
510    let mut out = Vec::with_capacity(total);
511
512    if options.nulls_first {
513        // Place null indices first
514        out.extend_from_slice(&nulls[..nulls.len().min(out_limit)]);
515        let rem = out_limit - out.len();
516        out.extend(valids.iter().map(|&(i, _)| i).take(rem));
517    } else {
518        // Place non-null indices first
519        out.extend(valids.iter().map(|&(i, _)| i).take(out_limit));
520        let rem = out_limit - out.len();
521        out.extend_from_slice(&nulls[..rem]);
522    }
523
524    out.into()
525}
526
527fn sort_fixed_size_binary(
528    values: &FixedSizeBinaryArray,
529    value_indices: Vec<u32>,
530    nulls: Vec<u32>,
531    options: SortOptions,
532    limit: Option<usize>,
533) -> UInt32Array {
534    let mut valids = value_indices
535        .iter()
536        .copied()
537        .map(|index| (index, values.value(index as usize)))
538        .collect::<Vec<(u32, &[u8])>>();
539    sort_impl(options, &mut valids, &nulls, limit, Ord::cmp).into()
540}
541
542fn sort_dictionary<K: ArrowDictionaryKeyType>(
543    dict: &DictionaryArray<K>,
544    value_indices: Vec<u32>,
545    null_indices: Vec<u32>,
546    options: SortOptions,
547    limit: Option<usize>,
548) -> Result<UInt32Array, ArrowError> {
549    let keys: &PrimitiveArray<K> = dict.keys();
550    let rank = child_rank(dict.values().as_ref(), options)?;
551
552    // create tuples that are used for sorting
553    let mut valids = value_indices
554        .into_iter()
555        .map(|index| {
556            let key: K::Native = keys.value(index as usize);
557            (index, rank[key.as_usize()])
558        })
559        .collect::<Vec<(u32, u32)>>();
560
561    Ok(sort_impl(options, &mut valids, &null_indices, limit, |a, b| a.cmp(&b)).into())
562}
563
564fn sort_list<O: OffsetSizeTrait>(
565    array: &GenericListArray<O>,
566    value_indices: Vec<u32>,
567    null_indices: Vec<u32>,
568    options: SortOptions,
569    limit: Option<usize>,
570) -> Result<UInt32Array, ArrowError> {
571    let rank = child_rank(array.values().as_ref(), options)?;
572    let offsets = array.value_offsets();
573    let mut valids = value_indices
574        .into_iter()
575        .map(|index| {
576            let end = offsets[index as usize + 1].as_usize();
577            let start = offsets[index as usize].as_usize();
578            (index, &rank[start..end])
579        })
580        .collect::<Vec<(u32, &[u32])>>();
581    Ok(sort_impl(options, &mut valids, &null_indices, limit, Ord::cmp).into())
582}
583
584fn sort_list_view<O: OffsetSizeTrait>(
585    array: &GenericListViewArray<O>,
586    value_indices: Vec<u32>,
587    null_indices: Vec<u32>,
588    options: SortOptions,
589    limit: Option<usize>,
590) -> Result<UInt32Array, ArrowError> {
591    let rank = child_rank(array.values().as_ref(), options)?;
592    let offsets = array.offsets();
593    let sizes = array.sizes();
594    let mut valids = value_indices
595        .into_iter()
596        .map(|index| {
597            let start = offsets[index as usize].as_usize();
598            let size = sizes[index as usize].as_usize();
599            let end = start + size;
600            (index, &rank[start..end])
601        })
602        .collect::<Vec<(u32, &[u32])>>();
603    Ok(sort_impl(options, &mut valids, &null_indices, limit, Ord::cmp).into())
604}
605
606fn sort_fixed_size_list(
607    array: &FixedSizeListArray,
608    value_indices: Vec<u32>,
609    null_indices: Vec<u32>,
610    options: SortOptions,
611    limit: Option<usize>,
612) -> Result<UInt32Array, ArrowError> {
613    let rank = child_rank(array.values().as_ref(), options)?;
614    let size = array.value_length() as usize;
615    let mut valids = value_indices
616        .into_iter()
617        .map(|index| {
618            let start = index as usize * size;
619            (index, &rank[start..start + size])
620        })
621        .collect::<Vec<(u32, &[u32])>>();
622    Ok(sort_impl(options, &mut valids, &null_indices, limit, Ord::cmp).into())
623}
624
625#[inline(never)]
626fn sort_impl<T: Copy>(
627    options: SortOptions,
628    valids: &mut [(u32, T)],
629    nulls: &[u32],
630    limit: Option<usize>,
631    mut cmp: impl FnMut(T, T) -> Ordering,
632) -> Vec<u32> {
633    let v_limit = match (limit, options.nulls_first) {
634        (Some(l), true) => l.saturating_sub(nulls.len()).min(valids.len()),
635        _ => valids.len(),
636    };
637
638    match options.descending {
639        false => sort_unstable_by(valids, v_limit, |a, b| cmp(a.1, b.1)),
640        true => sort_unstable_by(valids, v_limit, |a, b| cmp(a.1, b.1).reverse()),
641    }
642
643    let len = valids.len() + nulls.len();
644    let limit = limit.unwrap_or(len).min(len);
645    let mut out = Vec::with_capacity(len);
646    match options.nulls_first {
647        true => {
648            out.extend_from_slice(&nulls[..nulls.len().min(limit)]);
649            let remaining = limit - out.len();
650            out.extend(valids.iter().map(|x| x.0).take(remaining));
651        }
652        false => {
653            out.extend(valids.iter().map(|x| x.0).take(limit));
654            let remaining = limit - out.len();
655            out.extend_from_slice(&nulls[..remaining])
656        }
657    }
658    out
659}
660
661/// Computes the rank for a set of child values
662fn child_rank(values: &dyn Array, options: SortOptions) -> Result<Vec<u32>, ArrowError> {
663    // If parent sort order is descending we need to invert the value of nulls_first so that
664    // when the parent is sorted based on the produced ranks, nulls are still ordered correctly
665    let value_options = Some(SortOptions {
666        descending: false,
667        nulls_first: options.nulls_first != options.descending,
668    });
669    rank(values, value_options)
670}
671
672// Sort run array and return sorted run array.
673// The output RunArray will be encoded at the same level as input run array.
674// For e.g. an input RunArray { run_ends = [2,4,6,8], values = [1,2,1,2] }
675// will result in output RunArray { run_ends = [2,4,6,8], values = [1,1,2,2] }
676// and not RunArray { run_ends = [4,8], values = [1,2] }
677fn sort_run(
678    values: &dyn Array,
679    options: Option<SortOptions>,
680    limit: Option<usize>,
681) -> Result<ArrayRef, ArrowError> {
682    match values.data_type() {
683        DataType::RunEndEncoded(run_ends_field, _) => match run_ends_field.data_type() {
684            DataType::Int16 => sort_run_downcasted::<Int16Type>(values, options, limit),
685            DataType::Int32 => sort_run_downcasted::<Int32Type>(values, options, limit),
686            DataType::Int64 => sort_run_downcasted::<Int64Type>(values, options, limit),
687            dt => unreachable!("Not valid run ends data type {dt}"),
688        },
689        dt => Err(ArrowError::InvalidArgumentError(format!(
690            "Input is not a run encoded array. Input data type {dt}"
691        ))),
692    }
693}
694
695fn sort_run_downcasted<R: RunEndIndexType>(
696    values: &dyn Array,
697    options: Option<SortOptions>,
698    limit: Option<usize>,
699) -> Result<ArrayRef, ArrowError> {
700    let run_array = values.as_any().downcast_ref::<RunArray<R>>().unwrap();
701
702    // Determine the length of output run array.
703    let output_len = if let Some(limit) = limit {
704        limit.min(run_array.len())
705    } else {
706        run_array.len()
707    };
708
709    let run_ends = run_array.run_ends();
710
711    let mut new_run_ends = Vec::with_capacity(run_ends.len());
712    let mut new_run_end: usize = 0;
713    let mut new_physical_len: usize = 0;
714
715    let consume_runs = |run_length, _| {
716        new_run_end += run_length;
717        new_physical_len += 1;
718        new_run_ends.push(R::Native::from_usize(new_run_end).unwrap());
719    };
720
721    let (values_indices, run_values) = sort_run_inner(run_array, options, output_len, consume_runs);
722
723    let new_run_ends = unsafe {
724        // Safety:
725        // The function builds a valid run_ends array and hence need not be validated.
726        ArrayDataBuilder::new(R::DATA_TYPE)
727            .len(new_physical_len)
728            .add_buffer(new_run_ends.into())
729            .build_unchecked()
730    };
731
732    // slice the sorted value indices based on limit.
733    let new_values_indices: PrimitiveArray<UInt32Type> = values_indices
734        .slice(0, new_run_ends.len())
735        .into_data()
736        .into();
737
738    let new_values = take(&run_values, &new_values_indices, None)?;
739
740    // Build sorted run array
741    let builder = ArrayDataBuilder::new(run_array.data_type().clone())
742        .len(new_run_end)
743        .add_child_data(new_run_ends)
744        .add_child_data(new_values.into_data());
745    let array_data: RunArray<R> = unsafe {
746        // Safety:
747        //  This function builds a valid run array and hence can skip validation.
748        builder.build_unchecked().into()
749    };
750    Ok(Arc::new(array_data))
751}
752
753// Sort to indices for run encoded array.
754// This function will be slow for run array as it decodes the physical indices to
755// logical indices and to get the run array back, the logical indices has to be
756// encoded back to run array.
757fn sort_run_to_indices<R: RunEndIndexType>(
758    values: &dyn Array,
759    options: SortOptions,
760    limit: Option<usize>,
761) -> UInt32Array {
762    let run_array = values.as_any().downcast_ref::<RunArray<R>>().unwrap();
763    let output_len = if let Some(limit) = limit {
764        limit.min(run_array.len())
765    } else {
766        run_array.len()
767    };
768    let mut result: Vec<u32> = Vec::with_capacity(output_len);
769
770    //Add all logical indices belonging to a physical index to the output
771    let consume_runs = |run_length, logical_start| {
772        result.extend(logical_start as u32..(logical_start + run_length) as u32);
773    };
774    sort_run_inner(run_array, Some(options), output_len, consume_runs);
775
776    UInt32Array::from(result)
777}
778
779fn sort_run_inner<R: RunEndIndexType, F>(
780    run_array: &RunArray<R>,
781    options: Option<SortOptions>,
782    output_len: usize,
783    mut consume_runs: F,
784) -> (PrimitiveArray<UInt32Type>, ArrayRef)
785where
786    F: FnMut(usize, usize),
787{
788    // slice the run_array.values based on offset and length.
789    let start_physical_index = run_array.get_start_physical_index();
790    let end_physical_index = run_array.get_end_physical_index();
791    let physical_len = end_physical_index - start_physical_index + 1;
792    let run_values = run_array.values().slice(start_physical_index, physical_len);
793
794    // All the values have to be sorted irrespective of input limit.
795    let values_indices = sort_to_indices(&run_values, options, None).unwrap();
796
797    let mut remaining_len = output_len;
798
799    let run_ends = run_array.run_ends().values();
800
801    assert_eq!(
802        0,
803        values_indices.null_count(),
804        "The output of sort_to_indices should not have null values. Its values is {}",
805        values_indices.null_count()
806    );
807
808    // Calculate `run length` of sorted value indices.
809    // Find the `logical index` at which the run starts.
810    // Call the consumer using the run length and starting logical index.
811    for physical_index in values_indices.values() {
812        // As the values were sliced with offset = start_physical_index, it has to be added back
813        // before accessing `RunArray::run_ends`
814        let physical_index = *physical_index as usize + start_physical_index;
815
816        // calculate the run length and logical index of sorted values
817        let (run_length, logical_index_start) = unsafe {
818            // Safety:
819            // The index will be within bounds as its in bounds of start_physical_index
820            // and len, both of which are within bounds of run_array
821            if physical_index == start_physical_index {
822                (
823                    run_ends.get_unchecked(physical_index).as_usize() - run_array.offset(),
824                    0,
825                )
826            } else if physical_index == end_physical_index {
827                let prev_run_end = run_ends.get_unchecked(physical_index - 1).as_usize();
828                (
829                    run_array.offset() + run_array.len() - prev_run_end,
830                    prev_run_end - run_array.offset(),
831                )
832            } else {
833                let prev_run_end = run_ends.get_unchecked(physical_index - 1).as_usize();
834                (
835                    run_ends.get_unchecked(physical_index).as_usize() - prev_run_end,
836                    prev_run_end - run_array.offset(),
837                )
838            }
839        };
840        let new_run_length = run_length.min(remaining_len);
841        consume_runs(new_run_length, logical_index_start);
842        remaining_len -= new_run_length;
843
844        if remaining_len == 0 {
845            break;
846        }
847    }
848
849    if remaining_len > 0 {
850        panic!("Remaining length should be zero its values is {remaining_len}")
851    }
852    (values_indices, run_values)
853}
854
855/// One column to be used in lexicographical sort
856#[derive(Clone, Debug)]
857pub struct SortColumn {
858    /// The column to sort
859    pub values: ArrayRef,
860    /// Sort options for this column
861    pub options: Option<SortOptions>,
862}
863
864/// Sort a list of `ArrayRef` using `SortOptions` provided for each array.
865///
866/// Performs an unstable lexicographical sort on values and indices.
867///
868/// Returns an `ArrowError::ComputeError(String)` if any of the array type is either unsupported by
869/// `lexsort_to_indices` or `take`.
870///
871/// # Example:
872///
873/// ```
874/// # use std::convert::From;
875/// # use std::sync::Arc;
876/// # use arrow_array::{ArrayRef, StringArray, PrimitiveArray};
877/// # use arrow_array::types::Int64Type;
878/// # use arrow_array::cast::AsArray;
879/// # use arrow_ord::sort::{SortColumn, SortOptions, lexsort};
880/// let sorted_columns = lexsort(&vec![
881///     SortColumn {
882///         values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
883///             None,
884///             Some(-2),
885///             Some(89),
886///             Some(-64),
887///             Some(101),
888///         ])) as ArrayRef,
889///         options: None,
890///     },
891///     SortColumn {
892///         values: Arc::new(StringArray::from(vec![
893///             Some("hello"),
894///             Some("world"),
895///             Some(","),
896///             Some("foobar"),
897///             Some("!"),
898///         ])) as ArrayRef,
899///         options: Some(SortOptions {
900///             descending: true,
901///             nulls_first: false,
902///         }),
903///     },
904/// ], None).unwrap();
905///
906/// assert_eq!(sorted_columns[0].as_primitive::<Int64Type>().value(1), -64);
907/// assert!(sorted_columns[0].is_null(0));
908/// ```
909///
910/// Note: for multi-column sorts without a limit, using the [row format](https://docs.rs/arrow-row/latest/arrow_row/)
911/// may be significantly faster
912///
913pub fn lexsort(columns: &[SortColumn], limit: Option<usize>) -> Result<Vec<ArrayRef>, ArrowError> {
914    let indices = lexsort_to_indices(columns, limit)?;
915    columns
916        .iter()
917        .map(|c| take(c.values.as_ref(), &indices, None))
918        .collect()
919}
920
921/// Sort elements lexicographically from a list of `ArrayRef` into an unsigned integer
922/// (`UInt32Array`) of indices.
923///
924/// # Example
925///
926/// ```
927/// # use std::sync::Arc;
928/// # use arrow_array::{ArrayRef, UInt32Array, Int32Array, RecordBatch, StringArray};
929/// # use arrow_ord::sort::{lexsort_to_indices, SortColumn};
930/// # use arrow_select::take::take_record_batch;
931/// // Two columns (a, b). Values (2,x), (1, z), (1(a))
932/// let batch = RecordBatch::try_from_iter(vec![
933///     ("a", Arc::new(Int32Array::from(vec![2, 1, 1])) as ArrayRef),
934///     ("b", Arc::new(StringArray::from(vec!["x", "z", "a"])) as ArrayRef),
935/// ])
936/// .unwrap();
937///
938/// // Configure sort by (a, b)
939/// let sort_columns = vec![
940///     SortColumn {
941///         values: batch.column(0).clone(),
942///         options: None,
943///     },
944///     SortColumn {
945///         values: batch.column(1).clone(),
946///         options: None,
947///     },
948/// ];
949///
950/// // indices of the rows of (a,b), in lexicographic order
951/// let indices = lexsort_to_indices(&sort_columns, None).unwrap();
952/// assert_eq!(&indices, &UInt32Array::from(vec![2, 1, 0]));
953/// // Create new sorted RecordBatch by copying values at indices
954/// let sorted = take_record_batch(&batch, &indices).unwrap();
955///
956/// assert_eq!(sorted.column(0).as_ref(), &Int32Array::from(vec![1, 1, 2]));
957/// assert_eq!(sorted.column(1).as_ref(), &StringArray::from(vec!["a", "z", "x"]));
958/// ```
959///
960/// Note: for multi-column sorts without a limit, using the [row format](https://docs.rs/arrow-row/latest/arrow_row/)
961/// may be significantly faster
962pub fn lexsort_to_indices(
963    columns: &[SortColumn],
964    limit: Option<usize>,
965) -> Result<UInt32Array, ArrowError> {
966    if columns.is_empty() {
967        return Err(ArrowError::InvalidArgumentError(
968            "Sort requires at least one column".to_string(),
969        ));
970    }
971    if columns.len() == 1 && can_sort_to_indices(columns[0].values.data_type()) {
972        // fallback to non-lexical sort
973        let column = &columns[0];
974        return sort_to_indices(&column.values, column.options, limit);
975    }
976
977    let row_count = columns[0].values.len();
978    if columns.iter().any(|item| item.values.len() != row_count) {
979        return Err(ArrowError::ComputeError(
980            "lexical sort columns have different row counts".to_string(),
981        ));
982    }
983
984    let len = limit.unwrap_or(row_count).min(row_count);
985
986    if len == 0 {
987        return Ok(UInt32Array::from(Vec::<u32>::new()));
988    }
989
990    // The heap path avoids allocating and partially sorting all row indices
991    // when the requested limit is a small fraction of the input. For larger
992    // limits, the existing partial-sort path is preferred because heap
993    // maintenance costs grow with the requested limit.
994    let value_indices = match limit {
995        Some(limit) if limit <= row_count / 10 => match columns.len() {
996            2 => lexsort_topk_fixed::<2>(columns, row_count, len)?,
997            3 => lexsort_topk_fixed::<3>(columns, row_count, len)?,
998            4 => lexsort_topk_fixed::<4>(columns, row_count, len)?,
999            5 => lexsort_topk_fixed::<5>(columns, row_count, len)?,
1000            _ => {
1001                let lexicographical_comparator = LexicographicalComparator::try_new(columns)?;
1002                lexsort_topk(row_count, len, |a, b| {
1003                    lexicographical_comparator.compare(a, b)
1004                })
1005            }
1006        },
1007        _ => {
1008            let mut value_indices = (0..row_count).collect::<Vec<usize>>();
1009
1010            // Instantiate specialized versions of comparisons for small numbers
1011            // of columns as it helps the compiler generate better code.
1012            match columns.len() {
1013                2 => sort_fixed_column::<2>(columns, &mut value_indices, len)?,
1014                3 => sort_fixed_column::<3>(columns, &mut value_indices, len)?,
1015                4 => sort_fixed_column::<4>(columns, &mut value_indices, len)?,
1016                5 => sort_fixed_column::<5>(columns, &mut value_indices, len)?,
1017                _ => {
1018                    let lexicographical_comparator = LexicographicalComparator::try_new(columns)?;
1019                    sort_unstable_by(&mut value_indices, len, |a, b| {
1020                        lexicographical_comparator.compare(*a, *b)
1021                    });
1022                }
1023            }
1024
1025            value_indices.truncate(len);
1026            value_indices
1027        }
1028    };
1029
1030    Ok(UInt32Array::from(
1031        value_indices
1032            .into_iter()
1033            .map(|i| i as u32)
1034            .collect::<Vec<_>>(),
1035    ))
1036}
1037
1038// Sort a fixed number of columns using FixedLexicographicalComparator
1039fn sort_fixed_column<const N: usize>(
1040    columns: &[SortColumn],
1041    value_indices: &mut [usize],
1042    len: usize,
1043) -> Result<(), ArrowError> {
1044    let lexicographical_comparator = FixedLexicographicalComparator::<N>::try_new(columns)?;
1045    sort_unstable_by(value_indices, len, |a, b| {
1046        lexicographical_comparator.compare(*a, *b)
1047    });
1048    Ok(())
1049}
1050
1051// Uses the fixed-column comparator for the bounded heap path.
1052fn lexsort_topk_fixed<const N: usize>(
1053    columns: &[SortColumn],
1054    row_count: usize,
1055    limit: usize,
1056) -> Result<Vec<usize>, ArrowError> {
1057    let lexicographical_comparator = FixedLexicographicalComparator::<N>::try_new(columns)?;
1058    Ok(lexsort_topk(row_count, limit, |a, b| {
1059        lexicographical_comparator.compare(a, b)
1060    }))
1061}
1062
1063// Keeps the smallest `limit` indices in a bounded max-heap.
1064// The root is the largest retained index according to `compare`.
1065fn lexsort_topk(
1066    row_count: usize,
1067    limit: usize,
1068    mut compare: impl FnMut(usize, usize) -> Ordering,
1069) -> Vec<usize> {
1070    let mut heap = Vec::with_capacity(limit);
1071
1072    for idx in 0..row_count {
1073        if heap.len() < limit {
1074            heap.push(idx);
1075            let pos = heap.len() - 1;
1076            sift_up_worst_heap(&mut heap, pos, &mut compare);
1077        } else if compare(idx, heap[0]) == Ordering::Less {
1078            heap[0] = idx;
1079            sift_down_worst_heap(&mut heap, 0, &mut compare);
1080        }
1081    }
1082
1083    heap.sort_unstable_by(|a, b| compare(*a, *b));
1084    heap
1085}
1086
1087// Moves a newly inserted index toward the root while it is larger than its parent.
1088fn sift_up_worst_heap(
1089    heap: &mut [usize],
1090    mut pos: usize,
1091    compare: &mut impl FnMut(usize, usize) -> Ordering,
1092) {
1093    while pos > 0 {
1094        let parent = (pos - 1) / 2;
1095
1096        if compare(heap[parent], heap[pos]) != Ordering::Less {
1097            break;
1098        }
1099
1100        heap.swap(parent, pos);
1101        pos = parent;
1102    }
1103}
1104
1105// Moves the root down until both children are no larger than it.
1106// The larger child is selected at each step so the worst retained row remains
1107// at heap[0].
1108fn sift_down_worst_heap(
1109    heap: &mut [usize],
1110    mut pos: usize,
1111    compare: &mut impl FnMut(usize, usize) -> Ordering,
1112) {
1113    loop {
1114        let left = pos * 2 + 1;
1115        if left >= heap.len() {
1116            break;
1117        }
1118
1119        let right = left + 1;
1120        let worst = if right < heap.len() && compare(heap[left], heap[right]) == Ordering::Less {
1121            right
1122        } else {
1123            left
1124        };
1125
1126        if compare(heap[pos], heap[worst]) != Ordering::Less {
1127            break;
1128        }
1129
1130        heap.swap(pos, worst);
1131        pos = worst;
1132    }
1133}
1134
1135/// It's unstable_sort, may not preserve the order of equal elements
1136pub fn partial_sort<T, F>(v: &mut [T], limit: usize, mut is_less: F)
1137where
1138    F: FnMut(&T, &T) -> Ordering,
1139{
1140    if let Some(n) = limit.checked_sub(1) {
1141        let (before, _mid, _after) = v.select_nth_unstable_by(n, &mut is_less);
1142        before.sort_unstable_by(is_less);
1143    }
1144}
1145
1146/// A lexicographical comparator that wraps given array data (columns) and can lexicographically compare data
1147/// at given two indices. The lifetime is the same at the data wrapped.
1148pub struct LexicographicalComparator {
1149    compare_items: Vec<DynComparator>,
1150}
1151
1152impl LexicographicalComparator {
1153    /// lexicographically compare values at the wrapped columns with given indices.
1154    pub fn compare(&self, a_idx: usize, b_idx: usize) -> Ordering {
1155        for comparator in &self.compare_items {
1156            match comparator(a_idx, b_idx) {
1157                Ordering::Equal => {}
1158                r => return r,
1159            }
1160        }
1161        Ordering::Equal
1162    }
1163
1164    /// Create a new lex comparator that will wrap the given sort columns and give comparison
1165    /// results with two indices.
1166    pub fn try_new(columns: &[SortColumn]) -> Result<LexicographicalComparator, ArrowError> {
1167        let compare_items = columns
1168            .iter()
1169            .map(|c| {
1170                make_comparator(
1171                    c.values.as_ref(),
1172                    c.values.as_ref(),
1173                    c.options.unwrap_or_default(),
1174                )
1175            })
1176            .collect::<Result<Vec<_>, ArrowError>>()?;
1177        Ok(LexicographicalComparator { compare_items })
1178    }
1179}
1180
1181/// A lexicographical comparator that wraps given array data (columns) and can lexicographically compare data
1182/// at given two indices. This version of the comparator is for compile-time constant number of columns.
1183/// The lifetime is the same at the data wrapped.
1184pub struct FixedLexicographicalComparator<const N: usize> {
1185    compare_items: [DynComparator; N],
1186}
1187
1188impl<const N: usize> FixedLexicographicalComparator<N> {
1189    /// lexicographically compare values at the wrapped columns with given indices.
1190    pub fn compare(&self, a_idx: usize, b_idx: usize) -> Ordering {
1191        for comparator in &self.compare_items {
1192            match comparator(a_idx, b_idx) {
1193                Ordering::Equal => {}
1194                r => return r,
1195            }
1196        }
1197        Ordering::Equal
1198    }
1199
1200    /// Create a new lex comparator that will wrap the given sort columns and give comparison
1201    /// results with two indices.
1202    /// The number of columns should be equal to the compile-time constant N.
1203    pub fn try_new(
1204        columns: &[SortColumn],
1205    ) -> Result<FixedLexicographicalComparator<N>, ArrowError> {
1206        let compare_items = columns
1207            .iter()
1208            .map(|c| {
1209                make_comparator(
1210                    c.values.as_ref(),
1211                    c.values.as_ref(),
1212                    c.options.unwrap_or_default(),
1213                )
1214            })
1215            .collect::<Result<Vec<_>, ArrowError>>()?
1216            .try_into();
1217        let compare_items: [Box<dyn Fn(usize, usize) -> Ordering + Send + Sync + 'static>; N] =
1218            compare_items.map_err(|_| {
1219                ArrowError::ComputeError("Could not create fixed size array".to_string())
1220            })?;
1221        Ok(FixedLexicographicalComparator { compare_items })
1222    }
1223}
1224
1225#[cfg(test)]
1226mod tests {
1227    use super::*;
1228    use arrow_array::builder::{
1229        BooleanBuilder, FixedSizeListBuilder, GenericListBuilder, Int64Builder, ListBuilder,
1230        PrimitiveRunBuilder,
1231    };
1232    use arrow_buffer::{NullBuffer, i256};
1233    use arrow_schema::Field;
1234    use half::f16;
1235    use rand::rngs::StdRng;
1236    use rand::seq::SliceRandom;
1237    use rand::{Rng, RngExt, SeedableRng};
1238
1239    fn create_decimal_array<T: DecimalType>(
1240        data: Vec<Option<usize>>,
1241        precision: u8,
1242        scale: i8,
1243    ) -> PrimitiveArray<T> {
1244        data.into_iter()
1245            .map(|x| x.and_then(T::Native::from_usize))
1246            .collect::<PrimitiveArray<T>>()
1247            .with_precision_and_scale(precision, scale)
1248            .unwrap()
1249    }
1250
1251    fn create_decimal256_array(data: Vec<Option<i256>>) -> Decimal256Array {
1252        data.into_iter()
1253            .collect::<Decimal256Array>()
1254            .with_precision_and_scale(53, 6)
1255            .unwrap()
1256    }
1257
1258    fn test_sort_to_indices_decimal_array<T: DecimalType>(
1259        data: Vec<Option<usize>>,
1260        options: Option<SortOptions>,
1261        limit: Option<usize>,
1262        expected_data: Vec<u32>,
1263        precision: u8,
1264        scale: i8,
1265    ) {
1266        let output = create_decimal_array::<T>(data, precision, scale);
1267        let expected = UInt32Array::from(expected_data);
1268        let output = sort_to_indices(&output, options, limit).unwrap();
1269        assert_eq!(output, expected)
1270    }
1271
1272    fn test_sort_to_indices_decimal256_array(
1273        data: Vec<Option<i256>>,
1274        options: Option<SortOptions>,
1275        limit: Option<usize>,
1276        expected_data: Vec<u32>,
1277    ) {
1278        let output = create_decimal256_array(data);
1279        let expected = UInt32Array::from(expected_data);
1280        let output = sort_to_indices(&output, options, limit).unwrap();
1281        assert_eq!(output, expected)
1282    }
1283
1284    fn test_sort_decimal_array<T: DecimalType>(
1285        data: Vec<Option<usize>>,
1286        options: Option<SortOptions>,
1287        limit: Option<usize>,
1288        expected_data: Vec<Option<usize>>,
1289        p: u8,
1290        s: i8,
1291    ) {
1292        let output = create_decimal_array::<T>(data, p, s);
1293        let expected = Arc::new(create_decimal_array::<T>(expected_data, p, s)) as ArrayRef;
1294        let output = match limit {
1295            Some(_) => sort_limit(&output, options, limit).unwrap(),
1296            _ => sort(&output, options).unwrap(),
1297        };
1298        assert_eq!(&output, &expected)
1299    }
1300
1301    fn test_sort_decimal256_array(
1302        data: Vec<Option<i256>>,
1303        options: Option<SortOptions>,
1304        limit: Option<usize>,
1305        expected_data: Vec<Option<i256>>,
1306    ) {
1307        let output = create_decimal256_array(data);
1308        let expected = Arc::new(create_decimal256_array(expected_data)) as ArrayRef;
1309        let output = match limit {
1310            Some(_) => sort_limit(&output, options, limit).unwrap(),
1311            _ => sort(&output, options).unwrap(),
1312        };
1313        assert_eq!(&output, &expected)
1314    }
1315
1316    fn test_sort_to_indices_boolean_arrays(
1317        data: Vec<Option<bool>>,
1318        options: Option<SortOptions>,
1319        limit: Option<usize>,
1320        expected_data: Vec<u32>,
1321    ) {
1322        let output = BooleanArray::from(data);
1323        let expected = UInt32Array::from(expected_data);
1324        let output = sort_to_indices(&output, options, limit).unwrap();
1325        assert_eq!(output, expected)
1326    }
1327
1328    fn test_sort_to_indices_primitive_arrays<T>(
1329        data: Vec<Option<T::Native>>,
1330        options: Option<SortOptions>,
1331        limit: Option<usize>,
1332        expected_data: Vec<u32>,
1333    ) where
1334        T: ArrowPrimitiveType,
1335        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1336    {
1337        let output = PrimitiveArray::<T>::from(data);
1338        let expected = UInt32Array::from(expected_data);
1339        let output = sort_to_indices(&output, options, limit).unwrap();
1340        assert_eq!(output, expected)
1341    }
1342
1343    fn test_sort_primitive_arrays<T>(
1344        data: Vec<Option<T::Native>>,
1345        options: Option<SortOptions>,
1346        limit: Option<usize>,
1347        expected_data: Vec<Option<T::Native>>,
1348    ) where
1349        T: ArrowPrimitiveType,
1350        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1351    {
1352        let output = PrimitiveArray::<T>::from(data);
1353        let expected = Arc::new(PrimitiveArray::<T>::from(expected_data)) as ArrayRef;
1354        let output = match limit {
1355            Some(_) => sort_limit(&output, options, limit).unwrap(),
1356            _ => sort(&output, options).unwrap(),
1357        };
1358        assert_eq!(&output, &expected)
1359    }
1360
1361    fn test_sort_to_indices_string_arrays(
1362        data: Vec<Option<&str>>,
1363        options: Option<SortOptions>,
1364        limit: Option<usize>,
1365        expected_data: Vec<u32>,
1366    ) {
1367        let output = StringArray::from(data);
1368        let expected = UInt32Array::from(expected_data);
1369        let output = sort_to_indices(&output, options, limit).unwrap();
1370        assert_eq!(output, expected)
1371    }
1372
1373    /// Tests both Utf8 and LargeUtf8
1374    fn test_sort_string_arrays(
1375        data: Vec<Option<&str>>,
1376        options: Option<SortOptions>,
1377        limit: Option<usize>,
1378        expected_data: Vec<Option<&str>>,
1379    ) {
1380        let output = StringArray::from(data.clone());
1381        let expected = Arc::new(StringArray::from(expected_data.clone())) as ArrayRef;
1382        let output = match limit {
1383            Some(_) => sort_limit(&output, options, limit).unwrap(),
1384            _ => sort(&output, options).unwrap(),
1385        };
1386        assert_eq!(&output, &expected);
1387
1388        let output = LargeStringArray::from(data.clone());
1389        let expected = Arc::new(LargeStringArray::from(expected_data.clone())) as ArrayRef;
1390        let output = match limit {
1391            Some(_) => sort_limit(&output, options, limit).unwrap(),
1392            _ => sort(&output, options).unwrap(),
1393        };
1394        assert_eq!(&output, &expected);
1395
1396        let output = StringViewArray::from(data);
1397        let expected = Arc::new(StringViewArray::from(expected_data)) as ArrayRef;
1398        let output = match limit {
1399            Some(_) => sort_limit(&output, options, limit).unwrap(),
1400            _ => sort(&output, options).unwrap(),
1401        };
1402        assert_eq!(&output, &expected);
1403    }
1404
1405    fn test_sort_string_dict_arrays<T: ArrowDictionaryKeyType>(
1406        data: Vec<Option<&str>>,
1407        options: Option<SortOptions>,
1408        limit: Option<usize>,
1409        expected_data: Vec<Option<&str>>,
1410    ) {
1411        let array = data.into_iter().collect::<DictionaryArray<T>>();
1412        let array_values = array.values().clone();
1413        let dict = array_values
1414            .as_any()
1415            .downcast_ref::<StringArray>()
1416            .expect("Unable to get dictionary values");
1417
1418        let sorted = match limit {
1419            Some(_) => sort_limit(&array, options, limit).unwrap(),
1420            _ => sort(&array, options).unwrap(),
1421        };
1422        let sorted = sorted
1423            .as_any()
1424            .downcast_ref::<DictionaryArray<T>>()
1425            .unwrap();
1426        let sorted_values = sorted.values();
1427        let sorted_dict = sorted_values
1428            .as_any()
1429            .downcast_ref::<StringArray>()
1430            .expect("Unable to get dictionary values");
1431        let sorted_keys = sorted.keys();
1432
1433        assert_eq!(sorted_dict, dict);
1434
1435        let sorted_strings = StringArray::from_iter((0..sorted.len()).map(|i| {
1436            if sorted.is_valid(i) {
1437                Some(sorted_dict.value(sorted_keys.value(i).as_usize()))
1438            } else {
1439                None
1440            }
1441        }));
1442        let expected = StringArray::from(expected_data);
1443
1444        assert_eq!(sorted_strings, expected)
1445    }
1446
1447    fn test_sort_primitive_dict_arrays<K: ArrowDictionaryKeyType, T: ArrowPrimitiveType>(
1448        keys: PrimitiveArray<K>,
1449        values: PrimitiveArray<T>,
1450        options: Option<SortOptions>,
1451        limit: Option<usize>,
1452        expected_data: Vec<Option<T::Native>>,
1453    ) where
1454        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1455    {
1456        let array = DictionaryArray::<K>::new(keys, Arc::new(values));
1457        let array_values = array.values().clone();
1458        let dict = array_values.as_primitive::<T>();
1459
1460        let sorted = match limit {
1461            Some(_) => sort_limit(&array, options, limit).unwrap(),
1462            _ => sort(&array, options).unwrap(),
1463        };
1464        let sorted = sorted
1465            .as_any()
1466            .downcast_ref::<DictionaryArray<K>>()
1467            .unwrap();
1468        let sorted_values = sorted.values();
1469        let sorted_dict = sorted_values
1470            .as_any()
1471            .downcast_ref::<PrimitiveArray<T>>()
1472            .expect("Unable to get dictionary values");
1473        let sorted_keys = sorted.keys();
1474
1475        assert_eq!(sorted_dict, dict);
1476
1477        let sorted_values: PrimitiveArray<T> = From::<Vec<Option<T::Native>>>::from(
1478            (0..sorted.len())
1479                .map(|i| {
1480                    let key = sorted_keys.value(i).as_usize();
1481                    if sorted.is_valid(i) && sorted_dict.is_valid(key) {
1482                        Some(sorted_dict.value(key))
1483                    } else {
1484                        None
1485                    }
1486                })
1487                .collect::<Vec<Option<T::Native>>>(),
1488        );
1489        let expected: PrimitiveArray<T> = From::<Vec<Option<T::Native>>>::from(expected_data);
1490
1491        assert_eq!(sorted_values, expected)
1492    }
1493
1494    fn test_sort_list_arrays<T>(
1495        data: Vec<Option<Vec<Option<T::Native>>>>,
1496        options: Option<SortOptions>,
1497        limit: Option<usize>,
1498        expected_data: Vec<Option<Vec<Option<T::Native>>>>,
1499        fixed_length: Option<i32>,
1500    ) where
1501        T: ArrowPrimitiveType,
1502        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
1503    {
1504        // for FixedSizedList
1505        if let Some(length) = fixed_length {
1506            let input = Arc::new(FixedSizeListArray::from_iter_primitive::<T, _, _>(
1507                data.clone(),
1508                length,
1509            ));
1510            let sorted = match limit {
1511                Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1512                _ => sort(&(input as ArrayRef), options).unwrap(),
1513            };
1514            let expected = Arc::new(FixedSizeListArray::from_iter_primitive::<T, _, _>(
1515                expected_data.clone(),
1516                length,
1517            )) as ArrayRef;
1518
1519            assert_eq!(&sorted, &expected);
1520        }
1521
1522        // for List
1523        let input = Arc::new(ListArray::from_iter_primitive::<T, _, _>(data.clone()));
1524        let sorted = match limit {
1525            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1526            _ => sort(&(input as ArrayRef), options).unwrap(),
1527        };
1528        let expected = Arc::new(ListArray::from_iter_primitive::<T, _, _>(
1529            expected_data.clone(),
1530        )) as ArrayRef;
1531
1532        assert_eq!(&sorted, &expected);
1533
1534        // for ListView
1535        let input = Arc::new(ListViewArray::from_iter_primitive::<T, _, _>(data.clone()));
1536        let sorted = match limit {
1537            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1538            _ => sort(&(input as ArrayRef), options).unwrap(),
1539        };
1540        let expected = Arc::new(ListViewArray::from_iter_primitive::<T, _, _>(
1541            expected_data.clone(),
1542        )) as ArrayRef;
1543        assert_eq!(&sorted, &expected);
1544
1545        // for LargeList
1546        let input = Arc::new(LargeListArray::from_iter_primitive::<T, _, _>(data.clone()));
1547        let sorted = match limit {
1548            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1549            _ => sort(&(input as ArrayRef), options).unwrap(),
1550        };
1551        let expected = Arc::new(LargeListArray::from_iter_primitive::<T, _, _>(
1552            expected_data.clone(),
1553        )) as ArrayRef;
1554        assert_eq!(&sorted, &expected);
1555
1556        // for LargeListView
1557        let input = Arc::new(LargeListViewArray::from_iter_primitive::<T, _, _>(data));
1558        let sorted = match limit {
1559            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1560            _ => sort(&(input as ArrayRef), options).unwrap(),
1561        };
1562        let expected = Arc::new(LargeListViewArray::from_iter_primitive::<T, _, _>(
1563            expected_data,
1564        )) as ArrayRef;
1565        assert_eq!(&sorted, &expected);
1566    }
1567
1568    fn test_lex_sort_arrays(
1569        input: Vec<SortColumn>,
1570        expected_output: Vec<ArrayRef>,
1571        limit: Option<usize>,
1572    ) {
1573        let sorted = lexsort(&input, limit).unwrap();
1574
1575        for (result, expected) in sorted.iter().zip(expected_output.iter()) {
1576            assert_eq!(result, expected);
1577        }
1578    }
1579
1580    /// slice all arrays in expected_output to offset/length
1581    fn slice_arrays(expected_output: Vec<ArrayRef>, offset: usize, length: usize) -> Vec<ArrayRef> {
1582        expected_output
1583            .into_iter()
1584            .map(|array| array.slice(offset, length))
1585            .collect()
1586    }
1587
1588    fn test_sort_binary_arrays(
1589        data: Vec<Option<Vec<u8>>>,
1590        options: Option<SortOptions>,
1591        limit: Option<usize>,
1592        expected_data: Vec<Option<Vec<u8>>>,
1593        fixed_length: Option<i32>,
1594    ) {
1595        // Fixed size binary array
1596        if let Some(length) = fixed_length {
1597            let input = Arc::new(
1598                FixedSizeBinaryArray::try_from_sparse_iter_with_size(data.iter().cloned(), length)
1599                    .unwrap(),
1600            );
1601            let sorted = match limit {
1602                Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1603                None => sort(&(input as ArrayRef), options).unwrap(),
1604            };
1605            let expected = Arc::new(
1606                FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1607                    expected_data.iter().cloned(),
1608                    length,
1609                )
1610                .unwrap(),
1611            ) as ArrayRef;
1612
1613            assert_eq!(&sorted, &expected);
1614        }
1615
1616        // Generic size binary array
1617        fn make_generic_binary_array<S: OffsetSizeTrait>(
1618            data: &[Option<Vec<u8>>],
1619        ) -> Arc<GenericBinaryArray<S>> {
1620            Arc::new(GenericBinaryArray::<S>::from_opt_vec(
1621                data.iter()
1622                    .map(|binary| binary.as_ref().map(Vec::as_slice))
1623                    .collect(),
1624            ))
1625        }
1626
1627        // BinaryArray
1628        let input = make_generic_binary_array::<i32>(&data);
1629        let sorted = match limit {
1630            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1631            None => sort(&(input as ArrayRef), options).unwrap(),
1632        };
1633        let expected = make_generic_binary_array::<i32>(&expected_data) as ArrayRef;
1634        assert_eq!(&sorted, &expected);
1635
1636        // LargeBinaryArray
1637        let input = make_generic_binary_array::<i64>(&data);
1638        let sorted = match limit {
1639            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
1640            None => sort(&(input as ArrayRef), options).unwrap(),
1641        };
1642        let expected = make_generic_binary_array::<i64>(&expected_data) as ArrayRef;
1643        assert_eq!(&sorted, &expected);
1644    }
1645
1646    #[test]
1647    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
1648    fn test_sort_to_indices_primitives() {
1649        test_sort_to_indices_primitive_arrays::<Int8Type>(
1650            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1651            None,
1652            None,
1653            vec![0, 5, 3, 1, 4, 2],
1654        );
1655        test_sort_to_indices_primitive_arrays::<Int16Type>(
1656            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1657            None,
1658            None,
1659            vec![0, 5, 3, 1, 4, 2],
1660        );
1661        test_sort_to_indices_primitive_arrays::<Int32Type>(
1662            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1663            None,
1664            None,
1665            vec![0, 5, 3, 1, 4, 2],
1666        );
1667        test_sort_to_indices_primitive_arrays::<Int64Type>(
1668            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1669            None,
1670            None,
1671            vec![0, 5, 3, 1, 4, 2],
1672        );
1673        test_sort_to_indices_primitive_arrays::<Float16Type>(
1674            vec![
1675                None,
1676                Some(f16::from_f32(-0.05)),
1677                Some(f16::from_f32(2.225)),
1678                Some(f16::from_f32(-1.01)),
1679                Some(f16::from_f32(-0.05)),
1680                None,
1681            ],
1682            None,
1683            None,
1684            vec![0, 5, 3, 1, 4, 2],
1685        );
1686        test_sort_to_indices_primitive_arrays::<Float32Type>(
1687            vec![
1688                None,
1689                Some(-0.05),
1690                Some(2.225),
1691                Some(-1.01),
1692                Some(-0.05),
1693                None,
1694            ],
1695            None,
1696            None,
1697            vec![0, 5, 3, 1, 4, 2],
1698        );
1699        test_sort_to_indices_primitive_arrays::<Float64Type>(
1700            vec![
1701                None,
1702                Some(-0.05),
1703                Some(2.225),
1704                Some(-1.01),
1705                Some(-0.05),
1706                None,
1707            ],
1708            None,
1709            None,
1710            vec![0, 5, 3, 1, 4, 2],
1711        );
1712
1713        // descending
1714        test_sort_to_indices_primitive_arrays::<Int8Type>(
1715            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1716            Some(SortOptions {
1717                descending: true,
1718                nulls_first: false,
1719            }),
1720            None,
1721            vec![2, 1, 4, 3, 0, 5],
1722        );
1723
1724        test_sort_to_indices_primitive_arrays::<Int16Type>(
1725            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1726            Some(SortOptions {
1727                descending: true,
1728                nulls_first: false,
1729            }),
1730            None,
1731            vec![2, 1, 4, 3, 0, 5],
1732        );
1733
1734        test_sort_to_indices_primitive_arrays::<Int32Type>(
1735            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1736            Some(SortOptions {
1737                descending: true,
1738                nulls_first: false,
1739            }),
1740            None,
1741            vec![2, 1, 4, 3, 0, 5],
1742        );
1743
1744        test_sort_to_indices_primitive_arrays::<Int64Type>(
1745            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1746            Some(SortOptions {
1747                descending: true,
1748                nulls_first: false,
1749            }),
1750            None,
1751            vec![2, 1, 4, 3, 0, 5],
1752        );
1753
1754        test_sort_to_indices_primitive_arrays::<Float16Type>(
1755            vec![
1756                None,
1757                Some(f16::from_f32(0.005)),
1758                Some(f16::from_f32(20.22)),
1759                Some(f16::from_f32(-10.3)),
1760                Some(f16::from_f32(0.005)),
1761                None,
1762            ],
1763            Some(SortOptions {
1764                descending: true,
1765                nulls_first: false,
1766            }),
1767            None,
1768            vec![2, 1, 4, 3, 0, 5],
1769        );
1770
1771        test_sort_to_indices_primitive_arrays::<Float32Type>(
1772            vec![
1773                None,
1774                Some(0.005),
1775                Some(20.22),
1776                Some(-10.3),
1777                Some(0.005),
1778                None,
1779            ],
1780            Some(SortOptions {
1781                descending: true,
1782                nulls_first: false,
1783            }),
1784            None,
1785            vec![2, 1, 4, 3, 0, 5],
1786        );
1787
1788        test_sort_to_indices_primitive_arrays::<Float64Type>(
1789            vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(0.0), None],
1790            Some(SortOptions {
1791                descending: true,
1792                nulls_first: false,
1793            }),
1794            None,
1795            vec![2, 1, 4, 3, 0, 5],
1796        );
1797
1798        // descending, nulls first
1799        test_sort_to_indices_primitive_arrays::<Int8Type>(
1800            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1801            Some(SortOptions {
1802                descending: true,
1803                nulls_first: true,
1804            }),
1805            None,
1806            vec![0, 5, 2, 1, 4, 3], // [5, 0, 2, 4, 1, 3]
1807        );
1808
1809        test_sort_to_indices_primitive_arrays::<Int16Type>(
1810            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1811            Some(SortOptions {
1812                descending: true,
1813                nulls_first: true,
1814            }),
1815            None,
1816            vec![0, 5, 2, 1, 4, 3], // [5, 0, 2, 4, 1, 3]
1817        );
1818
1819        test_sort_to_indices_primitive_arrays::<Int32Type>(
1820            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1821            Some(SortOptions {
1822                descending: true,
1823                nulls_first: true,
1824            }),
1825            None,
1826            vec![0, 5, 2, 1, 4, 3],
1827        );
1828
1829        test_sort_to_indices_primitive_arrays::<Int64Type>(
1830            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
1831            Some(SortOptions {
1832                descending: true,
1833                nulls_first: true,
1834            }),
1835            None,
1836            vec![0, 5, 2, 1, 4, 3],
1837        );
1838
1839        test_sort_to_indices_primitive_arrays::<Float16Type>(
1840            vec![
1841                None,
1842                Some(f16::from_f32(0.1)),
1843                Some(f16::from_f32(0.2)),
1844                Some(f16::from_f32(-1.3)),
1845                Some(f16::from_f32(0.01)),
1846                None,
1847            ],
1848            Some(SortOptions {
1849                descending: true,
1850                nulls_first: true,
1851            }),
1852            None,
1853            vec![0, 5, 2, 1, 4, 3],
1854        );
1855
1856        test_sort_to_indices_primitive_arrays::<Float32Type>(
1857            vec![None, Some(0.1), Some(0.2), Some(-1.3), Some(0.01), None],
1858            Some(SortOptions {
1859                descending: true,
1860                nulls_first: true,
1861            }),
1862            None,
1863            vec![0, 5, 2, 1, 4, 3],
1864        );
1865
1866        test_sort_to_indices_primitive_arrays::<Float64Type>(
1867            vec![None, Some(10.1), Some(100.2), Some(-1.3), Some(10.01), None],
1868            Some(SortOptions {
1869                descending: true,
1870                nulls_first: true,
1871            }),
1872            None,
1873            vec![0, 5, 2, 1, 4, 3],
1874        );
1875
1876        // valid values less than limit with extra nulls
1877        test_sort_to_indices_primitive_arrays::<Float64Type>(
1878            vec![Some(2.0), None, None, Some(1.0)],
1879            Some(SortOptions {
1880                descending: false,
1881                nulls_first: false,
1882            }),
1883            Some(3),
1884            vec![3, 0, 1],
1885        );
1886
1887        test_sort_to_indices_primitive_arrays::<Float64Type>(
1888            vec![Some(2.0), None, None, Some(1.0)],
1889            Some(SortOptions {
1890                descending: false,
1891                nulls_first: true,
1892            }),
1893            Some(3),
1894            vec![1, 2, 3],
1895        );
1896
1897        // more nulls than limit
1898        test_sort_to_indices_primitive_arrays::<Float64Type>(
1899            vec![Some(1.0), None, None, None],
1900            Some(SortOptions {
1901                descending: false,
1902                nulls_first: true,
1903            }),
1904            Some(2),
1905            vec![1, 2],
1906        );
1907
1908        test_sort_to_indices_primitive_arrays::<Float64Type>(
1909            vec![Some(1.0), None, None, None],
1910            Some(SortOptions {
1911                descending: false,
1912                nulls_first: false,
1913            }),
1914            Some(2),
1915            vec![0, 1],
1916        );
1917    }
1918
1919    #[test]
1920    fn test_sort_to_indices_primitive_more_nulls_than_limit() {
1921        test_sort_to_indices_primitive_arrays::<Int32Type>(
1922            vec![None, None, Some(3), None, Some(1), None, Some(2)],
1923            Some(SortOptions {
1924                descending: false,
1925                nulls_first: false,
1926            }),
1927            Some(2),
1928            vec![4, 6],
1929        );
1930    }
1931
1932    #[test]
1933    fn test_sort_boolean() {
1934        // boolean
1935        test_sort_to_indices_boolean_arrays(
1936            vec![None, Some(false), Some(true), Some(true), Some(false), None],
1937            None,
1938            None,
1939            vec![0, 5, 1, 4, 2, 3],
1940        );
1941
1942        // boolean, descending
1943        test_sort_to_indices_boolean_arrays(
1944            vec![None, Some(false), Some(true), Some(true), Some(false), None],
1945            Some(SortOptions {
1946                descending: true,
1947                nulls_first: false,
1948            }),
1949            None,
1950            vec![2, 3, 1, 4, 0, 5],
1951        );
1952
1953        // boolean, descending, nulls first
1954        test_sort_to_indices_boolean_arrays(
1955            vec![None, Some(false), Some(true), Some(true), Some(false), None],
1956            Some(SortOptions {
1957                descending: true,
1958                nulls_first: true,
1959            }),
1960            None,
1961            vec![0, 5, 2, 3, 1, 4],
1962        );
1963
1964        // boolean, descending, nulls first, limit
1965        test_sort_to_indices_boolean_arrays(
1966            vec![None, Some(false), Some(true), Some(true), Some(false), None],
1967            Some(SortOptions {
1968                descending: true,
1969                nulls_first: true,
1970            }),
1971            Some(3),
1972            vec![0, 5, 2],
1973        );
1974
1975        // valid values less than limit with extra nulls
1976        test_sort_to_indices_boolean_arrays(
1977            vec![Some(true), None, None, Some(false)],
1978            Some(SortOptions {
1979                descending: false,
1980                nulls_first: false,
1981            }),
1982            Some(3),
1983            vec![3, 0, 1],
1984        );
1985
1986        test_sort_to_indices_boolean_arrays(
1987            vec![Some(true), None, None, Some(false)],
1988            Some(SortOptions {
1989                descending: false,
1990                nulls_first: true,
1991            }),
1992            Some(3),
1993            vec![1, 2, 3],
1994        );
1995
1996        // more nulls than limit
1997        test_sort_to_indices_boolean_arrays(
1998            vec![Some(true), None, None, None],
1999            Some(SortOptions {
2000                descending: false,
2001                nulls_first: true,
2002            }),
2003            Some(2),
2004            vec![1, 2],
2005        );
2006
2007        test_sort_to_indices_boolean_arrays(
2008            vec![Some(true), None, None, None],
2009            Some(SortOptions {
2010                descending: false,
2011                nulls_first: false,
2012            }),
2013            Some(2),
2014            vec![0, 1],
2015        );
2016    }
2017
2018    /// Test sort boolean on each permutation of with/without limit and GenericListArray/FixedSizeListArray
2019    ///
2020    /// The input data must have the same length for all list items so that we can test FixedSizeListArray
2021    ///
2022    fn test_every_config_sort_boolean_list_arrays(
2023        data: Vec<Option<Vec<Option<bool>>>>,
2024        options: Option<SortOptions>,
2025        expected_data: Vec<Option<Vec<Option<bool>>>>,
2026    ) {
2027        let first_length = data
2028            .iter()
2029            .find_map(|x| x.as_ref().map(|x| x.len()))
2030            .unwrap_or(0);
2031        let first_non_match_length = data
2032            .iter()
2033            .map(|x| x.as_ref().map(|x| x.len()).unwrap_or(first_length))
2034            .position(|x| x != first_length);
2035
2036        assert_eq!(
2037            first_non_match_length, None,
2038            "All list items should have the same length {first_length}, input data is invalid"
2039        );
2040
2041        let first_non_match_length = expected_data
2042            .iter()
2043            .map(|x| x.as_ref().map(|x| x.len()).unwrap_or(first_length))
2044            .position(|x| x != first_length);
2045
2046        assert_eq!(
2047            first_non_match_length, None,
2048            "All list items should have the same length {first_length}, expected data is invalid"
2049        );
2050
2051        let limit = expected_data.len().saturating_div(2);
2052
2053        for &with_limit in &[false, true] {
2054            let (limit, expected_data) = if with_limit {
2055                (
2056                    Some(limit),
2057                    expected_data.iter().take(limit).cloned().collect(),
2058                )
2059            } else {
2060                (None, expected_data.clone())
2061            };
2062
2063            for &fixed_length in &[None, Some(first_length as i32)] {
2064                test_sort_boolean_list_arrays(
2065                    data.clone(),
2066                    options,
2067                    limit,
2068                    expected_data.clone(),
2069                    fixed_length,
2070                );
2071            }
2072        }
2073    }
2074
2075    fn test_sort_boolean_list_arrays(
2076        data: Vec<Option<Vec<Option<bool>>>>,
2077        options: Option<SortOptions>,
2078        limit: Option<usize>,
2079        expected_data: Vec<Option<Vec<Option<bool>>>>,
2080        fixed_length: Option<i32>,
2081    ) {
2082        fn build_fixed_boolean_list_array(
2083            data: Vec<Option<Vec<Option<bool>>>>,
2084            fixed_length: i32,
2085        ) -> ArrayRef {
2086            let mut builder = FixedSizeListBuilder::new(
2087                BooleanBuilder::with_capacity(fixed_length as usize),
2088                fixed_length,
2089            );
2090            for sublist in data {
2091                match sublist {
2092                    Some(sublist) => {
2093                        builder.values().extend(sublist);
2094                        builder.append(true);
2095                    }
2096                    None => {
2097                        builder
2098                            .values()
2099                            .extend(std::iter::repeat_n(None, fixed_length as usize));
2100                        builder.append(false);
2101                    }
2102                }
2103            }
2104            Arc::new(builder.finish()) as ArrayRef
2105        }
2106
2107        fn build_generic_boolean_list_array<OffsetSize: OffsetSizeTrait>(
2108            data: Vec<Option<Vec<Option<bool>>>>,
2109        ) -> ArrayRef {
2110            let mut builder = GenericListBuilder::<OffsetSize, _>::new(BooleanBuilder::new());
2111            builder.extend(data);
2112            Arc::new(builder.finish()) as ArrayRef
2113        }
2114
2115        // for FixedSizedList
2116        if let Some(length) = fixed_length {
2117            let input = build_fixed_boolean_list_array(data.clone(), length);
2118            let sorted = match limit {
2119                Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
2120                _ => sort(&(input as ArrayRef), options).unwrap(),
2121            };
2122            let expected = build_fixed_boolean_list_array(expected_data.clone(), length);
2123
2124            assert_eq!(&sorted, &expected);
2125        }
2126
2127        // for List
2128        let input = build_generic_boolean_list_array::<i32>(data.clone());
2129        let sorted = match limit {
2130            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
2131            _ => sort(&(input as ArrayRef), options).unwrap(),
2132        };
2133        let expected = build_generic_boolean_list_array::<i32>(expected_data.clone());
2134
2135        assert_eq!(&sorted, &expected);
2136
2137        // for LargeList
2138        let input = build_generic_boolean_list_array::<i64>(data.clone());
2139        let sorted = match limit {
2140            Some(_) => sort_limit(&(input as ArrayRef), options, limit).unwrap(),
2141            _ => sort(&(input as ArrayRef), options).unwrap(),
2142        };
2143        let expected = build_generic_boolean_list_array::<i64>(expected_data.clone());
2144
2145        assert_eq!(&sorted, &expected);
2146    }
2147
2148    #[test]
2149    fn test_sort_list_of_booleans() {
2150        // These are all the possible combinations of boolean values
2151        // There are 3^3 + 1 = 28 possible combinations (3 values to permutate - [true, false, null] and 1 None value)
2152        #[rustfmt::skip]
2153        let mut cases = vec![
2154            Some(vec![Some(true),  Some(true),  Some(true)]),
2155            Some(vec![Some(true),  Some(true),  Some(false)]),
2156            Some(vec![Some(true),  Some(true),  None]),
2157
2158            Some(vec![Some(true),  Some(false), Some(true)]),
2159            Some(vec![Some(true),  Some(false), Some(false)]),
2160            Some(vec![Some(true),  Some(false), None]),
2161
2162            Some(vec![Some(true),  None,        Some(true)]),
2163            Some(vec![Some(true),  None,        Some(false)]),
2164            Some(vec![Some(true),  None,        None]),
2165
2166            Some(vec![Some(false), Some(true),  Some(true)]),
2167            Some(vec![Some(false), Some(true),  Some(false)]),
2168            Some(vec![Some(false), Some(true),  None]),
2169
2170            Some(vec![Some(false), Some(false), Some(true)]),
2171            Some(vec![Some(false), Some(false), Some(false)]),
2172            Some(vec![Some(false), Some(false), None]),
2173
2174            Some(vec![Some(false), None,        Some(true)]),
2175            Some(vec![Some(false), None,        Some(false)]),
2176            Some(vec![Some(false), None,        None]),
2177
2178            Some(vec![None,        Some(true),  Some(true)]),
2179            Some(vec![None,        Some(true),  Some(false)]),
2180            Some(vec![None,        Some(true),  None]),
2181
2182            Some(vec![None,        Some(false), Some(true)]),
2183            Some(vec![None,        Some(false), Some(false)]),
2184            Some(vec![None,        Some(false), None]),
2185
2186            Some(vec![None,        None,        Some(true)]),
2187            Some(vec![None,        None,        Some(false)]),
2188            Some(vec![None,        None,        None]),
2189            None,
2190        ];
2191
2192        cases.shuffle(&mut StdRng::seed_from_u64(42));
2193
2194        // The order is false, true, null
2195        #[rustfmt::skip]
2196        let expected_descending_false_nulls_first_false = vec![
2197            Some(vec![Some(false), Some(false), Some(false)]),
2198            Some(vec![Some(false), Some(false), Some(true)]),
2199            Some(vec![Some(false), Some(false), None]),
2200
2201            Some(vec![Some(false), Some(true),  Some(false)]),
2202            Some(vec![Some(false), Some(true),  Some(true)]),
2203            Some(vec![Some(false), Some(true),  None]),
2204
2205            Some(vec![Some(false), None,        Some(false)]),
2206            Some(vec![Some(false), None,        Some(true)]),
2207            Some(vec![Some(false), None,        None]),
2208
2209            Some(vec![Some(true),  Some(false), Some(false)]),
2210            Some(vec![Some(true),  Some(false), Some(true)]),
2211            Some(vec![Some(true),  Some(false), None]),
2212
2213            Some(vec![Some(true),  Some(true),  Some(false)]),
2214            Some(vec![Some(true),  Some(true),  Some(true)]),
2215            Some(vec![Some(true),  Some(true),  None]),
2216
2217            Some(vec![Some(true),  None,        Some(false)]),
2218            Some(vec![Some(true),  None,        Some(true)]),
2219            Some(vec![Some(true),  None,        None]),
2220
2221            Some(vec![None,        Some(false), Some(false)]),
2222            Some(vec![None,        Some(false), Some(true)]),
2223            Some(vec![None,        Some(false), None]),
2224
2225            Some(vec![None,        Some(true),  Some(false)]),
2226            Some(vec![None,        Some(true),  Some(true)]),
2227            Some(vec![None,        Some(true),  None]),
2228
2229            Some(vec![None,        None,        Some(false)]),
2230            Some(vec![None,        None,        Some(true)]),
2231            Some(vec![None,        None,        None]),
2232            None,
2233        ];
2234        test_every_config_sort_boolean_list_arrays(
2235            cases.clone(),
2236            Some(SortOptions {
2237                descending: false,
2238                nulls_first: false,
2239            }),
2240            expected_descending_false_nulls_first_false,
2241        );
2242
2243        // The order is null, false, true
2244        #[rustfmt::skip]
2245        let expected_descending_false_nulls_first_true = vec![
2246            None,
2247
2248            Some(vec![None,        None,        None]),
2249            Some(vec![None,        None,        Some(false)]),
2250            Some(vec![None,        None,        Some(true)]),
2251
2252            Some(vec![None,        Some(false), None]),
2253            Some(vec![None,        Some(false), Some(false)]),
2254            Some(vec![None,        Some(false), Some(true)]),
2255
2256            Some(vec![None,        Some(true),  None]),
2257            Some(vec![None,        Some(true),  Some(false)]),
2258            Some(vec![None,        Some(true),  Some(true)]),
2259
2260            Some(vec![Some(false), None,        None]),
2261            Some(vec![Some(false), None,        Some(false)]),
2262            Some(vec![Some(false), None,        Some(true)]),
2263
2264            Some(vec![Some(false), Some(false), None]),
2265            Some(vec![Some(false), Some(false), Some(false)]),
2266            Some(vec![Some(false), Some(false), Some(true)]),
2267
2268            Some(vec![Some(false), Some(true),  None]),
2269            Some(vec![Some(false), Some(true),  Some(false)]),
2270            Some(vec![Some(false), Some(true),  Some(true)]),
2271
2272            Some(vec![Some(true),  None,        None]),
2273            Some(vec![Some(true),  None,        Some(false)]),
2274            Some(vec![Some(true),  None,        Some(true)]),
2275
2276            Some(vec![Some(true),  Some(false), None]),
2277            Some(vec![Some(true),  Some(false), Some(false)]),
2278            Some(vec![Some(true),  Some(false), Some(true)]),
2279
2280            Some(vec![Some(true),  Some(true),  None]),
2281            Some(vec![Some(true),  Some(true),  Some(false)]),
2282            Some(vec![Some(true),  Some(true),  Some(true)]),
2283        ];
2284
2285        test_every_config_sort_boolean_list_arrays(
2286            cases.clone(),
2287            Some(SortOptions {
2288                descending: false,
2289                nulls_first: true,
2290            }),
2291            expected_descending_false_nulls_first_true,
2292        );
2293
2294        // The order is true, false, null
2295        #[rustfmt::skip]
2296        let expected_descending_true_nulls_first_false = vec![
2297            Some(vec![Some(true),  Some(true),  Some(true)]),
2298            Some(vec![Some(true),  Some(true),  Some(false)]),
2299            Some(vec![Some(true),  Some(true),  None]),
2300
2301            Some(vec![Some(true),  Some(false), Some(true)]),
2302            Some(vec![Some(true),  Some(false), Some(false)]),
2303            Some(vec![Some(true),  Some(false), None]),
2304
2305            Some(vec![Some(true),  None,        Some(true)]),
2306            Some(vec![Some(true),  None,        Some(false)]),
2307            Some(vec![Some(true),  None,        None]),
2308
2309            Some(vec![Some(false), Some(true),  Some(true)]),
2310            Some(vec![Some(false), Some(true),  Some(false)]),
2311            Some(vec![Some(false), Some(true),  None]),
2312
2313            Some(vec![Some(false), Some(false), Some(true)]),
2314            Some(vec![Some(false), Some(false), Some(false)]),
2315            Some(vec![Some(false), Some(false), None]),
2316
2317            Some(vec![Some(false), None,        Some(true)]),
2318            Some(vec![Some(false), None,        Some(false)]),
2319            Some(vec![Some(false), None,        None]),
2320
2321            Some(vec![None,        Some(true),  Some(true)]),
2322            Some(vec![None,        Some(true),  Some(false)]),
2323            Some(vec![None,        Some(true),  None]),
2324
2325            Some(vec![None,        Some(false), Some(true)]),
2326            Some(vec![None,        Some(false), Some(false)]),
2327            Some(vec![None,        Some(false), None]),
2328
2329            Some(vec![None,        None,        Some(true)]),
2330            Some(vec![None,        None,        Some(false)]),
2331            Some(vec![None,        None,        None]),
2332
2333            None,
2334        ];
2335        test_every_config_sort_boolean_list_arrays(
2336            cases.clone(),
2337            Some(SortOptions {
2338                descending: true,
2339                nulls_first: false,
2340            }),
2341            expected_descending_true_nulls_first_false,
2342        );
2343
2344        // The order is null, true, false
2345        #[rustfmt::skip]
2346        let expected_descending_true_nulls_first_true = vec![
2347            None,
2348
2349            Some(vec![None,        None,        None]),
2350            Some(vec![None,        None,        Some(true)]),
2351            Some(vec![None,        None,        Some(false)]),
2352
2353            Some(vec![None,        Some(true),  None]),
2354            Some(vec![None,        Some(true),  Some(true)]),
2355            Some(vec![None,        Some(true),  Some(false)]),
2356
2357            Some(vec![None,        Some(false), None]),
2358            Some(vec![None,        Some(false), Some(true)]),
2359            Some(vec![None,        Some(false), Some(false)]),
2360
2361            Some(vec![Some(true),  None,        None]),
2362            Some(vec![Some(true),  None,        Some(true)]),
2363            Some(vec![Some(true),  None,        Some(false)]),
2364
2365            Some(vec![Some(true),  Some(true),  None]),
2366            Some(vec![Some(true),  Some(true),  Some(true)]),
2367            Some(vec![Some(true),  Some(true),  Some(false)]),
2368
2369            Some(vec![Some(true),  Some(false), None]),
2370            Some(vec![Some(true),  Some(false), Some(true)]),
2371            Some(vec![Some(true),  Some(false), Some(false)]),
2372
2373            Some(vec![Some(false), None,        None]),
2374            Some(vec![Some(false), None,        Some(true)]),
2375            Some(vec![Some(false), None,        Some(false)]),
2376
2377            Some(vec![Some(false), Some(true),  None]),
2378            Some(vec![Some(false), Some(true),  Some(true)]),
2379            Some(vec![Some(false), Some(true),  Some(false)]),
2380
2381            Some(vec![Some(false), Some(false), None]),
2382            Some(vec![Some(false), Some(false), Some(true)]),
2383            Some(vec![Some(false), Some(false), Some(false)]),
2384        ];
2385        // Testing with limit false and fixed_length None
2386        test_every_config_sort_boolean_list_arrays(
2387            cases.clone(),
2388            Some(SortOptions {
2389                descending: true,
2390                nulls_first: true,
2391            }),
2392            expected_descending_true_nulls_first_true,
2393        );
2394    }
2395
2396    fn test_sort_indices_decimal<T: DecimalType>(precision: u8, scale: i8) {
2397        // decimal default
2398        test_sort_to_indices_decimal_array::<T>(
2399            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2400            None,
2401            None,
2402            vec![0, 6, 4, 2, 3, 5, 1],
2403            precision,
2404            scale,
2405        );
2406        // decimal descending
2407        test_sort_to_indices_decimal_array::<T>(
2408            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2409            Some(SortOptions {
2410                descending: true,
2411                nulls_first: false,
2412            }),
2413            None,
2414            vec![1, 5, 3, 2, 4, 0, 6],
2415            precision,
2416            scale,
2417        );
2418        // decimal null_first and descending
2419        test_sort_to_indices_decimal_array::<T>(
2420            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2421            Some(SortOptions {
2422                descending: true,
2423                nulls_first: true,
2424            }),
2425            None,
2426            vec![0, 6, 1, 5, 3, 2, 4],
2427            precision,
2428            scale,
2429        );
2430        // decimal null_first
2431        test_sort_to_indices_decimal_array::<T>(
2432            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2433            Some(SortOptions {
2434                descending: false,
2435                nulls_first: true,
2436            }),
2437            None,
2438            vec![0, 6, 4, 2, 3, 5, 1],
2439            precision,
2440            scale,
2441        );
2442        // limit
2443        test_sort_to_indices_decimal_array::<T>(
2444            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2445            None,
2446            Some(3),
2447            vec![0, 6, 4],
2448            precision,
2449            scale,
2450        );
2451        // limit descending
2452        test_sort_to_indices_decimal_array::<T>(
2453            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2454            Some(SortOptions {
2455                descending: true,
2456                nulls_first: false,
2457            }),
2458            Some(3),
2459            vec![1, 5, 3],
2460            precision,
2461            scale,
2462        );
2463        // limit descending null_first
2464        test_sort_to_indices_decimal_array::<T>(
2465            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2466            Some(SortOptions {
2467                descending: true,
2468                nulls_first: true,
2469            }),
2470            Some(3),
2471            vec![0, 6, 1],
2472            precision,
2473            scale,
2474        );
2475        // limit null_first
2476        test_sort_to_indices_decimal_array::<T>(
2477            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2478            Some(SortOptions {
2479                descending: false,
2480                nulls_first: true,
2481            }),
2482            Some(3),
2483            vec![0, 6, 4],
2484            precision,
2485            scale,
2486        );
2487    }
2488
2489    #[test]
2490    fn test_sort_indices_decimal32() {
2491        test_sort_indices_decimal::<Decimal32Type>(8, 3);
2492    }
2493
2494    #[test]
2495    fn test_sort_indices_decimal64() {
2496        test_sort_indices_decimal::<Decimal64Type>(17, 5);
2497    }
2498
2499    #[test]
2500    fn test_sort_indices_decimal128() {
2501        test_sort_indices_decimal::<Decimal128Type>(23, 6);
2502    }
2503
2504    #[test]
2505    fn test_sort_indices_decimal256() {
2506        test_sort_indices_decimal::<Decimal256Type>(53, 6);
2507    }
2508
2509    #[test]
2510    fn test_sort_indices_decimal256_max_min() {
2511        let data = vec![
2512            None,
2513            Some(i256::MIN),
2514            Some(i256::from_i128(1)),
2515            Some(i256::MAX),
2516            Some(i256::from_i128(-1)),
2517        ];
2518        test_sort_to_indices_decimal256_array(
2519            data.clone(),
2520            Some(SortOptions {
2521                descending: false,
2522                nulls_first: true,
2523            }),
2524            None,
2525            vec![0, 1, 4, 2, 3],
2526        );
2527
2528        test_sort_to_indices_decimal256_array(
2529            data.clone(),
2530            Some(SortOptions {
2531                descending: true,
2532                nulls_first: true,
2533            }),
2534            None,
2535            vec![0, 3, 2, 4, 1],
2536        );
2537
2538        test_sort_to_indices_decimal256_array(
2539            data.clone(),
2540            Some(SortOptions {
2541                descending: false,
2542                nulls_first: true,
2543            }),
2544            Some(4),
2545            vec![0, 1, 4, 2],
2546        );
2547
2548        test_sort_to_indices_decimal256_array(
2549            data.clone(),
2550            Some(SortOptions {
2551                descending: true,
2552                nulls_first: true,
2553            }),
2554            Some(4),
2555            vec![0, 3, 2, 4],
2556        );
2557    }
2558
2559    fn test_sort_decimal<T: DecimalType>(precision: u8, scale: i8) {
2560        // decimal default
2561        test_sort_decimal_array::<T>(
2562            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2563            None,
2564            None,
2565            vec![None, None, Some(1), Some(2), Some(3), Some(4), Some(5)],
2566            precision,
2567            scale,
2568        );
2569        // decimal descending
2570        test_sort_decimal_array::<T>(
2571            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2572            Some(SortOptions {
2573                descending: true,
2574                nulls_first: false,
2575            }),
2576            None,
2577            vec![Some(5), Some(4), Some(3), Some(2), Some(1), None, None],
2578            precision,
2579            scale,
2580        );
2581        // decimal null_first and descending
2582        test_sort_decimal_array::<T>(
2583            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2584            Some(SortOptions {
2585                descending: true,
2586                nulls_first: true,
2587            }),
2588            None,
2589            vec![None, None, Some(5), Some(4), Some(3), Some(2), Some(1)],
2590            precision,
2591            scale,
2592        );
2593        // decimal null_first
2594        test_sort_decimal_array::<T>(
2595            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2596            Some(SortOptions {
2597                descending: false,
2598                nulls_first: true,
2599            }),
2600            None,
2601            vec![None, None, Some(1), Some(2), Some(3), Some(4), Some(5)],
2602            precision,
2603            scale,
2604        );
2605        // limit
2606        test_sort_decimal_array::<T>(
2607            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2608            None,
2609            Some(3),
2610            vec![None, None, Some(1)],
2611            precision,
2612            scale,
2613        );
2614        // limit descending
2615        test_sort_decimal_array::<T>(
2616            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2617            Some(SortOptions {
2618                descending: true,
2619                nulls_first: false,
2620            }),
2621            Some(3),
2622            vec![Some(5), Some(4), Some(3)],
2623            precision,
2624            scale,
2625        );
2626        // limit descending null_first
2627        test_sort_decimal_array::<T>(
2628            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2629            Some(SortOptions {
2630                descending: true,
2631                nulls_first: true,
2632            }),
2633            Some(3),
2634            vec![None, None, Some(5)],
2635            precision,
2636            scale,
2637        );
2638        // limit null_first
2639        test_sort_decimal_array::<T>(
2640            vec![None, Some(5), Some(2), Some(3), Some(1), Some(4), None],
2641            Some(SortOptions {
2642                descending: false,
2643                nulls_first: true,
2644            }),
2645            Some(3),
2646            vec![None, None, Some(1)],
2647            precision,
2648            scale,
2649        );
2650    }
2651
2652    #[test]
2653    fn test_sort_decimal32() {
2654        test_sort_decimal::<Decimal32Type>(8, 3);
2655    }
2656
2657    #[test]
2658    fn test_sort_decimal64() {
2659        test_sort_decimal::<Decimal64Type>(17, 5);
2660    }
2661
2662    #[test]
2663    fn test_sort_decimal128() {
2664        test_sort_decimal::<Decimal128Type>(23, 6);
2665    }
2666
2667    #[test]
2668    fn test_sort_decimal256() {
2669        test_sort_decimal::<Decimal256Type>(53, 6);
2670    }
2671
2672    #[test]
2673    fn test_sort_decimal256_max_min() {
2674        test_sort_decimal256_array(
2675            vec![
2676                None,
2677                Some(i256::MIN),
2678                Some(i256::from_i128(1)),
2679                Some(i256::MAX),
2680                Some(i256::from_i128(-1)),
2681                None,
2682            ],
2683            Some(SortOptions {
2684                descending: false,
2685                nulls_first: true,
2686            }),
2687            None,
2688            vec![
2689                None,
2690                None,
2691                Some(i256::MIN),
2692                Some(i256::from_i128(-1)),
2693                Some(i256::from_i128(1)),
2694                Some(i256::MAX),
2695            ],
2696        );
2697
2698        test_sort_decimal256_array(
2699            vec![
2700                None,
2701                Some(i256::MIN),
2702                Some(i256::from_i128(1)),
2703                Some(i256::MAX),
2704                Some(i256::from_i128(-1)),
2705                None,
2706            ],
2707            Some(SortOptions {
2708                descending: true,
2709                nulls_first: true,
2710            }),
2711            None,
2712            vec![
2713                None,
2714                None,
2715                Some(i256::MAX),
2716                Some(i256::from_i128(1)),
2717                Some(i256::from_i128(-1)),
2718                Some(i256::MIN),
2719            ],
2720        );
2721
2722        test_sort_decimal256_array(
2723            vec![
2724                None,
2725                Some(i256::MIN),
2726                Some(i256::from_i128(1)),
2727                Some(i256::MAX),
2728                Some(i256::from_i128(-1)),
2729                None,
2730            ],
2731            Some(SortOptions {
2732                descending: false,
2733                nulls_first: true,
2734            }),
2735            Some(4),
2736            vec![None, None, Some(i256::MIN), Some(i256::from_i128(-1))],
2737        );
2738
2739        test_sort_decimal256_array(
2740            vec![
2741                None,
2742                Some(i256::MIN),
2743                Some(i256::from_i128(1)),
2744                Some(i256::MAX),
2745                Some(i256::from_i128(-1)),
2746                None,
2747            ],
2748            Some(SortOptions {
2749                descending: true,
2750                nulls_first: true,
2751            }),
2752            Some(4),
2753            vec![None, None, Some(i256::MAX), Some(i256::from_i128(1))],
2754        );
2755    }
2756
2757    #[test]
2758    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
2759    fn test_sort_primitives() {
2760        // default case
2761        test_sort_primitive_arrays::<UInt8Type>(
2762            vec![None, Some(3), Some(5), Some(2), Some(3), None],
2763            None,
2764            None,
2765            vec![None, None, Some(2), Some(3), Some(3), Some(5)],
2766        );
2767        test_sort_primitive_arrays::<UInt16Type>(
2768            vec![None, Some(3), Some(5), Some(2), Some(3), None],
2769            None,
2770            None,
2771            vec![None, None, Some(2), Some(3), Some(3), Some(5)],
2772        );
2773        test_sort_primitive_arrays::<UInt32Type>(
2774            vec![None, Some(3), Some(5), Some(2), Some(3), None],
2775            None,
2776            None,
2777            vec![None, None, Some(2), Some(3), Some(3), Some(5)],
2778        );
2779        test_sort_primitive_arrays::<UInt64Type>(
2780            vec![None, Some(3), Some(5), Some(2), Some(3), None],
2781            None,
2782            None,
2783            vec![None, None, Some(2), Some(3), Some(3), Some(5)],
2784        );
2785
2786        // descending
2787        test_sort_primitive_arrays::<Int8Type>(
2788            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2789            Some(SortOptions {
2790                descending: true,
2791                nulls_first: false,
2792            }),
2793            None,
2794            vec![Some(2), Some(0), Some(0), Some(-1), None, None],
2795        );
2796        test_sort_primitive_arrays::<Int16Type>(
2797            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2798            Some(SortOptions {
2799                descending: true,
2800                nulls_first: false,
2801            }),
2802            None,
2803            vec![Some(2), Some(0), Some(0), Some(-1), None, None],
2804        );
2805        test_sort_primitive_arrays::<Int32Type>(
2806            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2807            Some(SortOptions {
2808                descending: true,
2809                nulls_first: false,
2810            }),
2811            None,
2812            vec![Some(2), Some(0), Some(0), Some(-1), None, None],
2813        );
2814        test_sort_primitive_arrays::<Int16Type>(
2815            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2816            Some(SortOptions {
2817                descending: true,
2818                nulls_first: false,
2819            }),
2820            None,
2821            vec![Some(2), Some(0), Some(0), Some(-1), None, None],
2822        );
2823
2824        // descending, nulls first
2825        test_sort_primitive_arrays::<Int8Type>(
2826            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2827            Some(SortOptions {
2828                descending: true,
2829                nulls_first: true,
2830            }),
2831            None,
2832            vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
2833        );
2834        test_sort_primitive_arrays::<Int16Type>(
2835            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2836            Some(SortOptions {
2837                descending: true,
2838                nulls_first: true,
2839            }),
2840            None,
2841            vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
2842        );
2843        test_sort_primitive_arrays::<Int32Type>(
2844            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2845            Some(SortOptions {
2846                descending: true,
2847                nulls_first: true,
2848            }),
2849            None,
2850            vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
2851        );
2852        test_sort_primitive_arrays::<Int64Type>(
2853            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2854            Some(SortOptions {
2855                descending: true,
2856                nulls_first: true,
2857            }),
2858            None,
2859            vec![None, None, Some(2), Some(0), Some(0), Some(-1)],
2860        );
2861
2862        test_sort_primitive_arrays::<Int64Type>(
2863            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2864            Some(SortOptions {
2865                descending: true,
2866                nulls_first: true,
2867            }),
2868            Some(3),
2869            vec![None, None, Some(2)],
2870        );
2871
2872        test_sort_primitive_arrays::<Float16Type>(
2873            vec![
2874                None,
2875                Some(f16::from_f32(0.0)),
2876                Some(f16::from_f32(2.0)),
2877                Some(f16::from_f32(-1.0)),
2878                Some(f16::from_f32(0.0)),
2879                None,
2880            ],
2881            Some(SortOptions {
2882                descending: true,
2883                nulls_first: true,
2884            }),
2885            None,
2886            vec![
2887                None,
2888                None,
2889                Some(f16::from_f32(2.0)),
2890                Some(f16::from_f32(0.0)),
2891                Some(f16::from_f32(0.0)),
2892                Some(f16::from_f32(-1.0)),
2893            ],
2894        );
2895
2896        test_sort_primitive_arrays::<Float32Type>(
2897            vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(0.0), None],
2898            Some(SortOptions {
2899                descending: true,
2900                nulls_first: true,
2901            }),
2902            None,
2903            vec![None, None, Some(2.0), Some(0.0), Some(0.0), Some(-1.0)],
2904        );
2905        test_sort_primitive_arrays::<Float64Type>(
2906            vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None],
2907            Some(SortOptions {
2908                descending: true,
2909                nulls_first: true,
2910            }),
2911            None,
2912            vec![None, None, Some(f64::NAN), Some(2.0), Some(0.0), Some(-1.0)],
2913        );
2914        test_sort_primitive_arrays::<Float64Type>(
2915            vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
2916            Some(SortOptions {
2917                descending: true,
2918                nulls_first: true,
2919            }),
2920            None,
2921            vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
2922        );
2923
2924        // int8 nulls first
2925        test_sort_primitive_arrays::<Int8Type>(
2926            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2927            Some(SortOptions {
2928                descending: false,
2929                nulls_first: true,
2930            }),
2931            None,
2932            vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
2933        );
2934        test_sort_primitive_arrays::<Int16Type>(
2935            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2936            Some(SortOptions {
2937                descending: false,
2938                nulls_first: true,
2939            }),
2940            None,
2941            vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
2942        );
2943        test_sort_primitive_arrays::<Int32Type>(
2944            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2945            Some(SortOptions {
2946                descending: false,
2947                nulls_first: true,
2948            }),
2949            None,
2950            vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
2951        );
2952        test_sort_primitive_arrays::<Int64Type>(
2953            vec![None, Some(0), Some(2), Some(-1), Some(0), None],
2954            Some(SortOptions {
2955                descending: false,
2956                nulls_first: true,
2957            }),
2958            None,
2959            vec![None, None, Some(-1), Some(0), Some(0), Some(2)],
2960        );
2961        test_sort_primitive_arrays::<Float16Type>(
2962            vec![
2963                None,
2964                Some(f16::from_f32(0.0)),
2965                Some(f16::from_f32(2.0)),
2966                Some(f16::from_f32(-1.0)),
2967                Some(f16::from_f32(0.0)),
2968                None,
2969            ],
2970            Some(SortOptions {
2971                descending: false,
2972                nulls_first: true,
2973            }),
2974            None,
2975            vec![
2976                None,
2977                None,
2978                Some(f16::from_f32(-1.0)),
2979                Some(f16::from_f32(0.0)),
2980                Some(f16::from_f32(0.0)),
2981                Some(f16::from_f32(2.0)),
2982            ],
2983        );
2984        test_sort_primitive_arrays::<Float32Type>(
2985            vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(0.0), None],
2986            Some(SortOptions {
2987                descending: false,
2988                nulls_first: true,
2989            }),
2990            None,
2991            vec![None, None, Some(-1.0), Some(0.0), Some(0.0), Some(2.0)],
2992        );
2993        test_sort_primitive_arrays::<Float64Type>(
2994            vec![None, Some(0.0), Some(2.0), Some(-1.0), Some(f64::NAN), None],
2995            Some(SortOptions {
2996                descending: false,
2997                nulls_first: true,
2998            }),
2999            None,
3000            vec![None, None, Some(-1.0), Some(0.0), Some(2.0), Some(f64::NAN)],
3001        );
3002        test_sort_primitive_arrays::<Float64Type>(
3003            vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
3004            Some(SortOptions {
3005                descending: false,
3006                nulls_first: true,
3007            }),
3008            None,
3009            vec![Some(1.0), Some(f64::NAN), Some(f64::NAN), Some(f64::NAN)],
3010        );
3011
3012        // limit
3013        test_sort_primitive_arrays::<Float64Type>(
3014            vec![Some(f64::NAN), Some(f64::NAN), Some(f64::NAN), Some(1.0)],
3015            Some(SortOptions {
3016                descending: false,
3017                nulls_first: true,
3018            }),
3019            Some(2),
3020            vec![Some(1.0), Some(f64::NAN)],
3021        );
3022
3023        // limit with actual value
3024        test_sort_primitive_arrays::<Float64Type>(
3025            vec![Some(2.0), Some(4.0), Some(3.0), Some(1.0)],
3026            Some(SortOptions {
3027                descending: false,
3028                nulls_first: true,
3029            }),
3030            Some(3),
3031            vec![Some(1.0), Some(2.0), Some(3.0)],
3032        );
3033
3034        // valid values less than limit with extra nulls
3035        test_sort_primitive_arrays::<Float64Type>(
3036            vec![Some(2.0), None, None, Some(1.0)],
3037            Some(SortOptions {
3038                descending: false,
3039                nulls_first: false,
3040            }),
3041            Some(3),
3042            vec![Some(1.0), Some(2.0), None],
3043        );
3044
3045        test_sort_primitive_arrays::<Float64Type>(
3046            vec![Some(2.0), None, None, Some(1.0)],
3047            Some(SortOptions {
3048                descending: false,
3049                nulls_first: true,
3050            }),
3051            Some(3),
3052            vec![None, None, Some(1.0)],
3053        );
3054
3055        // more nulls than limit
3056        test_sort_primitive_arrays::<Float64Type>(
3057            vec![Some(2.0), None, None, None],
3058            Some(SortOptions {
3059                descending: false,
3060                nulls_first: true,
3061            }),
3062            Some(2),
3063            vec![None, None],
3064        );
3065
3066        test_sort_primitive_arrays::<Float64Type>(
3067            vec![Some(2.0), None, None, None],
3068            Some(SortOptions {
3069                descending: false,
3070                nulls_first: false,
3071            }),
3072            Some(2),
3073            vec![Some(2.0), None],
3074        );
3075    }
3076
3077    #[test]
3078    fn test_sort_to_indices_strings() {
3079        test_sort_to_indices_string_arrays(
3080            vec![
3081                None,
3082                Some("bad"),
3083                Some("sad"),
3084                None,
3085                Some("glad"),
3086                Some("-ad"),
3087            ],
3088            None,
3089            None,
3090            vec![0, 3, 5, 1, 4, 2],
3091        );
3092
3093        test_sort_to_indices_string_arrays(
3094            vec![
3095                None,
3096                Some("bad"),
3097                Some("sad"),
3098                None,
3099                Some("glad"),
3100                Some("-ad"),
3101            ],
3102            Some(SortOptions {
3103                descending: true,
3104                nulls_first: false,
3105            }),
3106            None,
3107            vec![2, 4, 1, 5, 0, 3],
3108        );
3109
3110        test_sort_to_indices_string_arrays(
3111            vec![
3112                None,
3113                Some("bad"),
3114                Some("sad"),
3115                None,
3116                Some("glad"),
3117                Some("-ad"),
3118            ],
3119            Some(SortOptions {
3120                descending: false,
3121                nulls_first: true,
3122            }),
3123            None,
3124            vec![0, 3, 5, 1, 4, 2],
3125        );
3126
3127        test_sort_to_indices_string_arrays(
3128            vec![
3129                None,
3130                Some("bad"),
3131                Some("sad"),
3132                None,
3133                Some("glad"),
3134                Some("-ad"),
3135            ],
3136            Some(SortOptions {
3137                descending: true,
3138                nulls_first: true,
3139            }),
3140            None,
3141            vec![0, 3, 2, 4, 1, 5],
3142        );
3143
3144        test_sort_to_indices_string_arrays(
3145            vec![
3146                None,
3147                Some("bad"),
3148                Some("sad"),
3149                None,
3150                Some("glad"),
3151                Some("-ad"),
3152            ],
3153            Some(SortOptions {
3154                descending: true,
3155                nulls_first: true,
3156            }),
3157            Some(3),
3158            vec![0, 3, 2],
3159        );
3160
3161        // valid values less than limit with extra nulls
3162        test_sort_to_indices_string_arrays(
3163            vec![Some("def"), None, None, Some("abc")],
3164            Some(SortOptions {
3165                descending: false,
3166                nulls_first: false,
3167            }),
3168            Some(3),
3169            vec![3, 0, 1],
3170        );
3171
3172        test_sort_to_indices_string_arrays(
3173            vec![Some("def"), None, None, Some("abc")],
3174            Some(SortOptions {
3175                descending: false,
3176                nulls_first: true,
3177            }),
3178            Some(3),
3179            vec![1, 2, 3],
3180        );
3181
3182        // more nulls than limit
3183        test_sort_to_indices_string_arrays(
3184            vec![Some("def"), None, None, None],
3185            Some(SortOptions {
3186                descending: false,
3187                nulls_first: true,
3188            }),
3189            Some(2),
3190            vec![1, 2],
3191        );
3192
3193        test_sort_to_indices_string_arrays(
3194            vec![Some("def"), None, None, None],
3195            Some(SortOptions {
3196                descending: false,
3197                nulls_first: false,
3198            }),
3199            Some(2),
3200            vec![0, 1],
3201        );
3202    }
3203
3204    #[test]
3205    fn test_sort_strings() {
3206        test_sort_string_arrays(
3207            vec![
3208                None,
3209                Some("bad"),
3210                Some("sad"),
3211                Some("long string longer than 12 bytes"),
3212                None,
3213                Some("glad"),
3214                Some("lang string longer than 12 bytes"),
3215                Some("-ad"),
3216            ],
3217            None,
3218            None,
3219            vec![
3220                None,
3221                None,
3222                Some("-ad"),
3223                Some("bad"),
3224                Some("glad"),
3225                Some("lang string longer than 12 bytes"),
3226                Some("long string longer than 12 bytes"),
3227                Some("sad"),
3228            ],
3229        );
3230
3231        test_sort_string_arrays(
3232            vec![
3233                None,
3234                Some("bad"),
3235                Some("sad"),
3236                Some("long string longer than 12 bytes"),
3237                None,
3238                Some("glad"),
3239                Some("lang string longer than 12 bytes"),
3240                Some("-ad"),
3241            ],
3242            Some(SortOptions {
3243                descending: true,
3244                nulls_first: false,
3245            }),
3246            None,
3247            vec![
3248                Some("sad"),
3249                Some("long string longer than 12 bytes"),
3250                Some("lang string longer than 12 bytes"),
3251                Some("glad"),
3252                Some("bad"),
3253                Some("-ad"),
3254                None,
3255                None,
3256            ],
3257        );
3258
3259        test_sort_string_arrays(
3260            vec![
3261                None,
3262                Some("bad"),
3263                Some("long string longer than 12 bytes"),
3264                Some("sad"),
3265                None,
3266                Some("glad"),
3267                Some("lang string longer than 12 bytes"),
3268                Some("-ad"),
3269            ],
3270            Some(SortOptions {
3271                descending: false,
3272                nulls_first: true,
3273            }),
3274            None,
3275            vec![
3276                None,
3277                None,
3278                Some("-ad"),
3279                Some("bad"),
3280                Some("glad"),
3281                Some("lang string longer than 12 bytes"),
3282                Some("long string longer than 12 bytes"),
3283                Some("sad"),
3284            ],
3285        );
3286
3287        test_sort_string_arrays(
3288            vec![
3289                None,
3290                Some("bad"),
3291                Some("long string longer than 12 bytes"),
3292                Some("sad"),
3293                None,
3294                Some("glad"),
3295                Some("lang string longer than 12 bytes"),
3296                Some("-ad"),
3297            ],
3298            Some(SortOptions {
3299                descending: true,
3300                nulls_first: true,
3301            }),
3302            None,
3303            vec![
3304                None,
3305                None,
3306                Some("sad"),
3307                Some("long string longer than 12 bytes"),
3308                Some("lang string longer than 12 bytes"),
3309                Some("glad"),
3310                Some("bad"),
3311                Some("-ad"),
3312            ],
3313        );
3314
3315        test_sort_string_arrays(
3316            vec![
3317                None,
3318                Some("bad"),
3319                Some("long string longer than 12 bytes"),
3320                Some("sad"),
3321                None,
3322                Some("glad"),
3323                Some("lang string longer than 12 bytes"),
3324                Some("-ad"),
3325            ],
3326            Some(SortOptions {
3327                descending: true,
3328                nulls_first: true,
3329            }),
3330            Some(3),
3331            vec![None, None, Some("sad")],
3332        );
3333
3334        // valid values less than limit with extra nulls
3335        test_sort_string_arrays(
3336            vec![
3337                Some("def long string longer than 12"),
3338                None,
3339                None,
3340                Some("abc"),
3341            ],
3342            Some(SortOptions {
3343                descending: false,
3344                nulls_first: false,
3345            }),
3346            Some(3),
3347            vec![Some("abc"), Some("def long string longer than 12"), None],
3348        );
3349
3350        test_sort_string_arrays(
3351            vec![
3352                Some("def long string longer than 12"),
3353                None,
3354                None,
3355                Some("abc"),
3356            ],
3357            Some(SortOptions {
3358                descending: false,
3359                nulls_first: true,
3360            }),
3361            Some(3),
3362            vec![None, None, Some("abc")],
3363        );
3364
3365        // more nulls than limit
3366        test_sort_string_arrays(
3367            vec![Some("def long string longer than 12"), None, None, None],
3368            Some(SortOptions {
3369                descending: false,
3370                nulls_first: true,
3371            }),
3372            Some(2),
3373            vec![None, None],
3374        );
3375
3376        test_sort_string_arrays(
3377            vec![Some("def long string longer than 12"), None, None, None],
3378            Some(SortOptions {
3379                descending: false,
3380                nulls_first: false,
3381            }),
3382            Some(2),
3383            vec![Some("def long string longer than 12"), None],
3384        );
3385    }
3386
3387    #[test]
3388    fn test_sort_run_to_run() {
3389        test_sort_run_inner(|array, sort_options, limit| sort_run(array, sort_options, limit));
3390    }
3391
3392    #[test]
3393    fn test_sort_run_to_indices() {
3394        test_sort_run_inner(|array, sort_options, limit| {
3395            let indices = sort_to_indices(array, sort_options, limit).unwrap();
3396            take(array, &indices, None)
3397        });
3398    }
3399
3400    fn test_sort_run_inner<F>(sort_fn: F)
3401    where
3402        F: Fn(&dyn Array, Option<SortOptions>, Option<usize>) -> Result<ArrayRef, ArrowError>,
3403    {
3404        // Create an input array for testing
3405        let total_len = 80;
3406        let vals: Vec<Option<i32>> = vec![Some(1), None, Some(2), Some(3), Some(4), None, Some(5)];
3407        let repeats: Vec<usize> = vec![1, 3, 2, 4];
3408        let mut input_array: Vec<Option<i32>> = Vec::with_capacity(total_len);
3409        for ix in 0_usize..32 {
3410            let repeat: usize = repeats[ix % repeats.len()];
3411            let val: Option<i32> = vals[ix % vals.len()];
3412            input_array.resize(input_array.len() + repeat, val);
3413        }
3414
3415        // create run array using input_array
3416        // Encode the input_array to run array
3417        let mut builder =
3418            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
3419        builder.extend(input_array.iter().copied());
3420        let run_array = builder.finish();
3421
3422        // slice lengths that are tested
3423        let slice_lens = [
3424            1, 2, 3, 4, 5, 6, 7, 37, 38, 39, 40, 41, 42, 43, 74, 75, 76, 77, 78, 79, 80,
3425        ];
3426        for slice_len in slice_lens {
3427            test_sort_run_inner2(
3428                input_array.as_slice(),
3429                &run_array,
3430                0,
3431                slice_len,
3432                None,
3433                &sort_fn,
3434            );
3435            test_sort_run_inner2(
3436                input_array.as_slice(),
3437                &run_array,
3438                total_len - slice_len,
3439                slice_len,
3440                None,
3441                &sort_fn,
3442            );
3443            // Test with non zero limit
3444            if slice_len > 1 {
3445                test_sort_run_inner2(
3446                    input_array.as_slice(),
3447                    &run_array,
3448                    0,
3449                    slice_len,
3450                    Some(slice_len / 2),
3451                    &sort_fn,
3452                );
3453                test_sort_run_inner2(
3454                    input_array.as_slice(),
3455                    &run_array,
3456                    total_len - slice_len,
3457                    slice_len,
3458                    Some(slice_len / 2),
3459                    &sort_fn,
3460                );
3461            }
3462        }
3463    }
3464
3465    fn test_sort_run_inner2<F>(
3466        input_array: &[Option<i32>],
3467        run_array: &RunArray<Int16Type>,
3468        offset: usize,
3469        length: usize,
3470        limit: Option<usize>,
3471        sort_fn: &F,
3472    ) where
3473        F: Fn(&dyn Array, Option<SortOptions>, Option<usize>) -> Result<ArrayRef, ArrowError>,
3474    {
3475        // Run the sort and build actual result
3476        let sliced_array = run_array.slice(offset, length);
3477        let sorted_sliced_array = sort_fn(&sliced_array, None, limit).unwrap();
3478        let sorted_run_array = sorted_sliced_array
3479            .as_any()
3480            .downcast_ref::<RunArray<Int16Type>>()
3481            .unwrap();
3482        let typed_run_array = sorted_run_array
3483            .downcast::<PrimitiveArray<Int32Type>>()
3484            .unwrap();
3485        let actual: Vec<Option<i32>> = typed_run_array.into_iter().collect();
3486
3487        // build expected result.
3488        let mut sliced_input = input_array[offset..(offset + length)].to_owned();
3489        sliced_input.sort();
3490        let expected = if let Some(limit) = limit {
3491            sliced_input.iter().take(limit).copied().collect()
3492        } else {
3493            sliced_input
3494        };
3495
3496        assert_eq!(expected, actual)
3497    }
3498
3499    #[test]
3500    fn test_sort_string_dicts() {
3501        test_sort_string_dict_arrays::<Int8Type>(
3502            vec![
3503                None,
3504                Some("bad"),
3505                Some("sad"),
3506                None,
3507                Some("glad"),
3508                Some("-ad"),
3509            ],
3510            None,
3511            None,
3512            vec![
3513                None,
3514                None,
3515                Some("-ad"),
3516                Some("bad"),
3517                Some("glad"),
3518                Some("sad"),
3519            ],
3520        );
3521
3522        test_sort_string_dict_arrays::<Int16Type>(
3523            vec![
3524                None,
3525                Some("bad"),
3526                Some("sad"),
3527                None,
3528                Some("glad"),
3529                Some("-ad"),
3530            ],
3531            Some(SortOptions {
3532                descending: true,
3533                nulls_first: false,
3534            }),
3535            None,
3536            vec![
3537                Some("sad"),
3538                Some("glad"),
3539                Some("bad"),
3540                Some("-ad"),
3541                None,
3542                None,
3543            ],
3544        );
3545
3546        test_sort_string_dict_arrays::<Int32Type>(
3547            vec![
3548                None,
3549                Some("bad"),
3550                Some("sad"),
3551                None,
3552                Some("glad"),
3553                Some("-ad"),
3554            ],
3555            Some(SortOptions {
3556                descending: false,
3557                nulls_first: true,
3558            }),
3559            None,
3560            vec![
3561                None,
3562                None,
3563                Some("-ad"),
3564                Some("bad"),
3565                Some("glad"),
3566                Some("sad"),
3567            ],
3568        );
3569
3570        test_sort_string_dict_arrays::<Int16Type>(
3571            vec![
3572                None,
3573                Some("bad"),
3574                Some("sad"),
3575                None,
3576                Some("glad"),
3577                Some("-ad"),
3578            ],
3579            Some(SortOptions {
3580                descending: true,
3581                nulls_first: true,
3582            }),
3583            None,
3584            vec![
3585                None,
3586                None,
3587                Some("sad"),
3588                Some("glad"),
3589                Some("bad"),
3590                Some("-ad"),
3591            ],
3592        );
3593
3594        test_sort_string_dict_arrays::<Int16Type>(
3595            vec![
3596                None,
3597                Some("bad"),
3598                Some("sad"),
3599                None,
3600                Some("glad"),
3601                Some("-ad"),
3602            ],
3603            Some(SortOptions {
3604                descending: true,
3605                nulls_first: true,
3606            }),
3607            Some(3),
3608            vec![None, None, Some("sad")],
3609        );
3610
3611        // valid values less than limit with extra nulls
3612        test_sort_string_dict_arrays::<Int16Type>(
3613            vec![Some("def"), None, None, Some("abc")],
3614            Some(SortOptions {
3615                descending: false,
3616                nulls_first: false,
3617            }),
3618            Some(3),
3619            vec![Some("abc"), Some("def"), None],
3620        );
3621
3622        test_sort_string_dict_arrays::<Int16Type>(
3623            vec![Some("def"), None, None, Some("abc")],
3624            Some(SortOptions {
3625                descending: false,
3626                nulls_first: true,
3627            }),
3628            Some(3),
3629            vec![None, None, Some("abc")],
3630        );
3631
3632        // more nulls than limit
3633        test_sort_string_dict_arrays::<Int16Type>(
3634            vec![Some("def"), None, None, None],
3635            Some(SortOptions {
3636                descending: false,
3637                nulls_first: true,
3638            }),
3639            Some(2),
3640            vec![None, None],
3641        );
3642
3643        test_sort_string_dict_arrays::<Int16Type>(
3644            vec![Some("def"), None, None, None],
3645            Some(SortOptions {
3646                descending: false,
3647                nulls_first: false,
3648            }),
3649            Some(2),
3650            vec![Some("def"), None],
3651        );
3652    }
3653
3654    #[test]
3655    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
3656    fn test_sort_list() {
3657        test_sort_list_arrays::<Int8Type>(
3658            vec![
3659                Some(vec![Some(1)]),
3660                Some(vec![Some(4)]),
3661                Some(vec![Some(2)]),
3662                Some(vec![Some(3)]),
3663            ],
3664            Some(SortOptions {
3665                descending: false,
3666                nulls_first: false,
3667            }),
3668            None,
3669            vec![
3670                Some(vec![Some(1)]),
3671                Some(vec![Some(2)]),
3672                Some(vec![Some(3)]),
3673                Some(vec![Some(4)]),
3674            ],
3675            Some(1),
3676        );
3677
3678        test_sort_list_arrays::<Float16Type>(
3679            vec![
3680                Some(vec![Some(f16::from_f32(1.0)), Some(f16::from_f32(0.0))]),
3681                Some(vec![
3682                    Some(f16::from_f32(4.0)),
3683                    Some(f16::from_f32(3.0)),
3684                    Some(f16::from_f32(2.0)),
3685                    Some(f16::from_f32(1.0)),
3686                ]),
3687                Some(vec![
3688                    Some(f16::from_f32(2.0)),
3689                    Some(f16::from_f32(3.0)),
3690                    Some(f16::from_f32(4.0)),
3691                ]),
3692                Some(vec![
3693                    Some(f16::from_f32(3.0)),
3694                    Some(f16::from_f32(3.0)),
3695                    Some(f16::from_f32(3.0)),
3696                    Some(f16::from_f32(3.0)),
3697                ]),
3698                Some(vec![Some(f16::from_f32(1.0)), Some(f16::from_f32(1.0))]),
3699            ],
3700            Some(SortOptions {
3701                descending: false,
3702                nulls_first: false,
3703            }),
3704            None,
3705            vec![
3706                Some(vec![Some(f16::from_f32(1.0)), Some(f16::from_f32(0.0))]),
3707                Some(vec![Some(f16::from_f32(1.0)), Some(f16::from_f32(1.0))]),
3708                Some(vec![
3709                    Some(f16::from_f32(2.0)),
3710                    Some(f16::from_f32(3.0)),
3711                    Some(f16::from_f32(4.0)),
3712                ]),
3713                Some(vec![
3714                    Some(f16::from_f32(3.0)),
3715                    Some(f16::from_f32(3.0)),
3716                    Some(f16::from_f32(3.0)),
3717                    Some(f16::from_f32(3.0)),
3718                ]),
3719                Some(vec![
3720                    Some(f16::from_f32(4.0)),
3721                    Some(f16::from_f32(3.0)),
3722                    Some(f16::from_f32(2.0)),
3723                    Some(f16::from_f32(1.0)),
3724                ]),
3725            ],
3726            None,
3727        );
3728
3729        test_sort_list_arrays::<Float32Type>(
3730            vec![
3731                Some(vec![Some(1.0), Some(0.0)]),
3732                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
3733                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
3734                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
3735                Some(vec![Some(1.0), Some(1.0)]),
3736            ],
3737            Some(SortOptions {
3738                descending: false,
3739                nulls_first: false,
3740            }),
3741            None,
3742            vec![
3743                Some(vec![Some(1.0), Some(0.0)]),
3744                Some(vec![Some(1.0), Some(1.0)]),
3745                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
3746                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
3747                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
3748            ],
3749            None,
3750        );
3751
3752        test_sort_list_arrays::<Float64Type>(
3753            vec![
3754                Some(vec![Some(1.0), Some(0.0)]),
3755                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
3756                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
3757                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
3758                Some(vec![Some(1.0), Some(1.0)]),
3759            ],
3760            Some(SortOptions {
3761                descending: false,
3762                nulls_first: false,
3763            }),
3764            None,
3765            vec![
3766                Some(vec![Some(1.0), Some(0.0)]),
3767                Some(vec![Some(1.0), Some(1.0)]),
3768                Some(vec![Some(2.0), Some(3.0), Some(4.0)]),
3769                Some(vec![Some(3.0), Some(3.0), Some(3.0), Some(3.0)]),
3770                Some(vec![Some(4.0), Some(3.0), Some(2.0), Some(1.0)]),
3771            ],
3772            None,
3773        );
3774
3775        test_sort_list_arrays::<Int32Type>(
3776            vec![
3777                Some(vec![Some(1), Some(0)]),
3778                Some(vec![Some(4), Some(3), Some(2), Some(1)]),
3779                Some(vec![Some(2), Some(3), Some(4)]),
3780                Some(vec![Some(3), Some(3), Some(3), Some(3)]),
3781                Some(vec![Some(1), Some(1)]),
3782            ],
3783            Some(SortOptions {
3784                descending: false,
3785                nulls_first: false,
3786            }),
3787            None,
3788            vec![
3789                Some(vec![Some(1), Some(0)]),
3790                Some(vec![Some(1), Some(1)]),
3791                Some(vec![Some(2), Some(3), Some(4)]),
3792                Some(vec![Some(3), Some(3), Some(3), Some(3)]),
3793                Some(vec![Some(4), Some(3), Some(2), Some(1)]),
3794            ],
3795            None,
3796        );
3797
3798        test_sort_list_arrays::<Int32Type>(
3799            vec![
3800                None,
3801                Some(vec![Some(4), None, Some(2)]),
3802                Some(vec![Some(2), Some(3), Some(4)]),
3803                None,
3804                Some(vec![Some(3), Some(3), None]),
3805            ],
3806            Some(SortOptions {
3807                descending: false,
3808                nulls_first: false,
3809            }),
3810            None,
3811            vec![
3812                Some(vec![Some(2), Some(3), Some(4)]),
3813                Some(vec![Some(3), Some(3), None]),
3814                Some(vec![Some(4), None, Some(2)]),
3815                None,
3816                None,
3817            ],
3818            Some(3),
3819        );
3820
3821        test_sort_list_arrays::<Int32Type>(
3822            vec![
3823                Some(vec![Some(1), Some(0)]),
3824                Some(vec![Some(4), Some(3), Some(2), Some(1)]),
3825                Some(vec![Some(2), Some(3), Some(4)]),
3826                Some(vec![Some(3), Some(3), Some(3), Some(3)]),
3827                Some(vec![Some(1), Some(1)]),
3828            ],
3829            Some(SortOptions {
3830                descending: false,
3831                nulls_first: false,
3832            }),
3833            Some(2),
3834            vec![Some(vec![Some(1), Some(0)]), Some(vec![Some(1), Some(1)])],
3835            None,
3836        );
3837
3838        // valid values less than limit with extra nulls
3839        test_sort_list_arrays::<Int32Type>(
3840            vec![Some(vec![Some(1)]), None, None, Some(vec![Some(2)])],
3841            Some(SortOptions {
3842                descending: false,
3843                nulls_first: false,
3844            }),
3845            Some(3),
3846            vec![Some(vec![Some(1)]), Some(vec![Some(2)]), None],
3847            None,
3848        );
3849
3850        test_sort_list_arrays::<Int32Type>(
3851            vec![Some(vec![Some(1)]), None, None, Some(vec![Some(2)])],
3852            Some(SortOptions {
3853                descending: false,
3854                nulls_first: true,
3855            }),
3856            Some(3),
3857            vec![None, None, Some(vec![Some(1)])],
3858            None,
3859        );
3860
3861        // more nulls than limit
3862        test_sort_list_arrays::<Int32Type>(
3863            vec![Some(vec![Some(1)]), None, None, None],
3864            Some(SortOptions {
3865                descending: false,
3866                nulls_first: true,
3867            }),
3868            Some(2),
3869            vec![None, None],
3870            None,
3871        );
3872
3873        test_sort_list_arrays::<Int32Type>(
3874            vec![Some(vec![Some(1)]), None, None, None],
3875            Some(SortOptions {
3876                descending: false,
3877                nulls_first: false,
3878            }),
3879            Some(2),
3880            vec![Some(vec![Some(1)]), None],
3881            None,
3882        );
3883    }
3884
3885    #[test]
3886    fn test_sort_binary() {
3887        test_sort_binary_arrays(
3888            vec![
3889                Some(vec![0, 0, 0]),
3890                Some(vec![0, 0, 5]),
3891                Some(vec![0, 0, 3]),
3892                Some(vec![0, 0, 7]),
3893                Some(vec![0, 0, 1]),
3894            ],
3895            Some(SortOptions {
3896                descending: false,
3897                nulls_first: false,
3898            }),
3899            None,
3900            vec![
3901                Some(vec![0, 0, 0]),
3902                Some(vec![0, 0, 1]),
3903                Some(vec![0, 0, 3]),
3904                Some(vec![0, 0, 5]),
3905                Some(vec![0, 0, 7]),
3906            ],
3907            Some(3),
3908        );
3909
3910        // with nulls
3911        test_sort_binary_arrays(
3912            vec![
3913                Some(vec![0, 0, 0]),
3914                None,
3915                Some(vec![0, 0, 3]),
3916                Some(vec![0, 0, 7]),
3917                Some(vec![0, 0, 1]),
3918                None,
3919            ],
3920            Some(SortOptions {
3921                descending: false,
3922                nulls_first: false,
3923            }),
3924            None,
3925            vec![
3926                Some(vec![0, 0, 0]),
3927                Some(vec![0, 0, 1]),
3928                Some(vec![0, 0, 3]),
3929                Some(vec![0, 0, 7]),
3930                None,
3931                None,
3932            ],
3933            Some(3),
3934        );
3935
3936        test_sort_binary_arrays(
3937            vec![
3938                Some(vec![3, 5, 7]),
3939                None,
3940                Some(vec![1, 7, 1]),
3941                Some(vec![2, 7, 3]),
3942                None,
3943                Some(vec![1, 4, 3]),
3944            ],
3945            Some(SortOptions {
3946                descending: false,
3947                nulls_first: false,
3948            }),
3949            None,
3950            vec![
3951                Some(vec![1, 4, 3]),
3952                Some(vec![1, 7, 1]),
3953                Some(vec![2, 7, 3]),
3954                Some(vec![3, 5, 7]),
3955                None,
3956                None,
3957            ],
3958            Some(3),
3959        );
3960
3961        // descending
3962        test_sort_binary_arrays(
3963            vec![
3964                Some(vec![0, 0, 0]),
3965                None,
3966                Some(vec![0, 0, 3]),
3967                Some(vec![0, 0, 7]),
3968                Some(vec![0, 0, 1]),
3969                None,
3970            ],
3971            Some(SortOptions {
3972                descending: true,
3973                nulls_first: false,
3974            }),
3975            None,
3976            vec![
3977                Some(vec![0, 0, 7]),
3978                Some(vec![0, 0, 3]),
3979                Some(vec![0, 0, 1]),
3980                Some(vec![0, 0, 0]),
3981                None,
3982                None,
3983            ],
3984            Some(3),
3985        );
3986
3987        // nulls first
3988        test_sort_binary_arrays(
3989            vec![
3990                Some(vec![0, 0, 0]),
3991                None,
3992                Some(vec![0, 0, 3]),
3993                Some(vec![0, 0, 7]),
3994                Some(vec![0, 0, 1]),
3995                None,
3996            ],
3997            Some(SortOptions {
3998                descending: false,
3999                nulls_first: true,
4000            }),
4001            None,
4002            vec![
4003                None,
4004                None,
4005                Some(vec![0, 0, 0]),
4006                Some(vec![0, 0, 1]),
4007                Some(vec![0, 0, 3]),
4008                Some(vec![0, 0, 7]),
4009            ],
4010            Some(3),
4011        );
4012
4013        // limit
4014        test_sort_binary_arrays(
4015            vec![
4016                Some(vec![0, 0, 0]),
4017                None,
4018                Some(vec![0, 0, 3]),
4019                Some(vec![0, 0, 7]),
4020                Some(vec![0, 0, 1]),
4021                None,
4022            ],
4023            Some(SortOptions {
4024                descending: false,
4025                nulls_first: true,
4026            }),
4027            Some(4),
4028            vec![None, None, Some(vec![0, 0, 0]), Some(vec![0, 0, 1])],
4029            Some(3),
4030        );
4031
4032        // var length
4033        test_sort_binary_arrays(
4034            vec![
4035                Some(b"Hello".to_vec()),
4036                None,
4037                Some(b"from".to_vec()),
4038                Some(b"Apache".to_vec()),
4039                Some(b"Arrow-rs".to_vec()),
4040                None,
4041            ],
4042            Some(SortOptions {
4043                descending: false,
4044                nulls_first: false,
4045            }),
4046            None,
4047            vec![
4048                Some(b"Apache".to_vec()),
4049                Some(b"Arrow-rs".to_vec()),
4050                Some(b"Hello".to_vec()),
4051                Some(b"from".to_vec()),
4052                None,
4053                None,
4054            ],
4055            None,
4056        );
4057
4058        // limit
4059        test_sort_binary_arrays(
4060            vec![
4061                Some(b"Hello".to_vec()),
4062                None,
4063                Some(b"from".to_vec()),
4064                Some(b"Apache".to_vec()),
4065                Some(b"Arrow-rs".to_vec()),
4066                None,
4067            ],
4068            Some(SortOptions {
4069                descending: false,
4070                nulls_first: true,
4071            }),
4072            Some(4),
4073            vec![
4074                None,
4075                None,
4076                Some(b"Apache".to_vec()),
4077                Some(b"Arrow-rs".to_vec()),
4078            ],
4079            None,
4080        );
4081    }
4082
4083    #[test]
4084    fn test_lex_sort_single_column() {
4085        let input = vec![SortColumn {
4086            values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4087                Some(17),
4088                Some(2),
4089                Some(-1),
4090                Some(0),
4091            ])) as ArrayRef,
4092            options: None,
4093        }];
4094        let expected = vec![Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4095            Some(-1),
4096            Some(0),
4097            Some(2),
4098            Some(17),
4099        ])) as ArrayRef];
4100        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4101        test_lex_sort_arrays(input.clone(), slice_arrays(expected, 0, 2), Some(2));
4102
4103        // Explicitly test a limit on the sort as a demonstration
4104        let expected = vec![Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4105            Some(-1),
4106            Some(0),
4107            Some(2),
4108        ])) as ArrayRef];
4109        test_lex_sort_arrays(input, expected, Some(3));
4110    }
4111
4112    #[test]
4113    fn test_lex_sort_unaligned_rows() {
4114        let input = vec![
4115            SortColumn {
4116                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![None, Some(-1)]))
4117                    as ArrayRef,
4118                options: None,
4119            },
4120            SortColumn {
4121                values: Arc::new(StringArray::from(vec![Some("foo")])) as ArrayRef,
4122                options: None,
4123            },
4124        ];
4125        assert!(
4126            lexsort(&input, None).is_err(),
4127            "lexsort should reject columns with different row counts"
4128        );
4129    }
4130
4131    #[test]
4132    fn test_lex_sort_mixed_types() {
4133        let input = vec![
4134            SortColumn {
4135                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4136                    Some(0),
4137                    Some(2),
4138                    Some(-1),
4139                    Some(0),
4140                ])) as ArrayRef,
4141                options: None,
4142            },
4143            SortColumn {
4144                values: Arc::new(PrimitiveArray::<UInt32Type>::from(vec![
4145                    Some(101),
4146                    Some(8),
4147                    Some(7),
4148                    Some(102),
4149                ])) as ArrayRef,
4150                options: None,
4151            },
4152            SortColumn {
4153                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4154                    Some(-1),
4155                    Some(-2),
4156                    Some(-3),
4157                    Some(-4),
4158                ])) as ArrayRef,
4159                options: None,
4160            },
4161        ];
4162        let expected = vec![
4163            Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4164                Some(-1),
4165                Some(0),
4166                Some(0),
4167                Some(2),
4168            ])) as ArrayRef,
4169            Arc::new(PrimitiveArray::<UInt32Type>::from(vec![
4170                Some(7),
4171                Some(101),
4172                Some(102),
4173                Some(8),
4174            ])) as ArrayRef,
4175            Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4176                Some(-3),
4177                Some(-1),
4178                Some(-4),
4179                Some(-2),
4180            ])) as ArrayRef,
4181        ];
4182        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4183        test_lex_sort_arrays(input, slice_arrays(expected, 0, 2), Some(2));
4184
4185        // test mix of string and in64 with option
4186        let input = vec![
4187            SortColumn {
4188                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4189                    Some(0),
4190                    Some(2),
4191                    Some(-1),
4192                    Some(0),
4193                ])) as ArrayRef,
4194                options: Some(SortOptions {
4195                    descending: true,
4196                    nulls_first: true,
4197                }),
4198            },
4199            SortColumn {
4200                values: Arc::new(StringArray::from(vec![
4201                    Some("foo"),
4202                    Some("9"),
4203                    Some("7"),
4204                    Some("bar"),
4205                ])) as ArrayRef,
4206                options: Some(SortOptions {
4207                    descending: true,
4208                    nulls_first: true,
4209                }),
4210            },
4211        ];
4212        let expected = vec![
4213            Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4214                Some(2),
4215                Some(0),
4216                Some(0),
4217                Some(-1),
4218            ])) as ArrayRef,
4219            Arc::new(StringArray::from(vec![
4220                Some("9"),
4221                Some("foo"),
4222                Some("bar"),
4223                Some("7"),
4224            ])) as ArrayRef,
4225        ];
4226        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4227        test_lex_sort_arrays(input, slice_arrays(expected, 0, 3), Some(3));
4228
4229        // test sort with nulls first
4230        let input = vec![
4231            SortColumn {
4232                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4233                    None,
4234                    Some(-1),
4235                    Some(2),
4236                    None,
4237                ])) as ArrayRef,
4238                options: Some(SortOptions {
4239                    descending: true,
4240                    nulls_first: true,
4241                }),
4242            },
4243            SortColumn {
4244                values: Arc::new(StringArray::from(vec![
4245                    Some("foo"),
4246                    Some("world"),
4247                    Some("hello"),
4248                    None,
4249                ])) as ArrayRef,
4250                options: Some(SortOptions {
4251                    descending: true,
4252                    nulls_first: true,
4253                }),
4254            },
4255        ];
4256        let expected = vec![
4257            Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4258                None,
4259                None,
4260                Some(2),
4261                Some(-1),
4262            ])) as ArrayRef,
4263            Arc::new(StringArray::from(vec![
4264                None,
4265                Some("foo"),
4266                Some("hello"),
4267                Some("world"),
4268            ])) as ArrayRef,
4269        ];
4270        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4271        test_lex_sort_arrays(input, slice_arrays(expected, 0, 1), Some(1));
4272
4273        // test sort with nulls last
4274        let input = vec![
4275            SortColumn {
4276                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4277                    None,
4278                    Some(-1),
4279                    Some(2),
4280                    None,
4281                ])) as ArrayRef,
4282                options: Some(SortOptions {
4283                    descending: true,
4284                    nulls_first: false,
4285                }),
4286            },
4287            SortColumn {
4288                values: Arc::new(StringArray::from(vec![
4289                    Some("foo"),
4290                    Some("world"),
4291                    Some("hello"),
4292                    None,
4293                ])) as ArrayRef,
4294                options: Some(SortOptions {
4295                    descending: true,
4296                    nulls_first: false,
4297                }),
4298            },
4299        ];
4300        let expected = vec![
4301            Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4302                Some(2),
4303                Some(-1),
4304                None,
4305                None,
4306            ])) as ArrayRef,
4307            Arc::new(StringArray::from(vec![
4308                Some("hello"),
4309                Some("world"),
4310                Some("foo"),
4311                None,
4312            ])) as ArrayRef,
4313        ];
4314        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4315        test_lex_sort_arrays(input, slice_arrays(expected, 0, 2), Some(2));
4316
4317        // test sort with opposite options
4318        let input = vec![
4319            SortColumn {
4320                values: Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4321                    None,
4322                    Some(-1),
4323                    Some(2),
4324                    Some(-1),
4325                    None,
4326                ])) as ArrayRef,
4327                options: Some(SortOptions {
4328                    descending: false,
4329                    nulls_first: false,
4330                }),
4331            },
4332            SortColumn {
4333                values: Arc::new(StringArray::from(vec![
4334                    Some("foo"),
4335                    Some("bar"),
4336                    Some("world"),
4337                    Some("hello"),
4338                    None,
4339                ])) as ArrayRef,
4340                options: Some(SortOptions {
4341                    descending: true,
4342                    nulls_first: true,
4343                }),
4344            },
4345        ];
4346        let expected = vec![
4347            Arc::new(PrimitiveArray::<Int64Type>::from(vec![
4348                Some(-1),
4349                Some(-1),
4350                Some(2),
4351                None,
4352                None,
4353            ])) as ArrayRef,
4354            Arc::new(StringArray::from(vec![
4355                Some("hello"),
4356                Some("bar"),
4357                Some("world"),
4358                None,
4359                Some("foo"),
4360            ])) as ArrayRef,
4361        ];
4362        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4363        test_lex_sort_arrays(input.clone(), slice_arrays(expected.clone(), 0, 5), Some(5));
4364
4365        // Limiting by more rows than present is ok
4366        test_lex_sort_arrays(input, slice_arrays(expected, 0, 5), Some(10));
4367
4368        // test with FixedSizeListArray, arrays order: [UInt32, FixedSizeList(UInt32, 1)]
4369
4370        // case1
4371        let primitive_array_data = vec![
4372            Some(2),
4373            Some(3),
4374            Some(2),
4375            Some(0),
4376            None,
4377            Some(2),
4378            Some(1),
4379            Some(2),
4380        ];
4381        let list_array_data = vec![
4382            None,
4383            Some(vec![Some(4)]),
4384            Some(vec![Some(3)]),
4385            Some(vec![Some(1)]),
4386            Some(vec![Some(5)]),
4387            Some(vec![Some(0)]),
4388            Some(vec![Some(2)]),
4389            Some(vec![None]),
4390        ];
4391
4392        let expected_primitive_array_data = vec![
4393            None,
4394            Some(0),
4395            Some(1),
4396            Some(2),
4397            Some(2),
4398            Some(2),
4399            Some(2),
4400            Some(3),
4401        ];
4402        let expected_list_array_data = vec![
4403            Some(vec![Some(5)]),
4404            Some(vec![Some(1)]),
4405            Some(vec![Some(2)]),
4406            None, // <-
4407            Some(vec![None]),
4408            Some(vec![Some(0)]),
4409            Some(vec![Some(3)]), // <-
4410            Some(vec![Some(4)]),
4411        ];
4412        test_lex_sort_mixed_types_with_fixed_size_list::<Int32Type>(
4413            primitive_array_data.clone(),
4414            list_array_data.clone(),
4415            expected_primitive_array_data.clone(),
4416            expected_list_array_data,
4417            None,
4418            None,
4419        );
4420
4421        // case2
4422        let primitive_array_options = SortOptions {
4423            descending: false,
4424            nulls_first: true,
4425        };
4426        let list_array_options = SortOptions {
4427            descending: false,
4428            nulls_first: false, // has been modified
4429        };
4430        let expected_list_array_data = vec![
4431            Some(vec![Some(5)]),
4432            Some(vec![Some(1)]),
4433            Some(vec![Some(2)]),
4434            Some(vec![Some(0)]), // <-
4435            Some(vec![Some(3)]),
4436            Some(vec![None]),
4437            None, // <-
4438            Some(vec![Some(4)]),
4439        ];
4440        test_lex_sort_mixed_types_with_fixed_size_list::<Int32Type>(
4441            primitive_array_data.clone(),
4442            list_array_data.clone(),
4443            expected_primitive_array_data.clone(),
4444            expected_list_array_data,
4445            Some(primitive_array_options),
4446            Some(list_array_options),
4447        );
4448
4449        // case3
4450        let primitive_array_options = SortOptions {
4451            descending: false,
4452            nulls_first: true,
4453        };
4454        let list_array_options = SortOptions {
4455            descending: true, // has been modified
4456            nulls_first: true,
4457        };
4458        let expected_list_array_data = vec![
4459            Some(vec![Some(5)]),
4460            Some(vec![Some(1)]),
4461            Some(vec![Some(2)]),
4462            None, // <-
4463            Some(vec![None]),
4464            Some(vec![Some(3)]),
4465            Some(vec![Some(0)]), // <-
4466            Some(vec![Some(4)]),
4467        ];
4468        test_lex_sort_mixed_types_with_fixed_size_list::<Int32Type>(
4469            primitive_array_data.clone(),
4470            list_array_data.clone(),
4471            expected_primitive_array_data,
4472            expected_list_array_data,
4473            Some(primitive_array_options),
4474            Some(list_array_options),
4475        );
4476
4477        // test with ListArray/LargeListArray, arrays order: [List<UInt32>/LargeList<UInt32>, UInt32]
4478
4479        let list_array_data = vec![
4480            Some(vec![Some(2), Some(1)]), // 0
4481            None,                         // 10
4482            Some(vec![Some(3)]),          // 1
4483            Some(vec![Some(2), Some(0)]), // 2
4484            Some(vec![None, Some(2)]),    // 3
4485            Some(vec![Some(0)]),          // none
4486            None,                         // 11
4487            Some(vec![Some(2), None]),    // 4
4488            Some(vec![None]),             // 5
4489            Some(vec![Some(2), Some(1)]), // 6
4490        ];
4491        let primitive_array_data = vec![
4492            Some(0),
4493            Some(10),
4494            Some(1),
4495            Some(2),
4496            Some(3),
4497            None,
4498            Some(11),
4499            Some(4),
4500            Some(5),
4501            Some(6),
4502        ];
4503        let expected_list_array_data = vec![
4504            None,
4505            None,
4506            Some(vec![None]),
4507            Some(vec![None, Some(2)]),
4508            Some(vec![Some(0)]),
4509            Some(vec![Some(2), None]),
4510            Some(vec![Some(2), Some(0)]),
4511            Some(vec![Some(2), Some(1)]),
4512            Some(vec![Some(2), Some(1)]),
4513            Some(vec![Some(3)]),
4514        ];
4515        let expected_primitive_array_data = vec![
4516            Some(10),
4517            Some(11),
4518            Some(5),
4519            Some(3),
4520            None,
4521            Some(4),
4522            Some(2),
4523            Some(0),
4524            Some(6),
4525            Some(1),
4526        ];
4527        test_lex_sort_mixed_types_with_list::<Int32Type>(
4528            list_array_data.clone(),
4529            primitive_array_data.clone(),
4530            expected_list_array_data,
4531            expected_primitive_array_data,
4532            None,
4533            None,
4534        );
4535    }
4536
4537    fn test_lex_sort_mixed_types_with_fixed_size_list<T>(
4538        primitive_array_data: Vec<Option<T::Native>>,
4539        list_array_data: Vec<Option<Vec<Option<T::Native>>>>,
4540        expected_primitive_array_data: Vec<Option<T::Native>>,
4541        expected_list_array_data: Vec<Option<Vec<Option<T::Native>>>>,
4542        primitive_array_options: Option<SortOptions>,
4543        list_array_options: Option<SortOptions>,
4544    ) where
4545        T: ArrowPrimitiveType,
4546        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
4547    {
4548        let input = vec![
4549            SortColumn {
4550                values: Arc::new(PrimitiveArray::<T>::from(primitive_array_data.clone()))
4551                    as ArrayRef,
4552                options: primitive_array_options,
4553            },
4554            SortColumn {
4555                values: Arc::new(FixedSizeListArray::from_iter_primitive::<T, _, _>(
4556                    list_array_data.clone(),
4557                    1,
4558                )) as ArrayRef,
4559                options: list_array_options,
4560            },
4561        ];
4562
4563        let expected = vec![
4564            Arc::new(PrimitiveArray::<T>::from(
4565                expected_primitive_array_data.clone(),
4566            )) as ArrayRef,
4567            Arc::new(FixedSizeListArray::from_iter_primitive::<T, _, _>(
4568                expected_list_array_data.clone(),
4569                1,
4570            )) as ArrayRef,
4571        ];
4572
4573        test_lex_sort_arrays(input.clone(), expected.clone(), None);
4574        test_lex_sort_arrays(input.clone(), slice_arrays(expected.clone(), 0, 5), Some(5));
4575    }
4576
4577    fn test_lex_sort_mixed_types_with_list<T>(
4578        list_array_data: Vec<Option<Vec<Option<T::Native>>>>,
4579        primitive_array_data: Vec<Option<T::Native>>,
4580        expected_list_array_data: Vec<Option<Vec<Option<T::Native>>>>,
4581        expected_primitive_array_data: Vec<Option<T::Native>>,
4582        list_array_options: Option<SortOptions>,
4583        primitive_array_options: Option<SortOptions>,
4584    ) where
4585        T: ArrowPrimitiveType,
4586        PrimitiveArray<T>: From<Vec<Option<T::Native>>>,
4587    {
4588        macro_rules! run_test {
4589            ($ARRAY_TYPE:ident) => {
4590                let input = vec![
4591                    SortColumn {
4592                        values: Arc::new(<$ARRAY_TYPE>::from_iter_primitive::<T, _, _>(
4593                            list_array_data.clone(),
4594                        )) as ArrayRef,
4595                        options: list_array_options.clone(),
4596                    },
4597                    SortColumn {
4598                        values: Arc::new(PrimitiveArray::<T>::from(primitive_array_data.clone()))
4599                            as ArrayRef,
4600                        options: primitive_array_options.clone(),
4601                    },
4602                ];
4603
4604                let expected = vec![
4605                    Arc::new(<$ARRAY_TYPE>::from_iter_primitive::<T, _, _>(
4606                        expected_list_array_data.clone(),
4607                    )) as ArrayRef,
4608                    Arc::new(PrimitiveArray::<T>::from(
4609                        expected_primitive_array_data.clone(),
4610                    )) as ArrayRef,
4611                ];
4612
4613                test_lex_sort_arrays(input.clone(), expected.clone(), None);
4614                test_lex_sort_arrays(input.clone(), slice_arrays(expected.clone(), 0, 5), Some(5));
4615            };
4616        }
4617        run_test!(ListArray);
4618        run_test!(LargeListArray);
4619    }
4620
4621    #[test]
4622    fn test_lex_sort_limit_paths() {
4623        let input = vec![
4624            SortColumn {
4625                values: Arc::new(Int32Array::from_iter_values((0..100).rev())) as ArrayRef,
4626                options: None,
4627            },
4628            SortColumn {
4629                values: Arc::new(Int32Array::from_iter_values(0..100)) as ArrayRef,
4630                options: None,
4631            },
4632        ];
4633
4634        // Exercise the bounded heap path.
4635        let indices = lexsort_to_indices(&input, Some(10)).unwrap();
4636        let expected = UInt32Array::from_iter_values((90..100).rev());
4637        assert_eq!(indices, expected);
4638
4639        // Exercise the existing partial-sort path.
4640        let indices = lexsort_to_indices(&input, Some(11)).unwrap();
4641        let expected = UInt32Array::from_iter_values((89..100).rev());
4642        assert_eq!(indices, expected);
4643    }
4644
4645    #[test]
4646    fn test_partial_sort() {
4647        let mut before: Vec<&str> = vec![
4648            "a", "cat", "mat", "on", "sat", "the", "xxx", "xxxx", "fdadfdsf",
4649        ];
4650        let mut d = before.clone();
4651        d.sort_unstable();
4652
4653        for last in 0..before.len() {
4654            partial_sort(&mut before, last, |a, b| a.cmp(b));
4655            assert_eq!(&d[0..last], &before.as_slice()[0..last]);
4656        }
4657    }
4658
4659    #[test]
4660    fn test_partial_rand_sort() {
4661        let size = 1000u32;
4662        let mut rng = StdRng::seed_from_u64(42);
4663        let mut before: Vec<u32> = (0..size).map(|_| rng.random::<u32>()).collect();
4664        let mut d = before.clone();
4665        let last = (rng.next_u32() % size) as usize;
4666        d.sort_unstable();
4667
4668        partial_sort(&mut before, last, |a, b| a.cmp(b));
4669        assert_eq!(&d[0..last], &before[0..last]);
4670    }
4671
4672    #[test]
4673    fn test_sort_int8_dicts() {
4674        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4675        let values = Int8Array::from(vec![1, 3, 5]);
4676        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4677            keys,
4678            values,
4679            None,
4680            None,
4681            vec![None, None, Some(1), Some(3), Some(5), Some(5)],
4682        );
4683
4684        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4685        let values = Int8Array::from(vec![1, 3, 5]);
4686        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4687            keys,
4688            values,
4689            Some(SortOptions {
4690                descending: true,
4691                nulls_first: false,
4692            }),
4693            None,
4694            vec![Some(5), Some(5), Some(3), Some(1), None, None],
4695        );
4696
4697        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4698        let values = Int8Array::from(vec![1, 3, 5]);
4699        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4700            keys,
4701            values,
4702            Some(SortOptions {
4703                descending: false,
4704                nulls_first: false,
4705            }),
4706            None,
4707            vec![Some(1), Some(3), Some(5), Some(5), None, None],
4708        );
4709
4710        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4711        let values = Int8Array::from(vec![1, 3, 5]);
4712        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4713            keys,
4714            values,
4715            Some(SortOptions {
4716                descending: true,
4717                nulls_first: true,
4718            }),
4719            Some(3),
4720            vec![None, None, Some(5)],
4721        );
4722
4723        // Values have `None`.
4724        let keys = Int8Array::from(vec![
4725            Some(1_i8),
4726            None,
4727            Some(3),
4728            None,
4729            Some(2),
4730            Some(3),
4731            Some(0),
4732        ]);
4733        let values = Int8Array::from(vec![Some(1), Some(3), None, Some(5)]);
4734        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4735            keys,
4736            values,
4737            None,
4738            None,
4739            vec![None, None, None, Some(1), Some(3), Some(5), Some(5)],
4740        );
4741
4742        let keys = Int8Array::from(vec![
4743            Some(1_i8),
4744            None,
4745            Some(3),
4746            None,
4747            Some(2),
4748            Some(3),
4749            Some(0),
4750        ]);
4751        let values = Int8Array::from(vec![Some(1), Some(3), None, Some(5)]);
4752        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4753            keys,
4754            values,
4755            Some(SortOptions {
4756                descending: false,
4757                nulls_first: false,
4758            }),
4759            None,
4760            vec![Some(1), Some(3), Some(5), Some(5), None, None, None],
4761        );
4762
4763        let keys = Int8Array::from(vec![
4764            Some(1_i8),
4765            None,
4766            Some(3),
4767            None,
4768            Some(2),
4769            Some(3),
4770            Some(0),
4771        ]);
4772        let values = Int8Array::from(vec![Some(1), Some(3), None, Some(5)]);
4773        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4774            keys,
4775            values,
4776            Some(SortOptions {
4777                descending: true,
4778                nulls_first: false,
4779            }),
4780            None,
4781            vec![Some(5), Some(5), Some(3), Some(1), None, None, None],
4782        );
4783
4784        let keys = Int8Array::from(vec![
4785            Some(1_i8),
4786            None,
4787            Some(3),
4788            None,
4789            Some(2),
4790            Some(3),
4791            Some(0),
4792        ]);
4793        let values = Int8Array::from(vec![Some(1), Some(3), None, Some(5)]);
4794        test_sort_primitive_dict_arrays::<Int8Type, Int8Type>(
4795            keys,
4796            values,
4797            Some(SortOptions {
4798                descending: true,
4799                nulls_first: true,
4800            }),
4801            None,
4802            vec![None, None, None, Some(5), Some(5), Some(3), Some(1)],
4803        );
4804    }
4805
4806    #[test]
4807    fn test_sort_f32_dicts() {
4808        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4809        let values = Float32Array::from(vec![1.2, 3.0, 5.1]);
4810        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4811            keys,
4812            values,
4813            None,
4814            None,
4815            vec![None, None, Some(1.2), Some(3.0), Some(5.1), Some(5.1)],
4816        );
4817
4818        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4819        let values = Float32Array::from(vec![1.2, 3.0, 5.1]);
4820        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4821            keys,
4822            values,
4823            Some(SortOptions {
4824                descending: true,
4825                nulls_first: false,
4826            }),
4827            None,
4828            vec![Some(5.1), Some(5.1), Some(3.0), Some(1.2), None, None],
4829        );
4830
4831        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4832        let values = Float32Array::from(vec![1.2, 3.0, 5.1]);
4833        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4834            keys,
4835            values,
4836            Some(SortOptions {
4837                descending: false,
4838                nulls_first: false,
4839            }),
4840            None,
4841            vec![Some(1.2), Some(3.0), Some(5.1), Some(5.1), None, None],
4842        );
4843
4844        let keys = Int8Array::from(vec![Some(1_i8), None, Some(2), None, Some(2), Some(0)]);
4845        let values = Float32Array::from(vec![1.2, 3.0, 5.1]);
4846        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4847            keys,
4848            values,
4849            Some(SortOptions {
4850                descending: true,
4851                nulls_first: true,
4852            }),
4853            Some(3),
4854            vec![None, None, Some(5.1)],
4855        );
4856
4857        // Values have `None`.
4858        let keys = Int8Array::from(vec![
4859            Some(1_i8),
4860            None,
4861            Some(3),
4862            None,
4863            Some(2),
4864            Some(3),
4865            Some(0),
4866        ]);
4867        let values = Float32Array::from(vec![Some(1.2), Some(3.0), None, Some(5.1)]);
4868        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4869            keys,
4870            values,
4871            None,
4872            None,
4873            vec![None, None, None, Some(1.2), Some(3.0), Some(5.1), Some(5.1)],
4874        );
4875
4876        let keys = Int8Array::from(vec![
4877            Some(1_i8),
4878            None,
4879            Some(3),
4880            None,
4881            Some(2),
4882            Some(3),
4883            Some(0),
4884        ]);
4885        let values = Float32Array::from(vec![Some(1.2), Some(3.0), None, Some(5.1)]);
4886        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4887            keys,
4888            values,
4889            Some(SortOptions {
4890                descending: false,
4891                nulls_first: false,
4892            }),
4893            None,
4894            vec![Some(1.2), Some(3.0), Some(5.1), Some(5.1), None, None, None],
4895        );
4896
4897        let keys = Int8Array::from(vec![
4898            Some(1_i8),
4899            None,
4900            Some(3),
4901            None,
4902            Some(2),
4903            Some(3),
4904            Some(0),
4905        ]);
4906        let values = Float32Array::from(vec![Some(1.2), Some(3.0), None, Some(5.1)]);
4907        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4908            keys,
4909            values,
4910            Some(SortOptions {
4911                descending: true,
4912                nulls_first: false,
4913            }),
4914            None,
4915            vec![Some(5.1), Some(5.1), Some(3.0), Some(1.2), None, None, None],
4916        );
4917
4918        let keys = Int8Array::from(vec![
4919            Some(1_i8),
4920            None,
4921            Some(3),
4922            None,
4923            Some(2),
4924            Some(3),
4925            Some(0),
4926        ]);
4927        let values = Float32Array::from(vec![Some(1.2), Some(3.0), None, Some(5.1)]);
4928        test_sort_primitive_dict_arrays::<Int8Type, Float32Type>(
4929            keys,
4930            values,
4931            Some(SortOptions {
4932                descending: true,
4933                nulls_first: true,
4934            }),
4935            None,
4936            vec![None, None, None, Some(5.1), Some(5.1), Some(3.0), Some(1.2)],
4937        );
4938    }
4939
4940    #[test]
4941    fn test_lexicographic_comparator_null_dict_values() {
4942        let values = Int32Array::new(
4943            vec![1, 2, 3, 4].into(),
4944            Some(NullBuffer::from(vec![true, false, false, true])),
4945        );
4946        let keys = Int32Array::new(
4947            vec![0, 1, 53, 3].into(),
4948            Some(NullBuffer::from(vec![true, true, false, true])),
4949        );
4950        // [1, NULL, NULL, 4]
4951        let dict = DictionaryArray::new(keys, Arc::new(values));
4952
4953        let comparator = LexicographicalComparator::try_new(&[SortColumn {
4954            values: Arc::new(dict),
4955            options: None,
4956        }])
4957        .unwrap();
4958        // 1.cmp(NULL)
4959        assert_eq!(comparator.compare(0, 1), Ordering::Greater);
4960        // NULL.cmp(NULL)
4961        assert_eq!(comparator.compare(2, 1), Ordering::Equal);
4962        // NULL.cmp(4)
4963        assert_eq!(comparator.compare(2, 3), Ordering::Less);
4964    }
4965
4966    #[test]
4967    fn sort_list_equal() {
4968        let a = {
4969            let mut builder = FixedSizeListBuilder::new(Int64Builder::new(), 2);
4970            for value in [[1, 5], [0, 3], [1, 3]] {
4971                builder.values().append_slice(&value);
4972                builder.append(true);
4973            }
4974            builder.finish()
4975        };
4976
4977        let sort_indices = sort_to_indices(&a, None, None).unwrap();
4978        assert_eq!(sort_indices.values(), &[1, 2, 0]);
4979
4980        let a = {
4981            let mut builder = ListBuilder::new(Int64Builder::new());
4982            for value in [[1, 5], [0, 3], [1, 3]] {
4983                builder.values().append_slice(&value);
4984                builder.append(true);
4985            }
4986            builder.finish()
4987        };
4988
4989        let sort_indices = sort_to_indices(&a, None, None).unwrap();
4990        assert_eq!(sort_indices.values(), &[1, 2, 0]);
4991    }
4992
4993    #[test]
4994    fn sort_struct_fallback_to_lexsort() {
4995        let float = Arc::new(Float32Array::from(vec![1.0, -0.1, 3.5, 1.0]));
4996        let int = Arc::new(Int32Array::from(vec![42, 28, 19, 31]));
4997
4998        let struct_array = StructArray::from(vec![
4999            (
5000                Arc::new(Field::new("b", DataType::Float32, false)),
5001                float.clone() as ArrayRef,
5002            ),
5003            (
5004                Arc::new(Field::new("c", DataType::Int32, false)),
5005                int.clone() as ArrayRef,
5006            ),
5007        ]);
5008
5009        assert!(!can_sort_to_indices(struct_array.data_type()));
5010        assert!(
5011            sort_to_indices(&struct_array, None, None)
5012                .err()
5013                .unwrap()
5014                .to_string()
5015                .contains("Sort not supported for data type")
5016        );
5017
5018        let sort_columns = vec![SortColumn {
5019            values: Arc::new(struct_array.clone()) as ArrayRef,
5020            options: None,
5021        }];
5022        let sorted = lexsort(&sort_columns, None).unwrap();
5023
5024        let expected_struct_array = Arc::new(StructArray::from(vec![
5025            (
5026                Arc::new(Field::new("b", DataType::Float32, false)),
5027                Arc::new(Float32Array::from(vec![-0.1, 1.0, 1.0, 3.5])) as ArrayRef,
5028            ),
5029            (
5030                Arc::new(Field::new("c", DataType::Int32, false)),
5031                Arc::new(Int32Array::from(vec![28, 31, 42, 19])) as ArrayRef,
5032            ),
5033        ])) as ArrayRef;
5034
5035        assert_eq!(&sorted[0], &expected_struct_array);
5036    }
5037
5038    /// A simple, correct but slower reference implementation.
5039    fn naive_partition(array: &BooleanArray) -> (Vec<u32>, Vec<u32>) {
5040        let len = array.len();
5041        let mut valid = Vec::with_capacity(len);
5042        let mut nulls = Vec::with_capacity(len);
5043        for i in 0..len {
5044            if array.is_valid(i) {
5045                valid.push(i as u32);
5046            } else {
5047                nulls.push(i as u32);
5048            }
5049        }
5050        (valid, nulls)
5051    }
5052
5053    #[test]
5054    #[cfg_attr(miri, ignore)] // Takes too long
5055    fn fuzz_partition_validity() {
5056        let mut rng = StdRng::seed_from_u64(0xF00D_CAFE);
5057        for _ in 0..1_000 {
5058            // build a random BooleanArray with some nulls
5059            let len = rng.random_range(0..512);
5060            let mut builder = BooleanBuilder::new();
5061            for _ in 0..len {
5062                if rng.random_bool(0.2) {
5063                    builder.append_null();
5064                } else {
5065                    builder.append_value(rng.random_bool(0.5));
5066                }
5067            }
5068            let array = builder.finish();
5069
5070            // Test both implementations on the full array
5071            let (v1, n1) = partition_validity(&array);
5072            let (v2, n2) = naive_partition(&array);
5073            assert_eq!(v1, v2, "valid mismatch on full array");
5074            assert_eq!(n1, n2, "null  mismatch on full array");
5075
5076            if len >= 8 {
5077                // 1) Random slice within the array
5078                let max_offset = len - 4;
5079                let offset = rng.random_range(0..=max_offset);
5080                let max_slice_len = len - offset;
5081                let slice_len = rng.random_range(1..=max_slice_len);
5082
5083                // Bind the sliced ArrayRef to keep it alive
5084                let sliced = array.slice(offset, slice_len);
5085                let slice = sliced
5086                    .as_any()
5087                    .downcast_ref::<BooleanArray>()
5088                    .expect("slice should be a BooleanArray");
5089
5090                let (sv1, sn1) = partition_validity(slice);
5091                let (sv2, sn2) = naive_partition(slice);
5092                assert_eq!(
5093                    sv1, sv2,
5094                    "valid mismatch on random slice at offset {offset} length {slice_len}",
5095                );
5096                assert_eq!(
5097                    sn1, sn2,
5098                    "null mismatch on random slice at offset {offset} length {slice_len}",
5099                );
5100
5101                // 2) Ensure we test slices that start beyond one 64-bit chunk boundary
5102                if len > 68 {
5103                    let offset2 = rng.random_range(65..(len - 3));
5104                    let len2 = rng.random_range(1..=(len - offset2));
5105
5106                    let sliced2 = array.slice(offset2, len2);
5107                    let slice2 = sliced2
5108                        .as_any()
5109                        .downcast_ref::<BooleanArray>()
5110                        .expect("slice2 should be a BooleanArray");
5111
5112                    let (sv3, sn3) = partition_validity(slice2);
5113                    let (sv4, sn4) = naive_partition(slice2);
5114                    assert_eq!(
5115                        sv3, sv4,
5116                        "valid mismatch on chunk-crossing slice at offset {offset2} length {len2}",
5117                    );
5118                    assert_eq!(
5119                        sn3, sn4,
5120                        "null mismatch on chunk-crossing slice at offset {offset2} length {len2}",
5121                    );
5122                }
5123            }
5124        }
5125    }
5126
5127    // A few small deterministic checks
5128    #[test]
5129    fn test_partition_edge_cases() {
5130        // all valid
5131        let array = BooleanArray::from(vec![Some(true), Some(false), Some(true)]);
5132        let (valid, nulls) = partition_validity(&array);
5133        assert_eq!(valid, vec![0, 1, 2]);
5134        assert!(nulls.is_empty());
5135
5136        // all null
5137        let array = BooleanArray::from(vec![None, None, None]);
5138        let (valid, nulls) = partition_validity(&array);
5139        assert!(valid.is_empty());
5140        assert_eq!(nulls, vec![0, 1, 2]);
5141
5142        // alternating
5143        let array = BooleanArray::from(vec![Some(true), None, Some(true), None]);
5144        let (valid, nulls) = partition_validity(&array);
5145        assert_eq!(valid, vec![0, 2]);
5146        assert_eq!(nulls, vec![1, 3]);
5147    }
5148
5149    // Test specific edge case strings that exercise the 4-byte prefix logic
5150    #[test]
5151    fn test_specific_edge_cases() {
5152        let test_cases = vec![
5153            // Key test cases for lengths 1-4 that test prefix padding
5154            "a", "ab", "ba", "baa", "abba", "abbc", "abc", "cda",
5155            // Test cases where first 4 bytes are same but subsequent bytes differ
5156            "abcd", "abcde", "abcdf", "abcdaaa", "abcdbbb",
5157            // Test cases with length < 4 that require padding
5158            "z", "za", "zaa", "zaaa", "zaaab", // Empty string
5159            "",      // Test various length combinations with same prefix
5160            "test", "test1", "test12", "test123", "test1234",
5161        ];
5162
5163        // Use standard library sort as reference
5164        let mut expected = test_cases.clone();
5165        expected.sort_unstable();
5166
5167        // Use our sorting algorithm
5168        let string_array = StringArray::from(test_cases.clone());
5169        let indices: Vec<u32> = (0..test_cases.len() as u32).collect();
5170        let result = sort_bytes(
5171            &string_array,
5172            indices,
5173            vec![], // no nulls
5174            SortOptions::default(),
5175            None,
5176        );
5177
5178        // Verify results
5179        let sorted_strings: Vec<&str> = result
5180            .values()
5181            .iter()
5182            .map(|&idx| test_cases[idx as usize])
5183            .collect();
5184
5185        assert_eq!(sorted_strings, expected);
5186    }
5187
5188    // Test sorting correctness for different length combinations
5189    #[test]
5190    fn test_length_combinations() {
5191        let test_cases = vec![
5192            // Focus on testing strings of length 1-4, as these affect padding logic
5193            ("", 0),
5194            ("a", 1),
5195            ("ab", 2),
5196            ("abc", 3),
5197            ("abcd", 4),
5198            ("abcde", 5),
5199            ("b", 1),
5200            ("ba", 2),
5201            ("bab", 3),
5202            ("babc", 4),
5203            ("babcd", 5),
5204            // Test same prefix with different lengths
5205            ("test", 4),
5206            ("test1", 5),
5207            ("test12", 6),
5208            ("test123", 7),
5209        ];
5210
5211        let strings: Vec<&str> = test_cases.iter().map(|(s, _)| *s).collect();
5212        let mut expected = strings.clone();
5213        expected.sort_unstable();
5214
5215        let string_array = StringArray::from(strings.clone());
5216        let indices: Vec<u32> = (0..strings.len() as u32).collect();
5217        let result = sort_bytes(&string_array, indices, vec![], SortOptions::default(), None);
5218
5219        let sorted_strings: Vec<&str> = result
5220            .values()
5221            .iter()
5222            .map(|&idx| strings[idx as usize])
5223            .collect();
5224
5225        assert_eq!(sorted_strings, expected);
5226    }
5227
5228    // Test UTF-8 string handling
5229    #[test]
5230    fn test_utf8_strings() {
5231        let test_cases = vec![
5232            "a",
5233            "你",       // 3-byte UTF-8 character
5234            "你好",     // 6 bytes
5235            "你好世界", // 12 bytes
5236            "🎉",       // 4-byte emoji
5237            "🎉🎊",     // 8 bytes
5238            "café",     // Contains accent character
5239            "naïve",
5240            "Москва", // Cyrillic script
5241            "東京",   // Japanese kanji
5242            "한국",   // Korean
5243        ];
5244
5245        let mut expected = test_cases.clone();
5246        expected.sort_unstable();
5247
5248        let string_array = StringArray::from(test_cases.clone());
5249        let indices: Vec<u32> = (0..test_cases.len() as u32).collect();
5250        let result = sort_bytes(&string_array, indices, vec![], SortOptions::default(), None);
5251
5252        let sorted_strings: Vec<&str> = result
5253            .values()
5254            .iter()
5255            .map(|&idx| test_cases[idx as usize])
5256            .collect();
5257
5258        assert_eq!(sorted_strings, expected);
5259    }
5260
5261    // Fuzz testing: generate random UTF-8 strings and verify sort correctness
5262    #[test]
5263    #[cfg_attr(miri, ignore)] // Takes too long
5264    fn test_fuzz_random_strings() {
5265        let mut rng = StdRng::seed_from_u64(42); // Fixed seed for reproducibility
5266
5267        for _ in 0..100 {
5268            // Run 100 rounds of fuzz testing
5269            let mut test_strings = Vec::new();
5270
5271            // Generate 20-50 random strings
5272            let num_strings = rng.random_range(20..=50);
5273
5274            for _ in 0..num_strings {
5275                let string = generate_random_string(&mut rng);
5276                test_strings.push(string);
5277            }
5278
5279            // Use standard library sort as reference
5280            let mut expected = test_strings.clone();
5281            expected.sort();
5282
5283            // Use our sorting algorithm
5284            let string_array = StringArray::from(test_strings.clone());
5285            let indices: Vec<u32> = (0..test_strings.len() as u32).collect();
5286            let result = sort_bytes(&string_array, indices, vec![], SortOptions::default(), None);
5287
5288            let sorted_strings: Vec<String> = result
5289                .values()
5290                .iter()
5291                .map(|&idx| test_strings[idx as usize].clone())
5292                .collect();
5293
5294            assert_eq!(
5295                sorted_strings, expected,
5296                "Fuzz test failed with input: {test_strings:?}"
5297            );
5298        }
5299    }
5300
5301    // Helper function to generate random UTF-8 strings
5302    fn generate_random_string(rng: &mut StdRng) -> String {
5303        // Bias towards generating short strings, especially length 1-4
5304        let length = if rng.random_bool(0.6) {
5305            rng.random_range(0..=4) // 60% probability for 0-4 length strings
5306        } else {
5307            rng.random_range(5..=20) // 40% probability for longer strings
5308        };
5309
5310        if length == 0 {
5311            return String::new();
5312        }
5313
5314        let mut result = String::new();
5315        let mut current_len = 0;
5316
5317        while current_len < length {
5318            let c = generate_random_char(rng);
5319            let char_len = c.len_utf8();
5320
5321            // Ensure we don't exceed target length
5322            if current_len + char_len <= length {
5323                result.push(c);
5324                current_len += char_len;
5325            } else {
5326                // If adding this character would exceed length, fill with ASCII
5327                let remaining = length - current_len;
5328                for _ in 0..remaining {
5329                    result.push(rng.random_range('a'..='z'));
5330                }
5331                break;
5332            }
5333        }
5334
5335        result
5336    }
5337
5338    // Generate random characters (including various UTF-8 characters)
5339    fn generate_random_char(rng: &mut StdRng) -> char {
5340        match rng.random_range(0..10) {
5341            0..=5 => rng.random_range('a'..='z'), // 60% ASCII lowercase
5342            6 => rng.random_range('A'..='Z'),     // 10% ASCII uppercase
5343            7 => rng.random_range('0'..='9'),     // 10% digits
5344            8 => {
5345                // 10% Chinese characters
5346                let chinese_chars = ['你', '好', '世', '界', '测', '试', '中', '文'];
5347                chinese_chars[rng.random_range(0..chinese_chars.len())]
5348            }
5349            9 => {
5350                // 10% other Unicode characters (single `char`s)
5351                let special_chars = ['é', 'ï', '🎉', '🎊', 'α', 'β', 'γ'];
5352                special_chars[rng.random_range(0..special_chars.len())]
5353            }
5354            _ => unreachable!(),
5355        }
5356    }
5357
5358    // Test descending sort order
5359    #[test]
5360    fn test_descending_sort() {
5361        let test_cases = vec!["a", "ab", "ba", "baa", "abba", "abbc", "abc", "cda"];
5362
5363        let mut expected = test_cases.clone();
5364        expected.sort_unstable();
5365        expected.reverse(); // Descending order
5366
5367        let string_array = StringArray::from(test_cases.clone());
5368        let indices: Vec<u32> = (0..test_cases.len() as u32).collect();
5369        let result = sort_bytes(
5370            &string_array,
5371            indices,
5372            vec![],
5373            SortOptions {
5374                descending: true,
5375                nulls_first: false,
5376            },
5377            None,
5378        );
5379
5380        let sorted_strings: Vec<&str> = result
5381            .values()
5382            .iter()
5383            .map(|&idx| test_cases[idx as usize])
5384            .collect();
5385
5386        assert_eq!(sorted_strings, expected);
5387    }
5388
5389    // Stress test: large number of strings with same prefix
5390    #[test]
5391    fn test_same_prefix_stress() {
5392        let mut test_cases = Vec::new();
5393        let prefix = "same";
5394
5395        // Generate many strings with the same prefix
5396        for i in 0..1000 {
5397            test_cases.push(format!("{prefix}{i:04}"));
5398        }
5399
5400        let mut expected = test_cases.clone();
5401        expected.sort();
5402
5403        let string_array = StringArray::from(test_cases.clone());
5404        let indices: Vec<u32> = (0..test_cases.len() as u32).collect();
5405        let result = sort_bytes(&string_array, indices, vec![], SortOptions::default(), None);
5406
5407        let sorted_strings: Vec<String> = result
5408            .values()
5409            .iter()
5410            .map(|&idx| test_cases[idx as usize].clone())
5411            .collect();
5412
5413        assert_eq!(sorted_strings, expected);
5414    }
5415
5416    // Test limit parameter
5417    #[test]
5418    fn test_with_limit() {
5419        let test_cases = vec!["z", "y", "x", "w", "v", "u", "t", "s"];
5420        let limit = 3;
5421
5422        let mut expected = test_cases.clone();
5423        expected.sort_unstable();
5424        expected.truncate(limit);
5425
5426        let string_array = StringArray::from(test_cases.clone());
5427        let indices: Vec<u32> = (0..test_cases.len() as u32).collect();
5428        let result = sort_bytes(
5429            &string_array,
5430            indices,
5431            vec![],
5432            SortOptions::default(),
5433            Some(limit),
5434        );
5435
5436        let sorted_strings: Vec<&str> = result
5437            .values()
5438            .iter()
5439            .map(|&idx| test_cases[idx as usize])
5440            .collect();
5441
5442        assert_eq!(sorted_strings, expected);
5443        assert_eq!(sorted_strings.len(), limit);
5444    }
5445
5446    #[test]
5447    fn test_empty_run() {
5448        let run = RunArray::try_new(
5449            &Int16Array::from(vec![1, 2, 3]),
5450            &Int32Array::from(vec![1, 5, 2]),
5451        )
5452        .unwrap();
5453
5454        let sorted = sort(&run.slice(1, 0), None).unwrap();
5455        assert!(sorted.is_empty());
5456        // ensure output run array upholds safety invariants
5457        sorted.into_data().validate_full().unwrap();
5458
5459        let sorted = sort_limit(&run, None, Some(0)).unwrap();
5460        assert!(sorted.is_empty());
5461        sorted.into_data().validate_full().unwrap();
5462
5463        let indices = sort_to_indices(&run, None, Some(0)).unwrap();
5464        assert!(indices.is_empty());
5465    }
5466}
arrow_ord/sort.rs

arrow_ord/
sort.rs
