1use super::*;
15use crate::catalog_cmds::CatalogCommand;
16use crate::jobs::{run_build_publish, BuildPublishJob, JobContext, JobError, JobKind, JobTarget};
17use crate::retention::{PinGuard, PinSource};
18use crate::schema::IndexDef;
19use crate::wal::DdlOp;
20
21#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
23enum IndexBuildKind {
24 Create {
25 definition: IndexDef,
26 },
27 Replace {
28 expected_old_name: String,
29 new_definition: IndexDef,
30 },
31}
32
33impl IndexBuildKind {
34 fn definition(&self) -> &IndexDef {
35 match self {
36 Self::Create { definition } => definition,
37 Self::Replace { new_definition, .. } => new_definition,
38 }
39 }
40}
41
42const INDEX_BUILD_SPEC_VERSION: u16 = 1;
43
44#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
47#[serde(deny_unknown_fields)]
48struct IndexBuildSpec {
49 version: u16,
50 table: String,
51 expected_schema_sequence: u64,
52 kind: IndexBuildKind,
53}
54
55impl IndexBuildSpec {
56 fn encode(&self) -> Result<Vec<u8>> {
57 serde_json::to_vec(self)
58 .map_err(|error| MongrelError::Other(format!("serialize index-build spec: {error}")))
59 }
60
61 fn decode(bytes: &[u8]) -> Result<Self> {
62 let spec: Self = serde_json::from_slice(bytes)
63 .map_err(|error| MongrelError::Other(format!("decode index-build spec: {error}")))?;
64 if spec.version != INDEX_BUILD_SPEC_VERSION {
65 return Err(MongrelError::Other(format!(
66 "unsupported index-build spec version {}",
67 spec.version
68 )));
69 }
70 Ok(spec)
71 }
72}
73
74struct IndexBuildJob<'a> {
76 db: &'a Database,
77 table: String,
78 kind: IndexBuildKind,
79 expected_schema_sequence: u64,
80 snapshot_epoch: Option<Epoch>,
82 pin: Option<PinGuard>,
84 snapshot_row_count: u64,
86 snapshot_rows: std::collections::BTreeMap<RowId, crate::memtable::Row>,
88 artifact: Option<crate::engine::SecondaryIndexArtifact>,
90 memory_reservation: Option<crate::memory::Reservation>,
92 built_through: Option<Epoch>,
94 built_data_generation: Option<u64>,
96 published: bool,
98}
99
100impl IndexBuildJob<'_> {
101 fn catalog_command(&self) -> CatalogCommand {
102 match &self.kind {
103 IndexBuildKind::Create { definition } => CatalogCommand::AddIndex {
104 table: self.table.clone(),
105 index: definition.clone(),
106 },
107 IndexBuildKind::Replace {
108 expected_old_name,
109 new_definition,
110 } => CatalogCommand::ReplaceIndex {
111 table: self.table.clone(),
112 expected_schema_sequence: self.expected_schema_sequence,
113 expected_old_name: expected_old_name.clone(),
114 new_definition: new_definition.clone(),
115 },
116 }
117 }
118
119 fn release_pin(&mut self) {
120 self.pin = None;
121 }
122
123 fn reserve_hidden_memory(
124 &mut self,
125 schema: &Schema,
126 rows: &[crate::memtable::Row],
127 ) -> std::result::Result<(), JobError> {
128 let row_bytes = rows.iter().fold(0u64, |total, row| {
129 total.saturating_add(row.estimated_bytes().saturating_add(64))
130 });
131 let mut estimate = row_bytes.saturating_mul(2);
132 if self.definition().kind == crate::schema::IndexKind::Ann {
133 let dim = schema
134 .columns
135 .iter()
136 .find(|column| column.id == self.definition().column_id)
137 .and_then(|column| match column.ty {
138 crate::schema::TypeId::Embedding { dim } => Some(dim as u64),
139 _ => None,
140 })
141 .ok_or_else(|| {
142 JobError::Phase(format!(
143 "ANN index {} has no embedding column",
144 self.definition().name
145 ))
146 })?;
147 let options = self.definition().options.ann.clone().unwrap_or_default();
148 let vector_bytes = match options.quantization {
150 crate::schema::AnnQuantization::BinarySign => dim.div_ceil(8),
151 crate::schema::AnnQuantization::Dense => dim.saturating_mul(4),
152 crate::schema::AnnQuantization::Product { num_subvectors, .. } => {
155 dim.saturating_mul(4).max(num_subvectors as u64)
156 }
157 };
158 let (per_node_struct, transient) = match options.algorithm {
160 crate::schema::AnnAlgorithm::Hnsw => {
161 let levels = (crate::index::hnsw::MAX_HNSW_LEVEL as u64).saturating_add(1);
162 let adjacency_slots = (options.m as u64).saturating_mul(2).saturating_add(
163 (options.m as u64).saturating_mul(levels.saturating_sub(1)),
164 );
165 let per_node = vector_bytes
166 .saturating_add(std::mem::size_of::<RowId>() as u64)
167 .saturating_add(
168 levels.saturating_mul(std::mem::size_of::<Vec<usize>>() as u64),
169 )
170 .saturating_add(
171 adjacency_slots.saturating_mul(std::mem::size_of::<usize>() as u64),
172 );
173 let transient = (options.ef_construction as u64)
174 .saturating_mul(dim.saturating_add(32))
175 .saturating_mul(8);
176 (per_node, transient)
177 }
178 crate::schema::AnnAlgorithm::DiskAnn => {
179 let r = options.diskann.as_ref().map(|d| d.r as u64).unwrap_or(64);
181 let per_node = vector_bytes
182 .saturating_add(std::mem::size_of::<RowId>() as u64)
183 .saturating_add(r.saturating_mul(std::mem::size_of::<usize>() as u64));
184 let transient = options
185 .diskann
186 .as_ref()
187 .map(|d| {
188 (d.l as u64)
189 .saturating_mul(dim.saturating_add(32))
190 .saturating_mul(8)
191 })
192 .unwrap_or(0);
193 (per_node, transient)
194 }
195 crate::schema::AnnAlgorithm::Ivf => {
196 let nlist = options.ivf.as_ref().map(|o| o.nlist as u64).unwrap_or(256);
199 let per_node = vector_bytes.saturating_add(std::mem::size_of::<RowId>() as u64);
200 let transient = nlist.saturating_mul(dim.saturating_mul(4));
201 (per_node, transient)
202 }
203 };
204 estimate = estimate.saturating_add(
206 per_node_struct
207 .saturating_mul(rows.len() as u64)
208 .saturating_mul(2),
209 );
210 estimate = estimate.saturating_add(transient);
211 }
212
213 let requested = usize::try_from(estimate).unwrap_or(usize::MAX);
214 let governor = self.db.memory_governor();
215 let result = match self.memory_reservation.as_mut() {
216 Some(reservation) => reservation.resize(estimate),
217 None => governor
218 .try_reserve(estimate, crate::memory::MemoryClass::Compaction)
219 .map(|reservation| self.memory_reservation = Some(reservation)),
220 };
221 result.map_err(|error| {
222 let limit = match error {
223 crate::memory::MemoryError::Exhausted { available, .. } => {
224 usize::try_from(available).unwrap_or(usize::MAX)
225 }
226 crate::memory::MemoryError::LowPriorityRejected { .. } => 0,
227 crate::memory::MemoryError::InvalidConfig(_) => {
228 usize::try_from(governor.class_budget(crate::memory::MemoryClass::Compaction))
229 .unwrap_or(usize::MAX)
230 }
231 };
232 JobError::ResourceLimitExceeded {
233 resource: "index build memory",
234 requested,
235 limit,
236 }
237 })
238 }
239
240 fn definition(&self) -> &IndexDef {
241 self.kind.definition()
242 }
243
244 fn publication_already_applied(&self, schema: &Schema) -> bool {
245 schema.schema_id == self.expected_schema_sequence.saturating_add(1)
246 && schema
247 .indexes
248 .iter()
249 .any(|index| index == self.definition())
250 && match &self.kind {
251 IndexBuildKind::Create { .. } => true,
252 IndexBuildKind::Replace {
253 expected_old_name,
254 new_definition,
255 } => {
256 expected_old_name == &new_definition.name
257 || !schema
258 .indexes
259 .iter()
260 .any(|index| index.name == *expected_old_name)
261 }
262 }
263 }
264
265 fn refresh_hidden(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
270 context.check_cancelled()?;
271 let handle = self.db.table(&self.table).map_err(job_phase_err)?;
272 let (generation, snapshot) = handle
273 .read_generation_with_context(None)
274 .map_err(job_phase_err)?;
275 if self.pin.is_none() {
276 self.pin = Some(
277 Arc::clone(generation.pin_registry())
278 .pin(PinSource::OnlineIndexBuild, snapshot.epoch),
279 );
280 self.snapshot_epoch.get_or_insert(snapshot.epoch);
281 }
282
283 let rows = generation.visible_rows(snapshot).map_err(job_phase_err)?;
284 self.reserve_hidden_memory(generation.schema(), &rows)?;
285 let next: std::collections::BTreeMap<_, _> =
286 rows.into_iter().map(|row| (row.row_id, row)).collect();
287 self.snapshot_rows
288 .retain(|row_id, _| next.contains_key(row_id));
289 for (row_id, row) in next {
290 self.snapshot_rows.insert(row_id, row);
291 }
292
293 let rows: Vec<_> = self.snapshot_rows.values().cloned().collect();
294 let total = rows.len();
295 let definition = self.definition().clone();
296 let mut last_reported = None;
297 let artifact = generation
298 .build_secondary_index_artifact(&definition, &rows, |done, total| {
299 context.check_cancelled().map_err(MongrelError::from)?;
300 if total > 0
301 && last_reported != Some(done)
302 && (done.is_multiple_of(256) || done == total)
303 {
304 context
305 .report_progress(done as u64, total as u64)
306 .map_err(MongrelError::from)?;
307 last_reported = Some(done);
308 }
309 Ok(())
310 })
311 .map_err(job_phase_err)?;
312 self.snapshot_row_count = total as u64;
313 self.built_through = Some(snapshot.epoch);
314 self.built_data_generation = Some(generation.data_generation());
315 self.artifact = Some(artifact);
316 Ok(())
317 }
318}
319
320impl BuildPublishJob for IndexBuildJob<'_> {
321 fn checkpoint_state(&self) -> Vec<u8> {
322 let epoch = self.snapshot_epoch.map(|e| e.0).unwrap_or(0);
324 let built_through = self.built_through.map(|epoch| epoch.0).unwrap_or(0);
325 let built_data_generation = self.built_data_generation.unwrap_or(0);
326 let published = u8::from(self.published);
327 let mut out = Vec::with_capacity(33);
328 out.extend_from_slice(&epoch.to_le_bytes());
329 out.extend_from_slice(&self.snapshot_row_count.to_le_bytes());
330 out.extend_from_slice(&built_through.to_le_bytes());
331 out.extend_from_slice(&built_data_generation.to_le_bytes());
332 out.push(published);
333 out
334 }
335
336 fn restore_checkpoint(&mut self, state: &[u8]) -> std::result::Result<(), JobError> {
337 if state.len() < 17 {
338 return Err(JobError::Phase(
339 "index-build checkpoint is truncated".into(),
340 ));
341 }
342 let mut epoch_bytes = [0u8; 8];
343 epoch_bytes.copy_from_slice(&state[0..8]);
344 let epoch = u64::from_le_bytes(epoch_bytes);
345 if epoch != 0 {
346 self.snapshot_epoch = Some(Epoch(epoch));
347 }
348 let mut count_bytes = [0u8; 8];
349 count_bytes.copy_from_slice(&state[8..16]);
350 self.snapshot_row_count = u64::from_le_bytes(count_bytes);
351 if state.len() >= 33 {
352 let mut built_bytes = [0u8; 8];
353 built_bytes.copy_from_slice(&state[16..24]);
354 let built = u64::from_le_bytes(built_bytes);
355 if built != 0 {
356 self.built_through = Some(Epoch(built));
357 }
358 let mut generation_bytes = [0u8; 8];
359 generation_bytes.copy_from_slice(&state[24..32]);
360 self.built_data_generation = Some(u64::from_le_bytes(generation_bytes));
361 self.published = state[32] != 0;
362 } else if state.len() >= 25 {
363 let mut built_bytes = [0u8; 8];
364 built_bytes.copy_from_slice(&state[16..24]);
365 let built = u64::from_le_bytes(built_bytes);
366 if built != 0 {
367 self.built_through = Some(Epoch(built));
368 }
369 self.published = state[24] != 0;
370 } else {
371 self.published = state[16] != 0;
372 }
373 Ok(())
374 }
375
376 fn record_pending(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
377 context.check_cancelled()?;
378 let catalog = self.db.catalog.read();
381 let entry = catalog
382 .live(&self.table)
383 .ok_or_else(|| JobError::Phase(format!("table {:?} not found", self.table)))?;
384 if entry.schema.schema_id != self.expected_schema_sequence {
385 return Err(JobError::Phase(format!(
386 "schema sequence changed before index build started on {:?}",
387 self.table
388 )));
389 }
390 let command = self.catalog_command();
391 crate::catalog_cmds::apply(&catalog, &command).map_err(job_phase_err)?;
392 Ok(())
393 }
394
395 fn pin_snapshot(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
396 context.check_cancelled()?;
397 if self.pin.is_some() {
398 return Ok(());
399 }
400 let handle = self.db.table(&self.table).map_err(job_phase_err)?;
401 let table = handle.lock();
402 let epoch = table.current_epoch();
403 let pin = Arc::clone(table.pin_registry()).pin(PinSource::OnlineIndexBuild, epoch);
404 self.snapshot_epoch = Some(epoch);
405 self.pin = Some(pin);
406 Ok(())
407 }
408
409 fn build_hidden(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
410 self.refresh_hidden(context)
411 }
412
413 fn catch_up(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
414 self.refresh_hidden(context)
415 }
416
417 fn validate(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
418 context.check_cancelled()?;
419 let catalog = self.db.catalog.read();
420 let entry = catalog
421 .live(&self.table)
422 .ok_or_else(|| JobError::Phase(format!("table {:?} not found", self.table)))?;
423 if entry.schema.schema_id != self.expected_schema_sequence {
424 return Err(JobError::Phase(format!(
425 "schema sequence conflict during index build on {:?}: expected {}, found {}",
426 self.table, self.expected_schema_sequence, entry.schema.schema_id
427 )));
428 }
429 crate::catalog_cmds::apply(&catalog, &self.catalog_command()).map_err(job_phase_err)?;
430 Ok(())
431 }
432
433 fn publish(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
434 context.check_cancelled()?;
435 if self.published {
436 return Ok(());
437 }
438 match self.db.publish_index_build(self, context) {
439 Ok(()) => {
440 self.published = true;
441 Ok(())
442 }
443 Err(error @ MongrelError::DurableCommit { .. }) => {
444 let applied = self
445 .db
446 .catalog
447 .read()
448 .live(&self.table)
449 .is_some_and(|entry| self.publication_already_applied(&entry.schema));
450 if applied {
451 self.published = true;
452 Ok(())
453 } else {
454 Err(job_phase_err(error))
455 }
456 }
457 Err(error) => Err(job_phase_err(error)),
458 }
459 }
460
461 fn release_old(&mut self, context: &JobContext) -> std::result::Result<(), JobError> {
462 context.check_cancelled()?;
463 self.release_pin();
466 self.artifact = None;
467 self.memory_reservation = None;
468 self.snapshot_rows.clear();
469 Ok(())
470 }
471
472 fn rollback(&mut self) -> std::result::Result<(), JobError> {
473 self.release_pin();
476 self.artifact = None;
477 self.memory_reservation = None;
478 self.snapshot_rows.clear();
479 Ok(())
480 }
481}
482
483fn job_phase_err(error: MongrelError) -> JobError {
484 match error {
485 MongrelError::Cancelled => JobError::Cancelled,
486 MongrelError::Conflict(message) => JobError::Phase(message),
487 other => JobError::Phase(other.to_string()),
488 }
489}
490
491impl Database {
492 pub fn index_build_target_definition(&self, job_id: u64) -> Result<IndexDef> {
499 let record = self
500 .job_registry
501 .get(job_id)
502 .ok_or_else(|| MongrelError::NotFound(format!("job {job_id} not found")))?;
503 if record.kind != JobKind::IndexBuild {
504 return Err(MongrelError::InvalidArgument(format!(
505 "job {job_id} is not an index build"
506 )));
507 }
508 let definition = record.definition.as_deref().ok_or_else(|| {
509 MongrelError::Other(format!(
510 "index-build job {job_id} has no durable definition"
511 ))
512 })?;
513 Ok(IndexBuildSpec::decode(definition)?
514 .kind
515 .definition()
516 .clone())
517 }
518
519 pub fn create_index(&self, table: &str, definition: IndexDef) -> Result<u64> {
524 self.run_index_build(
525 table,
526 IndexBuildKind::Create {
527 definition: definition.clone(),
528 },
529 Some(definition),
530 None,
531 )
532 }
533
534 pub fn submit_create_index(&self, table: &str, definition: IndexDef) -> Result<u64> {
540 let (job_id, _) = self.submit_index_build(
541 table,
542 IndexBuildKind::Create {
543 definition: definition.clone(),
544 },
545 Some(definition),
546 None,
547 )?;
548 Ok(job_id)
549 }
550
551 pub fn start_create_index(self: &Arc<Self>, table: &str, definition: IndexDef) -> Result<u64> {
553 let (job_id, spec) = self.submit_index_build(
554 table,
555 IndexBuildKind::Create {
556 definition: definition.clone(),
557 },
558 Some(definition),
559 None,
560 )?;
561 self.spawn_index_build(job_id, spec)?;
562 Ok(job_id)
563 }
564
565 pub fn replace_index(
571 &self,
572 table: &str,
573 expected_old_name: &str,
574 new_definition: IndexDef,
575 ) -> Result<u64> {
576 self.run_index_build(
577 table,
578 IndexBuildKind::Replace {
579 expected_old_name: expected_old_name.to_string(),
580 new_definition: new_definition.clone(),
581 },
582 Some(new_definition),
583 Some(expected_old_name),
584 )
585 }
586
587 pub fn submit_replace_index(
589 &self,
590 table: &str,
591 expected_old_name: &str,
592 new_definition: IndexDef,
593 ) -> Result<u64> {
594 let (job_id, _) = self.submit_index_build(
595 table,
596 IndexBuildKind::Replace {
597 expected_old_name: expected_old_name.to_string(),
598 new_definition: new_definition.clone(),
599 },
600 Some(new_definition),
601 Some(expected_old_name),
602 )?;
603 Ok(job_id)
604 }
605
606 pub fn start_replace_index(
608 self: &Arc<Self>,
609 table: &str,
610 expected_old_name: &str,
611 new_definition: IndexDef,
612 ) -> Result<u64> {
613 let (job_id, spec) = self.submit_index_build(
614 table,
615 IndexBuildKind::Replace {
616 expected_old_name: expected_old_name.to_string(),
617 new_definition: new_definition.clone(),
618 },
619 Some(new_definition),
620 Some(expected_old_name),
621 )?;
622 self.spawn_index_build(job_id, spec)?;
623 Ok(job_id)
624 }
625
626 pub fn resume_index_build(self: &Arc<Self>, job_id: u64) -> Result<()> {
629 let record = self
630 .job_registry
631 .get(job_id)
632 .ok_or_else(|| MongrelError::NotFound(format!("job {job_id}")))?;
633 if record.kind != JobKind::IndexBuild {
634 return Err(MongrelError::InvalidArgument(format!(
635 "job {job_id} is not an index build"
636 )));
637 }
638 let definition = record.definition.as_deref().ok_or_else(|| {
639 MongrelError::Other(format!(
640 "index build job {job_id} has no reconstructible definition"
641 ))
642 })?;
643 let spec = IndexBuildSpec::decode(definition)?;
644 match record.state {
645 crate::jobs::JobState::Paused => self
646 .job_registry
647 .resume(job_id)
648 .map_err(MongrelError::from)?,
649 crate::jobs::JobState::Pending => {}
650 state => {
651 return Err(MongrelError::Conflict(format!(
652 "index build job {job_id} cannot resume from {state:?}"
653 )));
654 }
655 }
656 self.spawn_index_build(job_id, spec)
657 }
658
659 pub fn drop_index(&self, table: &str, name: &str) -> Result<()> {
662 use std::sync::atomic::Ordering;
663
664 self.require(&crate::auth::Permission::Ddl)?;
665 if self.poisoned.load(Ordering::Relaxed) {
666 return Err(MongrelError::Other(
667 "database poisoned by fsync error".into(),
668 ));
669 }
670 let _operation = self.admit_operation()?;
671 let _schema_barrier = self.acquire_schema_barrier_exclusive()?;
672 let _ddl = self.ddl_lock.lock();
673 let _security_write = self.security_write()?;
674 self.require(&crate::auth::Permission::Ddl)?;
675
676 let table_id = {
677 let catalog = self.catalog.read();
678 catalog
679 .live(table)
680 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not found")))?
681 .table_id
682 };
683 let handle = self
684 .tables
685 .read()
686 .get(&table_id)
687 .cloned()
688 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not mounted")))?;
689
690 let command = CatalogCommand::RemoveIndex {
691 table: table.to_string(),
692 name: name.to_string(),
693 };
694 let prepared_schema = match crate::catalog_cmds::apply(&self.catalog.read(), &command)? {
696 crate::catalog_cmds::CatalogDelta::SchemaReplaced { schema, .. } => schema,
697 crate::catalog_cmds::CatalogDelta::NoOp => return Ok(()),
698 other => {
699 return Err(MongrelError::Other(format!(
700 "unexpected catalog delta for drop_index: {other:?}"
701 )));
702 }
703 };
704
705 crate::catalog::inject_hook("index.publish.before")?;
706
707 let durable_epoch = std::cell::Cell::new(None);
708 let result: Result<()> = (|| {
709 let mut table_guard = handle.lock();
710 let commit_lock = Arc::clone(&self.commit_lock);
711 let _commit = commit_lock.lock();
712 let epoch = self.epoch.bump_assigned();
713 let mut epoch_guard = EpochGuard::new(self.epoch.as_ref(), epoch);
714 let txn_id = self.alloc_txn_id()?;
715 let mut next_catalog = self.catalog.read().clone();
716 let catalog_entry_index = next_catalog
717 .tables
718 .iter()
719 .position(|entry| entry.table_id == table_id)
720 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not found")))?;
721 self.apply_catalog_command_to(&mut next_catalog, command)?;
722 next_catalog.tables[catalog_entry_index].schema = prepared_schema.clone();
723 next_catalog.db_epoch = next_catalog.db_epoch.max(epoch.0);
724
725 let commit_seq = {
726 let mut wal = self.shared_wal.lock();
727 let append: Result<u64> = (|| {
728 append_catalog_snapshot(&mut wal, txn_id, &next_catalog)?;
729 wal.append_commit(txn_id, epoch, &[])
730 })();
731 append.map_err(|error| self.commit_outcome_unknown(epoch, error))?
732 };
733 let receipt = self.await_durable_commit(txn_id, commit_seq, epoch)?;
734 durable_epoch.set(Some(epoch));
735
736 table_guard.publish_index_drop(prepared_schema)?;
737 let schema = table_guard.schema().clone();
738 drop(table_guard);
739 next_catalog.tables[catalog_entry_index].schema = schema;
740 let catalog_result =
741 catalog::write_atomic(&self.root, &next_catalog, self.meta_dek.as_ref());
742 *self.catalog.write() = next_catalog;
743 self.publish_committed(&receipt, epoch)?;
744 epoch_guard.disarm();
745 if let Err(error) = catalog_result {
746 self.poisoned.store(true, Ordering::Relaxed);
747 self.lifecycle.poison();
748 return Err(MongrelError::DurableCommit {
749 epoch: epoch.0,
750 message: error.to_string(),
751 });
752 }
753 crate::catalog::inject_hook("index.publish.after")?;
754 Ok(())
755 })();
756 result.map_err(|error| match (durable_epoch.get(), error) {
757 (_, error @ MongrelError::DurableCommit { .. }) => error,
758 (Some(epoch), error) => MongrelError::DurableCommit {
759 epoch: epoch.0,
760 message: error.to_string(),
761 },
762 (None, error) => error,
763 })
764 }
765
766 fn run_index_build(
767 &self,
768 table: &str,
769 kind: IndexBuildKind,
770 definition: Option<IndexDef>,
771 expected_old_name: Option<&str>,
772 ) -> Result<u64> {
773 let (job_id, spec) = self.submit_index_build(table, kind, definition, expected_old_name)?;
774 self.drive_index_build(job_id, spec)
775 }
776
777 fn submit_index_build(
778 &self,
779 table: &str,
780 kind: IndexBuildKind,
781 definition: Option<IndexDef>,
782 expected_old_name: Option<&str>,
783 ) -> Result<(u64, IndexBuildSpec)> {
784 self.require(&crate::auth::Permission::Ddl)?;
785 if self.poisoned.load(Ordering::Relaxed) {
786 return Err(MongrelError::Other(
787 "database poisoned by fsync error".into(),
788 ));
789 }
790 if table.is_empty() {
791 return Err(MongrelError::InvalidArgument(
792 "index DDL requires a non-empty table name".into(),
793 ));
794 }
795
796 let (expected_schema_sequence, index_name) = {
798 let catalog = self.catalog.read();
799 let entry = catalog
800 .live(table)
801 .ok_or_else(|| MongrelError::NotFound(format!("table {table:?} not found")))?;
802 if let Some(old) = expected_old_name {
803 if !entry.schema.indexes.iter().any(|index| index.name == old) {
804 return Err(MongrelError::NotFound(format!(
805 "index {old} does not exist on {table}"
806 )));
807 }
808 }
809 if let Some(def) = &definition {
810 def.validate_options()?;
811 if let crate::schema::AnnQuantization::Product { num_subvectors, .. } = def
817 .options
818 .ann
819 .as_ref()
820 .map(|o| o.quantization)
821 .unwrap_or_default()
822 {
823 let dim = entry
824 .schema
825 .columns
826 .iter()
827 .find(|c| c.id == def.column_id)
828 .and_then(|c| match c.ty {
829 crate::schema::TypeId::Embedding { dim } => Some(dim as usize),
830 _ => None,
831 })
832 .ok_or_else(|| {
833 JobError::Phase(format!(
834 "ANN index {} references a non-embedding column",
835 def.name
836 ))
837 })?;
838 if dim == 0 || dim % num_subvectors as usize != 0 {
839 return Err(MongrelError::Schema(format!(
840 "product quantization num_subvectors ({num_subvectors}) must evenly divide the column dimension ({dim})"
841 )));
842 }
843 }
844 }
845 let command = match &kind {
848 IndexBuildKind::Create { definition } => CatalogCommand::AddIndex {
849 table: table.to_string(),
850 index: definition.clone(),
851 },
852 IndexBuildKind::Replace {
853 expected_old_name,
854 new_definition,
855 } => CatalogCommand::ReplaceIndex {
856 table: table.to_string(),
857 expected_schema_sequence: entry.schema.schema_id,
858 expected_old_name: expected_old_name.clone(),
859 new_definition: new_definition.clone(),
860 },
861 };
862 crate::catalog_cmds::apply(&catalog, &command)?;
863 let name = match &kind {
864 IndexBuildKind::Create { definition } => definition.name.clone(),
865 IndexBuildKind::Replace { new_definition, .. } => new_definition.name.clone(),
866 };
867 (entry.schema.schema_id, name)
868 };
869
870 let _ = self.table(table)?;
872
873 let spec = IndexBuildSpec {
874 version: INDEX_BUILD_SPEC_VERSION,
875 table: table.to_string(),
876 expected_schema_sequence,
877 kind,
878 };
879 let job_id = self.job_registry.submit_with_definition(
880 JobKind::IndexBuild,
881 JobTarget {
882 table: table.to_string(),
883 index: Some(index_name),
884 },
885 Some(spec.encode()?),
886 )?;
887 Ok((job_id, spec))
888 }
889
890 fn spawn_index_build(self: &Arc<Self>, job_id: u64, spec: IndexBuildSpec) -> Result<()> {
891 let db = Arc::clone(self);
892 std::thread::Builder::new()
893 .name(format!("mongreldb-index-{job_id}"))
894 .spawn(move || {
895 let _ = db.drive_index_build(job_id, spec);
897 })
898 .map_err(|error| {
899 MongrelError::Io(std::io::Error::new(
900 error.kind(),
901 format!("spawn index build job {job_id}: {error}"),
902 ))
903 })?;
904 Ok(())
905 }
906
907 fn drive_index_build(&self, job_id: u64, spec: IndexBuildSpec) -> Result<u64> {
908 let mut job = IndexBuildJob {
909 db: self,
910 table: spec.table,
911 kind: spec.kind,
912 expected_schema_sequence: spec.expected_schema_sequence,
913 snapshot_epoch: None,
914 pin: None,
915 snapshot_row_count: 0,
916 snapshot_rows: std::collections::BTreeMap::new(),
917 artifact: None,
918 memory_reservation: None,
919 built_through: None,
920 built_data_generation: None,
921 published: false,
922 };
923
924 match run_build_publish(self.job_registry.as_ref(), job_id, &mut job) {
925 Ok(()) => {
926 let record = self.job_registry.get(job_id);
927 match record.map(|r| r.state) {
928 Some(crate::jobs::JobState::Succeeded) => Ok(job_id),
929 Some(crate::jobs::JobState::Paused) => Err(MongrelError::Other(format!(
930 "index build job {job_id} parked mid-run; resume and redrive"
931 ))),
932 Some(state) => Err(MongrelError::Other(format!(
933 "index build job {job_id} ended in unexpected state {state:?}"
934 ))),
935 None => Err(MongrelError::NotFound(format!("job {job_id}"))),
936 }
937 }
938 Err(JobError::Cancelled) => {
939 if job.published
943 || self
944 .job_registry
945 .get(job_id)
946 .is_some_and(|r| r.state == crate::jobs::JobState::Succeeded)
947 {
948 Ok(job_id)
949 } else {
950 Err(MongrelError::Cancelled)
951 }
952 }
953 Err(error) => {
954 if job.published {
955 Ok(job_id)
957 } else {
958 Err(error.into())
959 }
960 }
961 }
962 }
963
964 fn publish_index_build(&self, job: &mut IndexBuildJob<'_>, context: &JobContext) -> Result<()> {
969 use std::sync::atomic::Ordering;
970
971 if self.poisoned.load(Ordering::Relaxed) {
972 return Err(MongrelError::Other(
973 "database poisoned by fsync error".into(),
974 ));
975 }
976 self.require(&crate::auth::Permission::Ddl)?;
978
979 let table_name = job.table.clone();
980 let table_id = {
981 let catalog = self.catalog.read();
982 let entry = catalog
983 .live(&table_name)
984 .ok_or_else(|| MongrelError::NotFound(format!("table {table_name:?} not found")))?;
985 if entry.schema.schema_id != job.expected_schema_sequence {
986 if job.publication_already_applied(&entry.schema) {
987 job.published = true;
988 return Ok(());
989 }
990 return Err(MongrelError::Conflict(format!(
991 "index publish on {table_name}: expected schema sequence {}, found {}",
992 job.expected_schema_sequence, entry.schema.schema_id
993 )));
994 }
995 entry.table_id
996 };
997 let handle =
998 self.tables.read().get(&table_id).cloned().ok_or_else(|| {
999 MongrelError::NotFound(format!("table {table_name:?} not mounted"))
1000 })?;
1001
1002 loop {
1003 context.check_cancelled().map_err(MongrelError::from)?;
1004 if job.artifact.is_none() || job.built_data_generation.is_none() {
1005 job.refresh_hidden(context).map_err(MongrelError::from)?;
1006 }
1007 let built_data_generation = job.built_data_generation.ok_or_else(|| {
1008 MongrelError::Other("index build has no content watermark".into())
1009 })?;
1010 let command = job.catalog_command();
1011 let prepared_schema = match crate::catalog_cmds::apply(&self.catalog.read(), &command)?
1012 {
1013 crate::catalog_cmds::CatalogDelta::SchemaReplaced { schema, .. } => schema,
1014 other => {
1015 return Err(MongrelError::Other(format!(
1016 "unexpected catalog delta for index publish: {other:?}"
1017 )));
1018 }
1019 };
1020 let durable_epoch = std::cell::Cell::new(None);
1021
1022 let _schema_barrier = self.acquire_schema_barrier_exclusive()?;
1025 let _ddl = self.ddl_lock.lock();
1026 let _security_write = self.security_write()?;
1027 self.require(&crate::auth::Permission::Ddl)?;
1028
1029 {
1031 let catalog = self.catalog.read();
1032 let entry = catalog.live(&table_name).ok_or_else(|| {
1033 MongrelError::NotFound(format!("table {table_name:?} not found"))
1034 })?;
1035 if entry.schema.schema_id != job.expected_schema_sequence {
1036 if job.publication_already_applied(&entry.schema) {
1037 job.published = true;
1038 return Ok(());
1039 }
1040 return Err(MongrelError::Conflict(format!(
1041 "index publish on {table_name}: expected schema sequence {}, found {}",
1042 job.expected_schema_sequence, entry.schema.schema_id
1043 )));
1044 }
1045 }
1046
1047 let commit_lock = Arc::clone(&self.commit_lock);
1049 let _commit = commit_lock.lock();
1050 let mut table_guard = handle.lock();
1051 if table_guard.data_generation() != built_data_generation {
1052 drop(table_guard);
1053 drop(_commit);
1054 drop(_security_write);
1055 drop(_ddl);
1056 drop(_schema_barrier);
1057 job.refresh_hidden(context).map_err(MongrelError::from)?;
1058 continue;
1059 }
1060
1061 let result: Result<()> = (|| {
1062 crate::catalog::inject_hook("index.publish.before")?;
1065
1066 let epoch = self.epoch.bump_assigned();
1067 let mut epoch_guard = EpochGuard::new(self.epoch.as_ref(), epoch);
1068 let txn_id = self.alloc_txn_id()?;
1069 let mut next_catalog = self.catalog.read().clone();
1070 let catalog_entry_index = next_catalog
1071 .tables
1072 .iter()
1073 .position(|entry| entry.table_id == table_id)
1074 .ok_or_else(|| {
1075 MongrelError::NotFound(format!("table {table_name:?} not found"))
1076 })?;
1077 self.apply_catalog_command_to(&mut next_catalog, command)?;
1078 next_catalog.tables[catalog_entry_index].schema = prepared_schema.clone();
1079 next_catalog.db_epoch = next_catalog.db_epoch.max(epoch.0);
1080
1081 let commit_seq = {
1082 let mut wal = self.shared_wal.lock();
1083 let append: Result<u64> = (|| {
1084 let _ = DdlOp::encode_schema(&prepared_schema)?;
1086 append_catalog_snapshot(&mut wal, txn_id, &next_catalog)?;
1087 wal.append_commit(txn_id, epoch, &[])
1088 })();
1089 append.map_err(|error| self.commit_outcome_unknown(epoch, error))?
1090 };
1091 let receipt = self.await_durable_commit(txn_id, commit_seq, epoch)?;
1092 durable_epoch.set(Some(epoch));
1093
1094 let artifact = job.artifact.take().ok_or_else(|| {
1095 MongrelError::Other("index build lost its hidden artifact".into())
1096 })?;
1097 table_guard.publish_index_schema_change(prepared_schema, artifact)?;
1098 let schema = table_guard.schema().clone();
1099 drop(table_guard);
1100
1101 next_catalog.tables[catalog_entry_index].schema = schema;
1102 let catalog_result =
1103 catalog::write_atomic(&self.root, &next_catalog, self.meta_dek.as_ref());
1104 *self.catalog.write() = next_catalog;
1105 self.publish_committed(&receipt, epoch)?;
1106 epoch_guard.disarm();
1107 if let Err(error) = catalog_result {
1108 self.poisoned.store(true, Ordering::Relaxed);
1109 self.lifecycle.poison();
1110 return Err(MongrelError::DurableCommit {
1111 epoch: epoch.0,
1112 message: error.to_string(),
1113 });
1114 }
1115 crate::catalog::inject_hook("index.publish.after")?;
1116 Ok(())
1117 })();
1118 return result.map_err(|error| match (durable_epoch.get(), error) {
1119 (_, error @ MongrelError::DurableCommit { .. }) => error,
1120 (Some(epoch), error) => MongrelError::DurableCommit {
1121 epoch: epoch.0,
1122 message: error.to_string(),
1123 },
1124 (None, error) => error,
1125 });
1126 }
1127 }
1128}
1129
1130#[cfg(test)]
1131mod tests {
1132 use super::*;
1133 use crate::query::{Retriever, RetrieverScore};
1134 use crate::schema::{AnnOptions, AnnQuantization, IndexKind, IndexOptions, TypeId};
1135 use std::sync::Mutex;
1136
1137 static INDEX_DDL_TEST_LOCK: Mutex<()> = Mutex::new(());
1140
1141 fn lock_tests() -> std::sync::MutexGuard<'static, ()> {
1142 INDEX_DDL_TEST_LOCK
1143 .lock()
1144 .unwrap_or_else(|poisoned| poisoned.into_inner())
1145 }
1146
1147 fn embedding_schema(quantization: AnnQuantization) -> Schema {
1148 Schema {
1149 columns: vec![
1150 ColumnDef {
1151 id: 1,
1152 name: "id".into(),
1153 ty: TypeId::Int64,
1154 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
1155 default_value: None,
1156 embedding_source: None,
1157 },
1158 ColumnDef {
1159 id: 2,
1160 name: "embedding".into(),
1161 ty: TypeId::Embedding { dim: 4 },
1162 flags: ColumnFlags::empty(),
1163 default_value: None,
1164 embedding_source: None,
1165 },
1166 ],
1167 indexes: vec![IndexDef {
1168 name: "idx_embed_ann".into(),
1169 column_id: 2,
1170 kind: IndexKind::Ann,
1171 predicate: None,
1172 options: IndexOptions {
1173 ann: Some(AnnOptions {
1174 m: 8,
1175 ef_construction: 32,
1176 ef_search: 16,
1177 quantization,
1178 ..AnnOptions::default()
1179 }),
1180 ..IndexOptions::default()
1181 },
1182 }],
1183 ..Schema::default()
1184 }
1185 }
1186
1187 fn plain_embedding_schema() -> Schema {
1188 Schema {
1189 columns: vec![
1190 ColumnDef {
1191 id: 1,
1192 name: "id".into(),
1193 ty: TypeId::Int64,
1194 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
1195 default_value: None,
1196 embedding_source: None,
1197 },
1198 ColumnDef {
1199 id: 2,
1200 name: "embedding".into(),
1201 ty: TypeId::Embedding { dim: 4 },
1202 flags: ColumnFlags::empty(),
1203 default_value: None,
1204 embedding_source: None,
1205 },
1206 ],
1207 ..Schema::default()
1208 }
1209 }
1210
1211 fn put_row(db: &Database, table: &str, id: i64, embedding: Vec<f32>) {
1212 let mut txn = db.begin();
1213 txn.put(
1214 table,
1215 vec![(1, Value::Int64(id)), (2, Value::Embedding(embedding))],
1216 )
1217 .unwrap();
1218 txn.commit().unwrap();
1219 }
1220
1221 fn row_id_for_pk(db: &Database, table: &str, id: i64) -> RowId {
1222 let handle = db.table(table).unwrap();
1223 let table = handle.lock();
1224 table
1225 .visible_rows(db.snapshot_for_epoch(table.current_epoch()))
1226 .unwrap()
1227 .into_iter()
1228 .find(|row| row.columns.get(&1) == Some(&Value::Int64(id)))
1229 .map(|row| row.row_id)
1230 .unwrap()
1231 }
1232
1233 fn dense_definition() -> IndexDef {
1234 IndexDef {
1235 name: "idx_embed_ann".into(),
1236 column_id: 2,
1237 kind: IndexKind::Ann,
1238 predicate: None,
1239 options: IndexOptions {
1240 ann: Some(AnnOptions {
1241 quantization: AnnQuantization::Dense,
1242 ..AnnOptions::default()
1243 }),
1244 ..IndexOptions::default()
1245 },
1246 }
1247 }
1248
1249 #[test]
1250 fn create_index_indexes_snapshot_rows() {
1251 let _lock = lock_tests();
1252 let dir = tempfile::tempdir().unwrap();
1253 let db = Database::create(dir.path()).unwrap();
1254 db.create_table("docs", plain_embedding_schema()).unwrap();
1255 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1256 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1257
1258 let definition = IndexDef {
1259 name: "idx_embed_ann".into(),
1260 column_id: 2,
1261 kind: IndexKind::Ann,
1262 predicate: None,
1263 options: IndexOptions {
1264 ann: Some(AnnOptions {
1265 m: 8,
1266 ef_construction: 32,
1267 ef_search: 16,
1268 quantization: AnnQuantization::Dense,
1269 ..AnnOptions::default()
1270 }),
1271 ..IndexOptions::default()
1272 },
1273 };
1274 let job_id = db.create_index("docs", definition).unwrap();
1275 let record = db.job_registry().get(job_id).unwrap();
1276 assert_eq!(record.state, crate::jobs::JobState::Succeeded);
1277
1278 let handle = db.table("docs").unwrap();
1279 let mut table = handle.lock();
1280 let hits = table
1281 .retrieve(&Retriever::Ann {
1282 column_id: 2,
1283 query: vec![1.0, 0.0, 0.0, 0.0],
1284 k: 2,
1285 })
1286 .unwrap();
1287 assert_eq!(hits.len(), 2);
1288 assert!(matches!(
1289 hits[0].score,
1290 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-5
1291 ));
1292 assert_eq!(table.schema().schema_id, table.table_id().saturating_add(1));
1295 }
1296
1297 #[test]
1298 fn replace_binary_sign_with_dense_search_works() {
1299 let _lock = lock_tests();
1300 let dir = tempfile::tempdir().unwrap();
1301 let db = Database::create(dir.path()).unwrap();
1302 db.create_table("docs", embedding_schema(AnnQuantization::BinarySign))
1303 .unwrap();
1304 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1305 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1306
1307 let new_def = IndexDef {
1308 name: "idx_embed_ann".into(),
1309 column_id: 2,
1310 kind: IndexKind::Ann,
1311 predicate: None,
1312 options: IndexOptions {
1313 ann: Some(AnnOptions {
1314 m: 8,
1315 ef_construction: 32,
1316 ef_search: 16,
1317 quantization: AnnQuantization::Dense,
1318 ..AnnOptions::default()
1319 }),
1320 ..IndexOptions::default()
1321 },
1322 };
1323 db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1324
1325 let handle = db.table("docs").unwrap();
1326 let mut table = handle.lock();
1327 let hits = table
1328 .retrieve(&Retriever::Ann {
1329 column_id: 2,
1330 query: vec![1.0, 0.0, 0.0, 0.0],
1331 k: 1,
1332 })
1333 .unwrap();
1334 assert_eq!(hits.len(), 1);
1335 assert!(matches!(
1336 hits[0].score,
1337 RetrieverScore::AnnCosineDistance(_)
1338 ));
1339 assert!(!hits
1340 .iter()
1341 .any(|hit| matches!(hit.score, RetrieverScore::AnnHammingDistance(_))));
1342 }
1343
1344 #[test]
1345 fn drop_index_removes_schema_entry_without_rewrite() {
1346 let _lock = lock_tests();
1347 let dir = tempfile::tempdir().unwrap();
1348 let db = Database::create(dir.path()).unwrap();
1349 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1350 .unwrap();
1351 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1352 db.drop_index("docs", "idx_embed_ann").unwrap();
1353 let handle = db.table("docs").unwrap();
1354 let table = handle.lock();
1355 assert!(table.schema().indexes.is_empty());
1356 let epoch = table.current_epoch();
1358 drop(table);
1359 let rows = handle
1360 .lock()
1361 .visible_rows(db.snapshot_for_epoch(epoch))
1362 .unwrap();
1363 assert_eq!(rows.len(), 1);
1364 }
1365
1366 #[test]
1367 fn concurrent_writes_progress_during_create() {
1368 use std::sync::Barrier;
1369 let _lock = lock_tests();
1370 let dir = tempfile::tempdir().unwrap();
1371 let db = Arc::new(Database::create(dir.path()).unwrap());
1372 db.create_table("docs", plain_embedding_schema()).unwrap();
1373 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1374
1375 let start = Arc::new(Barrier::new(2));
1376 let writer = {
1377 let db = Arc::clone(&db);
1378 let start = Arc::clone(&start);
1379 std::thread::spawn(move || {
1380 start.wait();
1381 for id in 100..120 {
1382 put_row(&db, "docs", id, vec![0.0, 1.0, 0.0, 0.0]);
1383 }
1384 })
1385 };
1386 start.wait();
1387 let definition = IndexDef {
1388 name: "idx_embed_ann".into(),
1389 column_id: 2,
1390 kind: IndexKind::Ann,
1391 predicate: None,
1392 options: IndexOptions {
1393 ann: Some(AnnOptions {
1394 quantization: AnnQuantization::Dense,
1395 ..AnnOptions::default()
1396 }),
1397 ..IndexOptions::default()
1398 },
1399 };
1400 db.create_index("docs", definition).unwrap();
1401 writer.join().unwrap();
1402 let handle = db.table("docs").unwrap();
1404 let table = handle.lock();
1405 assert!(table
1406 .schema()
1407 .indexes
1408 .iter()
1409 .any(|index| index.name == "idx_embed_ann"));
1410 }
1411
1412 #[test]
1413 fn asynchronous_build_is_observable_and_replays_all_row_deltas() {
1414 use std::sync::Barrier;
1415 use std::time::Duration;
1416
1417 let _lock = lock_tests();
1418 let dir = tempfile::tempdir().unwrap();
1419 let db = Arc::new(Database::create(dir.path()).unwrap());
1420 db.create_table("docs", plain_embedding_schema()).unwrap();
1421 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1422 put_row(&db, "docs", 3, vec![0.0, 0.0, 1.0, 0.0]);
1423 let deleted_row = row_id_for_pk(&db, "docs", 3);
1424 let baseline = db
1425 .memory_governor()
1426 .usage(crate::memory::MemoryClass::Compaction);
1427
1428 let entered = Arc::new(Barrier::new(2));
1429 let release = Arc::new(Barrier::new(2));
1430 let callback_entered = Arc::clone(&entered);
1431 let callback_release = Arc::clone(&release);
1432 let _guard = mongreldb_fault::ScopedGuard::new(
1433 "job.build_hidden.after",
1434 mongreldb_fault::Action::Callback(Arc::new(move |_| {
1435 callback_entered.wait();
1436 callback_release.wait();
1437 })),
1438 );
1439
1440 let job_id = db.start_create_index("docs", dense_definition()).unwrap();
1441 entered.wait();
1442
1443 let running = db.job_registry().get(job_id).unwrap();
1444 assert_eq!(running.state, crate::jobs::JobState::Running);
1445 let persisted = IndexBuildSpec::decode(running.definition.as_deref().unwrap()).unwrap();
1446 assert_eq!(persisted.table, "docs");
1447 assert_eq!(persisted.kind.definition(), &dense_definition());
1448 assert_eq!(
1449 db.index_build_target_definition(job_id).unwrap(),
1450 dense_definition()
1451 );
1452 assert!(
1453 db.memory_governor()
1454 .usage(crate::memory::MemoryClass::Compaction)
1455 > baseline,
1456 "hidden generation must hold a governor reservation"
1457 );
1458
1459 put_row(&db, "docs", 1, vec![0.0, 1.0, 0.0, 0.0]);
1462 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1463 db.transaction(|txn| txn.delete("docs", deleted_row))
1464 .unwrap();
1465
1466 release.wait();
1467 let terminal = db
1468 .job_registry()
1469 .wait_terminal(job_id, Duration::from_secs(30))
1470 .unwrap();
1471 assert_eq!(terminal.state, crate::jobs::JobState::Succeeded);
1472 assert_eq!(
1473 db.memory_governor()
1474 .usage(crate::memory::MemoryClass::Compaction),
1475 baseline,
1476 "successful build must release its reservation"
1477 );
1478
1479 let expected: std::collections::HashSet<_> = [1, 2]
1480 .into_iter()
1481 .map(|id| row_id_for_pk(&db, "docs", id))
1482 .collect();
1483 let handle = db.table("docs").unwrap();
1484 let mut table = handle.lock();
1485 let actual: std::collections::HashSet<_> = table
1486 .retrieve(&Retriever::Ann {
1487 column_id: 2,
1488 query: vec![0.0, 1.0, 0.0, 0.0],
1489 k: 10,
1490 })
1491 .unwrap()
1492 .into_iter()
1493 .map(|hit| hit.row_id)
1494 .collect();
1495 assert_eq!(actual, expected);
1496 }
1497
1498 #[test]
1499 fn cancelling_real_hidden_build_rolls_back_pin_memory_and_schema() {
1500 use std::sync::Barrier;
1501 use std::time::Duration;
1502
1503 let _lock = lock_tests();
1504 let dir = tempfile::tempdir().unwrap();
1505 let db = Arc::new(Database::create(dir.path()).unwrap());
1506 db.create_table("docs", plain_embedding_schema()).unwrap();
1507 for id in 0..32 {
1508 put_row(&db, "docs", id, vec![1.0, id as f32, 0.0, 0.0]);
1509 }
1510 let baseline = db
1511 .memory_governor()
1512 .usage(crate::memory::MemoryClass::Compaction);
1513
1514 let entered = Arc::new(Barrier::new(2));
1515 let release = Arc::new(Barrier::new(2));
1516 let callback_entered = Arc::clone(&entered);
1517 let callback_release = Arc::clone(&release);
1518 let _guard = mongreldb_fault::ScopedGuard::new(
1519 "job.build_hidden.after",
1520 mongreldb_fault::Action::Callback(Arc::new(move |_| {
1521 callback_entered.wait();
1522 callback_release.wait();
1523 })),
1524 );
1525
1526 let job_id = db.start_create_index("docs", dense_definition()).unwrap();
1527 entered.wait();
1528 db.job_registry().cancel(job_id).unwrap();
1529 release.wait();
1530 let terminal = db
1531 .job_registry()
1532 .wait_terminal(job_id, Duration::from_secs(30))
1533 .unwrap();
1534 assert_eq!(terminal.state, crate::jobs::JobState::Failed);
1535 assert!(terminal
1536 .error
1537 .as_deref()
1538 .is_some_and(|error| error.contains("cancel")));
1539 assert_eq!(
1540 db.memory_governor()
1541 .usage(crate::memory::MemoryClass::Compaction),
1542 baseline
1543 );
1544 assert!(
1545 db.table("docs").unwrap().lock().schema().indexes.is_empty(),
1546 "pre-publication cancellation must keep the old schema"
1547 );
1548 }
1549
1550 #[test]
1551 fn pending_index_job_reconstructs_and_resumes_after_reopen() {
1552 use std::time::Duration;
1553
1554 let _lock = lock_tests();
1555 let dir = tempfile::tempdir().unwrap();
1556 let path = dir.path().to_path_buf();
1557 let job_id = {
1558 let db = Database::create(&path).unwrap();
1559 db.create_table("docs", plain_embedding_schema()).unwrap();
1560 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1561 let job_id = db.submit_create_index("docs", dense_definition()).unwrap();
1562 let pending = db.job_registry().get(job_id).unwrap();
1563 assert_eq!(pending.state, crate::jobs::JobState::Pending);
1564 assert!(pending.definition.is_some());
1565 job_id
1566 };
1567
1568 let db = Arc::new(Database::open(&path).unwrap());
1569 db.resume_index_build(job_id).unwrap();
1570 let terminal = db
1571 .job_registry()
1572 .wait_terminal(job_id, Duration::from_secs(30))
1573 .unwrap();
1574 assert_eq!(terminal.state, crate::jobs::JobState::Succeeded);
1575 assert_eq!(
1576 db.table("docs")
1577 .unwrap()
1578 .lock()
1579 .schema()
1580 .indexes
1581 .iter()
1582 .find(|index| index.name == "idx_embed_ann")
1583 .and_then(|index| index.options.ann.as_ref())
1584 .map(|options| options.quantization),
1585 Some(AnnQuantization::Dense)
1586 );
1587 }
1588
1589 #[test]
1590 fn dense_build_is_rejected_by_memory_admission_before_graph_allocation() {
1591 let _lock = lock_tests();
1592 let dir = tempfile::tempdir().unwrap();
1593 let path = dir.path().to_path_buf();
1594 let dim = 4_096u32;
1595 {
1596 let db = Database::create(&path).unwrap();
1597 let mut schema = plain_embedding_schema();
1598 schema.columns[1].ty = TypeId::Embedding { dim };
1599 db.create_table("docs", schema).unwrap();
1600 for id in 0..64 {
1601 let mut vector = vec![0.0; dim as usize];
1602 vector[id as usize] = 1.0;
1603 put_row(&db, "docs", id, vector);
1604 }
1605 }
1606
1607 let db = Database::open_with_options(
1608 &path,
1609 crate::database::OpenOptions::default().with_memory_budget_bytes(2 * 1024 * 1024),
1610 )
1611 .unwrap();
1612 let baseline = db
1613 .memory_governor()
1614 .usage(crate::memory::MemoryClass::Compaction);
1615 let error = db
1616 .create_index("docs", dense_definition())
1617 .expect_err("oversized Dense graph must be denied");
1618 assert!(
1619 matches!(
1620 error,
1621 MongrelError::ResourceLimitExceeded {
1622 resource: "index build memory",
1623 ..
1624 }
1625 ),
1626 "unexpected admission error: {error:?}"
1627 );
1628 assert_eq!(
1629 db.memory_governor()
1630 .usage(crate::memory::MemoryClass::Compaction),
1631 baseline,
1632 "rejected build must not leak a reservation"
1633 );
1634 assert!(db.table("docs").unwrap().lock().schema().indexes.is_empty());
1635 }
1636
1637 #[test]
1638 fn ddl_auth_fail_closed() {
1639 let _lock = lock_tests();
1640 let dir = tempfile::tempdir().unwrap();
1641 let path = dir.path().to_path_buf();
1642 {
1643 let db = Database::create_with_credentials(&path, "admin", "admin-pw").unwrap();
1644 db.create_table("docs", plain_embedding_schema()).unwrap();
1645 db.create_user("alice", "alice-pw").unwrap();
1646 db.create_role("reader").unwrap();
1647 db.grant_permission(
1648 "reader",
1649 crate::auth::Permission::Select {
1650 table: "docs".into(),
1651 },
1652 )
1653 .unwrap();
1654 db.grant_role("alice", "reader").unwrap();
1655 }
1656 let db = Database::open_with_credentials(&path, "alice", "alice-pw").unwrap();
1657 let definition = IndexDef {
1658 name: "idx".into(),
1659 column_id: 2,
1660 kind: IndexKind::Ann,
1661 predicate: None,
1662 options: IndexOptions {
1663 ann: Some(AnnOptions {
1664 quantization: AnnQuantization::Dense,
1665 ..AnnOptions::default()
1666 }),
1667 ..IndexOptions::default()
1668 },
1669 };
1670 let err = db.create_index("docs", definition).unwrap_err();
1671 assert!(
1672 matches!(err, MongrelError::PermissionDenied { .. }),
1673 "unexpected error: {err:?}"
1674 );
1675 }
1676
1677 #[test]
1678 fn cluster_replica_rejects_local_index_ddl() {
1679 let _lock = lock_tests();
1680 use mongreldb_types::ids::{ClusterId, DatabaseId, NodeId};
1681 let dir = tempfile::tempdir().unwrap();
1682 let cluster_id = ClusterId::from_bytes([1; 16]);
1683 let node_id = NodeId::from_bytes([2; 16]);
1684 let database_id = DatabaseId::from_bytes([3; 16]);
1685 let db =
1686 Database::create_cluster_replica(dir.path(), cluster_id, node_id, database_id).unwrap();
1687 assert!(db.is_read_only_replica());
1688 let definition = IndexDef {
1689 name: "idx".into(),
1690 column_id: 2,
1691 kind: IndexKind::Ann,
1692 predicate: None,
1693 options: Default::default(),
1694 };
1695 let err = db.create_index("docs", definition).unwrap_err();
1696 assert!(
1697 matches!(err, MongrelError::ReadOnlyReplica),
1698 "cluster replica must reject local index DDL: {err:?}"
1699 );
1700 }
1701
1702 #[test]
1703 fn fault_before_publish_leaves_old_schema() {
1704 let _lock = lock_tests();
1705 let _guard = mongreldb_fault::ScopedGuard::new(
1708 "index.publish.before",
1709 mongreldb_fault::Action::Fail,
1710 );
1711 let dir = tempfile::tempdir().unwrap();
1712 let db = Database::create(dir.path()).unwrap();
1713 db.create_table("docs", plain_embedding_schema()).unwrap();
1714 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1715
1716 let definition = IndexDef {
1717 name: "idx_embed_ann".into(),
1718 column_id: 2,
1719 kind: IndexKind::Ann,
1720 predicate: None,
1721 options: IndexOptions {
1722 ann: Some(AnnOptions {
1723 quantization: AnnQuantization::Dense,
1724 ..AnnOptions::default()
1725 }),
1726 ..IndexOptions::default()
1727 },
1728 };
1729 let err = db.create_index("docs", definition).unwrap_err();
1730 assert!(
1731 !matches!(err, MongrelError::NotFound(_)),
1732 "should fail at publish, not earlier: {err:?}"
1733 );
1734 let handle = db.table("docs").unwrap();
1735 let table = handle.lock();
1736 assert!(
1737 table.schema().indexes.is_empty(),
1738 "pre-publication fault must leave the old schema active"
1739 );
1740 }
1741
1742 #[test]
1743 fn cancel_pre_publish_leaves_old_index() {
1744 let _lock = lock_tests();
1745 let dir = tempfile::tempdir().unwrap();
1747 let db = Database::create(dir.path()).unwrap();
1748 db.create_table("docs", embedding_schema(AnnQuantization::BinarySign))
1749 .unwrap();
1750 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1751
1752 let job_id = db
1754 .job_registry()
1755 .submit(
1756 JobKind::IndexBuild,
1757 JobTarget {
1758 table: "docs".into(),
1759 index: Some("idx_embed_ann".into()),
1760 },
1761 )
1762 .unwrap();
1763 db.job_registry().cancel(job_id).unwrap();
1764 let record = db.job_registry().get(job_id).unwrap();
1765 assert_eq!(record.state, crate::jobs::JobState::Failed);
1766
1767 let handle = db.table("docs").unwrap();
1769 let table = handle.lock();
1770 let index = table
1771 .schema()
1772 .indexes
1773 .iter()
1774 .find(|index| index.name == "idx_embed_ann")
1775 .unwrap();
1776 assert_eq!(
1777 index
1778 .options
1779 .ann
1780 .as_ref()
1781 .map(|options| options.quantization),
1782 Some(AnnQuantization::BinarySign)
1783 );
1784 }
1785
1786 #[test]
1787 fn replace_dense_with_binary_sign_search_works() {
1788 let _lock = lock_tests();
1789 let dir = tempfile::tempdir().unwrap();
1790 let db = Database::create(dir.path()).unwrap();
1791 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1792 .unwrap();
1793 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1794 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1795
1796 let new_def = IndexDef {
1797 name: "idx_embed_ann".into(),
1798 column_id: 2,
1799 kind: IndexKind::Ann,
1800 predicate: None,
1801 options: IndexOptions {
1802 ann: Some(AnnOptions {
1803 m: 8,
1804 ef_construction: 32,
1805 ef_search: 16,
1806 quantization: AnnQuantization::BinarySign,
1807 ..AnnOptions::default()
1808 }),
1809 ..IndexOptions::default()
1810 },
1811 };
1812 db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1813
1814 let handle = db.table("docs").unwrap();
1815 let mut table = handle.lock();
1816 let hits = table
1817 .retrieve(&Retriever::Ann {
1818 column_id: 2,
1819 query: vec![1.0, 0.0, 0.0, 0.0],
1820 k: 1,
1821 })
1822 .unwrap();
1823 assert_eq!(hits.len(), 1);
1824 assert!(matches!(
1825 hits[0].score,
1826 RetrieverScore::AnnHammingDistance(_)
1827 ));
1828 assert!(!hits
1829 .iter()
1830 .any(|hit| matches!(hit.score, RetrieverScore::AnnCosineDistance(_))));
1831 }
1832
1833 #[test]
1834 fn replace_hnsw_with_diskann_changes_algorithm_online() {
1835 let _lock = lock_tests();
1839 let dir = tempfile::tempdir().unwrap();
1840 let db = Database::create(dir.path()).unwrap();
1841 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1842 .unwrap();
1843 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1844 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1845
1846 let new_def = IndexDef {
1847 name: "idx_embed_ann".into(),
1848 column_id: 2,
1849 kind: IndexKind::Ann,
1850 predicate: None,
1851 options: IndexOptions {
1852 ann: Some(AnnOptions {
1853 m: 8,
1854 ef_construction: 32,
1855 ef_search: 16,
1856 quantization: AnnQuantization::Dense,
1857 algorithm: crate::schema::AnnAlgorithm::DiskAnn,
1858 diskann: Some(crate::schema::DiskAnnOptions {
1859 r: 8,
1860 l: 16,
1861 beam_width: 4,
1862 alpha: 120,
1863 }),
1864 ..AnnOptions::default()
1865 }),
1866 ..IndexOptions::default()
1867 },
1868 };
1869 let job_id = db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1870 let record = db.job_registry().get(job_id).unwrap();
1871 assert_eq!(record.state, crate::jobs::JobState::Succeeded);
1872
1873 let handle = db.table("docs").unwrap();
1874 let mut table = handle.lock();
1875 let hits = table
1876 .retrieve(&Retriever::Ann {
1877 column_id: 2,
1878 query: vec![1.0, 0.0, 0.0, 0.0],
1879 k: 2,
1880 })
1881 .unwrap();
1882 assert_eq!(hits.len(), 2);
1883 assert!(matches!(
1885 hits[0].score,
1886 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-5
1887 ));
1888 }
1889
1890 #[test]
1891 fn product_quantization_dimension_divisibility_rejected_at_ddl() {
1892 let _lock = lock_tests();
1895 let dir = tempfile::tempdir().unwrap();
1896 let db = Database::create(dir.path()).unwrap();
1897 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1898 .unwrap();
1899 let bad_def = IndexDef {
1901 name: "idx_bad_pq".into(),
1902 column_id: 2,
1903 kind: IndexKind::Ann,
1904 predicate: None,
1905 options: IndexOptions {
1906 ann: Some(AnnOptions {
1907 quantization: AnnQuantization::Product {
1908 num_subvectors: 3,
1909 bits: 8,
1910 },
1911 product: Some(crate::schema::ProductQuantizerOptions::default()),
1912 ..AnnOptions::default()
1913 }),
1914 ..IndexOptions::default()
1915 },
1916 };
1917 let err = db.create_index("docs", bad_def).unwrap_err();
1918 assert!(
1919 err.to_string()
1920 .contains("must evenly divide the column dimension"),
1921 "expected divisibility error, got: {err}"
1922 );
1923 }
1924
1925 #[test]
1926 fn fault_after_publish_keeps_new_index() {
1927 let _lock = lock_tests();
1928 let _guard =
1931 mongreldb_fault::ScopedGuard::new("index.publish.after", mongreldb_fault::Action::Fail);
1932 let dir = tempfile::tempdir().unwrap();
1933 let db = Database::create(dir.path()).unwrap();
1934 db.create_table("docs", plain_embedding_schema()).unwrap();
1935 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1936
1937 let definition = IndexDef {
1938 name: "idx_embed_ann".into(),
1939 column_id: 2,
1940 kind: IndexKind::Ann,
1941 predicate: None,
1942 options: IndexOptions {
1943 ann: Some(AnnOptions {
1944 quantization: AnnQuantization::Dense,
1945 ..AnnOptions::default()
1946 }),
1947 ..IndexOptions::default()
1948 },
1949 };
1950 let _ = db.create_index("docs", definition);
1953 let handle = db.table("docs").unwrap();
1954 let table = handle.lock();
1955 let index = table
1956 .schema()
1957 .indexes
1958 .iter()
1959 .find(|index| index.name == "idx_embed_ann");
1960 assert!(
1961 index.is_some(),
1962 "post-publish fault must keep the new index active"
1963 );
1964 assert_eq!(
1965 index
1966 .unwrap()
1967 .options
1968 .ann
1969 .as_ref()
1970 .map(|options| options.quantization),
1971 Some(AnnQuantization::Dense)
1972 );
1973 }
1974
1975 #[test]
1976 fn dense_plaintext_close_reopen_preserves_search() {
1977 let _lock = lock_tests();
1978 let dir = tempfile::tempdir().unwrap();
1979 let path = dir.path().to_path_buf();
1980 {
1981 let db = Database::create(&path).unwrap();
1982 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1983 .unwrap();
1984 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1985 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1986 let handle = db.table("docs").unwrap();
1988 let mut table = handle.lock();
1989 table.flush().unwrap();
1990 let _ = table.publish_read_generation().unwrap();
1991 }
1992 let db = Database::open(&path).unwrap();
1993 let handle = db.table("docs").unwrap();
1994 let mut table = handle.lock();
1995 let hits = table
1996 .retrieve(&Retriever::Ann {
1997 column_id: 2,
1998 query: vec![1.0, 0.0, 0.0, 0.0],
1999 k: 1,
2000 })
2001 .unwrap();
2002 assert_eq!(hits.len(), 1);
2003 assert!(matches!(
2004 hits[0].score,
2005 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-4
2006 ));
2007 }
2008}