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(Snapshot::at(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 rows = table
1358 .visible_rows(Snapshot::at(table.current_epoch()))
1359 .unwrap();
1360 assert_eq!(rows.len(), 1);
1361 }
1362
1363 #[test]
1364 fn concurrent_writes_progress_during_create() {
1365 use std::sync::Barrier;
1366 let _lock = lock_tests();
1367 let dir = tempfile::tempdir().unwrap();
1368 let db = Arc::new(Database::create(dir.path()).unwrap());
1369 db.create_table("docs", plain_embedding_schema()).unwrap();
1370 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1371
1372 let start = Arc::new(Barrier::new(2));
1373 let writer = {
1374 let db = Arc::clone(&db);
1375 let start = Arc::clone(&start);
1376 std::thread::spawn(move || {
1377 start.wait();
1378 for id in 100..120 {
1379 put_row(&db, "docs", id, vec![0.0, 1.0, 0.0, 0.0]);
1380 }
1381 })
1382 };
1383 start.wait();
1384 let definition = IndexDef {
1385 name: "idx_embed_ann".into(),
1386 column_id: 2,
1387 kind: IndexKind::Ann,
1388 predicate: None,
1389 options: IndexOptions {
1390 ann: Some(AnnOptions {
1391 quantization: AnnQuantization::Dense,
1392 ..AnnOptions::default()
1393 }),
1394 ..IndexOptions::default()
1395 },
1396 };
1397 db.create_index("docs", definition).unwrap();
1398 writer.join().unwrap();
1399 let handle = db.table("docs").unwrap();
1401 let table = handle.lock();
1402 assert!(table
1403 .schema()
1404 .indexes
1405 .iter()
1406 .any(|index| index.name == "idx_embed_ann"));
1407 }
1408
1409 #[test]
1410 fn asynchronous_build_is_observable_and_replays_all_row_deltas() {
1411 use std::sync::Barrier;
1412 use std::time::Duration;
1413
1414 let _lock = lock_tests();
1415 let dir = tempfile::tempdir().unwrap();
1416 let db = Arc::new(Database::create(dir.path()).unwrap());
1417 db.create_table("docs", plain_embedding_schema()).unwrap();
1418 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1419 put_row(&db, "docs", 3, vec![0.0, 0.0, 1.0, 0.0]);
1420 let deleted_row = row_id_for_pk(&db, "docs", 3);
1421 let baseline = db
1422 .memory_governor()
1423 .usage(crate::memory::MemoryClass::Compaction);
1424
1425 let entered = Arc::new(Barrier::new(2));
1426 let release = Arc::new(Barrier::new(2));
1427 let callback_entered = Arc::clone(&entered);
1428 let callback_release = Arc::clone(&release);
1429 let _guard = mongreldb_fault::ScopedGuard::new(
1430 "job.build_hidden.after",
1431 mongreldb_fault::Action::Callback(Arc::new(move |_| {
1432 callback_entered.wait();
1433 callback_release.wait();
1434 })),
1435 );
1436
1437 let job_id = db.start_create_index("docs", dense_definition()).unwrap();
1438 entered.wait();
1439
1440 let running = db.job_registry().get(job_id).unwrap();
1441 assert_eq!(running.state, crate::jobs::JobState::Running);
1442 let persisted = IndexBuildSpec::decode(running.definition.as_deref().unwrap()).unwrap();
1443 assert_eq!(persisted.table, "docs");
1444 assert_eq!(persisted.kind.definition(), &dense_definition());
1445 assert_eq!(
1446 db.index_build_target_definition(job_id).unwrap(),
1447 dense_definition()
1448 );
1449 assert!(
1450 db.memory_governor()
1451 .usage(crate::memory::MemoryClass::Compaction)
1452 > baseline,
1453 "hidden generation must hold a governor reservation"
1454 );
1455
1456 put_row(&db, "docs", 1, vec![0.0, 1.0, 0.0, 0.0]);
1459 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1460 db.transaction(|txn| txn.delete("docs", deleted_row))
1461 .unwrap();
1462
1463 release.wait();
1464 let terminal = db
1465 .job_registry()
1466 .wait_terminal(job_id, Duration::from_secs(30))
1467 .unwrap();
1468 assert_eq!(terminal.state, crate::jobs::JobState::Succeeded);
1469 assert_eq!(
1470 db.memory_governor()
1471 .usage(crate::memory::MemoryClass::Compaction),
1472 baseline,
1473 "successful build must release its reservation"
1474 );
1475
1476 let expected: std::collections::HashSet<_> = [1, 2]
1477 .into_iter()
1478 .map(|id| row_id_for_pk(&db, "docs", id))
1479 .collect();
1480 let handle = db.table("docs").unwrap();
1481 let mut table = handle.lock();
1482 let actual: std::collections::HashSet<_> = table
1483 .retrieve(&Retriever::Ann {
1484 column_id: 2,
1485 query: vec![0.0, 1.0, 0.0, 0.0],
1486 k: 10,
1487 })
1488 .unwrap()
1489 .into_iter()
1490 .map(|hit| hit.row_id)
1491 .collect();
1492 assert_eq!(actual, expected);
1493 }
1494
1495 #[test]
1496 fn cancelling_real_hidden_build_rolls_back_pin_memory_and_schema() {
1497 use std::sync::Barrier;
1498 use std::time::Duration;
1499
1500 let _lock = lock_tests();
1501 let dir = tempfile::tempdir().unwrap();
1502 let db = Arc::new(Database::create(dir.path()).unwrap());
1503 db.create_table("docs", plain_embedding_schema()).unwrap();
1504 for id in 0..32 {
1505 put_row(&db, "docs", id, vec![1.0, id as f32, 0.0, 0.0]);
1506 }
1507 let baseline = db
1508 .memory_governor()
1509 .usage(crate::memory::MemoryClass::Compaction);
1510
1511 let entered = Arc::new(Barrier::new(2));
1512 let release = Arc::new(Barrier::new(2));
1513 let callback_entered = Arc::clone(&entered);
1514 let callback_release = Arc::clone(&release);
1515 let _guard = mongreldb_fault::ScopedGuard::new(
1516 "job.build_hidden.after",
1517 mongreldb_fault::Action::Callback(Arc::new(move |_| {
1518 callback_entered.wait();
1519 callback_release.wait();
1520 })),
1521 );
1522
1523 let job_id = db.start_create_index("docs", dense_definition()).unwrap();
1524 entered.wait();
1525 db.job_registry().cancel(job_id).unwrap();
1526 release.wait();
1527 let terminal = db
1528 .job_registry()
1529 .wait_terminal(job_id, Duration::from_secs(30))
1530 .unwrap();
1531 assert_eq!(terminal.state, crate::jobs::JobState::Failed);
1532 assert!(terminal
1533 .error
1534 .as_deref()
1535 .is_some_and(|error| error.contains("cancel")));
1536 assert_eq!(
1537 db.memory_governor()
1538 .usage(crate::memory::MemoryClass::Compaction),
1539 baseline
1540 );
1541 assert!(
1542 db.table("docs").unwrap().lock().schema().indexes.is_empty(),
1543 "pre-publication cancellation must keep the old schema"
1544 );
1545 }
1546
1547 #[test]
1548 fn pending_index_job_reconstructs_and_resumes_after_reopen() {
1549 use std::time::Duration;
1550
1551 let _lock = lock_tests();
1552 let dir = tempfile::tempdir().unwrap();
1553 let path = dir.path().to_path_buf();
1554 let job_id = {
1555 let db = Database::create(&path).unwrap();
1556 db.create_table("docs", plain_embedding_schema()).unwrap();
1557 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1558 let job_id = db.submit_create_index("docs", dense_definition()).unwrap();
1559 let pending = db.job_registry().get(job_id).unwrap();
1560 assert_eq!(pending.state, crate::jobs::JobState::Pending);
1561 assert!(pending.definition.is_some());
1562 job_id
1563 };
1564
1565 let db = Arc::new(Database::open(&path).unwrap());
1566 db.resume_index_build(job_id).unwrap();
1567 let terminal = db
1568 .job_registry()
1569 .wait_terminal(job_id, Duration::from_secs(30))
1570 .unwrap();
1571 assert_eq!(terminal.state, crate::jobs::JobState::Succeeded);
1572 assert_eq!(
1573 db.table("docs")
1574 .unwrap()
1575 .lock()
1576 .schema()
1577 .indexes
1578 .iter()
1579 .find(|index| index.name == "idx_embed_ann")
1580 .and_then(|index| index.options.ann.as_ref())
1581 .map(|options| options.quantization),
1582 Some(AnnQuantization::Dense)
1583 );
1584 }
1585
1586 #[test]
1587 fn dense_build_is_rejected_by_memory_admission_before_graph_allocation() {
1588 let _lock = lock_tests();
1589 let dir = tempfile::tempdir().unwrap();
1590 let path = dir.path().to_path_buf();
1591 let dim = 4_096u32;
1592 {
1593 let db = Database::create(&path).unwrap();
1594 let mut schema = plain_embedding_schema();
1595 schema.columns[1].ty = TypeId::Embedding { dim };
1596 db.create_table("docs", schema).unwrap();
1597 for id in 0..64 {
1598 let mut vector = vec![0.0; dim as usize];
1599 vector[id as usize] = 1.0;
1600 put_row(&db, "docs", id, vector);
1601 }
1602 }
1603
1604 let db = Database::open_with_options(
1605 &path,
1606 crate::database::OpenOptions::default().with_memory_budget_bytes(2 * 1024 * 1024),
1607 )
1608 .unwrap();
1609 let baseline = db
1610 .memory_governor()
1611 .usage(crate::memory::MemoryClass::Compaction);
1612 let error = db
1613 .create_index("docs", dense_definition())
1614 .expect_err("oversized Dense graph must be denied");
1615 assert!(
1616 matches!(
1617 error,
1618 MongrelError::ResourceLimitExceeded {
1619 resource: "index build memory",
1620 ..
1621 }
1622 ),
1623 "unexpected admission error: {error:?}"
1624 );
1625 assert_eq!(
1626 db.memory_governor()
1627 .usage(crate::memory::MemoryClass::Compaction),
1628 baseline,
1629 "rejected build must not leak a reservation"
1630 );
1631 assert!(db.table("docs").unwrap().lock().schema().indexes.is_empty());
1632 }
1633
1634 #[test]
1635 fn ddl_auth_fail_closed() {
1636 let _lock = lock_tests();
1637 let dir = tempfile::tempdir().unwrap();
1638 let path = dir.path().to_path_buf();
1639 {
1640 let db = Database::create_with_credentials(&path, "admin", "admin-pw").unwrap();
1641 db.create_table("docs", plain_embedding_schema()).unwrap();
1642 db.create_user("alice", "alice-pw").unwrap();
1643 db.create_role("reader").unwrap();
1644 db.grant_permission(
1645 "reader",
1646 crate::auth::Permission::Select {
1647 table: "docs".into(),
1648 },
1649 )
1650 .unwrap();
1651 db.grant_role("alice", "reader").unwrap();
1652 }
1653 let db = Database::open_with_credentials(&path, "alice", "alice-pw").unwrap();
1654 let definition = IndexDef {
1655 name: "idx".into(),
1656 column_id: 2,
1657 kind: IndexKind::Ann,
1658 predicate: None,
1659 options: IndexOptions {
1660 ann: Some(AnnOptions {
1661 quantization: AnnQuantization::Dense,
1662 ..AnnOptions::default()
1663 }),
1664 ..IndexOptions::default()
1665 },
1666 };
1667 let err = db.create_index("docs", definition).unwrap_err();
1668 assert!(
1669 matches!(err, MongrelError::PermissionDenied { .. }),
1670 "unexpected error: {err:?}"
1671 );
1672 }
1673
1674 #[test]
1675 fn cluster_replica_rejects_local_index_ddl() {
1676 let _lock = lock_tests();
1677 use mongreldb_types::ids::{ClusterId, DatabaseId, NodeId};
1678 let dir = tempfile::tempdir().unwrap();
1679 let cluster_id = ClusterId::from_bytes([1; 16]);
1680 let node_id = NodeId::from_bytes([2; 16]);
1681 let database_id = DatabaseId::from_bytes([3; 16]);
1682 let db =
1683 Database::create_cluster_replica(dir.path(), cluster_id, node_id, database_id).unwrap();
1684 assert!(db.is_read_only_replica());
1685 let definition = IndexDef {
1686 name: "idx".into(),
1687 column_id: 2,
1688 kind: IndexKind::Ann,
1689 predicate: None,
1690 options: Default::default(),
1691 };
1692 let err = db.create_index("docs", definition).unwrap_err();
1693 assert!(
1694 matches!(err, MongrelError::ReadOnlyReplica),
1695 "cluster replica must reject local index DDL: {err:?}"
1696 );
1697 }
1698
1699 #[test]
1700 fn fault_before_publish_leaves_old_schema() {
1701 let _lock = lock_tests();
1702 let _guard = mongreldb_fault::ScopedGuard::new(
1705 "index.publish.before",
1706 mongreldb_fault::Action::Fail,
1707 );
1708 let dir = tempfile::tempdir().unwrap();
1709 let db = Database::create(dir.path()).unwrap();
1710 db.create_table("docs", plain_embedding_schema()).unwrap();
1711 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1712
1713 let definition = IndexDef {
1714 name: "idx_embed_ann".into(),
1715 column_id: 2,
1716 kind: IndexKind::Ann,
1717 predicate: None,
1718 options: IndexOptions {
1719 ann: Some(AnnOptions {
1720 quantization: AnnQuantization::Dense,
1721 ..AnnOptions::default()
1722 }),
1723 ..IndexOptions::default()
1724 },
1725 };
1726 let err = db.create_index("docs", definition).unwrap_err();
1727 assert!(
1728 !matches!(err, MongrelError::NotFound(_)),
1729 "should fail at publish, not earlier: {err:?}"
1730 );
1731 let handle = db.table("docs").unwrap();
1732 let table = handle.lock();
1733 assert!(
1734 table.schema().indexes.is_empty(),
1735 "pre-publication fault must leave the old schema active"
1736 );
1737 }
1738
1739 #[test]
1740 fn cancel_pre_publish_leaves_old_index() {
1741 let _lock = lock_tests();
1742 let dir = tempfile::tempdir().unwrap();
1744 let db = Database::create(dir.path()).unwrap();
1745 db.create_table("docs", embedding_schema(AnnQuantization::BinarySign))
1746 .unwrap();
1747 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1748
1749 let job_id = db
1751 .job_registry()
1752 .submit(
1753 JobKind::IndexBuild,
1754 JobTarget {
1755 table: "docs".into(),
1756 index: Some("idx_embed_ann".into()),
1757 },
1758 )
1759 .unwrap();
1760 db.job_registry().cancel(job_id).unwrap();
1761 let record = db.job_registry().get(job_id).unwrap();
1762 assert_eq!(record.state, crate::jobs::JobState::Failed);
1763
1764 let handle = db.table("docs").unwrap();
1766 let table = handle.lock();
1767 let index = table
1768 .schema()
1769 .indexes
1770 .iter()
1771 .find(|index| index.name == "idx_embed_ann")
1772 .unwrap();
1773 assert_eq!(
1774 index
1775 .options
1776 .ann
1777 .as_ref()
1778 .map(|options| options.quantization),
1779 Some(AnnQuantization::BinarySign)
1780 );
1781 }
1782
1783 #[test]
1784 fn replace_dense_with_binary_sign_search_works() {
1785 let _lock = lock_tests();
1786 let dir = tempfile::tempdir().unwrap();
1787 let db = Database::create(dir.path()).unwrap();
1788 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1789 .unwrap();
1790 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1791 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1792
1793 let new_def = IndexDef {
1794 name: "idx_embed_ann".into(),
1795 column_id: 2,
1796 kind: IndexKind::Ann,
1797 predicate: None,
1798 options: IndexOptions {
1799 ann: Some(AnnOptions {
1800 m: 8,
1801 ef_construction: 32,
1802 ef_search: 16,
1803 quantization: AnnQuantization::BinarySign,
1804 ..AnnOptions::default()
1805 }),
1806 ..IndexOptions::default()
1807 },
1808 };
1809 db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1810
1811 let handle = db.table("docs").unwrap();
1812 let mut table = handle.lock();
1813 let hits = table
1814 .retrieve(&Retriever::Ann {
1815 column_id: 2,
1816 query: vec![1.0, 0.0, 0.0, 0.0],
1817 k: 1,
1818 })
1819 .unwrap();
1820 assert_eq!(hits.len(), 1);
1821 assert!(matches!(
1822 hits[0].score,
1823 RetrieverScore::AnnHammingDistance(_)
1824 ));
1825 assert!(!hits
1826 .iter()
1827 .any(|hit| matches!(hit.score, RetrieverScore::AnnCosineDistance(_))));
1828 }
1829
1830 #[test]
1831 fn replace_hnsw_with_diskann_changes_algorithm_online() {
1832 let _lock = lock_tests();
1836 let dir = tempfile::tempdir().unwrap();
1837 let db = Database::create(dir.path()).unwrap();
1838 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1839 .unwrap();
1840 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1841 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1842
1843 let new_def = IndexDef {
1844 name: "idx_embed_ann".into(),
1845 column_id: 2,
1846 kind: IndexKind::Ann,
1847 predicate: None,
1848 options: IndexOptions {
1849 ann: Some(AnnOptions {
1850 m: 8,
1851 ef_construction: 32,
1852 ef_search: 16,
1853 quantization: AnnQuantization::Dense,
1854 algorithm: crate::schema::AnnAlgorithm::DiskAnn,
1855 diskann: Some(crate::schema::DiskAnnOptions {
1856 r: 8,
1857 l: 16,
1858 beam_width: 4,
1859 alpha: 120,
1860 }),
1861 ..AnnOptions::default()
1862 }),
1863 ..IndexOptions::default()
1864 },
1865 };
1866 let job_id = db.replace_index("docs", "idx_embed_ann", new_def).unwrap();
1867 let record = db.job_registry().get(job_id).unwrap();
1868 assert_eq!(record.state, crate::jobs::JobState::Succeeded);
1869
1870 let handle = db.table("docs").unwrap();
1871 let mut table = handle.lock();
1872 let hits = table
1873 .retrieve(&Retriever::Ann {
1874 column_id: 2,
1875 query: vec![1.0, 0.0, 0.0, 0.0],
1876 k: 2,
1877 })
1878 .unwrap();
1879 assert_eq!(hits.len(), 2);
1880 assert!(matches!(
1882 hits[0].score,
1883 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-5
1884 ));
1885 }
1886
1887 #[test]
1888 fn product_quantization_dimension_divisibility_rejected_at_ddl() {
1889 let _lock = lock_tests();
1892 let dir = tempfile::tempdir().unwrap();
1893 let db = Database::create(dir.path()).unwrap();
1894 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1895 .unwrap();
1896 let bad_def = IndexDef {
1898 name: "idx_bad_pq".into(),
1899 column_id: 2,
1900 kind: IndexKind::Ann,
1901 predicate: None,
1902 options: IndexOptions {
1903 ann: Some(AnnOptions {
1904 quantization: AnnQuantization::Product {
1905 num_subvectors: 3,
1906 bits: 8,
1907 },
1908 product: Some(crate::schema::ProductQuantizerOptions::default()),
1909 ..AnnOptions::default()
1910 }),
1911 ..IndexOptions::default()
1912 },
1913 };
1914 let err = db.create_index("docs", bad_def).unwrap_err();
1915 assert!(
1916 err.to_string()
1917 .contains("must evenly divide the column dimension"),
1918 "expected divisibility error, got: {err}"
1919 );
1920 }
1921
1922 #[test]
1923 fn fault_after_publish_keeps_new_index() {
1924 let _lock = lock_tests();
1925 let _guard =
1928 mongreldb_fault::ScopedGuard::new("index.publish.after", mongreldb_fault::Action::Fail);
1929 let dir = tempfile::tempdir().unwrap();
1930 let db = Database::create(dir.path()).unwrap();
1931 db.create_table("docs", plain_embedding_schema()).unwrap();
1932 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1933
1934 let definition = IndexDef {
1935 name: "idx_embed_ann".into(),
1936 column_id: 2,
1937 kind: IndexKind::Ann,
1938 predicate: None,
1939 options: IndexOptions {
1940 ann: Some(AnnOptions {
1941 quantization: AnnQuantization::Dense,
1942 ..AnnOptions::default()
1943 }),
1944 ..IndexOptions::default()
1945 },
1946 };
1947 let _ = db.create_index("docs", definition);
1950 let handle = db.table("docs").unwrap();
1951 let table = handle.lock();
1952 let index = table
1953 .schema()
1954 .indexes
1955 .iter()
1956 .find(|index| index.name == "idx_embed_ann");
1957 assert!(
1958 index.is_some(),
1959 "post-publish fault must keep the new index active"
1960 );
1961 assert_eq!(
1962 index
1963 .unwrap()
1964 .options
1965 .ann
1966 .as_ref()
1967 .map(|options| options.quantization),
1968 Some(AnnQuantization::Dense)
1969 );
1970 }
1971
1972 #[test]
1973 fn dense_plaintext_close_reopen_preserves_search() {
1974 let _lock = lock_tests();
1975 let dir = tempfile::tempdir().unwrap();
1976 let path = dir.path().to_path_buf();
1977 {
1978 let db = Database::create(&path).unwrap();
1979 db.create_table("docs", embedding_schema(AnnQuantization::Dense))
1980 .unwrap();
1981 put_row(&db, "docs", 1, vec![1.0, 0.0, 0.0, 0.0]);
1982 put_row(&db, "docs", 2, vec![0.0, 1.0, 0.0, 0.0]);
1983 let handle = db.table("docs").unwrap();
1985 let mut table = handle.lock();
1986 table.flush().unwrap();
1987 let _ = table.publish_read_generation().unwrap();
1988 }
1989 let db = Database::open(&path).unwrap();
1990 let handle = db.table("docs").unwrap();
1991 let mut table = handle.lock();
1992 let hits = table
1993 .retrieve(&Retriever::Ann {
1994 column_id: 2,
1995 query: vec![1.0, 0.0, 0.0, 0.0],
1996 k: 1,
1997 })
1998 .unwrap();
1999 assert_eq!(hits.len(), 1);
2000 assert!(matches!(
2001 hits[0].score,
2002 RetrieverScore::AnnCosineDistance(d) if d.abs() < 1e-4
2003 ));
2004 }
2005}