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 TryFrom<&Fields> for StructType {
32    type Error = crate::error::Error;
33    fn try_from(value: &Fields) -> Result<Self, Self::Error> {
34        let fields = value
35            .iter()
36            .map(|field| {
37                Ok(StructField::new(
38                    field.name().clone(),
39                    ConcreteDataType::try_from(field.data_type())?,
40                    field.is_nullable(),
41                ))
42            })
43            .collect::<Result<Vec<StructField>, Self::Error>>()?;
44        Ok(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    pub fn fields(&self) -> Arc<Vec<StructField>> {
105        self.fields.clone()
106    }
107
108    pub fn as_arrow_fields(&self) -> Fields {
109        self.fields
110            .iter()
111            .map(|f| Field::new(f.name.clone(), f.data_type.as_arrow_type(), f.nullable))
112            .collect()
113    }
114}
115
116#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
117pub struct StructField {
118    name: String,
119    data_type: ConcreteDataType,
120    nullable: bool,
121    metadata: BTreeMap<String, String>,
122}
123
124impl StructField {
125    pub fn new(name: String, data_type: ConcreteDataType, nullable: bool) -> Self {
126        StructField {
127            name,
128            data_type,
129            nullable,
130            metadata: BTreeMap::new(),
131        }
132    }
133
134    pub fn name(&self) -> &str {
135        &self.name
136    }
137
138    pub fn take_name(self) -> String {
139        self.name
140    }
141
142    pub fn data_type(&self) -> &ConcreteDataType {
143        &self.data_type
144    }
145
146    pub fn is_nullable(&self) -> bool {
147        self.nullable
148    }
149
150    pub(crate) fn insert_metadata(&mut self, key: impl ToString, value: impl ToString) {
151        self.metadata.insert(key.to_string(), value.to_string());
152    }
153
154    #[expect(unused)]
155    pub(crate) fn metadata(&self, key: &str) -> Option<&str> {
156        self.metadata.get(key).map(String::as_str)
157    }
158
159    pub fn to_df_field(&self) -> Field {
160        let metadata = self
161            .metadata
162            .iter()
163            .map(|(k, v)| (k.clone(), v.clone()))
164            .collect();
165        Field::new(
166            self.name.clone(),
167            self.data_type.as_arrow_type(),
168            self.nullable,
169        )
170        .with_metadata(metadata)
171    }
172}
datatypes/types/struct_type.rs

datatypes/types/
struct_type.rs
