1use crate::logical_type::assign_fresh_field_ids;
2use crate::{DataField, FieldId, LogicalType, LogicalTypeKind, Result, TableError};
3use serde::{Deserialize, Serialize};
4use std::collections::{HashMap, HashSet};
5
6#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
8pub struct TableSchema {
9 pub fields: Vec<DataField>,
10 pub primary_key: Vec<FieldId>,
11 pub bucket_key: Vec<FieldId>,
12}
13
14impl TableSchema {
15 pub fn builder() -> TableSchemaBuilder {
17 TableSchemaBuilder::new()
18 }
19
20 pub fn new(
26 fields: Vec<DataField>,
27 primary_key: Vec<FieldId>,
28 bucket_key: Vec<FieldId>,
29 ) -> Result<Self> {
30 let schema = Self {
31 fields,
32 primary_key,
33 bucket_key,
34 };
35 schema.validate()?;
36 Ok(schema)
37 }
38
39 pub(crate) fn validate(&self) -> Result<()> {
40 if self.fields.is_empty() {
41 return Err(TableError::InvalidSchema(
42 "table must contain at least one field".to_string(),
43 ));
44 }
45 if self.primary_key.is_empty() {
46 return Err(TableError::InvalidSchema(
47 "primary key must not be empty".to_string(),
48 ));
49 }
50 if self.bucket_key.is_empty() {
51 return Err(TableError::InvalidSchema(
52 "bucket key must not be empty".to_string(),
53 ));
54 }
55 if self.bucket_key.len() > self.primary_key.len()
56 || self.primary_key[..self.bucket_key.len()] != self.bucket_key
57 {
58 return Err(TableError::InvalidSchema(
59 "bucket key must be a prefix of the primary key".to_string(),
60 ));
61 }
62
63 let mut field_ids = HashSet::new();
64 let mut field_names = HashSet::new();
65 for field in &self.fields {
66 validate_field_tree(field, &mut field_ids)?;
67 if !field_names.insert(field.name.as_str()) {
68 return Err(TableError::InvalidSchema(format!(
69 "duplicate top-level field name: {}",
70 field.name
71 )));
72 }
73 }
74
75 let top_level = self
76 .fields
77 .iter()
78 .map(|field| (field.id, field))
79 .collect::<HashMap<_, _>>();
80 let mut key_ids = HashSet::new();
81 for id in &self.primary_key {
82 let field = top_level.get(id).ok_or_else(|| {
83 TableError::InvalidSchema(format!("primary-key field {} does not exist", id.0))
84 })?;
85 if !key_ids.insert(*id) {
86 return Err(TableError::InvalidSchema(format!(
87 "duplicate primary-key field: {}",
88 id.0
89 )));
90 }
91 if !field.logical_type.is_key_compatible() {
92 return Err(TableError::InvalidSchema(format!(
93 "field '{}' cannot be used in a primary key",
94 field.name
95 )));
96 }
97 }
98 Ok(())
99 }
100}
101
102#[derive(Default)]
108pub struct TableSchemaBuilder {
109 fields: Vec<(String, LogicalType)>,
110 primary_key: Vec<String>,
111 bucket_key: Vec<String>,
112}
113
114impl TableSchemaBuilder {
115 pub fn new() -> Self {
116 Self::default()
117 }
118
119 pub fn field(mut self, name: impl Into<String>, logical_type: LogicalType) -> Self {
121 self.fields.push((name.into(), logical_type));
122 self
123 }
124
125 pub fn primary_key<I, S>(mut self, field_names: I) -> Self
127 where
128 I: IntoIterator<Item = S>,
129 S: AsRef<str>,
130 {
131 self.primary_key = field_names
132 .into_iter()
133 .map(|field_name| field_name.as_ref().to_string())
134 .collect();
135 self
136 }
137
138 pub fn bucket_key<I, S>(mut self, field_names: I) -> Self
141 where
142 I: IntoIterator<Item = S>,
143 S: AsRef<str>,
144 {
145 self.bucket_key = field_names
146 .into_iter()
147 .map(|field_name| field_name.as_ref().to_string())
148 .collect();
149 self
150 }
151
152 pub fn build(self) -> Result<TableSchema> {
154 let mut fields = self
155 .fields
156 .into_iter()
157 .map(|(name, logical_type)| DataField {
158 id: FieldId(0),
159 name,
160 logical_type,
161 })
162 .collect::<Vec<_>>();
163 let mut next_id = 0;
164 assign_fresh_field_ids(&mut fields, &mut next_id)?;
165 let primary_key = resolve_key_names(&fields, &self.primary_key, "primary key")?;
166 let bucket_key = resolve_key_names(&fields, &self.bucket_key, "bucket key")?;
167 TableSchema::new(fields, primary_key, bucket_key)
168 }
169}
170
171fn resolve_key_names(
172 fields: &[DataField],
173 field_names: &[String],
174 key_name: &str,
175) -> Result<Vec<FieldId>> {
176 field_names
177 .iter()
178 .map(|field_name| {
179 fields
180 .iter()
181 .find(|field| field.name == *field_name)
182 .map(|field| field.id)
183 .ok_or_else(|| {
184 TableError::InvalidSchema(format!(
185 "{key_name} field '{field_name}' does not exist"
186 ))
187 })
188 })
189 .collect()
190}
191
192fn validate_field_tree(field: &DataField, ids: &mut HashSet<FieldId>) -> Result<()> {
193 field.validate()?;
194 if !ids.insert(field.id) {
195 return Err(TableError::InvalidSchema(format!(
196 "duplicate field id: {}",
197 field.id.0
198 )));
199 }
200 validate_nested_fields(&field.logical_type, ids)
201}
202
203fn validate_nested_fields(logical_type: &LogicalType, ids: &mut HashSet<FieldId>) -> Result<()> {
204 match &logical_type.kind {
205 LogicalTypeKind::List { element_type } => validate_nested_fields(element_type, ids),
206 LogicalTypeKind::Map {
207 key_type,
208 value_type,
209 } => {
210 validate_nested_fields(key_type, ids)?;
211 validate_nested_fields(value_type, ids)
212 }
213 LogicalTypeKind::Struct { fields } => {
214 let mut names = HashSet::new();
215 for field in fields {
216 if !names.insert(field.name.as_str()) {
217 return Err(TableError::InvalidSchema(format!(
218 "duplicate nested field name: {}",
219 field.name
220 )));
221 }
222 validate_field_tree(field, ids)?;
223 }
224 Ok(())
225 }
226 LogicalTypeKind::Extension { extension } => {
227 validate_nested_fields(&extension.physical_type, ids)
228 }
229 _ => Ok(()),
230 }
231}