1use serde::{Deserialize, Serialize};
7use std::collections::{HashMap, HashSet};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum ColumnType {
13 Bool,
14 Int8,
15 Int16,
16 Int32,
17 Int64,
18 Float32,
19 Float64,
20 Text,
21 Bytes,
22 Json,
23 Date,
24 DateTime,
25 TimestampNanos,
26 Date64,
27 Time64,
28 Interval,
29 Decimal128,
30 Uuid,
32 JsonNative,
34 Array,
36 Embedding,
39 Sparse,
42}
43
44#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
46#[serde(rename_all = "snake_case")]
47pub enum DefaultKind {
48 Static(serde_json::Value),
50 Now,
52 Uuid,
54 Sequence(String),
56 CustomName(String),
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
67#[serde(rename_all = "snake_case")]
68pub enum EmbeddingSpecNormalization {
69 #[default]
70 None,
71 L2,
72}
73
74#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
75#[serde(rename_all = "snake_case")]
76pub enum EmbeddingWriteFailurePolicy {
77 #[default]
78 AbortWrite,
79}
80
81#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
82pub struct GeneratedEmbeddingSpec {
83 pub provider_id: String,
84 pub model_id: String,
85 pub model_version: String,
86 pub source_columns: Vec<u32>,
87 pub input_template: String,
88 pub dimension: u32,
89 pub normalization: EmbeddingSpecNormalization,
90 pub failure_policy: EmbeddingWriteFailurePolicy,
91}
92
93#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94#[serde(tag = "kind", rename_all = "snake_case")]
95pub enum EmbeddingSource {
96 SuppliedByApplication,
98 LocalModel {
100 model_path: String,
102 model_id: String,
104 },
105 ConfiguredModel {
107 provider_id: String,
108 model_id: String,
109 model_version: String,
110 },
111 GeneratedColumn {
113 provider: String,
115 },
116 GeneratedColumnSpec { spec: GeneratedEmbeddingSpec },
118}
119
120#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
122pub struct Column {
123 pub id: u32,
125 pub name: String,
127 pub storage_type: ColumnType,
129 pub application_type: ColumnType,
131 pub nullable: bool,
133 pub primary_key: bool,
135 pub default: Option<DefaultKind>,
137 pub generated: bool,
139 #[serde(default, skip_serializing_if = "Option::is_none")]
141 pub enum_values: Option<Vec<String>>,
142 #[serde(default, skip_serializing_if = "Option::is_none")]
144 pub min: Option<f64>,
145 #[serde(default, skip_serializing_if = "Option::is_none")]
147 pub max: Option<f64>,
148 #[serde(default, skip_serializing_if = "Option::is_none")]
150 pub min_length: Option<usize>,
151 #[serde(default, skip_serializing_if = "Option::is_none")]
153 pub max_length: Option<usize>,
154 #[serde(default, skip_serializing_if = "Option::is_none")]
156 pub regex: Option<String>,
157 #[serde(default, skip_serializing_if = "Option::is_none")]
159 pub check_expr: Option<String>,
160 #[serde(default, skip_serializing_if = "Option::is_none")]
162 pub embedding_dim: Option<u32>,
163 #[serde(default, skip_serializing_if = "Option::is_none")]
166 pub embedding_source: Option<EmbeddingSource>,
167 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
169 pub encrypted: bool,
170 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
173 pub encrypted_indexable: bool,
174}
175
176impl Column {
177 pub fn new(id: u32, name: impl Into<String>, storage_type: ColumnType) -> Self {
179 Self {
180 id,
181 name: name.into(),
182 storage_type,
183 application_type: storage_type,
184 nullable: false,
185 primary_key: false,
186 default: None,
187 generated: false,
188 enum_values: None,
189 min: None,
190 max: None,
191 min_length: None,
192 max_length: None,
193 regex: None,
194 check_expr: None,
195 embedding_dim: None,
196 embedding_source: None,
197 encrypted: false,
198 encrypted_indexable: false,
199 }
200 }
201}
202
203#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
205#[serde(rename_all = "snake_case")]
206pub enum IndexKind {
207 #[default]
209 Bitmap,
210 Fm,
212 Ann,
214 Sparse,
216 MinHash,
219 LearnedRange,
222}
223
224#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
226#[serde(rename_all = "snake_case")]
227pub enum AnnQuantization {
228 #[default]
230 BinarySign,
231 Dense,
233 Product { num_subvectors: u16, bits: u8 },
237}
238
239#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
243pub enum AnnAlgorithm {
244 #[default]
246 #[serde(rename = "hnsw")]
247 Hnsw,
248 #[serde(rename = "diskann")]
250 DiskAnn,
251 #[serde(rename = "ivf")]
253 Ivf,
254}
255
256#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
258pub struct Index {
259 pub name: String,
260 pub columns: Vec<String>,
261 pub unique: bool,
262 #[serde(default)]
265 pub kind: IndexKind,
266 #[serde(default)]
268 pub ann_quantization: AnnQuantization,
269 #[serde(default)]
272 pub ann_algorithm: AnnAlgorithm,
273 #[serde(default, skip_serializing_if = "Option::is_none")]
275 pub predicate: Option<String>,
276 #[serde(default, skip_serializing_if = "Option::is_none")]
278 pub ann_m: Option<usize>,
279 #[serde(default, skip_serializing_if = "Option::is_none")]
281 pub ann_ef_construction: Option<usize>,
282 #[serde(default, skip_serializing_if = "Option::is_none")]
284 pub ann_ef_search: Option<usize>,
285 #[serde(default, skip_serializing_if = "Option::is_none")]
287 pub ann_diskann_r: Option<usize>,
288 #[serde(default, skip_serializing_if = "Option::is_none")]
290 pub ann_diskann_l: Option<usize>,
291 #[serde(default, skip_serializing_if = "Option::is_none")]
293 pub ann_diskann_beam_width: Option<usize>,
294 #[serde(default, skip_serializing_if = "Option::is_none")]
296 pub ann_diskann_alpha: Option<u32>,
297 #[serde(default, skip_serializing_if = "Option::is_none")]
299 pub ann_ivf_nlist: Option<usize>,
300 #[serde(default, skip_serializing_if = "Option::is_none")]
302 pub ann_ivf_nprobe: Option<usize>,
303 #[serde(default, skip_serializing_if = "Option::is_none")]
305 pub ann_ivf_training_samples: Option<usize>,
306 #[serde(default, skip_serializing_if = "Option::is_none")]
308 pub ann_pq_training_samples: Option<usize>,
309 #[serde(default, skip_serializing_if = "Option::is_none")]
311 pub ann_pq_seed: Option<u64>,
312 #[serde(default, skip_serializing_if = "Option::is_none")]
314 pub ann_pq_rerank_factor: Option<usize>,
315 #[serde(default, skip_serializing_if = "Option::is_none")]
317 pub minhash_permutations: Option<usize>,
318 #[serde(default, skip_serializing_if = "Option::is_none")]
320 pub minhash_bands: Option<usize>,
321 #[serde(default, skip_serializing_if = "Option::is_none")]
323 pub learned_range_epsilon: Option<usize>,
324}
325
326#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
328pub struct UniqueConstraint {
329 pub name: String,
330 pub columns: Vec<String>,
331}
332
333#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
335pub struct ForeignKey {
336 pub name: String,
337 pub columns: Vec<String>,
338 pub references_table: String,
339 pub references_columns: Vec<String>,
340 #[serde(default)]
341 pub on_delete: ForeignKeyAction,
342}
343
344#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
346#[serde(rename_all = "snake_case")]
347pub enum ForeignKeyAction {
348 #[default]
349 Restrict,
350 Cascade,
351 SetNull,
352}
353
354#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
356pub struct CheckConstraint {
357 pub name: String,
358 pub expr: String,
359}
360
361#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
363pub struct Sequence {
364 pub name: String,
365 pub next_value: i64,
366}
367
368#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
370pub struct Table {
371 pub id: u32,
373 pub name: String,
374 pub columns: Vec<Column>,
375 pub primary_key: Vec<String>,
376 #[serde(default, skip_serializing_if = "Vec::is_empty")]
377 pub indexes: Vec<Index>,
378 #[serde(default, skip_serializing_if = "Vec::is_empty")]
379 pub foreign_keys: Vec<ForeignKey>,
380 #[serde(default, skip_serializing_if = "Vec::is_empty")]
381 pub unique_constraints: Vec<UniqueConstraint>,
382 #[serde(default, skip_serializing_if = "Vec::is_empty")]
383 pub check_constraints: Vec<CheckConstraint>,
384}
385
386impl Table {
387 pub fn column(&self, name: &str) -> Option<&Column> {
389 self.columns.iter().find(|c| c.name == name)
390 }
391
392 pub fn is_pk_column(&self, name: &str) -> bool {
394 self.primary_key.iter().any(|c| c == name)
395 }
396}
397
398#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
400pub enum SchemaError {
401 #[error("duplicate table name \"{0}\"")]
402 DuplicateTableName(String),
403 #[error("duplicate table id {0}")]
404 DuplicateTableId(u32),
405 #[error("duplicate column name \"{1}\" in table \"{0}\"")]
406 DuplicateColumnName(String, String),
407 #[error("duplicate column id {1} in table \"{0}\"")]
408 DuplicateColumnId(String, u32),
409 #[error("primary key column \"{1}\" not found in table \"{0}\"")]
410 MissingPrimaryKeyColumn(String, String),
411 #[error("index \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
412 MissingIndexColumn(String, String, String),
413 #[error("unique constraint \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
414 MissingUniqueColumn(String, String, String),
415 #[error("foreign key \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
416 MissingForeignKeyColumn(String, String, String),
417 #[error("foreign key \"{1}\" references unknown table \"{2}\"")]
418 MissingReferencedTable(String, String, String),
419 #[error("foreign key \"{1}\" references unknown column \"{2}\" on table \"{3}\"")]
420 MissingReferencedColumn(String, String, String, String),
421 #[error(
422 "column \"{1}\" on table \"{0}\" sets embedding_source but is not an embedding column"
423 )]
424 EmbeddingSourceOnNonEmbedding(String, String),
425 #[error(
426 "embedding column \"{1}\" on table \"{0}\" with LocalModel/GeneratedColumn source requires embedding_dim > 0"
427 )]
428 EmbeddingSourceMissingDim(String, String),
429 #[error("generated embedding column \"{1}\" on table \"{0}\" is invalid: {2}")]
430 InvalidGeneratedEmbeddingSpec(String, String, String),
431}
432
433#[derive(Debug, Clone, PartialEq, Serialize)]
435pub struct Schema {
436 pub tables: Vec<Table>,
437 by_name: HashMap<String, usize>,
438 by_id: HashMap<u32, usize>,
439}
440
441impl<'de> serde::Deserialize<'de> for Schema {
442 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
443 where
444 D: serde::Deserializer<'de>,
445 {
446 #[derive(serde::Deserialize)]
447 struct SchemaHelper {
448 tables: Vec<Table>,
449 }
450 let helper = SchemaHelper::deserialize(deserializer)?;
451 Schema::new(helper.tables).map_err(serde::de::Error::custom)
452 }
453}
454
455fn synthesize_unique_from_indexes(table: &mut Table) {
460 let mut synthesized: Vec<UniqueConstraint> = Vec::new();
461 for idx in &table.indexes {
462 if !idx.unique {
463 continue;
464 }
465 let covered = table
466 .unique_constraints
467 .iter()
468 .chain(synthesized.iter())
469 .any(|u| u.columns == idx.columns);
470 if !covered {
471 synthesized.push(UniqueConstraint {
472 name: idx.name.clone(),
473 columns: idx.columns.clone(),
474 });
475 }
476 }
477 table.unique_constraints.extend(synthesized);
478}
479
480impl Schema {
481 pub fn new(mut tables: Vec<Table>) -> Result<Self, SchemaError> {
483 for table in &mut tables {
484 synthesize_unique_from_indexes(table);
485 }
486
487 let mut by_name = HashMap::with_capacity(tables.len());
488 let mut by_id = HashMap::with_capacity(tables.len());
489
490 for (idx, table) in tables.iter().enumerate() {
491 if by_name.contains_key(&table.name) {
492 return Err(SchemaError::DuplicateTableName(table.name.clone()));
493 }
494 if by_id.contains_key(&table.id) {
495 return Err(SchemaError::DuplicateTableId(table.id));
496 }
497 by_name.insert(table.name.clone(), idx);
498 by_id.insert(table.id, idx);
499 }
500
501 for table in &tables {
502 Self::validate_table(table, &by_name)?;
503 }
504
505 Ok(Self {
506 tables,
507 by_name,
508 by_id,
509 })
510 }
511
512 fn validate_table(
513 table: &Table,
514 table_names: &HashMap<String, usize>,
515 ) -> Result<(), SchemaError> {
516 let mut column_names = HashMap::with_capacity(table.columns.len());
517 let mut column_ids = HashMap::with_capacity(table.columns.len());
518
519 for col in &table.columns {
520 if column_names.contains_key(&col.name) {
521 return Err(SchemaError::DuplicateColumnName(
522 table.name.clone(),
523 col.name.clone(),
524 ));
525 }
526 if column_ids.contains_key(&col.id) {
527 return Err(SchemaError::DuplicateColumnId(table.name.clone(), col.id));
528 }
529 if col.embedding_source.is_some() && col.storage_type != ColumnType::Embedding {
530 return Err(SchemaError::EmbeddingSourceOnNonEmbedding(
531 table.name.clone(),
532 col.name.clone(),
533 ));
534 }
535 if matches!(
536 col.embedding_source,
537 Some(
538 EmbeddingSource::LocalModel { .. }
539 | EmbeddingSource::ConfiguredModel { .. }
540 | EmbeddingSource::GeneratedColumn { .. }
541 | EmbeddingSource::GeneratedColumnSpec { .. }
542 )
543 ) && col.embedding_dim.unwrap_or(0) == 0
544 {
545 return Err(SchemaError::EmbeddingSourceMissingDim(
546 table.name.clone(),
547 col.name.clone(),
548 ));
549 }
550 column_names.insert(col.name.clone(), col.id);
551 column_ids.insert(col.id, col.name.clone());
552 }
553
554 for col in &table.columns {
555 if let Some(EmbeddingSource::ConfiguredModel {
556 provider_id,
557 model_id,
558 model_version,
559 }) = col.embedding_source.as_ref()
560 {
561 if provider_id.is_empty() || model_id.is_empty() || model_version.is_empty() {
562 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
563 table.name.clone(),
564 col.name.clone(),
565 "provider, model, and version are required".into(),
566 ));
567 }
568 }
569 }
570
571 for col in &table.columns {
572 let Some(EmbeddingSource::GeneratedColumnSpec { spec }) = col.embedding_source.as_ref()
573 else {
574 continue;
575 };
576 if spec.provider_id.is_empty()
577 || spec.model_id.is_empty()
578 || spec.model_version.is_empty()
579 || spec.source_columns.is_empty()
580 {
581 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
582 table.name.clone(),
583 col.name.clone(),
584 "provider, model, version, and source columns are required".into(),
585 ));
586 }
587 if col.embedding_dim != Some(spec.dimension) {
588 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
589 table.name.clone(),
590 col.name.clone(),
591 "spec dimension must match embedding_dim".into(),
592 ));
593 }
594 let mut seen = HashSet::new();
595 if spec.source_columns.iter().any(|source_id| {
596 *source_id == col.id
597 || !seen.insert(*source_id)
598 || !column_ids.contains_key(source_id)
599 }) {
600 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
601 table.name.clone(),
602 col.name.clone(),
603 "source columns must exist, be unique, and exclude the target".into(),
604 ));
605 }
606 }
607
608 for pk in &table.primary_key {
609 if !column_names.contains_key(pk) {
610 return Err(SchemaError::MissingPrimaryKeyColumn(
611 table.name.clone(),
612 pk.clone(),
613 ));
614 }
615 }
616
617 for idx in &table.indexes {
618 for col in &idx.columns {
619 if !column_names.contains_key(col) {
620 return Err(SchemaError::MissingIndexColumn(
621 table.name.clone(),
622 idx.name.clone(),
623 col.clone(),
624 ));
625 }
626 }
627 }
628
629 for uq in &table.unique_constraints {
630 for col in &uq.columns {
631 if !column_names.contains_key(col) {
632 return Err(SchemaError::MissingUniqueColumn(
633 table.name.clone(),
634 uq.name.clone(),
635 col.clone(),
636 ));
637 }
638 }
639 }
640
641 for fk in &table.foreign_keys {
642 for col in &fk.columns {
643 if !column_names.contains_key(col) {
644 return Err(SchemaError::MissingForeignKeyColumn(
645 table.name.clone(),
646 fk.name.clone(),
647 col.clone(),
648 ));
649 }
650 }
651 if !table_names.contains_key(&fk.references_table) {
652 return Err(SchemaError::MissingReferencedTable(
653 table.name.clone(),
654 fk.name.clone(),
655 fk.references_table.clone(),
656 ));
657 }
658 }
659
660 Ok(())
661 }
662
663 pub fn table(&self, name: &str) -> Option<&Table> {
665 self.by_name.get(name).map(|&idx| &self.tables[idx])
666 }
667
668 pub fn table_by_id(&self, id: u32) -> Option<&Table> {
670 self.by_id.get(&id).map(|&idx| &self.tables[idx])
671 }
672
673 pub fn has_table(&self, name: &str) -> bool {
675 self.by_name.contains_key(name)
676 }
677
678 pub fn rename_table(&mut self, from: &str, to: &str) -> bool {
683 if from == to {
684 return self.has_table(from);
685 }
686 if !self.has_table(from) || self.has_table(to) {
687 return false;
688 }
689 let idx = *self.by_name.get(from).unwrap();
690 self.tables[idx].name = to.to_string();
691 self.by_name.remove(from);
692 self.by_name.insert(to.to_string(), idx);
693 true
694 }
695}
696
697#[cfg(test)]
698mod tests {
699 use super::*;
700
701 fn make_table(name: &str, id: u32) -> Table {
702 Table {
703 id,
704 name: name.into(),
705 columns: vec![Column::new(1, "id", ColumnType::Int64)],
706 primary_key: vec!["id".into()],
707 indexes: vec![],
708 foreign_keys: vec![],
709 unique_constraints: vec![],
710 check_constraints: vec![],
711 }
712 }
713
714 #[test]
715 fn schema_rejects_duplicate_table_name() {
716 let err = Schema::new(vec![make_table("a", 1), make_table("a", 2)]).unwrap_err();
717 assert!(matches!(err, SchemaError::DuplicateTableName(n) if n == "a"));
718 }
719
720 #[test]
721 fn schema_rejects_duplicate_table_id() {
722 let err = Schema::new(vec![make_table("a", 1), make_table("b", 1)]).unwrap_err();
723 assert!(matches!(err, SchemaError::DuplicateTableId(1)));
724 }
725
726 #[test]
727 fn schema_rejects_missing_pk_column() {
728 let t = Table {
729 id: 1,
730 name: "t".into(),
731 columns: vec![Column::new(1, "x", ColumnType::Text)],
732 primary_key: vec!["id".into()],
733 indexes: vec![],
734 foreign_keys: vec![],
735 unique_constraints: vec![],
736 check_constraints: vec![],
737 };
738 let err = Schema::new(vec![t]).unwrap_err();
739 assert!(matches!(err, SchemaError::MissingPrimaryKeyColumn(_, _)));
740 }
741
742 #[test]
743 fn unique_index_synthesizes_unique_constraint() {
744 let schema = Schema::new(vec![Table {
745 id: 1,
746 name: "users".into(),
747 columns: vec![
748 Column::new(1, "id", ColumnType::Int64),
749 Column::new(2, "email", ColumnType::Text),
750 Column::new(3, "handle", ColumnType::Text),
751 ],
752 primary_key: vec!["id".into()],
753 indexes: vec![
754 Index {
755 name: "idx_email".into(),
756 columns: vec!["email".into()],
757 unique: true,
758 kind: Default::default(),
759 ann_quantization: Default::default(),
760 ..Default::default()
761 },
762 Index {
764 name: "idx_handle".into(),
765 columns: vec!["handle".into()],
766 unique: false,
767 kind: Default::default(),
768 ann_quantization: Default::default(),
769 ..Default::default()
770 },
771 ],
772 foreign_keys: vec![],
773 unique_constraints: vec![],
774 check_constraints: vec![],
775 }])
776 .unwrap();
777 let table = schema.table("users").unwrap();
778 assert_eq!(table.unique_constraints.len(), 1);
779 assert_eq!(
780 table.unique_constraints[0].columns,
781 vec!["email".to_string()]
782 );
783 }
784
785 #[test]
786 fn unique_index_does_not_duplicate_existing_constraint() {
787 let schema = Schema::new(vec![Table {
788 id: 1,
789 name: "users".into(),
790 columns: vec![
791 Column::new(1, "id", ColumnType::Int64),
792 Column::new(2, "email", ColumnType::Text),
793 ],
794 primary_key: vec!["id".into()],
795 indexes: vec![Index {
796 name: "idx_email".into(),
797 columns: vec!["email".into()],
798 unique: true,
799 kind: Default::default(),
800 ann_quantization: Default::default(),
801 ..Default::default()
802 }],
803 foreign_keys: vec![],
804 unique_constraints: vec![UniqueConstraint {
805 name: "uq_email".into(),
806 columns: vec!["email".into()],
807 }],
808 check_constraints: vec![],
809 }])
810 .unwrap();
811 let table = schema.table("users").unwrap();
813 assert_eq!(table.unique_constraints.len(), 1);
814 assert_eq!(table.unique_constraints[0].name, "uq_email");
815 }
816
817 #[test]
818 fn schema_roundtrips_json() {
819 let schema = Schema::new(vec![Table {
820 id: 1,
821 name: "users".into(),
822 columns: vec![
823 Column::new(1, "id", ColumnType::Int64),
824 Column {
825 nullable: true,
826 ..Column::new(2, "email", ColumnType::Text)
827 },
828 ],
829 primary_key: vec!["id".into()],
830 indexes: vec![Index {
831 name: "idx_email".into(),
832 columns: vec!["email".into()],
833 unique: true,
834 kind: Default::default(),
835 ann_quantization: Default::default(),
836 ..Default::default()
837 }],
838 foreign_keys: vec![],
839 unique_constraints: vec![],
840 check_constraints: vec![CheckConstraint {
841 name: "chk_id_positive".into(),
842 expr: "id > 0".into(),
843 }],
844 }])
845 .unwrap();
846
847 let json = serde_json::to_string(&schema).unwrap();
848 let decoded: Schema = serde_json::from_str(&json).unwrap();
849 assert_eq!(decoded.tables.len(), 1);
850 assert_eq!(decoded.table("users").unwrap().columns.len(), 2);
851 }
852
853 #[test]
854 fn dense_ann_quantization_roundtrips_and_old_json_defaults_binary() {
855 let dense = Index {
856 name: "idx_embedding".into(),
857 columns: vec!["embedding".into()],
858 unique: false,
859 kind: IndexKind::Ann,
860 ann_quantization: AnnQuantization::Dense,
861 predicate: Some("embedding IS NOT NULL".into()),
862 ann_m: Some(24),
863 ann_ef_construction: Some(96),
864 ann_ef_search: Some(48),
865 ..Default::default()
866 };
867 let json = serde_json::to_value(&dense).unwrap();
868 assert_eq!(json["ann_quantization"], "dense");
869 assert_eq!(json["ann_m"], 24);
870 assert_eq!(json["predicate"], "embedding IS NOT NULL");
871 assert_eq!(
872 serde_json::from_value::<Index>(json)
873 .unwrap()
874 .ann_quantization,
875 AnnQuantization::Dense
876 );
877
878 let old: Index = serde_json::from_value(serde_json::json!({
879 "name": "idx_embedding",
880 "columns": ["embedding"],
881 "unique": false,
882 "kind": "ann"
883 }))
884 .unwrap();
885 assert_eq!(old.ann_quantization, AnnQuantization::BinarySign);
886 }
887
888 #[test]
889 fn swappable_ann_algorithm_and_product_quantization_roundtrip() {
890 let diskann = Index {
891 name: "idx_diskann".into(),
892 columns: vec!["embedding".into()],
893 unique: false,
894 kind: IndexKind::Ann,
895 ann_algorithm: AnnAlgorithm::DiskAnn,
896 ann_quantization: AnnQuantization::Dense,
897 ann_diskann_r: Some(128),
898 ann_diskann_l: Some(256),
899 ann_diskann_beam_width: Some(4),
900 ann_diskann_alpha: Some(130),
901 ..Default::default()
902 };
903 let json = serde_json::to_value(&diskann).unwrap();
904 assert_eq!(json["ann_algorithm"], "diskann");
905 assert_eq!(json["ann_diskann_r"], 128);
906 assert_eq!(json["ann_diskann_alpha"], 130);
907 let decoded: Index = serde_json::from_value(json).unwrap();
908 assert_eq!(decoded.ann_algorithm, AnnAlgorithm::DiskAnn);
909 assert_eq!(decoded.ann_diskann_l, Some(256));
910
911 let ivf = Index {
912 name: "idx_ivf".into(),
913 columns: vec!["embedding".into()],
914 kind: IndexKind::Ann,
915 ann_algorithm: AnnAlgorithm::Ivf,
916 ann_quantization: AnnQuantization::Dense,
917 ann_ivf_nlist: Some(512),
918 ann_ivf_nprobe: Some(16),
919 ann_ivf_training_samples: Some(20_000),
920 ..Default::default()
921 };
922 let json = serde_json::to_value(&ivf).unwrap();
923 assert_eq!(json["ann_algorithm"], "ivf");
924 assert_eq!(json["ann_ivf_nlist"], 512);
925 assert_eq!(json["ann_ivf_training_samples"], 20_000);
926
927 let pq = Index {
928 name: "idx_pq".into(),
929 columns: vec!["embedding".into()],
930 kind: IndexKind::Ann,
931 ann_quantization: AnnQuantization::Product {
932 num_subvectors: 32,
933 bits: 8,
934 },
935 ann_pq_training_samples: Some(10_000),
936 ann_pq_seed: Some(42),
937 ann_pq_rerank_factor: Some(3),
938 ..Default::default()
939 };
940 let json = serde_json::to_value(&pq).unwrap();
941 assert_eq!(json["ann_quantization"]["product"]["num_subvectors"], 32);
942 assert_eq!(json["ann_quantization"]["product"]["bits"], 8);
943 assert_eq!(json["ann_pq_seed"], 42);
944 let decoded: Index = serde_json::from_value(json).unwrap();
945 assert_eq!(
946 decoded.ann_quantization,
947 AnnQuantization::Product {
948 num_subvectors: 32,
949 bits: 8
950 }
951 );
952
953 let hnsw = Index {
955 name: "idx_hnsw".into(),
956 columns: vec!["embedding".into()],
957 kind: IndexKind::Ann,
958 ann_quantization: AnnQuantization::Dense,
959 ..Default::default()
960 };
961 let json = serde_json::to_value(&hnsw).unwrap();
962 assert_eq!(json["ann_algorithm"], "hnsw");
963 }
964
965 #[test]
966 fn embedding_source_roundtrips_json() {
967 let mut emb = Column::new(2, "vec", ColumnType::Embedding);
968 emb.embedding_dim = Some(4);
969 emb.embedding_source = Some(EmbeddingSource::LocalModel {
970 model_path: "/models/kit-mini".into(),
971 model_id: "kit-mini".into(),
972 });
973 let schema = Schema::new(vec![Table {
974 id: 1,
975 name: "docs".into(),
976 columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
977 primary_key: vec!["id".into()],
978 indexes: vec![],
979 foreign_keys: vec![],
980 unique_constraints: vec![],
981 check_constraints: vec![],
982 }])
983 .unwrap();
984 let json = serde_json::to_string(&schema).unwrap();
985 assert!(json.contains("local_model"));
986 assert!(json.contains("kit-mini"));
987 let decoded: Schema = serde_json::from_str(&json).unwrap();
988 let col = decoded.table("docs").unwrap().column("vec").unwrap();
989 assert_eq!(
990 col.embedding_source,
991 Some(EmbeddingSource::LocalModel {
992 model_path: "/models/kit-mini".into(),
993 model_id: "kit-mini".into(),
994 })
995 );
996 }
997
998 #[test]
999 fn embedding_source_rejected_on_non_embedding() {
1000 let mut col = Column::new(2, "name", ColumnType::Text);
1001 col.embedding_source = Some(EmbeddingSource::SuppliedByApplication);
1002 let err = Schema::new(vec![Table {
1003 id: 1,
1004 name: "t".into(),
1005 columns: vec![Column::new(1, "id", ColumnType::Int64), col],
1006 primary_key: vec!["id".into()],
1007 indexes: vec![],
1008 foreign_keys: vec![],
1009 unique_constraints: vec![],
1010 check_constraints: vec![],
1011 }])
1012 .unwrap_err();
1013 assert!(matches!(
1014 err,
1015 SchemaError::EmbeddingSourceOnNonEmbedding(_, _)
1016 ));
1017 }
1018
1019 #[test]
1020 fn generated_embedding_requires_dim() {
1021 let mut emb = Column::new(2, "vec", ColumnType::Embedding);
1022 emb.embedding_source = Some(EmbeddingSource::GeneratedColumn {
1023 provider: "local-test".into(),
1024 });
1025 let err = Schema::new(vec![Table {
1026 id: 1,
1027 name: "t".into(),
1028 columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
1029 primary_key: vec!["id".into()],
1030 indexes: vec![],
1031 foreign_keys: vec![],
1032 unique_constraints: vec![],
1033 check_constraints: vec![],
1034 }])
1035 .unwrap_err();
1036 assert!(matches!(err, SchemaError::EmbeddingSourceMissingDim(_, _)));
1037 }
1038
1039 #[test]
1040 fn generated_embedding_spec_roundtrips_and_validates_sources() {
1041 let mut emb = Column::new(3, "vec", ColumnType::Embedding);
1042 emb.embedding_dim = Some(4);
1043 emb.embedding_source = Some(EmbeddingSource::GeneratedColumnSpec {
1044 spec: GeneratedEmbeddingSpec {
1045 provider_id: "provider".into(),
1046 model_id: "model".into(),
1047 model_version: "1".into(),
1048 source_columns: vec![2],
1049 input_template: "{body}".into(),
1050 dimension: 4,
1051 normalization: EmbeddingSpecNormalization::None,
1052 failure_policy: EmbeddingWriteFailurePolicy::AbortWrite,
1053 },
1054 });
1055 let schema = Schema::new(vec![Table {
1056 id: 1,
1057 name: "docs".into(),
1058 columns: vec![
1059 Column::new(1, "id", ColumnType::Int64),
1060 Column::new(2, "body", ColumnType::Text),
1061 emb,
1062 ],
1063 primary_key: vec!["id".into()],
1064 indexes: vec![],
1065 foreign_keys: vec![],
1066 unique_constraints: vec![],
1067 check_constraints: vec![],
1068 }])
1069 .unwrap();
1070 let json = serde_json::to_string(&schema).unwrap();
1071 assert!(json.contains("generated_column_spec"));
1072 assert_eq!(serde_json::from_str::<Schema>(&json).unwrap(), schema);
1073 }
1074}