datatypes/types/
struct_type.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::BTreeMap;
16use std::sync::Arc;
17
18use arrow::datatypes::{DataType as ArrowDataType, Field};
19use arrow_schema::Fields;
20use serde::{Deserialize, Serialize};
21
22use crate::error::Result;
23use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId};
24use crate::value::Value;
25use crate::vectors::StructVectorBuilder;
26
27#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
28pub struct StructType {
29    fields: Arc<Vec<StructField>>,
30}
31
32impl From<&Fields> for StructType {
33    fn from(value: &Fields) -> Self {
34        let fields = value
35            .iter()
36            .map(|field| {
37                StructField::new(
38                    field.name().clone(),
39                    ConcreteDataType::from_arrow_type(field.data_type()),
40                    field.is_nullable(),
41                )
42            })
43            .collect::<Vec<_>>();
44        StructType {
45            fields: Arc::new(fields),
46        }
47    }
48}
49
50impl<const N: usize> From<[StructField; N]> for StructType {
51    fn from(value: [StructField; N]) -> Self {
52        let value: Box<[StructField]> = Box::new(value);
53        Self {
54            fields: Arc::new(value.into_vec()),
55        }
56    }
57}
58
59impl DataType for StructType {
60    fn name(&self) -> String {
61        format!(
62            "Struct<{}>",
63            self.fields
64                .iter()
65                .map(|f| format!(r#""{}": {}"#, f.name(), f.data_type()))
66                .collect::<Vec<_>>()
67                .join(", ")
68        )
69    }
70
71    fn logical_type_id(&self) -> LogicalTypeId {
72        LogicalTypeId::Struct
73    }
74
75    fn default_value(&self) -> Value {
76        Value::Null
77    }
78
79    fn as_arrow_type(&self) -> ArrowDataType {
80        let fields = self.as_arrow_fields();
81        ArrowDataType::Struct(fields)
82    }
83
84    fn create_mutable_vector(&self, capacity: usize) -> Box<dyn crate::prelude::MutableVector> {
85        Box::new(StructVectorBuilder::with_type_and_capacity(
86            self.clone(),
87            capacity,
88        ))
89    }
90
91    fn try_cast(&self, _from: Value) -> Option<Value> {
92        // TODO(discord9): what is the meaning of casting from Value to StructFields?
93        None
94    }
95}
96
97impl StructType {
98    pub fn new(fields: Arc<Vec<StructField>>) -> Self {
99        StructType {
100            fields: fields.clone(),
101        }
102    }
103
104    /// Like the `From<&Fields>` impl, but returns an error for arrow field
105    /// types greptimedb cannot represent instead of panicking.
106    pub fn try_from_arrow_fields(fields: &Fields) -> Result<Self> {
107        let struct_fields = fields
108            .iter()
109            .map(|field| {
110                Ok(StructField::new(
111                    field.name().clone(),
112                    ConcreteDataType::try_from(field.data_type())?,
113                    field.is_nullable(),
114                ))
115            })
116            .collect::<Result<Vec<_>>>()?;
117        Ok(StructType {
118            fields: Arc::new(struct_fields),
119        })
120    }
121
122    pub fn fields(&self) -> Arc<Vec<StructField>> {
123        self.fields.clone()
124    }
125
126    pub fn as_arrow_fields(&self) -> Fields {
127        self.fields
128            .iter()
129            .map(|f| Field::new(f.name.clone(), f.data_type.as_arrow_type(), f.nullable))
130            .collect()
131    }
132}
133
134#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
135pub struct StructField {
136    name: String,
137    data_type: ConcreteDataType,
138    nullable: bool,
139    metadata: BTreeMap<String, String>,
140}
141
142impl StructField {
143    pub fn new<T: Into<String>>(name: T, data_type: ConcreteDataType, nullable: bool) -> Self {
144        StructField {
145            name: name.into(),
146            data_type,
147            nullable,
148            metadata: BTreeMap::new(),
149        }
150    }
151
152    pub fn name(&self) -> &str {
153        &self.name
154    }
155
156    pub fn take_name(self) -> String {
157        self.name
158    }
159
160    pub fn data_type(&self) -> &ConcreteDataType {
161        &self.data_type
162    }
163
164    pub fn is_nullable(&self) -> bool {
165        self.nullable
166    }
167
168    #[expect(unused)]
169    pub(crate) fn metadata(&self, key: &str) -> Option<&str> {
170        self.metadata.get(key).map(String::as_str)
171    }
172
173    pub fn to_df_field(&self) -> Field {
174        let metadata = self
175            .metadata
176            .iter()
177            .map(|(k, v)| (k.clone(), v.clone()))
178            .collect();
179        Field::new(
180            self.name.clone(),
181            self.data_type.as_arrow_type(),
182            self.nullable,
183        )
184        .with_metadata(metadata)
185    }
186}
datatypes/types/struct_type.rs

datatypes/types/
struct_type.rs
