Function general_array_slice

datafusion/functions-nested/src/extract.rs:613–690

(
    array: &GenericListArray<O>,
    from_array: &Int64Array,
    to_array: &Int64Array,
    stride: Option<&Int64Array>,
)

}

fn general_array_slice<O: OffsetSizeTrait>(
    array: &GenericListArray<O>,
    from_array: &Int64Array,
    to_array: &Int64Array,
    stride: Option<&Int64Array>,
) -> Result<ArrayRef>
where
    i64: TryInto<O>,
{
    let values = array.values();
    let original_data = values.to_data();
    let capacity = Capacities::Array(original_data.len());

    let mut mutable =
        MutableArrayData::with_capacities(vec![&original_data], true, capacity);

    // We have the slice syntax compatible with DuckDB v0.8.1.
    // The rule `adjusted_from_index` and `adjusted_to_index` follows the rule of array_slice in duckdb.

    let mut offsets = vec![O::usize_as(0)];

    let nulls = combine_input_nulls(array, from_array, to_array, stride);

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

        if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) {
            mutable.extend_nulls(1);
            offsets.push(offsets[row_index] + O::usize_as(1));
            continue;
        }

        // Empty arrays always return an empty array.
        if len == O::usize_as(0) {
            offsets.push(offsets[row_index]);
            continue;
        }

        let slice_plan = compute_slice_plan::<O>(
            len,
            from_array.value(row_index),
            to_array.value(row_index),
            stride.map(|s| s.value(row_index)),
        )?;

        match slice_plan {
            SlicePlan::Empty => offsets.push(offsets[row_index]),
            SlicePlan::Contiguous {
                start: rel_start,
                len: slice_len,
            } => {
                let start_index = (start + rel_start).to_usize().unwrap();
                let end_index = (start + rel_start + slice_len).to_usize().unwrap();
                mutable.extend(0, start_index, end_index);
                offsets.push(offsets[row_index] + slice_len);
            }