## 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/physical-expr/src/expressions/cast_column.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
    //! Physical expression for struct-aware casting of columns.
    19
    20
    use crate::physical_expr::PhysicalExpr;
    21
    use arrow::{
    22
        compute::CastOptions,
    23
        datatypes::{DataType, FieldRef, Schema},
    24
        record_batch::RecordBatch,
    25
    };
    26
    use datafusion_common::{
    27
        format::DEFAULT_CAST_OPTIONS, nested_struct::cast_column, Result, ScalarValue,
    28
    };
    29
    use datafusion_expr_common::columnar_value::ColumnarValue;
    30
    use std::{
    31
        any::Any,
    32
        fmt::{self, Display},
    33
        hash::Hash,
    34
        sync::Arc,
    35
    };
    36
    /// A physical expression that applies [`cast_column`] to its input.
    37
    ///
    38
    /// [`CastColumnExpr`] extends the regular [`CastExpr`](super::CastExpr) by
    39
    /// retaining schema metadata for both the input and output fields. This allows
    40
    /// the evaluator to perform struct-aware casts that honour nested field
    41
    /// ordering, preserve nullability, and fill missing fields with null values.
    42
    ///
    43
    /// This expression is intended for schema rewriting scenarios where the
    44
    /// planner already resolved the input column but needs to adapt its physical
    45
    /// representation to a new [`arrow::datatypes::Field`]. It mirrors the behaviour of the
    46
    /// [`datafusion_common::nested_struct::cast_column`] helper while integrating
    47
    /// with the `PhysicalExpr` trait so it can participate in the execution plan
    48
    /// like any other column expression.
    49
    #[derive(Debug, Clone, Eq)]
    50
    pub struct CastColumnExpr {
    51
        /// The physical expression producing the value to cast.
    52
        expr: Arc<dyn PhysicalExpr>,
    53
        /// The logical field of the input column.
    54
        input_field: FieldRef,
    55
        /// The field metadata describing the desired output column.
    56
        target_field: FieldRef,
    57
        /// Options forwarded to [`cast_column`].
    58
        cast_options: CastOptions<'static>,
    59
    }
    60
    61
    // Manually derive `PartialEq`/`Hash` as `Arc<dyn PhysicalExpr>` does not
    62
    // implement these traits by default for the trait object.
    63
    impl PartialEq for CastColumnExpr {
    64
    0
        fn eq(&self, other: &Self) -> bool {
    65
    0
            self.expr.eq(&other.expr)
    66
    0
                && self.input_field.eq(&other.input_field)
    67
    0
                && self.target_field.eq(&other.target_field)
    68
    0
                && self.cast_options.eq(&other.cast_options)
    69
    0
        }
    70
    }
    71
    72
    impl Hash for CastColumnExpr {
    73
    0
        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
    74
    0
            self.expr.hash(state);
    75
    0
            self.input_field.hash(state);
    76
    0
            self.target_field.hash(state);
    77
    0
            self.cast_options.hash(state);
    78
    0
        }
    79
    }
    80
    81
    impl CastColumnExpr {
    82
        /// Create a new [`CastColumnExpr`].
    83
    0
        pub fn new(
    84
    0
            expr: Arc<dyn PhysicalExpr>,
    85
    0
            input_field: FieldRef,
    86
    0
            target_field: FieldRef,
    87
    0
            cast_options: Option<CastOptions<'static>>,
    88
    0
        ) -> Self {
    89
    0
            Self {
    90
    0
                expr,
    91
    0
                input_field,
    92
    0
                target_field,
    93
    0
                cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS),
    94
    0
            }
    95
    0
        }
    96
    97
        /// The expression that produces the value to be cast.
    98
    0
        pub fn expr(&self) -> &Arc<dyn PhysicalExpr> {
    99
    0
            &self.expr
    100
    0
        }
    101
    102
        /// Field metadata describing the resolved input column.
    103
    0
        pub fn input_field(&self) -> &FieldRef {
    104
    0
            &self.input_field
    105
    0
        }
    106
    107
        /// Field metadata describing the output column after casting.
    108
    0
        pub fn target_field(&self) -> &FieldRef {
    109
    0
            &self.target_field
    110
    0
        }
    111
    }
    112
    113
    impl Display for CastColumnExpr {
    114
    0
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    115
    0
            write!(
    116
    0
                f,
    117
    0
                "CAST_COLUMN({} AS {:?})",
    118
                self.expr,
    119
    0
                self.target_field.data_type()
    120
            )
    121
    0
        }
    122
    }
    123
    124
    impl PhysicalExpr for CastColumnExpr {
    125
    0
        fn as_any(&self) -> &dyn Any {
    126
    0
            self
    127
    0
        }
    128
    129
    0
        fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
    130
    0
            Ok(self.target_field.data_type().clone())
    131
    0
        }
    132
    133
    0
        fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
    134
    0
            Ok(self.target_field.is_nullable())
    135
    0
        }
    136
    137
    0
        fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
    138
    0
            let value = self.expr.evaluate(batch)?;
    139
    0
            match value {
    140
    0
                ColumnarValue::Array(array) => {
    141
    0
                    let casted =
    142
    0
                        cast_column(&array, self.target_field.as_ref(), &self.cast_options)?;
    143
    0
                    Ok(ColumnarValue::Array(casted))
    144
                }
    145
    0
                ColumnarValue::Scalar(scalar) => {
    146
    0
                    let as_array = scalar.to_array_of_size(1)?;
    147
    0
                    let casted = cast_column(
    148
    0
                        &as_array,
    149
    0
                        self.target_field.as_ref(),
    150
    0
                        &self.cast_options,
    151
    0
                    )?;
    152
    0
                    let result = ScalarValue::try_from_array(casted.as_ref(), 0)?;
    153
    0
                    Ok(ColumnarValue::Scalar(result))
    154
                }
    155
            }
    156
    0
        }
    157
    158
    0
        fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
    159
    0
            Ok(Arc::clone(&self.target_field))
    160
    0
        }
    161
    162
    0
        fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
    163
    0
            vec![&self.expr]
    164
    0
        }
    165
    166
    0
        fn with_new_children(
    167
    0
            self: Arc<Self>,
    168
    0
            mut children: Vec<Arc<dyn PhysicalExpr>>,
    169
    0
        ) -> Result<Arc<dyn PhysicalExpr>> {
    170
    0
            assert_eq!(children.len(), 1);
    171
    0
            let child = children.pop().expect("CastColumnExpr child");
    172
    0
            Ok(Arc::new(Self::new(
    173
    0
                child,
    174
    0
                Arc::clone(&self.input_field),
    175
    0
                Arc::clone(&self.target_field),
    176
    0
                Some(self.cast_options.clone()),
    177
    0
            )))
    178
    0
        }
    179
    180
    0
        fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    181
    0
            Display::fmt(self, f)
    182
    0
        }
    183
    }
    184
    185
    #[cfg(test)]
    186
    mod tests {
    187
        use super::*;
    188
    189
        use crate::expressions::{Column, Literal};
    190
        use arrow::{
    191
            array::{Array, ArrayRef, BooleanArray, Int32Array, StringArray, StructArray},
    192
            datatypes::{DataType, Field, Fields, SchemaRef},
    193
        };
    194
        use datafusion_common::{
    195
            cast::{as_int64_array, as_string_array, as_struct_array, as_uint8_array},
    196
            Result as DFResult, ScalarValue,
    197
        };
    198
    199
    0
        fn make_schema(field: &Field) -> SchemaRef {
    200
    0
            Arc::new(Schema::new(vec![field.clone()]))
    201
    0
        }
    202
    203
    0
        fn make_struct_array(fields: Fields, arrays: Vec<ArrayRef>) -> StructArray {
    204
    0
            StructArray::new(fields, arrays, None)
    205
    0
        }
    206
    207
        #[test]
    208
    0
        fn cast_primitive_array() -> DFResult<()> {
    209
    0
            let input_field = Field::new("a", DataType::Int32, true);
    210
    0
            let target_field = Field::new("a", DataType::Int64, true);
    211
    0
            let schema = make_schema(&input_field);
    212
    213
    0
            let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]));
    214
    0
            let batch = RecordBatch::try_new(Arc::clone(&schema), vec![values])?;
    215
    216
    0
            let column = Arc::new(Column::new_with_schema("a", schema.as_ref())?);
    217
    0
            let expr = CastColumnExpr::new(
    218
    0
                column,
    219
    0
                Arc::new(input_field.clone()),
    220
    0
                Arc::new(target_field.clone()),
    221
    0
                None,
    222
            );
    223
    224
    0
            let result = expr.evaluate(&batch)?;
    225
    0
            let ColumnarValue::Array(array) = result else {
    226
    0
                panic!("expected array");
    227
            };
    228
    0
            let casted = as_int64_array(array.as_ref())?;
    229
    0
            assert_eq!(casted.value(0), 1);
    230
    0
            assert!(casted.is_null(1));
    231
    0
            assert_eq!(casted.value(2), 3);
    232
    0
            Ok(())
    233
    0
        }
    234
    235
        #[test]
    236
    0
        fn cast_struct_array_missing_child() -> DFResult<()> {
    237
    0
            let source_a = Field::new("a", DataType::Int32, true);
    238
    0
            let source_b = Field::new("b", DataType::Utf8, true);
    239
    0
            let input_field = Field::new(
    240
                "s",
    241
    0
                DataType::Struct(
    242
    0
                    vec![Arc::new(source_a.clone()), Arc::new(source_b.clone())].into(),
    243
    0
                ),
    244
                true,
    245
            );
    246
    0
            let target_a = Field::new("a", DataType::Int64, true);
    247
    0
            let target_c = Field::new("c", DataType::Utf8, true);
    248
    0
            let target_field = Field::new(
    249
                "s",
    250
    0
                DataType::Struct(
    251
    0
                    vec![Arc::new(target_a.clone()), Arc::new(target_c.clone())].into(),
    252
    0
                ),
    253
                true,
    254
            );
    255
    256
    0
            let schema = make_schema(&input_field);
    257
    0
            let struct_array = make_struct_array(
    258
    0
                vec![Arc::new(source_a.clone()), Arc::new(source_b.clone())].into(),
    259
    0
                vec![
    260
    0
                    Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef,
    261
    0
                    Arc::new(StringArray::from(vec![Some("alpha"), Some("beta")]))
    262
    0
                        as ArrayRef,
    263
                ],
    264
            );
    265
    0
            let batch = RecordBatch::try_new(
    266
    0
                Arc::clone(&schema),
    267
    0
                vec![Arc::new(struct_array) as Arc<_>],
    268
    0
            )?;
    269
    270
    0
            let column = Arc::new(Column::new_with_schema("s", schema.as_ref())?);
    271
    0
            let expr = CastColumnExpr::new(
    272
    0
                column,
    273
    0
                Arc::new(input_field.clone()),
    274
    0
                Arc::new(target_field.clone()),
    275
    0
                None,
    276
            );
    277
    278
    0
            let result = expr.evaluate(&batch)?;
    279
    0
            let ColumnarValue::Array(array) = result else {
    280
    0
                panic!("expected array");
    281
            };
    282
    0
            let struct_array = as_struct_array(array.as_ref())?;
    283
    0
            let cast_a = as_int64_array(struct_array.column_by_name("a").unwrap().as_ref())?;
    284
    0
            assert_eq!(cast_a.value(0), 1);
    285
    0
            assert!(cast_a.is_null(1));
    286
    287
    0
            let cast_c = as_string_array(struct_array.column_by_name("c").unwrap().as_ref())?;
    288
    0
            assert!(cast_c.is_null(0));
    289
    0
            assert!(cast_c.is_null(1));
    290
    0
            Ok(())
    291
    0
        }
    292
    293
        #[test]
    294
    0
        fn cast_nested_struct_array() -> DFResult<()> {
    295
    0
            let inner_source = Field::new(
    296
                "inner",
    297
    0
                DataType::Struct(
    298
    0
                    vec![Arc::new(Field::new("x", DataType::Int32, true))].into(),
    299
    0
                ),
    300
                true,
    301
            );
    302
    0
            let outer_field = Field::new(
    303
                "root",
    304
    0
                DataType::Struct(vec![Arc::new(inner_source.clone())].into()),
    305
                true,
    306
            );
    307
    308
    0
            let inner_target = Field::new(
    309
                "inner",
    310
    0
                DataType::Struct(
    311
    0
                    vec![
    312
    0
                        Arc::new(Field::new("x", DataType::Int64, true)),
    313
    0
                        Arc::new(Field::new("y", DataType::Boolean, true)),
    314
    0
                    ]
    315
    0
                    .into(),
    316
    0
                ),
    317
                true,
    318
            );
    319
    0
            let target_field = Field::new(
    320
                "root",
    321
    0
                DataType::Struct(vec![Arc::new(inner_target.clone())].into()),
    322
                true,
    323
            );
    324
    325
    0
            let schema = make_schema(&outer_field);
    326
    327
    0
            let inner_struct = make_struct_array(
    328
    0
                vec![Arc::new(Field::new("x", DataType::Int32, true))].into(),
    329
    0
                vec![Arc::new(Int32Array::from(vec![Some(7), None])) as ArrayRef],
    330
            );
    331
    0
            let outer_struct = make_struct_array(
    332
    0
                vec![Arc::new(inner_source.clone())].into(),
    333
    0
                vec![Arc::new(inner_struct) as ArrayRef],
    334
            );
    335
    0
            let batch = RecordBatch::try_new(
    336
    0
                Arc::clone(&schema),
    337
    0
                vec![Arc::new(outer_struct) as ArrayRef],
    338
    0
            )?;
    339
    340
    0
            let column = Arc::new(Column::new_with_schema("root", schema.as_ref())?);
    341
    0
            let expr = CastColumnExpr::new(
    342
    0
                column,
    343
    0
                Arc::new(outer_field.clone()),
    344
    0
                Arc::new(target_field.clone()),
    345
    0
                None,
    346
            );
    347
    348
    0
            let result = expr.evaluate(&batch)?;
    349
    0
            let ColumnarValue::Array(array) = result else {
    350
    0
                panic!("expected array");
    351
            };
    352
    0
            let struct_array = as_struct_array(array.as_ref())?;
    353
    0
            let inner =
    354
    0
                as_struct_array(struct_array.column_by_name("inner").unwrap().as_ref())?;
    355
    0
            let x = as_int64_array(inner.column_by_name("x").unwrap().as_ref())?;
    356
    0
            assert_eq!(x.value(0), 7);
    357
    0
            assert!(x.is_null(1));
    358
    0
            let y = inner.column_by_name("y").unwrap();
    359
    0
            let y = y
    360
    0
                .as_any()
    361
    0
                .downcast_ref::<BooleanArray>()
    362
    0
                .expect("boolean array");
    363
    0
            assert!(y.is_null(0));
    364
    0
            assert!(y.is_null(1));
    365
    0
            Ok(())
    366
    0
        }
    367
    368
        #[test]
    369
    0
        fn cast_struct_scalar() -> DFResult<()> {
    370
    0
            let source_field = Field::new("a", DataType::Int32, true);
    371
    0
            let input_field = Field::new(
    372
                "s",
    373
    0
                DataType::Struct(vec![Arc::new(source_field.clone())].into()),
    374
                true,
    375
            );
    376
    0
            let target_field = Field::new(
    377
                "s",
    378
    0
                DataType::Struct(
    379
    0
                    vec![Arc::new(Field::new("a", DataType::UInt8, true))].into(),
    380
    0
                ),
    381
                true,
    382
            );
    383
    384
    0
            let schema = make_schema(&input_field);
    385
    0
            let scalar_struct = StructArray::new(
    386
    0
                vec![Arc::new(source_field.clone())].into(),
    387
    0
                vec![Arc::new(Int32Array::from(vec![Some(9)])) as ArrayRef],
    388
    0
                None,
    389
            );
    390
    0
            let literal =
    391
    0
                Arc::new(Literal::new(ScalarValue::Struct(Arc::new(scalar_struct))));
    392
    0
            let expr = CastColumnExpr::new(
    393
    0
                literal,
    394
    0
                Arc::new(input_field.clone()),
    395
    0
                Arc::new(target_field.clone()),
    396
    0
                None,
    397
            );
    398
    399
    0
            let batch = RecordBatch::new_empty(Arc::clone(&schema));
    400
    0
            let result = expr.evaluate(&batch)?;
    401
    0
            let ColumnarValue::Scalar(ScalarValue::Struct(array)) = result else {
    402
    0
                panic!("expected struct scalar");
    403
            };
    404
    0
            let casted = array.column_by_name("a").unwrap();
    405
    0
            let casted = as_uint8_array(casted.as_ref())?;
    406
    0
            assert_eq!(casted.value(0), 9);
    407
    0
            Ok(())
    408
    0
        }
    409
    }