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::sync::Arc;
16
17use arrow::datatypes::{DataType as ArrowDataType, Field};
18use arrow_schema::Fields;
19use serde::{Deserialize, Serialize};
20
21use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId};
22use crate::value::Value;
23use crate::vectors::StructVectorBuilder;
24
25#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
26pub struct StructType {
27    fields: Arc<Vec<StructField>>,
28}
29
30impl TryFrom<&Fields> for StructType {
31    type Error = crate::error::Error;
32    fn try_from(value: &Fields) -> Result<Self, Self::Error> {
33        let fields = value
34            .iter()
35            .map(|field| {
36                Ok(StructField::new(
37                    field.name().clone(),
38                    ConcreteDataType::try_from(field.data_type())?,
39                    field.is_nullable(),
40                ))
41            })
42            .collect::<Result<Vec<StructField>, Self::Error>>()?;
43        Ok(StructType {
44            fields: Arc::new(fields),
45        })
46    }
47}
48
49impl DataType for StructType {
50    fn name(&self) -> String {
51        format!(
52            "Struct<{}>",
53            self.fields
54                .iter()
55                .map(|f| format!(r#""{}": {}"#, f.name(), f.data_type()))
56                .collect::<Vec<_>>()
57                .join(", ")
58        )
59    }
60
61    fn logical_type_id(&self) -> LogicalTypeId {
62        LogicalTypeId::Struct
63    }
64
65    fn default_value(&self) -> Value {
66        Value::Null
67    }
68
69    fn as_arrow_type(&self) -> ArrowDataType {
70        let fields = self.as_arrow_fields();
71        ArrowDataType::Struct(fields)
72    }
73
74    fn create_mutable_vector(&self, capacity: usize) -> Box<dyn crate::prelude::MutableVector> {
75        Box::new(StructVectorBuilder::with_type_and_capacity(
76            self.clone(),
77            capacity,
78        ))
79    }
80
81    fn try_cast(&self, _from: Value) -> Option<Value> {
82        // TODO(discord9): what is the meaning of casting from Value to StructFields?
83        None
84    }
85}
86
87impl StructType {
88    pub fn new(fields: Arc<Vec<StructField>>) -> Self {
89        StructType {
90            fields: fields.clone(),
91        }
92    }
93
94    pub fn fields(&self) -> Arc<Vec<StructField>> {
95        self.fields.clone()
96    }
97
98    pub fn as_arrow_fields(&self) -> Fields {
99        self.fields
100            .iter()
101            .map(|f| Field::new(f.name.clone(), f.data_type.as_arrow_type(), f.nullable))
102            .collect()
103    }
104}
105
106#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
107pub struct StructField {
108    name: String,
109    data_type: ConcreteDataType,
110    nullable: bool,
111}
112
113impl StructField {
114    pub fn new(name: String, data_type: ConcreteDataType, nullable: bool) -> Self {
115        StructField {
116            name,
117            data_type,
118            nullable,
119        }
120    }
121
122    pub fn name(&self) -> &str {
123        &self.name
124    }
125
126    pub fn take_name(self) -> String {
127        self.name
128    }
129
130    pub fn data_type(&self) -> &ConcreteDataType {
131        &self.data_type
132    }
133
134    pub fn is_nullable(&self) -> bool {
135        self.nullable
136    }
137
138    pub fn to_df_field(&self) -> Field {
139        Field::new(
140            self.name.clone(),
141            self.data_type.as_arrow_type(),
142            self.nullable,
143        )
144    }
145}