Skip to main content

mongreldb_core/
schema.rs

1use serde::{Deserialize, Serialize};
2
3use crate::constraint::TableConstraints;
4use crate::error::{MongrelError, Result};
5use crate::memtable::Value;
6
7/// Logical column types. The on-disk Arrow encoding is chosen at flush based on
8/// [`TypeId`] and run-time stats (e.g. low-cardinality strings → dictionary).
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10#[serde(tag = "kind", rename_all = "lowercase")]
11pub enum TypeId {
12    Bool,
13    Int8,
14    Int16,
15    Int32,
16    Int64,
17    UInt8,
18    UInt16,
19    UInt32,
20    UInt64,
21    Float32,
22    Float64,
23    TimestampNanos,
24    Date32,
25    /// Variable-length bytes (covers UTF-8 strings).
26    Bytes,
27    /// Fixed-size binary embedding of `dim` f32 components.
28    Embedding {
29        dim: u32,
30    },
31}
32
33impl TypeId {
34    /// Fixed size in bytes for fixed-width types, else `None`.
35    pub const fn fixed_size(self) -> Option<usize> {
36        match self {
37            TypeId::Bool => Some(1),
38            TypeId::Int8 | TypeId::UInt8 => Some(1),
39            TypeId::Int16 | TypeId::UInt16 => Some(2),
40            TypeId::Int32 | TypeId::UInt32 | TypeId::Float32 | TypeId::Date32 => Some(4),
41            TypeId::Int64 | TypeId::UInt64 | TypeId::Float64 | TypeId::TimestampNanos => Some(8),
42            TypeId::Bytes | TypeId::Embedding { .. } => None,
43        }
44    }
45}
46
47/// Per-column flags packed into a `u32`. Stored verbatim in the run header.
48#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
49pub struct ColumnFlags {
50    bits: u32,
51}
52
53impl ColumnFlags {
54    pub const NULLABLE: u32 = 1 << 0;
55    pub const PRIMARY_KEY: u32 = 1 << 1;
56    pub const ENCRYPTED: u32 = 1 << 2;
57    /// Store HMAC(value) for equality or OPE for range so indexes work without
58    /// decrypting.
59    pub const ENCRYPTED_INDEXABLE: u32 = 1 << 3;
60    /// Store 1 bit per dimension; similarity via popcount(XOR).
61    pub const EMBEDDING_BINARY_QUANTIZED: u32 = 1 << 4;
62    /// Engine-managed monotonic identity allocator. Valid only on a single
63    /// `Int64` primary-key column per table (see [`Schema::validate_auto_increment`]).
64    /// On insert, when the column is omitted or `Null`, the engine assigns the
65    /// next counter value; an explicit `Int64` value is honored and advances the
66    /// counter past it. Counters are 1-based, never reused, and independent of
67    /// the physical [`crate::rowid::RowId`].
68    pub const AUTO_INCREMENT: u32 = 1 << 5;
69
70    #[inline]
71    pub const fn empty() -> Self {
72        Self { bits: 0 }
73    }
74
75    #[inline]
76    pub const fn with(mut self, flag: u32) -> Self {
77        self.bits |= flag;
78        self
79    }
80
81    #[inline]
82    pub const fn without(mut self, flag: u32) -> Self {
83        self.bits &= !flag;
84        self
85    }
86
87    #[inline]
88    pub const fn contains(&self, flag: u32) -> bool {
89        self.bits & flag != 0
90    }
91
92    #[inline]
93    pub const fn bits(&self) -> u32 {
94        self.bits
95    }
96}
97
98/// A column definition. `id` is stable, monotonic, and never reused.
99#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
100pub struct ColumnDef {
101    pub id: u16,
102    pub name: String,
103    pub ty: TypeId,
104    pub flags: ColumnFlags,
105}
106
107/// Metadata updates supported by native ALTER COLUMN.
108#[derive(Debug, Clone, Default, Serialize, Deserialize)]
109pub struct AlterColumn {
110    pub name: Option<String>,
111    pub ty: Option<TypeId>,
112    pub flags: Option<ColumnFlags>,
113}
114
115impl AlterColumn {
116    pub fn rename(name: impl Into<String>) -> Self {
117        Self {
118            name: Some(name.into()),
119            ty: None,
120            flags: None,
121        }
122    }
123
124    pub fn set_type(ty: TypeId) -> Self {
125        Self {
126            name: None,
127            ty: Some(ty),
128            flags: None,
129        }
130    }
131
132    pub fn set_flags(flags: ColumnFlags) -> Self {
133        Self {
134            name: None,
135            ty: None,
136            flags: Some(flags),
137        }
138    }
139}
140
141/// The kind of secondary index to maintain for a column. The primary-key index
142/// (in-memory HOT + on-disk learned PGM) is implicit and not listed here.
143#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
144pub enum IndexKind {
145    /// Roaring bitmap (value → row-id set). Low-cardinality equality / IN.
146    Bitmap,
147    /// FM-index / wavelet tree for arbitrary substring + ranked access.
148    FmIndex,
149    /// Quantized-vector ANN (binary / PQ). For `Embedding` columns.
150    Ann,
151    /// Learned zonemap (PGM) for ordered range predicates.
152    LearnedRange,
153    /// MinHash/LSH set-similarity (AI dedup/join primitives).
154    MinHash,
155    /// Learned-sparse (SPLADE-style) retrieval over weighted token vectors.
156    Sparse,
157}
158
159#[derive(Debug, Clone, Serialize, Deserialize)]
160pub struct IndexDef {
161    pub name: String,
162    pub column_id: u16,
163    pub kind: IndexKind,
164}
165
166#[derive(Debug, Clone, Default, Serialize, Deserialize)]
167pub struct Schema {
168    pub schema_id: u64,
169    pub columns: Vec<ColumnDef>,
170    pub indexes: Vec<IndexDef>,
171    /// Phase 18.2: column co-location groups. Each inner Vec lists column IDs
172    /// that are always accessed together. The run writer writes their pages
173    /// adjacently so a scan touching those columns benefits from sequential
174    /// I/O and cache locality. Empty = no co-location (default).
175    #[serde(default)]
176    pub colocation: Vec<Vec<u16>>,
177    /// Engine-side declarative constraints (unique / FK / check). Empty by
178    /// default — legacy and Kit-managed tables carry no engine constraints and
179    /// behave exactly as before. When non-empty, the transaction layer enforces
180    /// them authoritatively at commit (see [`crate::database`]).
181    #[serde(default)]
182    pub constraints: TableConstraints,
183}
184
185impl Schema {
186    pub fn column(&self, name: &str) -> Option<&ColumnDef> {
187        self.columns.iter().find(|c| c.name == name)
188    }
189
190    pub fn primary_key(&self) -> Option<&ColumnDef> {
191        self.columns
192            .iter()
193            .find(|c| c.flags.contains(ColumnFlags::PRIMARY_KEY))
194    }
195
196    /// Return an error if any column that is not marked NULLABLE is either
197    /// missing from `columns` or present as `Value::Null`. A column carrying
198    /// [`ColumnFlags::AUTO_INCREMENT`] is exempt when omitted/`Null` because the
199    /// engine fills it in before this check runs.
200    pub fn validate_not_null(&self, columns: &[(u16, Value)]) -> Result<()> {
201        // Rows are short sparse `(id, value)` lists; a linear probe beats
202        // materializing a HashMap (and cloning every Value) per row.
203        let at = |id: u16| columns.iter().find(|(c, _)| *c == id).map(|(_, v)| v);
204        for col in &self.columns {
205            if col.flags.contains(ColumnFlags::NULLABLE) {
206                continue;
207            }
208            // The engine supplies the AUTO_INCREMENT value, so its absence is
209            // legal at this layer (filled in upstream of validation).
210            if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
211                match at(col.id) {
212                    None | Some(Value::Null) => continue,
213                    Some(_) => {}
214                }
215            }
216            match at(col.id) {
217                None => {
218                    return Err(MongrelError::InvalidArgument(format!(
219                        "column '{}' ({}) is NOT NULL but was omitted",
220                        col.name, col.id
221                    )));
222                }
223                Some(Value::Null) => {
224                    return Err(MongrelError::InvalidArgument(format!(
225                        "column '{}' ({}) is NOT NULL but got NULL",
226                        col.name, col.id
227                    )));
228                }
229                Some(_) => {}
230            }
231        }
232        Ok(())
233    }
234
235    /// Enforce the `AUTO_INCREMENT` column contract: at most one such column,
236    /// and it must be a non-nullable `Int64` primary key. Called at table
237    /// creation time so an invalid schema never reaches the insert path.
238    pub fn validate_auto_increment(&self) -> Result<()> {
239        let mut seen: Option<&ColumnDef> = None;
240        for col in &self.columns {
241            if !col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
242                continue;
243            }
244            if let Some(prev) = seen {
245                return Err(MongrelError::Schema(format!(
246                    "AUTO_INCREMENT may be set on at most one column; '{}' and '{}' both carry it",
247                    prev.name, col.name
248                )));
249            }
250            if col.ty != TypeId::Int64 {
251                return Err(MongrelError::Schema(format!(
252                    "AUTO_INCREMENT column '{}' must be Int64, is {:?}",
253                    col.name, col.ty
254                )));
255            }
256            if !col.flags.contains(ColumnFlags::PRIMARY_KEY) {
257                return Err(MongrelError::Schema(format!(
258                    "AUTO_INCREMENT column '{}' must also be the primary key",
259                    col.name
260                )));
261            }
262            if col.flags.contains(ColumnFlags::NULLABLE) {
263                return Err(MongrelError::Schema(format!(
264                    "AUTO_INCREMENT column '{}' must not be nullable",
265                    col.name
266                )));
267            }
268            seen = Some(col);
269        }
270        Ok(())
271    }
272
273    /// The single `AUTO_INCREMENT` column, if any.
274    pub fn auto_increment_column(&self) -> Option<&ColumnDef> {
275        self.columns
276            .iter()
277            .find(|c| c.flags.contains(ColumnFlags::AUTO_INCREMENT))
278    }
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    #[test]
286    fn flag_composition() {
287        let f = ColumnFlags::empty()
288            .with(ColumnFlags::PRIMARY_KEY)
289            .with(ColumnFlags::ENCRYPTED_INDEXABLE);
290        assert!(f.contains(ColumnFlags::PRIMARY_KEY));
291        assert!(f.contains(ColumnFlags::ENCRYPTED_INDEXABLE));
292        assert!(!f.contains(ColumnFlags::ENCRYPTED));
293    }
294
295    #[test]
296    fn fixed_size() {
297        assert_eq!(TypeId::Int64.fixed_size(), Some(8));
298        assert_eq!(TypeId::Bytes.fixed_size(), None);
299        assert_eq!(TypeId::Embedding { dim: 768 }.fixed_size(), None);
300    }
301
302    fn col(id: u16, name: &str, ty: TypeId, flags: ColumnFlags) -> ColumnDef {
303        ColumnDef {
304            id,
305            name: name.into(),
306            ty,
307            flags,
308        }
309    }
310
311    #[test]
312    fn auto_increment_validation_accepts_int64_pk() {
313        let s = Schema {
314            schema_id: 1,
315            columns: vec![col(
316                0,
317                "id",
318                TypeId::Int64,
319                ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
320            )],
321            indexes: vec![],
322            colocation: vec![],
323            constraints: Default::default(),
324        };
325        assert!(s.validate_auto_increment().is_ok());
326        assert_eq!(s.auto_increment_column().unwrap().id, 0);
327    }
328
329    #[test]
330    fn auto_increment_validation_rejects_non_pk() {
331        let s = Schema {
332            schema_id: 1,
333            columns: vec![
334                col(
335                    0,
336                    "id",
337                    TypeId::Int64,
338                    ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
339                ),
340                col(
341                    1,
342                    "seq",
343                    TypeId::Int64,
344                    ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
345                ),
346            ],
347            indexes: vec![],
348            colocation: vec![],
349            constraints: Default::default(),
350        };
351        assert!(s.validate_auto_increment().is_err());
352    }
353
354    #[test]
355    fn auto_increment_validation_rejects_non_int64() {
356        let s = Schema {
357            schema_id: 1,
358            columns: vec![col(
359                0,
360                "id",
361                TypeId::Bytes,
362                ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
363            )],
364            indexes: vec![],
365            colocation: vec![],
366            constraints: Default::default(),
367        };
368        assert!(s.validate_auto_increment().is_err());
369    }
370
371    #[test]
372    fn auto_increment_validation_rejects_two() {
373        let s = Schema {
374            schema_id: 1,
375            columns: vec![
376                col(
377                    0,
378                    "id",
379                    TypeId::Int64,
380                    ColumnFlags::empty()
381                        .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
382                ),
383                col(
384                    1,
385                    "id2",
386                    TypeId::Int64,
387                    ColumnFlags::empty().with(ColumnFlags::AUTO_INCREMENT),
388                ),
389            ],
390            indexes: vec![],
391            colocation: vec![],
392            constraints: Default::default(),
393        };
394        assert!(s.validate_auto_increment().is_err());
395    }
396
397    #[test]
398    fn auto_increment_exempt_from_not_null_when_omitted() {
399        let s = Schema {
400            schema_id: 1,
401            columns: vec![
402                col(
403                    0,
404                    "id",
405                    TypeId::Int64,
406                    ColumnFlags::empty()
407                        .with(ColumnFlags::PRIMARY_KEY | ColumnFlags::AUTO_INCREMENT),
408                ),
409                col(1, "name", TypeId::Bytes, ColumnFlags::empty()),
410            ],
411            indexes: vec![],
412            colocation: vec![],
413            constraints: Default::default(),
414        };
415        // Omitting the auto-inc column must not trip NOT NULL.
416        let cols = vec![(1u16, Value::Bytes(b"x".to_vec()))];
417        assert!(s.validate_not_null(&cols).is_ok());
418    }
419}