Skip to main content

cobble_table/
schema.rs

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/// User-visible semantic schema of a table.
7#[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    /// Start building a fresh schema with deterministic generated field ids.
16    pub fn builder() -> TableSchemaBuilder {
17        TableSchemaBuilder::new()
18    }
19
20    /// Construct a schema with explicit field ids.
21    ///
22    /// Prefer [`Self::builder`] for fresh schemas. This constructor preserves
23    /// ids exactly and is intended for restoration or advanced interoperability
24    /// with an already identified schema.
25    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/// Builder for a fresh table schema with generated, deterministic field ids.
103///
104/// All field ids, including those inside nested types, are assigned from zero
105/// in depth-first declaration order. Use [`TableSchema::new`] or load the
106/// persisted schema to preserve existing identities instead.
107#[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    /// Add a top-level field in schema order.
120    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    /// Set primary-key fields by their exact top-level names.
126    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    /// Set bucket-key fields by their exact top-level names.
139    /// This is required and must be a nonempty prefix of the primary key.
140    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    /// Assign preorder ids, resolve keys, and validate the completed schema.
153    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}