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_pq_training_samples: Option<usize>,
306 #[serde(default, skip_serializing_if = "Option::is_none")]
308 pub ann_pq_seed: Option<u64>,
309 #[serde(default, skip_serializing_if = "Option::is_none")]
311 pub ann_pq_rerank_factor: Option<usize>,
312 #[serde(default, skip_serializing_if = "Option::is_none")]
314 pub minhash_permutations: Option<usize>,
315 #[serde(default, skip_serializing_if = "Option::is_none")]
317 pub minhash_bands: Option<usize>,
318 #[serde(default, skip_serializing_if = "Option::is_none")]
320 pub learned_range_epsilon: Option<usize>,
321}
322
323#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
325pub struct UniqueConstraint {
326 pub name: String,
327 pub columns: Vec<String>,
328}
329
330#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
332pub struct ForeignKey {
333 pub name: String,
334 pub columns: Vec<String>,
335 pub references_table: String,
336 pub references_columns: Vec<String>,
337 #[serde(default)]
338 pub on_delete: ForeignKeyAction,
339}
340
341#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
343#[serde(rename_all = "snake_case")]
344pub enum ForeignKeyAction {
345 #[default]
346 Restrict,
347 Cascade,
348 SetNull,
349}
350
351#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
353pub struct CheckConstraint {
354 pub name: String,
355 pub expr: String,
356}
357
358#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
360pub struct Sequence {
361 pub name: String,
362 pub next_value: i64,
363}
364
365#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
367pub struct Table {
368 pub id: u32,
370 pub name: String,
371 pub columns: Vec<Column>,
372 pub primary_key: Vec<String>,
373 #[serde(default, skip_serializing_if = "Vec::is_empty")]
374 pub indexes: Vec<Index>,
375 #[serde(default, skip_serializing_if = "Vec::is_empty")]
376 pub foreign_keys: Vec<ForeignKey>,
377 #[serde(default, skip_serializing_if = "Vec::is_empty")]
378 pub unique_constraints: Vec<UniqueConstraint>,
379 #[serde(default, skip_serializing_if = "Vec::is_empty")]
380 pub check_constraints: Vec<CheckConstraint>,
381}
382
383impl Table {
384 pub fn column(&self, name: &str) -> Option<&Column> {
386 self.columns.iter().find(|c| c.name == name)
387 }
388
389 pub fn is_pk_column(&self, name: &str) -> bool {
391 self.primary_key.iter().any(|c| c == name)
392 }
393}
394
395#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
397pub enum SchemaError {
398 #[error("duplicate table name \"{0}\"")]
399 DuplicateTableName(String),
400 #[error("duplicate table id {0}")]
401 DuplicateTableId(u32),
402 #[error("duplicate column name \"{1}\" in table \"{0}\"")]
403 DuplicateColumnName(String, String),
404 #[error("duplicate column id {1} in table \"{0}\"")]
405 DuplicateColumnId(String, u32),
406 #[error("primary key column \"{1}\" not found in table \"{0}\"")]
407 MissingPrimaryKeyColumn(String, String),
408 #[error("index \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
409 MissingIndexColumn(String, String, String),
410 #[error("unique constraint \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
411 MissingUniqueColumn(String, String, String),
412 #[error("foreign key \"{1}\" references unknown column \"{2}\" in table \"{0}\"")]
413 MissingForeignKeyColumn(String, String, String),
414 #[error("foreign key \"{1}\" references unknown table \"{2}\"")]
415 MissingReferencedTable(String, String, String),
416 #[error("foreign key \"{1}\" references unknown column \"{2}\" on table \"{3}\"")]
417 MissingReferencedColumn(String, String, String, String),
418 #[error(
419 "column \"{1}\" on table \"{0}\" sets embedding_source but is not an embedding column"
420 )]
421 EmbeddingSourceOnNonEmbedding(String, String),
422 #[error(
423 "embedding column \"{1}\" on table \"{0}\" with LocalModel/GeneratedColumn source requires embedding_dim > 0"
424 )]
425 EmbeddingSourceMissingDim(String, String),
426 #[error("generated embedding column \"{1}\" on table \"{0}\" is invalid: {2}")]
427 InvalidGeneratedEmbeddingSpec(String, String, String),
428}
429
430#[derive(Debug, Clone, PartialEq, Serialize)]
432pub struct Schema {
433 pub tables: Vec<Table>,
434 by_name: HashMap<String, usize>,
435 by_id: HashMap<u32, usize>,
436}
437
438impl<'de> serde::Deserialize<'de> for Schema {
439 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
440 where
441 D: serde::Deserializer<'de>,
442 {
443 #[derive(serde::Deserialize)]
444 struct SchemaHelper {
445 tables: Vec<Table>,
446 }
447 let helper = SchemaHelper::deserialize(deserializer)?;
448 Schema::new(helper.tables).map_err(serde::de::Error::custom)
449 }
450}
451
452fn synthesize_unique_from_indexes(table: &mut Table) {
457 let mut synthesized: Vec<UniqueConstraint> = Vec::new();
458 for idx in &table.indexes {
459 if !idx.unique {
460 continue;
461 }
462 let covered = table
463 .unique_constraints
464 .iter()
465 .chain(synthesized.iter())
466 .any(|u| u.columns == idx.columns);
467 if !covered {
468 synthesized.push(UniqueConstraint {
469 name: idx.name.clone(),
470 columns: idx.columns.clone(),
471 });
472 }
473 }
474 table.unique_constraints.extend(synthesized);
475}
476
477impl Schema {
478 pub fn new(mut tables: Vec<Table>) -> Result<Self, SchemaError> {
480 for table in &mut tables {
481 synthesize_unique_from_indexes(table);
482 }
483
484 let mut by_name = HashMap::with_capacity(tables.len());
485 let mut by_id = HashMap::with_capacity(tables.len());
486
487 for (idx, table) in tables.iter().enumerate() {
488 if by_name.contains_key(&table.name) {
489 return Err(SchemaError::DuplicateTableName(table.name.clone()));
490 }
491 if by_id.contains_key(&table.id) {
492 return Err(SchemaError::DuplicateTableId(table.id));
493 }
494 by_name.insert(table.name.clone(), idx);
495 by_id.insert(table.id, idx);
496 }
497
498 for table in &tables {
499 Self::validate_table(table, &by_name)?;
500 }
501
502 Ok(Self {
503 tables,
504 by_name,
505 by_id,
506 })
507 }
508
509 fn validate_table(
510 table: &Table,
511 table_names: &HashMap<String, usize>,
512 ) -> Result<(), SchemaError> {
513 let mut column_names = HashMap::with_capacity(table.columns.len());
514 let mut column_ids = HashMap::with_capacity(table.columns.len());
515
516 for col in &table.columns {
517 if column_names.contains_key(&col.name) {
518 return Err(SchemaError::DuplicateColumnName(
519 table.name.clone(),
520 col.name.clone(),
521 ));
522 }
523 if column_ids.contains_key(&col.id) {
524 return Err(SchemaError::DuplicateColumnId(table.name.clone(), col.id));
525 }
526 if col.embedding_source.is_some() && col.storage_type != ColumnType::Embedding {
527 return Err(SchemaError::EmbeddingSourceOnNonEmbedding(
528 table.name.clone(),
529 col.name.clone(),
530 ));
531 }
532 if matches!(
533 col.embedding_source,
534 Some(
535 EmbeddingSource::LocalModel { .. }
536 | EmbeddingSource::ConfiguredModel { .. }
537 | EmbeddingSource::GeneratedColumn { .. }
538 | EmbeddingSource::GeneratedColumnSpec { .. }
539 )
540 ) && col.embedding_dim.unwrap_or(0) == 0
541 {
542 return Err(SchemaError::EmbeddingSourceMissingDim(
543 table.name.clone(),
544 col.name.clone(),
545 ));
546 }
547 column_names.insert(col.name.clone(), col.id);
548 column_ids.insert(col.id, col.name.clone());
549 }
550
551 for col in &table.columns {
552 if let Some(EmbeddingSource::ConfiguredModel {
553 provider_id,
554 model_id,
555 model_version,
556 }) = col.embedding_source.as_ref()
557 {
558 if provider_id.is_empty() || model_id.is_empty() || model_version.is_empty() {
559 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
560 table.name.clone(),
561 col.name.clone(),
562 "provider, model, and version are required".into(),
563 ));
564 }
565 }
566 }
567
568 for col in &table.columns {
569 let Some(EmbeddingSource::GeneratedColumnSpec { spec }) = col.embedding_source.as_ref()
570 else {
571 continue;
572 };
573 if spec.provider_id.is_empty()
574 || spec.model_id.is_empty()
575 || spec.model_version.is_empty()
576 || spec.source_columns.is_empty()
577 {
578 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
579 table.name.clone(),
580 col.name.clone(),
581 "provider, model, version, and source columns are required".into(),
582 ));
583 }
584 if col.embedding_dim != Some(spec.dimension) {
585 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
586 table.name.clone(),
587 col.name.clone(),
588 "spec dimension must match embedding_dim".into(),
589 ));
590 }
591 let mut seen = HashSet::new();
592 if spec.source_columns.iter().any(|source_id| {
593 *source_id == col.id
594 || !seen.insert(*source_id)
595 || !column_ids.contains_key(source_id)
596 }) {
597 return Err(SchemaError::InvalidGeneratedEmbeddingSpec(
598 table.name.clone(),
599 col.name.clone(),
600 "source columns must exist, be unique, and exclude the target".into(),
601 ));
602 }
603 }
604
605 for pk in &table.primary_key {
606 if !column_names.contains_key(pk) {
607 return Err(SchemaError::MissingPrimaryKeyColumn(
608 table.name.clone(),
609 pk.clone(),
610 ));
611 }
612 }
613
614 for idx in &table.indexes {
615 for col in &idx.columns {
616 if !column_names.contains_key(col) {
617 return Err(SchemaError::MissingIndexColumn(
618 table.name.clone(),
619 idx.name.clone(),
620 col.clone(),
621 ));
622 }
623 }
624 }
625
626 for uq in &table.unique_constraints {
627 for col in &uq.columns {
628 if !column_names.contains_key(col) {
629 return Err(SchemaError::MissingUniqueColumn(
630 table.name.clone(),
631 uq.name.clone(),
632 col.clone(),
633 ));
634 }
635 }
636 }
637
638 for fk in &table.foreign_keys {
639 for col in &fk.columns {
640 if !column_names.contains_key(col) {
641 return Err(SchemaError::MissingForeignKeyColumn(
642 table.name.clone(),
643 fk.name.clone(),
644 col.clone(),
645 ));
646 }
647 }
648 if !table_names.contains_key(&fk.references_table) {
649 return Err(SchemaError::MissingReferencedTable(
650 table.name.clone(),
651 fk.name.clone(),
652 fk.references_table.clone(),
653 ));
654 }
655 }
656
657 Ok(())
658 }
659
660 pub fn table(&self, name: &str) -> Option<&Table> {
662 self.by_name.get(name).map(|&idx| &self.tables[idx])
663 }
664
665 pub fn table_by_id(&self, id: u32) -> Option<&Table> {
667 self.by_id.get(&id).map(|&idx| &self.tables[idx])
668 }
669
670 pub fn has_table(&self, name: &str) -> bool {
672 self.by_name.contains_key(name)
673 }
674
675 pub fn rename_table(&mut self, from: &str, to: &str) -> bool {
680 if from == to {
681 return self.has_table(from);
682 }
683 if !self.has_table(from) || self.has_table(to) {
684 return false;
685 }
686 let idx = *self.by_name.get(from).unwrap();
687 self.tables[idx].name = to.to_string();
688 self.by_name.remove(from);
689 self.by_name.insert(to.to_string(), idx);
690 true
691 }
692}
693
694#[cfg(test)]
695mod tests {
696 use super::*;
697
698 fn make_table(name: &str, id: u32) -> Table {
699 Table {
700 id,
701 name: name.into(),
702 columns: vec![Column::new(1, "id", ColumnType::Int64)],
703 primary_key: vec!["id".into()],
704 indexes: vec![],
705 foreign_keys: vec![],
706 unique_constraints: vec![],
707 check_constraints: vec![],
708 }
709 }
710
711 #[test]
712 fn schema_rejects_duplicate_table_name() {
713 let err = Schema::new(vec![make_table("a", 1), make_table("a", 2)]).unwrap_err();
714 assert!(matches!(err, SchemaError::DuplicateTableName(n) if n == "a"));
715 }
716
717 #[test]
718 fn schema_rejects_duplicate_table_id() {
719 let err = Schema::new(vec![make_table("a", 1), make_table("b", 1)]).unwrap_err();
720 assert!(matches!(err, SchemaError::DuplicateTableId(1)));
721 }
722
723 #[test]
724 fn schema_rejects_missing_pk_column() {
725 let t = Table {
726 id: 1,
727 name: "t".into(),
728 columns: vec![Column::new(1, "x", ColumnType::Text)],
729 primary_key: vec!["id".into()],
730 indexes: vec![],
731 foreign_keys: vec![],
732 unique_constraints: vec![],
733 check_constraints: vec![],
734 };
735 let err = Schema::new(vec![t]).unwrap_err();
736 assert!(matches!(err, SchemaError::MissingPrimaryKeyColumn(_, _)));
737 }
738
739 #[test]
740 fn unique_index_synthesizes_unique_constraint() {
741 let schema = Schema::new(vec![Table {
742 id: 1,
743 name: "users".into(),
744 columns: vec![
745 Column::new(1, "id", ColumnType::Int64),
746 Column::new(2, "email", ColumnType::Text),
747 Column::new(3, "handle", ColumnType::Text),
748 ],
749 primary_key: vec!["id".into()],
750 indexes: vec![
751 Index {
752 name: "idx_email".into(),
753 columns: vec!["email".into()],
754 unique: true,
755 kind: Default::default(),
756 ann_quantization: Default::default(),
757 ..Default::default()
758 },
759 Index {
761 name: "idx_handle".into(),
762 columns: vec!["handle".into()],
763 unique: false,
764 kind: Default::default(),
765 ann_quantization: Default::default(),
766 ..Default::default()
767 },
768 ],
769 foreign_keys: vec![],
770 unique_constraints: vec![],
771 check_constraints: vec![],
772 }])
773 .unwrap();
774 let table = schema.table("users").unwrap();
775 assert_eq!(table.unique_constraints.len(), 1);
776 assert_eq!(
777 table.unique_constraints[0].columns,
778 vec!["email".to_string()]
779 );
780 }
781
782 #[test]
783 fn unique_index_does_not_duplicate_existing_constraint() {
784 let schema = Schema::new(vec![Table {
785 id: 1,
786 name: "users".into(),
787 columns: vec![
788 Column::new(1, "id", ColumnType::Int64),
789 Column::new(2, "email", ColumnType::Text),
790 ],
791 primary_key: vec!["id".into()],
792 indexes: vec![Index {
793 name: "idx_email".into(),
794 columns: vec!["email".into()],
795 unique: true,
796 kind: Default::default(),
797 ann_quantization: Default::default(),
798 ..Default::default()
799 }],
800 foreign_keys: vec![],
801 unique_constraints: vec![UniqueConstraint {
802 name: "uq_email".into(),
803 columns: vec!["email".into()],
804 }],
805 check_constraints: vec![],
806 }])
807 .unwrap();
808 let table = schema.table("users").unwrap();
810 assert_eq!(table.unique_constraints.len(), 1);
811 assert_eq!(table.unique_constraints[0].name, "uq_email");
812 }
813
814 #[test]
815 fn schema_roundtrips_json() {
816 let schema = Schema::new(vec![Table {
817 id: 1,
818 name: "users".into(),
819 columns: vec![
820 Column::new(1, "id", ColumnType::Int64),
821 Column {
822 nullable: true,
823 ..Column::new(2, "email", ColumnType::Text)
824 },
825 ],
826 primary_key: vec!["id".into()],
827 indexes: vec![Index {
828 name: "idx_email".into(),
829 columns: vec!["email".into()],
830 unique: true,
831 kind: Default::default(),
832 ann_quantization: Default::default(),
833 ..Default::default()
834 }],
835 foreign_keys: vec![],
836 unique_constraints: vec![],
837 check_constraints: vec![CheckConstraint {
838 name: "chk_id_positive".into(),
839 expr: "id > 0".into(),
840 }],
841 }])
842 .unwrap();
843
844 let json = serde_json::to_string(&schema).unwrap();
845 let decoded: Schema = serde_json::from_str(&json).unwrap();
846 assert_eq!(decoded.tables.len(), 1);
847 assert_eq!(decoded.table("users").unwrap().columns.len(), 2);
848 }
849
850 #[test]
851 fn dense_ann_quantization_roundtrips_and_old_json_defaults_binary() {
852 let dense = Index {
853 name: "idx_embedding".into(),
854 columns: vec!["embedding".into()],
855 unique: false,
856 kind: IndexKind::Ann,
857 ann_quantization: AnnQuantization::Dense,
858 predicate: Some("embedding IS NOT NULL".into()),
859 ann_m: Some(24),
860 ann_ef_construction: Some(96),
861 ann_ef_search: Some(48),
862 ..Default::default()
863 };
864 let json = serde_json::to_value(&dense).unwrap();
865 assert_eq!(json["ann_quantization"], "dense");
866 assert_eq!(json["ann_m"], 24);
867 assert_eq!(json["predicate"], "embedding IS NOT NULL");
868 assert_eq!(
869 serde_json::from_value::<Index>(json)
870 .unwrap()
871 .ann_quantization,
872 AnnQuantization::Dense
873 );
874
875 let old: Index = serde_json::from_value(serde_json::json!({
876 "name": "idx_embedding",
877 "columns": ["embedding"],
878 "unique": false,
879 "kind": "ann"
880 }))
881 .unwrap();
882 assert_eq!(old.ann_quantization, AnnQuantization::BinarySign);
883 }
884
885 #[test]
886 fn swappable_ann_algorithm_and_product_quantization_roundtrip() {
887 let diskann = Index {
888 name: "idx_diskann".into(),
889 columns: vec!["embedding".into()],
890 unique: false,
891 kind: IndexKind::Ann,
892 ann_algorithm: AnnAlgorithm::DiskAnn,
893 ann_quantization: AnnQuantization::Dense,
894 ann_diskann_r: Some(128),
895 ann_diskann_l: Some(256),
896 ann_diskann_beam_width: Some(4),
897 ann_diskann_alpha: Some(130),
898 ..Default::default()
899 };
900 let json = serde_json::to_value(&diskann).unwrap();
901 assert_eq!(json["ann_algorithm"], "diskann");
902 assert_eq!(json["ann_diskann_r"], 128);
903 assert_eq!(json["ann_diskann_alpha"], 130);
904 let decoded: Index = serde_json::from_value(json).unwrap();
905 assert_eq!(decoded.ann_algorithm, AnnAlgorithm::DiskAnn);
906 assert_eq!(decoded.ann_diskann_l, Some(256));
907
908 let ivf = Index {
909 name: "idx_ivf".into(),
910 columns: vec!["embedding".into()],
911 kind: IndexKind::Ann,
912 ann_algorithm: AnnAlgorithm::Ivf,
913 ann_quantization: AnnQuantization::Dense,
914 ann_ivf_nlist: Some(512),
915 ann_ivf_nprobe: Some(16),
916 ..Default::default()
917 };
918 let json = serde_json::to_value(&ivf).unwrap();
919 assert_eq!(json["ann_algorithm"], "ivf");
920 assert_eq!(json["ann_ivf_nlist"], 512);
921
922 let pq = Index {
923 name: "idx_pq".into(),
924 columns: vec!["embedding".into()],
925 kind: IndexKind::Ann,
926 ann_quantization: AnnQuantization::Product {
927 num_subvectors: 32,
928 bits: 8,
929 },
930 ann_pq_training_samples: Some(10_000),
931 ann_pq_seed: Some(42),
932 ann_pq_rerank_factor: Some(3),
933 ..Default::default()
934 };
935 let json = serde_json::to_value(&pq).unwrap();
936 assert_eq!(json["ann_quantization"]["product"]["num_subvectors"], 32);
937 assert_eq!(json["ann_quantization"]["product"]["bits"], 8);
938 assert_eq!(json["ann_pq_seed"], 42);
939 let decoded: Index = serde_json::from_value(json).unwrap();
940 assert_eq!(
941 decoded.ann_quantization,
942 AnnQuantization::Product {
943 num_subvectors: 32,
944 bits: 8
945 }
946 );
947
948 let hnsw = Index {
950 name: "idx_hnsw".into(),
951 columns: vec!["embedding".into()],
952 kind: IndexKind::Ann,
953 ann_quantization: AnnQuantization::Dense,
954 ..Default::default()
955 };
956 let json = serde_json::to_value(&hnsw).unwrap();
957 assert_eq!(json["ann_algorithm"], "hnsw");
958 }
959
960 #[test]
961 fn embedding_source_roundtrips_json() {
962 let mut emb = Column::new(2, "vec", ColumnType::Embedding);
963 emb.embedding_dim = Some(4);
964 emb.embedding_source = Some(EmbeddingSource::LocalModel {
965 model_path: "/models/kit-mini".into(),
966 model_id: "kit-mini".into(),
967 });
968 let schema = Schema::new(vec![Table {
969 id: 1,
970 name: "docs".into(),
971 columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
972 primary_key: vec!["id".into()],
973 indexes: vec![],
974 foreign_keys: vec![],
975 unique_constraints: vec![],
976 check_constraints: vec![],
977 }])
978 .unwrap();
979 let json = serde_json::to_string(&schema).unwrap();
980 assert!(json.contains("local_model"));
981 assert!(json.contains("kit-mini"));
982 let decoded: Schema = serde_json::from_str(&json).unwrap();
983 let col = decoded.table("docs").unwrap().column("vec").unwrap();
984 assert_eq!(
985 col.embedding_source,
986 Some(EmbeddingSource::LocalModel {
987 model_path: "/models/kit-mini".into(),
988 model_id: "kit-mini".into(),
989 })
990 );
991 }
992
993 #[test]
994 fn embedding_source_rejected_on_non_embedding() {
995 let mut col = Column::new(2, "name", ColumnType::Text);
996 col.embedding_source = Some(EmbeddingSource::SuppliedByApplication);
997 let err = Schema::new(vec![Table {
998 id: 1,
999 name: "t".into(),
1000 columns: vec![Column::new(1, "id", ColumnType::Int64), col],
1001 primary_key: vec!["id".into()],
1002 indexes: vec![],
1003 foreign_keys: vec![],
1004 unique_constraints: vec![],
1005 check_constraints: vec![],
1006 }])
1007 .unwrap_err();
1008 assert!(matches!(
1009 err,
1010 SchemaError::EmbeddingSourceOnNonEmbedding(_, _)
1011 ));
1012 }
1013
1014 #[test]
1015 fn generated_embedding_requires_dim() {
1016 let mut emb = Column::new(2, "vec", ColumnType::Embedding);
1017 emb.embedding_source = Some(EmbeddingSource::GeneratedColumn {
1018 provider: "local-test".into(),
1019 });
1020 let err = Schema::new(vec![Table {
1021 id: 1,
1022 name: "t".into(),
1023 columns: vec![Column::new(1, "id", ColumnType::Int64), emb],
1024 primary_key: vec!["id".into()],
1025 indexes: vec![],
1026 foreign_keys: vec![],
1027 unique_constraints: vec![],
1028 check_constraints: vec![],
1029 }])
1030 .unwrap_err();
1031 assert!(matches!(err, SchemaError::EmbeddingSourceMissingDim(_, _)));
1032 }
1033
1034 #[test]
1035 fn generated_embedding_spec_roundtrips_and_validates_sources() {
1036 let mut emb = Column::new(3, "vec", ColumnType::Embedding);
1037 emb.embedding_dim = Some(4);
1038 emb.embedding_source = Some(EmbeddingSource::GeneratedColumnSpec {
1039 spec: GeneratedEmbeddingSpec {
1040 provider_id: "provider".into(),
1041 model_id: "model".into(),
1042 model_version: "1".into(),
1043 source_columns: vec![2],
1044 input_template: "{body}".into(),
1045 dimension: 4,
1046 normalization: EmbeddingSpecNormalization::None,
1047 failure_policy: EmbeddingWriteFailurePolicy::AbortWrite,
1048 },
1049 });
1050 let schema = Schema::new(vec![Table {
1051 id: 1,
1052 name: "docs".into(),
1053 columns: vec![
1054 Column::new(1, "id", ColumnType::Int64),
1055 Column::new(2, "body", ColumnType::Text),
1056 emb,
1057 ],
1058 primary_key: vec!["id".into()],
1059 indexes: vec![],
1060 foreign_keys: vec![],
1061 unique_constraints: vec![],
1062 check_constraints: vec![],
1063 }])
1064 .unwrap();
1065 let json = serde_json::to_string(&schema).unwrap();
1066 assert!(json.contains("generated_column_spec"));
1067 assert_eq!(serde_json::from_str::<Schema>(&json).unwrap(), schema);
1068 }
1069}