Function array_sort_non_primitive

datafusion/functions-nested/src/sort.rs:384–444

Sort a non-pritive-typed ListArray by converting all rows at once using `RowConverter`, and then sort row indices by comparing encoded bytes (sort direction and null ordering are baked into the encoding), and materialize the result with a single `take()`.

(
    list_array: &GenericListArray<OffsetSize>,
    field: FieldRef,
    sort_options: Option<SortOptions>,
)

/// direction and null ordering are baked into the encoding), and materialize
/// the result with a single `take()`.
fn array_sort_non_primitive<OffsetSize: OffsetSizeTrait>(
    list_array: &GenericListArray<OffsetSize>,
    field: FieldRef,
    sort_options: Option<SortOptions>,
) -> Result<ArrayRef> {
    let row_count = list_array.len();
    let values = list_array.values();
    let offsets = list_array.offsets();
    let values_start = offsets[0].as_usize();
    let total_values = offsets[row_count].as_usize() - values_start;

    let converter = RowConverter::new(vec![SortField::new_with_options(
        values.data_type().clone(),
        sort_options.unwrap_or_default(),
    )])?;
    let values_sliced = values.slice(values_start, total_values);
    let rows = converter.convert_columns(&[Arc::clone(&values_sliced)])?;

    let mut indices: Vec<OffsetSize> = Vec::with_capacity(total_values);
    let mut new_offsets = Vec::with_capacity(row_count + 1);
    new_offsets.push(OffsetSize::usize_as(0));

    let mut sort_scratch: Vec<usize> = Vec::new();

    for (row_index, window) in offsets.windows(2).enumerate() {
        let start = window[0];
        let end = window[1];

        if list_array.is_null(row_index) {
            new_offsets.push(new_offsets[row_index]);
            continue;
        }

        let len = (end - start).as_usize();
        let local_start = start.as_usize() - values_start;

        if len <= 1 {
            indices.extend((local_start..local_start + len).map(OffsetSize::usize_as));
        } else {
            sort_scratch.clear();
            sort_scratch.extend(local_start..local_start + len);
            sort_scratch.sort_unstable_by(|&a, &b| rows.row(a).cmp(&rows.row(b)));
            indices.extend(sort_scratch.iter().map(|&i| OffsetSize::usize_as(i)));
        }

        new_offsets.push(new_offsets[row_index] + (end - start));
    }

    let sorted_values = if indices.is_empty() {
        new_empty_array(values.data_type())
    } else {
        take_by_indices(&values_sliced, indices)?
    };

    Ok(Arc::new(GenericListArray::<OffsetSize>::try_new(
        field,
        OffsetBuffer::<OffsetSize>::new(new_offsets.into()),
        sorted_values,