## 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_tag.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, AsArray, DictionaryArray, Int8Array, StringArray};
    19
    use arrow::datatypes::DataType;
    20
    use datafusion_common::utils::take_function_args;
    21
    use datafusion_common::{exec_datafusion_err, exec_err, Result, ScalarValue};
    22
    use datafusion_doc::Documentation;
    23
    use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
    24
    use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
    25
    use datafusion_macros::user_doc;
    26
    use std::sync::Arc;
    27
    28
    #[user_doc(
    29
        doc_section(label = "Union Functions"),
    30
        description = "Returns the name of the currently selected field in the union",
    31
        syntax_example = "union_tag(union_expression)",
    32
        sql_example = r#"```sql
    33
    ❯ select union_column, union_tag(union_column) from table_with_union;
    34
    +--------------+-------------------------+
    35
    | union_column | union_tag(union_column) |
    36
    +--------------+-------------------------+
    37
    | {a=1}        | a                       |
    38
    | {b=3.0}      | b                       |
    39
    | {a=4}        | a                       |
    40
    | {b=}         | b                       |
    41
    | {a=}         | a                       |
    42
    +--------------+-------------------------+
    43
    ```"#,
    44
        standard_argument(name = "union", prefix = "Union")
    45
    )]
    46
    #[derive(Debug, PartialEq, Eq, Hash)]
    47
    pub struct UnionTagFunc {
    48
        signature: Signature,
    49
    }
    50
    51
    impl Default for UnionTagFunc {
    52
    0
        fn default() -> Self {
    53
    0
            Self::new()
    54
    0
        }
    55
    }
    56
    57
    impl UnionTagFunc {
    58
    0
        pub fn new() -> Self {
    59
    0
            Self {
    60
    0
                signature: Signature::any(1, Volatility::Immutable),
    61
    0
            }
    62
    0
        }
    63
    }
    64
    65
    impl ScalarUDFImpl for UnionTagFunc {
    66
    0
        fn as_any(&self) -> &dyn std::any::Any {
    67
    0
            self
    68
    0
        }
    69
    70
    0
        fn name(&self) -> &str {
    71
    0
            "union_tag"
    72
    0
        }
    73
    74
    0
        fn signature(&self) -> &Signature {
    75
    0
            &self.signature
    76
    0
        }
    77
    78
    0
        fn return_type(&self, _: &[DataType]) -> Result<DataType> {
    79
    0
            Ok(DataType::Dictionary(
    80
    0
                Box::new(DataType::Int8),
    81
    0
                Box::new(DataType::Utf8),
    82
    0
            ))
    83
    0
        }
    84
    85
    0
        fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
    86
    0
            let [union_] = take_function_args("union_tag", args.args)?;
    87
    88
    0
            match union_ {
    89
    0
                ColumnarValue::Array(array)
    90
    0
                    if matches!(array.data_type(), DataType::Union(_, _)) =>
    91
                {
    92
    0
                    let union_array = array.as_union();
    93
    94
    0
                    let keys = Int8Array::try_new(union_array.type_ids().clone(), None)?;
    95
    96
    0
                    let fields = match union_array.data_type() {
    97
    0
                        DataType::Union(fields, _) => fields,
    98
    0
                        _ => unreachable!(),
    99
                    };
    100
    101
                    // Union fields type IDs only constraints are being unique and in the 0..128 range:
    102
                    // They may not start at 0, be sequential, or even contiguous.
    103
                    // Therefore, we allocate a values vector with a length equal to the highest type ID plus one,
    104
                    // ensuring that each field's name can be placed at the index corresponding to its type ID.
    105
    0
                    let values_len = fields
    106
    0
                        .iter()
    107
    0
                        .map(|(type_id, _)| type_id + 1)
    108
    0
                        .max()
    109
    0
                        .unwrap_or_default() as usize;
    110
    111
    0
                    let mut values = vec![""; values_len];
    112
    113
    0
                    for (type_id, field) in fields.iter() {
    114
    0
                        values[type_id as usize] = field.name().as_str()
    115
                    }
    116
    117
    0
                    let values = Arc::new(StringArray::from(values));
    118
    119
                    // SAFETY: union type_ids are validated to not be smaller than zero.
    120
                    // values len is the union biggest type id plus one.
    121
                    // keys is built from the union type_ids, which contains only valid type ids
    122
                    // therefore, `keys[i] >= values.len() || keys[i] < 0` never occurs
    123
    0
                    let dict = unsafe { DictionaryArray::new_unchecked(keys, values) };
    124
    125
    0
                    Ok(ColumnarValue::Array(Arc::new(dict)))
    126
                }
    127
    0
                ColumnarValue::Scalar(ScalarValue::Union(value, fields, _)) => match value {
    128
    0
                    Some((value_type_id, _)) => fields
    129
    0
                        .iter()
    130
    0
                        .find(|(type_id, _)| value_type_id == *type_id)
    131
    0
                        .map(|(_, field)| {
    132
    0
                            ColumnarValue::Scalar(ScalarValue::Dictionary(
    133
    0
                                Box::new(DataType::Int8),
    134
    0
                                Box::new(field.name().as_str().into()),
    135
    0
                            ))
    136
    0
                        })
    137
    0
                        .ok_or_else(|| {
    138
    0
                            exec_datafusion_err!(
    139
    0
                                "union_tag: union scalar with unknown type_id {value_type_id}"
    140
                            )
    141
    0
                        }),
    142
    0
                    None => Ok(ColumnarValue::Scalar(ScalarValue::try_new_null(
    143
    0
                        args.return_field.data_type(),
    144
    0
                    )?)),
    145
                },
    146
    0
                v => exec_err!("union_tag only support unions, got {:?}", v.data_type()),
    147
            }
    148
    0
        }
    149
    150
    0
        fn documentation(&self) -> Option<&Documentation> {
    151
    0
            self.doc()
    152
    0
        }
    153
    }
    154
    155
    #[cfg(test)]
    156
    mod tests {
    157
        use super::UnionTagFunc;
    158
        use arrow::datatypes::{DataType, Field, UnionFields, UnionMode};
    159
        use datafusion_common::config::ConfigOptions;
    160
        use datafusion_common::ScalarValue;
    161
        use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
    162
        use std::sync::Arc;
    163
    164
        // when it becomes possible to construct union scalars in SQL, this should go to sqllogictests
    165
        #[test]
    166
        fn union_scalar() {
    167
            let fields = [(0, Arc::new(Field::new("a", DataType::UInt32, false)))]
    168
                .into_iter()
    169
                .collect();
    170
    171
            let scalar = ScalarValue::Union(
    172
                Some((0, Box::new(ScalarValue::UInt32(Some(0))))),
    173
                fields,
    174
                UnionMode::Dense,
    175
            );
    176
    177
            let return_type =
    178
                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8));
    179
    180
            let result = UnionTagFunc::new()
    181
                .invoke_with_args(ScalarFunctionArgs {
    182
                    args: vec![ColumnarValue::Scalar(scalar)],
    183
                    number_rows: 1,
    184
                    return_field: Field::new("res", return_type, true).into(),
    185
                    arg_fields: vec![],
    186
                    config_options: Arc::new(ConfigOptions::default()),
    187
                })
    188
                .unwrap();
    189
    190
            assert_scalar(
    191
                result,
    192
                ScalarValue::Dictionary(Box::new(DataType::Int8), Box::new("a".into())),
    193
            );
    194
        }
    195
    196
        #[test]
    197
        fn union_scalar_empty() {
    198
            let scalar = ScalarValue::Union(None, UnionFields::empty(), UnionMode::Dense);
    199
    200
            let return_type =
    201
                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8));
    202
    203
            let result = UnionTagFunc::new()
    204
                .invoke_with_args(ScalarFunctionArgs {
    205
                    args: vec![ColumnarValue::Scalar(scalar)],
    206
                    number_rows: 1,
    207
                    return_field: Field::new("res", return_type, true).into(),
    208
                    arg_fields: vec![],
    209
                    config_options: Arc::new(ConfigOptions::default()),
    210
                })
    211
                .unwrap();
    212
    213
            assert_scalar(
    214
                result,
    215
                ScalarValue::Dictionary(
    216
                    Box::new(DataType::Int8),
    217
                    Box::new(ScalarValue::Utf8(None)),
    218
                ),
    219
            );
    220
        }
    221
    222
        fn assert_scalar(value: ColumnarValue, expected: ScalarValue) {
    223
            match value {
    224
                ColumnarValue::Array(array) => panic!("expected scalar got {array:?}"),
    225
                ColumnarValue::Scalar(scalar) => assert_eq!(scalar, expected),
    226
            }
    227
        }
    228
    }