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