Function general_remove

datafusion/functions-nested/src/remove.rs:377–469

For each element of `list_array[i]`, removed up to `arr_n[i]` occurrences of `element_array[i]`. The type of each **element** in `list_array` must be the same as the type of `element_array`. This function also handles nested arrays ([`arrow::array::ListArray`] of [`arrow::array::ListArray`]s) For example, when called to remove a list array (where each element is a list of int32s, the second argu

(
    list_array: &GenericListArray<OffsetSize>,
    element_array: &ArrayRef,
    arr_n: &[i64],
)

/// )
/// ```
fn general_remove<OffsetSize: OffsetSizeTrait>(
    list_array: &GenericListArray<OffsetSize>,
    element_array: &ArrayRef,
    arr_n: &[i64],
) -> Result<ArrayRef> {
    let list_field = match list_array.data_type() {
        DataType::List(field) | DataType::LargeList(field) => field,
        _ => {
            return exec_err!(
                "Expected List or LargeList data type, got {:?}",
                list_array.data_type()
            );
        }
    };
    let original_data = list_array.values().to_data();
    // Build up the offsets for the final output array
    let mut offsets = Vec::<OffsetSize>::with_capacity(arr_n.len() + 1);
    offsets.push(OffsetSize::zero());

    let mut mutable = MutableArrayData::with_capacities(
        vec![&original_data],
        false,
        Capacities::Array(original_data.len()),
    );

    // Pre-compute combined null bitmap
    let nulls = NullBuffer::union(list_array.nulls(), element_array.nulls());

    for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() {
        if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row_index)) {
            offsets.push(offsets[row_index]);
            continue;
        }

        let start = offset_window[0].to_usize().unwrap();
        let end = offset_window[1].to_usize().unwrap();
        // n is the number of elements to remove in this row
        let n = arr_n[row_index];

        // compare each element in the list, `false` means the element matches and should be removed
        let eq_array = utils::compare_element_to_list(
            &list_array.value(row_index),
            element_array,
            row_index,
            false,
        )?;

        let num_to_remove = eq_array.false_count();

        // Fast path: no elements to remove, copy entire row
        if num_to_remove == 0 {
            mutable.extend(0, start, end);
            offsets.push(offsets[row_index] + OffsetSize::usize_as(end - start));
            continue;
        }

        // Remove at most `n` matching elements
        let max_removals = n.min(num_to_remove as i64);