## 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/union_extract.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::Array;
    19
    use arrow::datatypes::{DataType, Field, FieldRef, UnionFields};
    20
    use datafusion_common::cast::as_union_array;
    21
    use datafusion_common::utils::take_function_args;
    22
    use datafusion_common::{
    23
        exec_datafusion_err, exec_err, internal_err, Result, ScalarValue,
    24
    };
    25
    use datafusion_doc::Documentation;
    26
    use datafusion_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs};
    27
    use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
    28
    use datafusion_macros::user_doc;
    29
    30
    #[user_doc(
    31
        doc_section(label = "Union Functions"),
    32
        description = "Returns the value of the given field in the union when selected, or NULL otherwise.",
    33
        syntax_example = "union_extract(union, field_name)",
    34
        sql_example = r#"```sql
    35
    ❯ select union_column, union_extract(union_column, 'a'), union_extract(union_column, 'b') from table_with_union;
    36
    +--------------+----------------------------------+----------------------------------+
    37
    | union_column | union_extract(union_column, 'a') | union_extract(union_column, 'b') |
    38
    +--------------+----------------------------------+----------------------------------+
    39
    | {a=1}        | 1                                |                                  |
    40
    | {b=3.0}      |                                  | 3.0                              |
    41
    | {a=4}        | 4                                |                                  |
    42
    | {b=}         |                                  |                                  |
    43
    | {a=}         |                                  |                                  |
    44
    +--------------+----------------------------------+----------------------------------+
    45
    ```"#,
    46
        standard_argument(name = "union", prefix = "Union"),
    47
        argument(
    48
            name = "field_name",
    49
            description = "String expression to operate on. Must be a constant."
    50
        )
    51
    )]
    52
    #[derive(Debug, PartialEq, Eq, Hash)]
    53
    pub struct UnionExtractFun {
    54
        signature: Signature,
    55
    }
    56
    57
    impl Default for UnionExtractFun {
    58
    0
        fn default() -> Self {
    59
    0
            Self::new()
    60
    0
        }
    61
    }
    62
    63
    impl UnionExtractFun {
    64
    0
        pub fn new() -> Self {
    65
    0
            Self {
    66
    0
                signature: Signature::any(2, Volatility::Immutable),
    67
    0
            }
    68
    0
        }
    69
    }
    70
    71
    impl ScalarUDFImpl for UnionExtractFun {
    72
    0
        fn as_any(&self) -> &dyn std::any::Any {
    73
    0
            self
    74
    0
        }
    75
    76
    0
        fn name(&self) -> &str {
    77
    0
            "union_extract"
    78
    0
        }
    79
    80
    0
        fn signature(&self) -> &Signature {
    81
    0
            &self.signature
    82
    0
        }
    83
    84
    0
        fn return_type(&self, _: &[DataType]) -> Result<DataType> {
    85
            // should be using return_field_from_args and not calling the default implementation
    86
    0
            internal_err!("union_extract should return type from args")
    87
    0
        }
    88
    89
    0
        fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
    90
    0
            if args.arg_fields.len() != 2 {
    91
    0
                return exec_err!(
    92
    0
                    "union_extract expects 2 arguments, got {} instead",
    93
    0
                    args.arg_fields.len()
    94
                );
    95
    0
            }
    96
    97
    0
            let DataType::Union(fields, _) = &args.arg_fields[0].data_type() else {
    98
    0
                return exec_err!(
    99
    0
                    "union_extract first argument must be a union, got {} instead",
    100
    0
                    args.arg_fields[0].data_type()
    101
                );
    102
            };
    103
    104
    0
            let Some(ScalarValue::Utf8(Some(field_name))) = &args.scalar_arguments[1] else {
    105
    0
                return exec_err!(
    106
    0
                    "union_extract second argument must be a non-null string literal, got {} instead",
    107
    0
                    args.arg_fields[1].data_type()
    108
                );
    109
            };
    110
    111
    0
            let field = find_field(fields, field_name)?.1;
    112
    113
    0
            Ok(Field::new(self.name(), field.data_type().clone(), true).into())
    114
    0
        }
    115
    116
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    117
    0
            let [array, target_name] = take_function_args("union_extract", args.args)?;
    118
    119
    0
            let target_name = match target_name {
    120
    0
                ColumnarValue::Scalar(ScalarValue::Utf8(Some(target_name))) => Ok(target_name),
    121
    0
                ColumnarValue::Scalar(ScalarValue::Utf8(None)) => exec_err!("union_extract second argument must be a non-null string literal, got a null instead"),
    122
    0
                _ => exec_err!("union_extract second argument must be a non-null string literal, got {} instead", target_name.data_type()),
    123
    0
            }?;
    124
    125
    0
            match array {
    126
    0
                ColumnarValue::Array(array) => {
    127
    0
                    let union_array = as_union_array(&array).map_err(|_| {
    128
    0
                        exec_datafusion_err!(
    129
    0
                            "union_extract first argument must be a union, got {} instead",
    130
    0
                            array.data_type()
    131
                        )
    132
    0
                    })?;
    133
    134
                    Ok(ColumnarValue::Array(
    135
    0
                        arrow::compute::kernels::union_extract::union_extract(
    136
    0
                            union_array,
    137
    0
                            &target_name,
    138
    0
                        )?,
    139
                    ))
    140
                }
    141
    0
                ColumnarValue::Scalar(ScalarValue::Union(value, fields, _)) => {
    142
    0
                    let (target_type_id, target) = find_field(&fields, &target_name)?;
    143
    144
    0
                    let result = match value {
    145
    0
                        Some((type_id, value)) if target_type_id == type_id => *value,
    146
    0
                        _ => ScalarValue::try_new_null(target.data_type())?,
    147
                    };
    148
    149
    0
                    Ok(ColumnarValue::Scalar(result))
    150
                }
    151
    0
                other => exec_err!(
    152
    0
                    "union_extract first argument must be a union, got {} instead",
    153
    0
                    other.data_type()
    154
                ),
    155
            }
    156
    0
        }
    157
    158
    0
        fn documentation(&self) -> Option<&Documentation> {
    159
    0
            self.doc()
    160
    0
        }
    161
    }
    162
    163
    0
    fn find_field<'a>(fields: &'a UnionFields, name: &str) -> Result<(i8, &'a FieldRef)> {
    164
    0
        fields
    165
    0
            .iter()
    166
    0
            .find(|field| field.1.name() == name)
    167
    0
            .ok_or_else(|| exec_datafusion_err!("field {name} not found on union"))
    168
    0
    }
    169
    170
    #[cfg(test)]
    171
    mod tests {
    172
        use arrow::datatypes::{DataType, Field, UnionFields, UnionMode};
    173
        use datafusion_common::config::ConfigOptions;
    174
        use datafusion_common::{Result, ScalarValue};
    175
        use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
    176
        use std::sync::Arc;
    177
    178
        use super::UnionExtractFun;
    179
    180
        // when it becomes possible to construct union scalars in SQL, this should go to sqllogictests
    181
        #[test]
    182
        fn test_scalar_value() -> Result<()> {
    183
            let fun = UnionExtractFun::new();
    184
    185
            let fields = UnionFields::new(
    186
                vec![1, 3],
    187
                vec![
    188
                    Field::new("str", DataType::Utf8, false),
    189
                    Field::new("int", DataType::Int32, false),
    190
                ],
    191
            );
    192
    193
            let args = vec![
    194
                ColumnarValue::Scalar(ScalarValue::Union(
    195
                    None,
    196
                    fields.clone(),
    197
                    UnionMode::Dense,
    198
                )),
    199
                ColumnarValue::Scalar(ScalarValue::new_utf8("str")),
    200
            ];
    201
            let arg_fields = args
    202
                .iter()
    203
                .map(|arg| Field::new("a", arg.data_type().clone(), true).into())
    204
                .collect::<Vec<_>>();
    205
    206
            let result = fun.invoke_with_args(ScalarFunctionArgs {
    207
                args,
    208
                arg_fields,
    209
                number_rows: 1,
    210
                return_field: Field::new("f", DataType::Utf8, true).into(),
    211
                config_options: Arc::new(ConfigOptions::default()),
    212
            })?;
    213
    214
            assert_scalar(result, ScalarValue::Utf8(None));
    215
    216
            let args = vec![
    217
                ColumnarValue::Scalar(ScalarValue::Union(
    218
                    Some((3, Box::new(ScalarValue::Int32(Some(42))))),
    219
                    fields.clone(),
    220
                    UnionMode::Dense,
    221
                )),
    222
                ColumnarValue::Scalar(ScalarValue::new_utf8("str")),
    223
            ];
    224
            let arg_fields = args
    225
                .iter()
    226
                .map(|arg| Field::new("a", arg.data_type().clone(), true).into())
    227
                .collect::<Vec<_>>();
    228
    229
            let result = fun.invoke_with_args(ScalarFunctionArgs {
    230
                args,
    231
                arg_fields,
    232
                number_rows: 1,
    233
                return_field: Field::new("f", DataType::Utf8, true).into(),
    234
                config_options: Arc::new(ConfigOptions::default()),
    235
            })?;
    236
    237
            assert_scalar(result, ScalarValue::Utf8(None));
    238
    239
            let args = vec![
    240
                ColumnarValue::Scalar(ScalarValue::Union(
    241
                    Some((1, Box::new(ScalarValue::new_utf8("42")))),
    242
                    fields.clone(),
    243
                    UnionMode::Dense,
    244
                )),
    245
                ColumnarValue::Scalar(ScalarValue::new_utf8("str")),
    246
            ];
    247
            let arg_fields = args
    248
                .iter()
    249
                .map(|arg| Field::new("a", arg.data_type().clone(), true).into())
    250
                .collect::<Vec<_>>();
    251
            let result = fun.invoke_with_args(ScalarFunctionArgs {
    252
                args,
    253
                arg_fields,
    254
                number_rows: 1,
    255
                return_field: Field::new("f", DataType::Utf8, true).into(),
    256
                config_options: Arc::new(ConfigOptions::default()),
    257
            })?;
    258
    259
            assert_scalar(result, ScalarValue::new_utf8("42"));
    260
    261
            Ok(())
    262
        }
    263
    264
        fn assert_scalar(value: ColumnarValue, expected: ScalarValue) {
    265
            match value {
    266
                ColumnarValue::Array(array) => panic!("expected scalar got {array:?}"),
    267
                ColumnarValue::Scalar(scalar) => assert_eq!(scalar, expected),
    268
            }
    269
        }
    270
    }