Function general_list_repeat

datafusion/functions-nested/src/repeat.rs:229–310

Handle List version of `general_repeat` For each element of `list_array[i]` repeat `count_array[i]` times. For example, ```text array_repeat( [[1, 2, 3], [4, 5], [6]], [2, 0, 1] => [[[1, 2, 3], [1, 2, 3]], [], [[6]]] ) ```

(
    list_array: &GenericListArray<O>,
    count_array: &Int64Array,
)

/// )
/// ```
fn general_list_repeat<O: OffsetSizeTrait>(
    list_array: &GenericListArray<O>,
    count_array: &Int64Array,
) -> Result<ArrayRef> {
    let list_offsets = list_array.value_offsets();

    // calculate capacities for pre-allocation
    let mut outer_total = 0usize;
    let mut inner_total = 0usize;
    for i in 0..count_array.len() {
        let count = get_count_with_validity(count_array, i);
        if count > 0 {
            outer_total += count;
            if list_array.is_valid(i) {
                let len = list_offsets[i + 1].to_usize().unwrap()
                    - list_offsets[i].to_usize().unwrap();
                inner_total += len * count;
            }
        }
    }

    // Build inner structures
    let mut inner_offsets = Vec::with_capacity(outer_total + 1);
    let mut take_indices = Vec::with_capacity(inner_total);
    let mut inner_nulls = BooleanBufferBuilder::new(outer_total);
    let mut inner_running = 0usize;
    inner_offsets.push(O::zero());

    for row_idx in 0..count_array.len() {
        let count = get_count_with_validity(count_array, row_idx);
        let list_is_valid = list_array.is_valid(row_idx);
        let start = list_offsets[row_idx].to_usize().unwrap();
        let end = list_offsets[row_idx + 1].to_usize().unwrap();
        let row_len = end - start;

        for _ in 0..count {
            inner_running = inner_running.checked_add(row_len).ok_or_else(|| {
                DataFusionError::Execution(
                    "array_repeat: inner offset overflowed usize".to_string(),
                )
            })?;
            let offset = O::from_usize(inner_running).ok_or_else(|| {
                DataFusionError::Execution(format!(
                    "array_repeat: offset {inner_running} exceeds the maximum value for offset type"
                ))
            })?;
            inner_offsets.push(offset);
            inner_nulls.append(list_is_valid);
            if list_is_valid {
                take_indices.extend(start as u64..end as u64);
            }
        }
    }

    // Build inner ListArray
    let inner_values = compute::take(
        list_array.values().as_ref(),
        &UInt64Array::from_iter_values(take_indices),