Skip to main content

mongreldb_core/
schema.rs

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