Function process_map_array

datafusion/functions/src/core/getfield.rs:104–150

Process a map array by finding matching keys and extracting corresponding values. This function handles both simple (scalar) and nested key types by using appropriate comparison strategies.

(
    array: &dyn Array,
    key_array: Arc<dyn Array>,
)

/// This function handles both simple (scalar) and nested key types by using
/// appropriate comparison strategies.
fn process_map_array(
    array: &dyn Array,
    key_array: Arc<dyn Array>,
) -> Result<ColumnarValue> {
    let map_array = as_map_array(array)?;
    let keys = if key_array.data_type().is_nested() {
        let comparator = make_comparator(
            map_array.keys().as_ref(),
            key_array.as_ref(),
            SortOptions::default(),
        )?;
        let len = map_array.keys().len().min(key_array.len());
        let values = (0..len).map(|i| comparator(i, i).is_eq()).collect();
        let nulls = NullBuffer::union(map_array.keys().nulls(), key_array.nulls());
        BooleanArray::new(values, nulls)
    } else {
        let be_compared = Scalar::new(key_array);
        arrow::compute::kernels::cmp::eq(&be_compared, map_array.keys())?
    };

    let original_data = map_array.entries().column(1).to_data();
    let capacity = Capacities::Array(original_data.len());
    let mut mutable =
        MutableArrayData::with_capacities(vec![&original_data], true, capacity);

    for entry in 0..map_array.len() {
        let start = map_array.value_offsets()[entry] as usize;
        let end = map_array.value_offsets()[entry + 1] as usize;

        let maybe_matched = keys
            .slice(start, end - start)
            .iter()
            .enumerate()
            .find(|(_, t)| t.unwrap());

        if maybe_matched.is_none() {
            mutable.extend_nulls(1);
            continue;
        }
        let (match_offset, _) = maybe_matched.unwrap();
        mutable.extend(0, start + match_offset, start + match_offset + 1);
    }

    let data = mutable.freeze();
    let data = make_array(data);
    Ok(ColumnarValue::Array(data))
}

/// Process a map array with a nested key type by iterating through entries
/// and using a comparator for key matching.