Function concat_internal

datafusion/functions-nested/src/concat.rs:384–450

(args: &[ArrayRef])

}

fn concat_internal<O: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
    let args = align_array_dimensions::<O>(args.to_vec())?;

    let list_arrays = args
        .iter()
        .map(|arg| as_generic_list_array::<O>(arg))
        .collect::<Result<Vec<_>>>()?;
    let row_count = list_arrays[0].len();

    // Extract underlying values ArrayData from each list array for MutableArrayData.
    let values_data: Vec<ArrayData> =
        list_arrays.iter().map(|la| la.values().to_data()).collect();
    let values_data_refs: Vec<&ArrayData> = values_data.iter().collect();

    // Estimate capacity as the sum of all values arrays' lengths.
    let total_capacity: usize = values_data.iter().map(|d| d.len()).sum();

    let mut mutable = MutableArrayData::with_capacities(
        values_data_refs,
        false,
        Capacities::Array(total_capacity),
    );
    let mut offsets: Vec<O> = Vec::with_capacity(row_count + 1);
    offsets.push(O::zero());

    // Compute the output null buffer: a row is null only if null in ALL input
    // arrays. This is the bitwise OR of validity bits (valid if valid in ANY
    // input). If any array has no null buffer (all valid), no output row can be
    // null.
    let nulls = list_arrays
        .iter()
        .filter_map(|la| la.nulls())
        .collect::<Vec<_>>();
    let valid = if nulls.len() == list_arrays.len() {
        nulls
            .iter()
            .map(|n| n.inner().clone())
            .reduce(|a, b| &a | &b)
            .map(NullBuffer::new)
    } else {
        None
    };

    for row_idx in 0..row_count {
        for (arr_idx, list_array) in list_arrays.iter().enumerate() {
            if list_array.is_null(row_idx) {
                continue;
            }
            let start = list_array.offsets()[row_idx].to_usize().unwrap();
            let end = list_array.offsets()[row_idx + 1].to_usize().unwrap();
            if start < end {
                mutable.extend(arr_idx, start, end);
            }
        }
        offsets.push(O::usize_as(mutable.len()));
    }

    let data_type = list_arrays[0].value_type();