Function extract_single_field

datafusion/functions/src/core/getfield.rs:194–250

Extract a single field from a struct or map array

(base: ColumnarValue, name: ScalarValue)

/// Extract a single field from a struct or map array
fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result<ColumnarValue> {
    let arrays = ColumnarValue::values_to_arrays(&[base])?;
    let array = Arc::clone(&arrays[0]);

    let string_value = name.try_as_str().flatten().map(|s| s.to_string());

    match (array.data_type(), name, string_value) {
        // Dictionary-encoded struct: extract the field from the dictionary's
        // values (the deduplicated struct array) and rebuild a dictionary with
        // the same keys. This preserves dictionary encoding without expanding.
        (DataType::Dictionary(_, value_type), _, Some(field_name))
            if matches!(value_type.as_ref(), DataType::Struct(_)) =>
        {
            let dict = array.as_any_dictionary();
            let values_struct = dict.values().as_struct();
            let field_col =
                values_struct.column_by_name(&field_name).ok_or_else(|| {
                    exec_datafusion_err!(
                        "Field {field_name} not found in dictionary struct"
                    )
                })?;
            Ok(ColumnarValue::Array(
                dict.with_values(Arc::clone(field_col)),
            ))
        }
        (DataType::Map(_, _), ScalarValue::List(arr), _) => {
            let key_array: Arc<dyn Array> = arr;
            process_map_array(&array, key_array)
        }
        (DataType::Map(_, _), ScalarValue::Struct(arr), _) => {
            process_map_array(&array, arr as Arc<dyn Array>)
        }
        (DataType::Map(_, _), other, _) => {
            let data_type = other.data_type();
            if data_type.is_nested() {
                process_map_with_nested_key(&array, &other.to_array()?)
            } else {
                process_map_array(&array, other.to_array()?)
            }
        }
        (DataType::Struct(_), _, Some(k)) => {
            let as_struct_array = as_struct_array(&array)?;
            match as_struct_array.column_by_name(&k) {
                None => exec_err!("Field {k} not found in struct"),
                Some(col) => Ok(ColumnarValue::Array(Arc::clone(col))),
            }
        }
        (DataType::Struct(_), name, _) => exec_err!(
            "get_field is only possible on struct with utf8 indexes. \
                         Received with {name:?} index"
        ),
        (DataType::Null, _, _) => Ok(ColumnarValue::Scalar(ScalarValue::Null)),
        (dt, name, _) => exec_err!(
            "get_field is only possible on maps or structs. Received {dt} with {name:?} index"
        ),
    }
}