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
}