Skip to main content

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::prelude::{ConcreteDataType, DataType, LogicalTypeId};
23use crate::value::Value;
24use crate::vectors::StructVectorBuilder;
25
26#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
27pub struct StructType {
28    fields: Arc<Vec<StructField>>,
29}
30
31impl From<&Fields> for StructType {
32    fn from(value: &Fields) -> Self {
33        let fields = value
34            .iter()
35            .map(|field| {
36                StructField::new(
37                    field.name().clone(),
38                    ConcreteDataType::from_arrow_type(field.data_type()),
39                    field.is_nullable(),
40                )
41            })
42            .collect::<Vec<_>>();
43        StructType {
44            fields: Arc::new(fields),
45        }
46    }
47}
48
49impl<const N: usize> From<[StructField; N]> for StructType {
50    fn from(value: [StructField; N]) -> Self {
51        let value: Box<[StructField]> = Box::new(value);
52        Self {
53            fields: Arc::new(value.into_vec()),
54        }
55    }
56}
57
58impl DataType for StructType {
59    fn name(&self) -> String {
60        format!(
61            "Struct<{}>",
62            self.fields
63                .iter()
64                .map(|f| format!(r#""{}": {}"#, f.name(), f.data_type()))
65                .collect::<Vec<_>>()
66                .join(", ")
67        )
68    }
69
70    fn logical_type_id(&self) -> LogicalTypeId {
71        LogicalTypeId::Struct
72    }
73
74    fn default_value(&self) -> Value {
75        Value::Null
76    }
77
78    fn as_arrow_type(&self) -> ArrowDataType {
79        let fields = self.as_arrow_fields();
80        ArrowDataType::Struct(fields)
81    }
82
83    fn create_mutable_vector(&self, capacity: usize) -> Box<dyn crate::prelude::MutableVector> {
84        Box::new(StructVectorBuilder::with_type_and_capacity(
85            self.clone(),
86            capacity,
87        ))
88    }
89
90    fn try_cast(&self, _from: Value) -> Option<Value> {
91        // TODO(discord9): what is the meaning of casting from Value to StructFields?
92        None
93    }
94}
95
96impl StructType {
97    pub fn new(fields: Arc<Vec<StructField>>) -> Self {
98        StructType {
99            fields: fields.clone(),
100        }
101    }
102
103    pub fn fields(&self) -> Arc<Vec<StructField>> {
104        self.fields.clone()
105    }
106
107    pub fn as_arrow_fields(&self) -> Fields {
108        self.fields
109            .iter()
110            .map(|f| Field::new(f.name.clone(), f.data_type.as_arrow_type(), f.nullable))
111            .collect()
112    }
113}
114
115#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
116pub struct StructField {
117    name: String,
118    data_type: ConcreteDataType,
119    nullable: bool,
120    metadata: BTreeMap<String, String>,
121}
122
123impl StructField {
124    pub fn new<T: Into<String>>(name: T, data_type: ConcreteDataType, nullable: bool) -> Self {
125        StructField {
126            name: name.into(),
127            data_type,
128            nullable,
129            metadata: BTreeMap::new(),
130        }
131    }
132
133    pub fn name(&self) -> &str {
134        &self.name
135    }
136
137    pub fn take_name(self) -> String {
138        self.name
139    }
140
141    pub fn data_type(&self) -> &ConcreteDataType {
142        &self.data_type
143    }
144
145    pub fn is_nullable(&self) -> bool {
146        self.nullable
147    }
148
149    pub(crate) fn insert_metadata(&mut self, key: impl ToString, value: impl ToString) {
150        self.metadata.insert(key.to_string(), value.to_string());
151    }
152
153    #[expect(unused)]
154    pub(crate) fn metadata(&self, key: &str) -> Option<&str> {
155        self.metadata.get(key).map(String::as_str)
156    }
157
158    pub fn to_df_field(&self) -> Field {
159        let metadata = self
160            .metadata
161            .iter()
162            .map(|(k, v)| (k.clone(), v.clone()))
163            .collect();
164        Field::new(
165            self.name.clone(),
166            self.data_type.as_arrow_type(),
167            self.nullable,
168        )
169        .with_metadata(metadata)
170    }
171}