## Coverage Report

#### Created: 2025-11-22 06:00

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /Users/andrewlamb/Software/datafusion/datafusion/functions/src/core/getfield.rs

    Line
    Count
    Source
    1
    // Licensed to the Apache Software Foundation (ASF) under one
    2
    // or more contributor license agreements.  See the NOTICE file
    3
    // distributed with this work for additional information
    4
    // regarding copyright ownership.  The ASF licenses this file
    5
    // to you under the Apache License, Version 2.0 (the
    6
    // "License"); you may not use this file except in compliance
    7
    // with the License.  You may obtain a copy of the License at
    8
    //
    9
    //   http://www.apache.org/licenses/LICENSE-2.0
    10
    //
    11
    // Unless required by applicable law or agreed to in writing,
    12
    // software distributed under the License is distributed on an
    13
    // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14
    // KIND, either express or implied.  See the License for the
    15
    // specific language governing permissions and limitations
    16
    // under the License.
    17
    18
    use arrow::array::{
    19
        make_array, make_comparator, Array, BooleanArray, Capacities, MutableArrayData,
    20
        Scalar,
    21
    };
    22
    use arrow::compute::SortOptions;
    23
    use arrow::datatypes::{DataType, Field, FieldRef};
    24
    use arrow_buffer::NullBuffer;
    25
    use datafusion_common::cast::{as_map_array, as_struct_array};
    26
    use datafusion_common::{
    27
        exec_err, internal_err, plan_datafusion_err, utils::take_function_args, Result,
    28
        ScalarValue,
    29
    };
    30
    use datafusion_expr::{
    31
        ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs,
    32
    };
    33
    use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
    34
    use datafusion_macros::user_doc;
    35
    use std::any::Any;
    36
    use std::sync::Arc;
    37
    38
    #[user_doc(
    39
        doc_section(label = "Other Functions"),
    40
        description = r#"Returns a field within a map or a struct with the given key.
    41
        Note: most users invoke `get_field` indirectly via field access
    42
        syntax such as `my_struct_col['field_name']` which results in a call to
    43
        `get_field(my_struct_col, 'field_name')`."#,
    44
        syntax_example = "get_field(expression1, expression2)",
    45
        sql_example = r#"```sql
    46
    > create table t (idx varchar, v varchar) as values ('data','fusion'), ('apache', 'arrow');
    47
    > select struct(idx, v) from t as c;
    48
    +-------------------------+
    49
    | struct(c.idx,c.v)       |
    50
    +-------------------------+
    51
    | {c0: data, c1: fusion}  |
    52
    | {c0: apache, c1: arrow} |
    53
    +-------------------------+
    54
    > select get_field((select struct(idx, v) from t), 'c0');
    55
    +-----------------------+
    56
    | struct(t.idx,t.v)[c0] |
    57
    +-----------------------+
    58
    | data                  |
    59
    | apache                |
    60
    +-----------------------+
    61
    > select get_field((select struct(idx, v) from t), 'c1');
    62
    +-----------------------+
    63
    | struct(t.idx,t.v)[c1] |
    64
    +-----------------------+
    65
    | fusion                |
    66
    | arrow                 |
    67
    +-----------------------+
    68
    ```"#,
    69
        argument(
    70
            name = "expression1",
    71
            description = "The map or struct to retrieve a field for."
    72
        ),
    73
        argument(
    74
            name = "expression2",
    75
            description = "The field name in the map or struct to retrieve data for. Must evaluate to a string."
    76
        )
    77
    )]
    78
    #[derive(Debug, PartialEq, Eq, Hash)]
    79
    pub struct GetFieldFunc {
    80
        signature: Signature,
    81
    }
    82
    83
    impl Default for GetFieldFunc {
    84
    0
        fn default() -> Self {
    85
    0
            Self::new()
    86
    0
        }
    87
    }
    88
    89
    impl GetFieldFunc {
    90
    0
        pub fn new() -> Self {
    91
    0
            Self {
    92
    0
                signature: Signature::any(2, Volatility::Immutable),
    93
    0
            }
    94
    0
        }
    95
    }
    96
    97
    // get_field(struct_array, field_name)
    98
    impl ScalarUDFImpl for GetFieldFunc {
    99
    0
        fn as_any(&self) -> &dyn Any {
    100
    0
            self
    101
    0
        }
    102
    103
    0
        fn name(&self) -> &str {
    104
    0
            "get_field"
    105
    0
        }
    106
    107
    0
        fn display_name(&self, args: &[Expr]) -> Result<String> {
    108
    0
            let [base, field_name] = take_function_args(self.name(), args)?;
    109
    110
    0
            let name = match field_name {
    111
    0
                Expr::Literal(name, _) => name.to_string(),
    112
    0
                other => other.schema_name().to_string(),
    113
            };
    114
    115
    0
            Ok(format!("{base}[{name}]"))
    116
    0
        }
    117
    118
    0
        fn schema_name(&self, args: &[Expr]) -> Result<String> {
    119
    0
            let [base, field_name] = take_function_args(self.name(), args)?;
    120
    0
            let name = match field_name {
    121
    0
                Expr::Literal(name, _) => name.to_string(),
    122
    0
                other => other.schema_name().to_string(),
    123
            };
    124
    125
    0
            Ok(format!("{}[{}]", base.schema_name(), name))
    126
    0
        }
    127
    128
    0
        fn signature(&self) -> &Signature {
    129
    0
            &self.signature
    130
    0
        }
    131
    132
    0
        fn return_type(&self, _: &[DataType]) -> Result<DataType> {
    133
    0
            internal_err!("return_field_from_args should be called instead")
    134
    0
        }
    135
    136
    0
        fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
    137
            // Length check handled in the signature
    138
    0
            debug_assert_eq!(args.scalar_arguments.len(), 2);
    139
    140
    0
            match (&args.arg_fields[0].data_type(), args.scalar_arguments[1].as_ref()) {
    141
    0
                (DataType::Map(fields, _), _) => {
    142
    0
                    match fields.data_type() {
    143
    0
                        DataType::Struct(fields) if fields.len() == 2 => {
    144
                            // Arrow's MapArray is essentially a ListArray of structs with two columns. They are
    145
                            // often named "key", and "value", but we don't require any specific naming here;
    146
                            // instead, we assume that the second column is the "value" column both here and in
    147
                            // execution.
    148
    0
                            let value_field = fields.get(1).expect("fields should have exactly two members");
    149
    150
    0
                            Ok(value_field.as_ref().clone().with_nullable(true).into())
    151
                        },
    152
    0
                        _ => exec_err!("Map fields must contain a Struct with exactly 2 fields"),
    153
                    }
    154
                }
    155
    0
                (DataType::Struct(fields),sv) => {
    156
    0
                    sv.and_then(|sv| sv.try_as_str().flatten().filter(|s| !s.is_empty()))
    157
    0
                    .map_or_else(
    158
    0
                        || exec_err!("Field name must be a non-empty string"),
    159
    0
                        |field_name| {
    160
    0
                        fields.iter().find(|f| f.name() == field_name)
    161
    0
                        .ok_or(plan_datafusion_err!("Field {field_name} not found in struct"))
    162
    0
                        .map(|f| {
    163
    0
                            let mut child_field = f.as_ref().clone();
    164
    165
                            // If the parent is nullable, then getting the child must be nullable,
    166
                            // so potentially override the return value
    167
    168
    0
                            if args.arg_fields[0].is_nullable() {
    169
    0
                                child_field = child_field.with_nullable(true);
    170
    0
                            }
    171
    0
                            Arc::new(child_field)
    172
    0
                        })
    173
    0
                    })
    174
                },
    175
    0
                (DataType::Null, _) => Ok(Field::new(self.name(), DataType::Null, true).into()),
    176
    0
                (other, _) => exec_err!("The expression to get an indexed field is only valid for `Struct`, `Map` or `Null` types, got {other}"),
    177
            }
    178
    0
        }
    179
    180
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    181
    0
            let [base, field_name] = take_function_args(self.name(), args.args)?;
    182
    183
    0
            if base.data_type().is_null() {
    184
    0
                return Ok(ColumnarValue::Scalar(ScalarValue::Null));
    185
    0
            }
    186
    187
    0
            let arrays =
    188
    0
                ColumnarValue::values_to_arrays(&[base.clone(), field_name.clone()])?;
    189
    0
            let array = Arc::clone(&arrays[0]);
    190
    0
            let name = match field_name {
    191
    0
                ColumnarValue::Scalar(name) => name,
    192
                _ => {
    193
    0
                    return exec_err!(
    194
    0
                        "get_field function requires the argument field_name to be a string"
    195
                    );
    196
                }
    197
            };
    198
    199
    0
            fn process_map_array(
    200
    0
                array: Arc<dyn Array>,
    201
    0
                key_array: Arc<dyn Array>,
    202
    0
            ) -> Result<ColumnarValue> {
    203
    0
                let map_array = as_map_array(array.as_ref())?;
    204
    0
                let keys = if key_array.data_type().is_nested() {
    205
    0
                    let comparator = make_comparator(
    206
    0
                        map_array.keys().as_ref(),
    207
    0
                        key_array.as_ref(),
    208
    0
                        SortOptions::default(),
    209
    0
                    )?;
    210
    0
                    let len = map_array.keys().len().min(key_array.len());
    211
    0
                    let values = (0..len).map(|i| comparator(i, i).is_eq()).collect();
    212
    0
                    let nulls =
    213
    0
                        NullBuffer::union(map_array.keys().nulls(), key_array.nulls());
    214
    0
                    BooleanArray::new(values, nulls)
    215
                } else {
    216
    0
                    let be_compared = Scalar::new(key_array);
    217
    0
                    arrow::compute::kernels::cmp::eq(&be_compared, map_array.keys())?
    218
                };
    219
    220
    0
                let original_data = map_array.entries().column(1).to_data();
    221
    0
                let capacity = Capacities::Array(original_data.len());
    222
    0
                let mut mutable =
    223
    0
                    MutableArrayData::with_capacities(vec![&original_data], true, capacity);
    224
    225
    0
                for entry in 0..map_array.len() {
    226
    0
                    let start = map_array.value_offsets()[entry] as usize;
    227
    0
                    let end = map_array.value_offsets()[entry + 1] as usize;
    228
    229
    0
                    let maybe_matched = keys
    230
    0
                        .slice(start, end - start)
    231
    0
                        .iter()
    232
    0
                        .enumerate()
    233
    0
                        .find(|(_, t)| t.unwrap());
    234
    235
    0
                    if maybe_matched.is_none() {
    236
    0
                        mutable.extend_nulls(1);
    237
    0
                        continue;
    238
    0
                    }
    239
    0
                    let (match_offset, _) = maybe_matched.unwrap();
    240
    0
                    mutable.extend(0, start + match_offset, start + match_offset + 1);
    241
                }
    242
    243
    0
                let data = mutable.freeze();
    244
    0
                let data = make_array(data);
    245
    0
                Ok(ColumnarValue::Array(data))
    246
    0
            }
    247
    248
    0
            fn process_map_with_nested_key(
    249
    0
                array: Arc<dyn Array>,
    250
    0
                key_array: Arc<dyn Array>,
    251
    0
            ) -> Result<ColumnarValue> {
    252
    0
                let map_array = as_map_array(array.as_ref())?;
    253
    254
    0
                let comparator = make_comparator(
    255
    0
                    map_array.keys().as_ref(),
    256
    0
                    key_array.as_ref(),
    257
    0
                    SortOptions::default(),
    258
    0
                )?;
    259
    260
    0
                let original_data = map_array.entries().column(1).to_data();
    261
    0
                let capacity = Capacities::Array(original_data.len());
    262
    0
                let mut mutable =
    263
    0
                    MutableArrayData::with_capacities(vec![&original_data], true, capacity);
    264
    265
    0
                for entry in 0..map_array.len() {
    266
    0
                    let start = map_array.value_offsets()[entry] as usize;
    267
    0
                    let end = map_array.value_offsets()[entry + 1] as usize;
    268
    269
    0
                    let mut found_match = false;
    270
    0
                    for i in start..end {
    271
    0
                        if comparator(i, 0).is_eq() {
    272
    0
                            mutable.extend(0, i, i + 1);
    273
    0
                            found_match = true;
    274
    0
                            break;
    275
    0
                        }
    276
                    }
    277
    278
    0
                    if !found_match {
    279
    0
                        mutable.extend_nulls(1);
    280
    0
                    }
    281
                }
    282
    283
    0
                let data = mutable.freeze();
    284
    0
                let data = make_array(data);
    285
    0
                Ok(ColumnarValue::Array(data))
    286
    0
            }
    287
    288
    0
            match (array.data_type(), name) {
    289
    0
                (DataType::Map(_, _), ScalarValue::List(arr)) => {
    290
    0
                    let key_array: Arc<dyn Array> = arr;
    291
    0
                    process_map_array(array, key_array)
    292
                }
    293
    0
                (DataType::Map(_, _), ScalarValue::Struct(arr)) => {
    294
    0
                    process_map_array(array, arr as Arc<dyn Array>)
    295
                }
    296
    0
                (DataType::Map(_, _), other) => {
    297
    0
                    let data_type = other.data_type();
    298
    0
                    if data_type.is_nested() {
    299
    0
                        process_map_with_nested_key(array, other.to_array()?)
    300
                    } else {
    301
    0
                        process_map_array(array, other.to_array()?)
    302
                    }
    303
                }
    304
    0
                (DataType::Struct(_), ScalarValue::Utf8(Some(k))) => {
    305
    0
                    let as_struct_array = as_struct_array(&array)?;
    306
    0
                    match as_struct_array.column_by_name(&k) {
    307
    0
                        None => exec_err!("get indexed field {k} not found in struct"),
    308
    0
                        Some(col) => Ok(ColumnarValue::Array(Arc::clone(col))),
    309
                    }
    310
                }
    311
    0
                (DataType::Struct(_), name) => exec_err!(
    312
    0
                    "get_field is only possible on struct with utf8 indexes. \
    313
    0
                                 Received with {name:?} index"
    314
                ),
    315
    0
                (DataType::Null, _) => Ok(ColumnarValue::Scalar(ScalarValue::Null)),
    316
    0
                (dt, name) => exec_err!(
    317
    0
                    "get_field is only possible on maps with utf8 indexes or struct \
    318
    0
                                             with utf8 indexes. Received {dt} with {name:?} index"
    319
                ),
    320
            }
    321
    0
        }
    322
    323
    0
        fn documentation(&self) -> Option<&Documentation> {
    324
    0
            self.doc()
    325
    0
        }
    326
    }