Function general_repeat

datafusion/functions-nested/src/repeat.rs:174–217

For each element of `array[i]` repeat `count_array[i]` times. Assumption for the input: 1. `count[i] >= 0` 2. `array.len() == count_array.len()` For example, ```text array_repeat( [1, 2, 3], [2, 0, 1] => [[1, 1], [], [3]] ) ```

(
    array: &ArrayRef,
    count_array: &Int64Array,
)

/// )
/// ```
fn general_repeat<O: OffsetSizeTrait>(
    array: &ArrayRef,
    count_array: &Int64Array,
) -> Result<ArrayRef> {
    let total_repeated_values: usize = (0..count_array.len())
        .map(|i| get_count_with_validity(count_array, i))
        .sum();

    let mut take_indices = Vec::with_capacity(total_repeated_values);
    let mut offsets = Vec::with_capacity(count_array.len() + 1);
    offsets.push(O::zero());
    let mut running_offset = 0usize;

    for idx in 0..count_array.len() {
        let count = get_count_with_validity(count_array, idx);
        running_offset = running_offset.checked_add(count).ok_or_else(|| {
            DataFusionError::Execution(
                "array_repeat: running_offset overflowed usize".to_string(),
            )
        })?;
        let offset = O::from_usize(running_offset).ok_or_else(|| {
            DataFusionError::Execution(format!(
                "array_repeat: offset {running_offset} exceeds the maximum value for offset type"
            ))
        })?;
        offsets.push(offset);
        take_indices.extend(std::iter::repeat_n(idx as u64, count));
    }

    // Build the flattened values
    let repeated_values = compute::take(
        array.as_ref(),
        &UInt64Array::from_iter_values(take_indices),
        None,
    )?;

    // Construct final ListArray
    Ok(Arc::new(GenericListArray::<O>::try_new(
        Arc::new(Field::new_list_field(array.data_type().to_owned(), true)),
        OffsetBuffer::new(offsets.into()),
        repeated_values,
        count_array.nulls().cloned(),
    )?))
}

/// Handle List version of `general_repeat`
///