## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/common/src/scalar/struct_builder.rs

    Line
    Count
    Source (jump to first uncovered line)
    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
    //! [`ScalarStructBuilder`] for building [`ScalarValue::Struct`]
    19
    20
    use crate::error::_internal_err;
    21
    use crate::{Result, ScalarValue};
    22
    use arrow::array::{ArrayRef, StructArray};
    23
    use arrow::datatypes::{DataType, FieldRef, Fields};
    24
    use arrow_schema::Field;
    25
    use std::sync::Arc;
    26
    27
    /// Builder for [`ScalarValue::Struct`].
    28
    ///
    29
    /// See examples on [`ScalarValue`]
    30
    #[derive(Debug, Default)]
    31
    pub struct ScalarStructBuilder {
    32
        fields: Vec<FieldRef>,
    33
        arrays: Vec<ArrayRef>,
    34
    }
    35
    36
    impl ScalarStructBuilder {
    37
        /// Create a new `ScalarStructBuilder`
    38
    0
        pub fn new() -> Self {
    39
    0
            Self::default()
    40
    0
        }
    41
    42
        /// Return a new [`ScalarValue::Struct`] with a single `null` value.
    43
        ///
    44
        /// Note this is different from a struct where each of the specified fields
    45
        /// are null (e.g. `{a: NULL}`)
    46
        ///
    47
        /// # Example
    48
        ///
    49
        /// ```rust
    50
        /// # use arrow::datatypes::{DataType, Field};
    51
        /// # use datafusion_common::scalar::ScalarStructBuilder;
    52
        /// let fields = vec![
    53
        ///    Field::new("a", DataType::Int32, false),
    54
        /// ];
    55
        /// let sv = ScalarStructBuilder::new_null(fields);
    56
        /// // Note this is `NULL`, not `{a: NULL}`
    57
        /// assert_eq!(format!("{sv}"), "NULL");
    58
        ///```
    59
        ///
    60
        /// To create a struct where the *fields* are null, use `Self::new()` and
    61
        /// pass null values for each field:
    62
        ///
    63
        /// ```rust
    64
        /// # use arrow::datatypes::{DataType, Field};
    65
        /// # use datafusion_common::scalar::{ScalarStructBuilder, ScalarValue};
    66
        /// // make a nullable field
    67
        /// let field = Field::new("a", DataType::Int32, true);
    68
        /// // add a null value for the "a" field
    69
        /// let sv = ScalarStructBuilder::new()
    70
        ///   .with_scalar(field, ScalarValue::Int32(None))
    71
        ///   .build()
    72
        ///   .unwrap();
    73
        /// // value is not null, but field is
    74
        /// assert_eq!(format!("{sv}"), "{a:}");
    75
        /// ```
    76
    0
        pub fn new_null(fields: impl IntoFields) -> ScalarValue {
    77
    0
            DataType::Struct(fields.into()).try_into().unwrap()
    78
    0
        }
    79
    80
        /// Add the specified field and [`ArrayRef`] to the struct.
    81
        ///
    82
        /// Note the array should have a single row.
    83
    0
        pub fn with_array(mut self, field: impl IntoFieldRef, value: ArrayRef) -> Self {
    84
    0
            self.fields.push(field.into_field_ref());
    85
    0
            self.arrays.push(value);
    86
    0
            self
    87
    0
        }
    88
    89
        /// Add the specified field and `ScalarValue` to the struct.
    90
    0
        pub fn with_scalar(self, field: impl IntoFieldRef, value: ScalarValue) -> Self {
    91
    0
            // valid scalar value should not fail
    92
    0
            let array = value.to_array().unwrap();
    93
    0
            self.with_array(field, array)
    94
    0
        }
    95
    96
        /// Add a field with the specified name and value to the struct.
    97
        /// the field is created with the specified data type and as non nullable
    98
    0
        pub fn with_name_and_scalar(self, name: &str, value: ScalarValue) -> Self {
    99
    0
            let field = Field::new(name, value.data_type(), false);
    100
    0
            self.with_scalar(field, value)
    101
    0
        }
    102
    103
        /// Return a [`ScalarValue::Struct`] with the fields and values added so far
    104
        ///
    105
        /// # Errors
    106
        ///
    107
        /// If the [`StructArray`] cannot be created (for example if there is a
    108
        /// mismatch between field types and arrays) or the arrays do not have
    109
        /// exactly one element.
    110
    0
        pub fn build(self) -> Result<ScalarValue> {
    111
    0
            let Self { fields, arrays } = self;
    112
    113
    0
            for array in &arrays {
    114
    0
                if array.len() != 1 {
    115
    0
                    return _internal_err!(
    116
    0
                        "Error building ScalarValue::Struct. \
    117
    0
                    Expected array with exactly one element, found array with {} elements",
    118
    0
                        array.len()
    119
    0
                    );
    120
    0
                }
    121
            }
    122
    123
    0
            let struct_array = StructArray::try_new(Fields::from(fields), arrays, None)?;
    124
    0
            Ok(ScalarValue::Struct(Arc::new(struct_array)))
    125
    0
        }
    126
    }
    127
    128
    /// Trait for converting a type into a [`FieldRef`]
    129
    ///
    130
    /// Used to avoid having to call `clone()` on a `FieldRef` when adding a field to
    131
    /// a `ScalarStructBuilder`.
    132
    ///
    133
    /// TODO potentially upstream this to arrow-rs so that we can
    134
    /// use impl `Into<FieldRef>` instead
    135
    pub trait IntoFieldRef {
    136
        fn into_field_ref(self) -> FieldRef;
    137
    }
    138
    139
    impl IntoFieldRef for FieldRef {
    140
    0
        fn into_field_ref(self) -> FieldRef {
    141
    0
            self
    142
    0
        }
    143
    }
    144
    145
    impl IntoFieldRef for &FieldRef {
    146
    0
        fn into_field_ref(self) -> FieldRef {
    147
    0
            Arc::clone(self)
    148
    0
        }
    149
    }
    150
    151
    impl IntoFieldRef for Field {
    152
    0
        fn into_field_ref(self) -> FieldRef {
    153
    0
            FieldRef::new(self)
    154
    0
        }
    155
    }
    156
    157
    /// Trait for converting a type into a [`Fields`]
    158
    ///
    159
    /// This avoids to avoid having to call clone() on an Arc'd `Fields` when adding
    160
    /// a field to a `ScalarStructBuilder`
    161
    ///
    162
    /// TODO potentially upstream this to arrow-rs so that we can
    163
    /// use impl `Into<Fields>` instead
    164
    pub trait IntoFields {
    165
        fn into(self) -> Fields;
    166
    }
    167
    168
    impl IntoFields for Fields {
    169
    0
        fn into(self) -> Fields {
    170
    0
            self
    171
    0
        }
    172
    }
    173
    174
    impl IntoFields for &Fields {
    175
    0
        fn into(self) -> Fields {
    176
    0
            self.clone()
    177
    0
        }
    178
    }
    179
    180
    impl IntoFields for Vec<Field> {
    181
        fn into(self) -> Fields {
    182
            Fields::from(self)
    183
        }
    184
    }