datatypes/types/
struct_type.rs1use 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 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}