1use std::collections::BTreeSet;
9use std::fs::{File, OpenOptions};
10use std::path::{Path, PathBuf};
11use std::sync::Arc;
12
13use arrow::array::{
14 Array, ArrayRef, FixedSizeBinaryArray, FixedSizeBinaryBuilder, FixedSizeListArray,
15 FixedSizeListBuilder, Float32Array, Float32Builder, TimestampMicrosecondArray,
16 TimestampMicrosecondBuilder,
17};
18use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit};
19use arrow::record_batch::RecordBatch;
20use parquet::arrow::ArrowWriter;
21use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
22
23use crate::{
24 PublishedSearchArtifact, SearchArtifactError, SearchArtifactKey, SearchCoordinationLimits,
25 SearchPublicationMode, SearchPublicationOutcome, SearchPublicationPlan, SearchSourceSnapshot,
26 SearchUpdateBuild, coordinate_search_update,
27};
28
29pub const VECTOR_DATA_FILE: &str = "vectors.parquet";
31pub const VECTOR_BACKEND_VERSION: &str = "arrow-parquet-58";
33pub const VECTOR_CONTRACT_VERSION: &str = "graphforge_vector_v1";
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq)]
38pub struct VectorStoreLimits {
39 pub dimensions: usize,
41 pub stored_vectors: usize,
43 pub vector_cells: usize,
45 pub eligible_nodes: usize,
47 pub results: usize,
49 pub parquet_bytes: u64,
51}
52
53impl Default for VectorStoreLimits {
54 fn default() -> Self {
55 Self {
56 dimensions: 4_096,
57 stored_vectors: 1_000_000,
58 vector_cells: 100_000_000,
59 eligible_nodes: 1_000_000,
60 results: 10_000,
61 parquet_bytes: 4 * 1024 * 1024 * 1024,
62 }
63 }
64}
65
66#[derive(Clone, Debug, PartialEq)]
68pub struct StoredVector {
69 pub node_uuid: [u8; 16],
71 pub vector: Vec<f32>,
73 pub updated_at_micros: i64,
75}
76
77#[derive(Clone, Copy, Debug, PartialEq, Eq)]
79pub enum VectorUpsertChange {
80 Unchanged,
82 Changed,
84}
85
86#[derive(Clone, Copy, Debug, PartialEq)]
88pub struct VectorSearchHit {
89 pub node_uuid: [u8; 16],
91 pub score: f64,
93}
94
95pub fn vector_schema(
101 dimension: usize,
102 limits: VectorStoreLimits,
103) -> Result<SchemaRef, SearchArtifactError> {
104 validate_dimension(dimension, limits)?;
105 let dimension =
106 i32::try_from(dimension).map_err(|_| exhausted("vector_dimensions", limits.dimensions))?;
107 Ok(Arc::new(Schema::new(vec![
108 Field::new("node_uuid", DataType::FixedSizeBinary(16), false),
109 Field::new(
110 "vector",
111 DataType::FixedSizeList(
112 Arc::new(Field::new("item", DataType::Float32, false)),
113 dimension,
114 ),
115 false,
116 ),
117 Field::new(
118 "updated_at",
119 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
120 false,
121 ),
122 ])))
123}
124
125pub fn validate_vector(
130 vector: &[f32],
131 limits: VectorStoreLimits,
132) -> Result<f64, SearchArtifactError> {
133 validate_dimension(vector.len(), limits)?;
134 let mut norm = 0.0_f64;
135 for &value in vector {
136 if !value.is_finite() {
137 return Err(invalid("vector", "values must be finite Float32"));
138 }
139 let value = f64::from(value);
140 norm = value.mul_add(value, norm);
141 if !norm.is_finite() {
142 return Err(invalid("vector", "norm exceeds supported range"));
143 }
144 }
145 if norm == 0.0 {
146 return Err(invalid("vector", "norm must be non-zero"));
147 }
148 Ok(norm)
149}
150
151pub fn apply_vector_upsert(
160 rows: &mut Vec<StoredVector>,
161 node_uuid: [u8; 16],
162 vector: &[f32],
163 updated_at_micros: i64,
164 limits: VectorStoreLimits,
165) -> Result<VectorUpsertChange, SearchArtifactError> {
166 validate_vector(vector, limits)?;
167 validate_rows(rows, Some(vector.len()), limits, Path::new("<memory>"))?;
168 match rows.binary_search_by_key(&node_uuid, |row| row.node_uuid) {
169 Ok(index) if rows[index].vector == vector => Ok(VectorUpsertChange::Unchanged),
170 Ok(index) => {
171 rows[index].vector.clear();
172 rows[index].vector.extend_from_slice(vector);
173 rows[index].updated_at_micros = updated_at_micros;
174 Ok(VectorUpsertChange::Changed)
175 }
176 Err(index) => {
177 if rows.len() >= limits.stored_vectors {
178 return Err(exhausted("stored_vectors", limits.stored_vectors));
179 }
180 checked_cells(rows.len() + 1, vector.len(), limits)?;
181 rows.insert(
182 index,
183 StoredVector {
184 node_uuid,
185 vector: vector.to_vec(),
186 updated_at_micros,
187 },
188 );
189 Ok(VectorUpsertChange::Changed)
190 }
191 }
192}
193
194#[allow(clippy::too_many_arguments)]
205pub fn upsert_published_vector<S, M, C>(
206 project_dir: &Path,
207 label: &str,
208 space: &str,
209 node_uuid: [u8; 16],
210 vector: &[f32],
211 updated_at_micros: i64,
212 limits: VectorStoreLimits,
213 coordination: SearchCoordinationLimits,
214 snapshot: S,
215 mut is_current_member: M,
216 checkpoint: C,
217) -> Result<SearchPublicationOutcome, SearchArtifactError>
218where
219 S: FnMut() -> Result<SearchSourceSnapshot, SearchArtifactError>,
220 M: FnMut([u8; 16]) -> Result<bool, SearchArtifactError>,
221 C: FnMut() -> Result<(), SearchArtifactError>,
222{
223 validate_vector(vector, limits)?;
224 let key = SearchArtifactKey::vector(label, space)?;
225 let dimension = u32::try_from(vector.len())
226 .map_err(|_| exhausted("vector_dimensions", limits.dimensions))?;
227 let plan = SearchPublicationPlan {
228 key: &key,
229 backend_version: VECTOR_BACKEND_VERSION,
230 contract_version: VECTOR_CONTRACT_VERSION,
231 dimension: Some(dimension),
232 mode: SearchPublicationMode::Replace,
233 };
234 coordinate_search_update(
235 project_dir,
236 plan,
237 coordination,
238 snapshot,
239 |current, build_dir, _, checkpoint| {
240 if !is_current_member(node_uuid)? {
241 return Err(invalid(
242 "node",
243 "UUID does not exist with the required label",
244 ));
245 }
246 let mut rows = match current {
247 Some(artifact) => {
248 validate_vector_manifest(artifact, &key, dimension)?;
249 read_vector_snapshot(&artifact.path, vector.len(), limits, &mut *checkpoint)?
250 }
251 None => Vec::new(),
252 };
253 match apply_vector_upsert(&mut rows, node_uuid, vector, updated_at_micros, limits)? {
254 VectorUpsertChange::Unchanged => Ok(SearchUpdateBuild::ReuseCurrent),
255 VectorUpsertChange::Changed => {
256 write_vector_snapshot(
257 build_dir,
258 &rows,
259 vector.len(),
260 limits,
261 &mut *checkpoint,
262 )?;
263 Ok(SearchUpdateBuild::Publish)
264 }
265 }
266 },
267 checkpoint,
268 )
269}
270
271pub fn write_vector_snapshot<C>(
277 build_dir: &Path,
278 rows: &[StoredVector],
279 dimension: usize,
280 limits: VectorStoreLimits,
281 mut checkpoint: C,
282) -> Result<PathBuf, SearchArtifactError>
283where
284 C: FnMut() -> Result<(), SearchArtifactError>,
285{
286 checkpoint()?;
287 validate_rows(rows, Some(dimension), limits, build_dir)?;
288 let schema = vector_schema(dimension, limits)?;
289 let batch = rows_to_batch(rows, schema.clone(), dimension, &mut checkpoint)?;
290 let path = build_dir.join(VECTOR_DATA_FILE);
291 let file = File::create(&path).map_err(|source| io("create vector snapshot", &path, source))?;
292 let mut writer = ArrowWriter::try_new(file, schema, None)
293 .map_err(|error| build(format!("create vector Parquet writer: {error}")))?;
294 writer
295 .write(&batch)
296 .map_err(|error| build(format!("write vector Parquet batch: {error}")))?;
297 checkpoint()?;
298 writer
299 .close()
300 .map_err(|error| build(format!("close vector Parquet writer: {error}")))?;
301 sync_vector_snapshot(&path)?;
302 let bytes = std::fs::metadata(&path)
303 .map_err(|source| io("inspect vector snapshot", &path, source))?
304 .len();
305 if bytes > limits.parquet_bytes {
306 return Err(SearchArtifactError::ResourceExhausted {
307 resource: "vector_parquet_bytes",
308 limit: limits.parquet_bytes,
309 });
310 }
311 Ok(path)
312}
313
314fn sync_vector_snapshot(path: &Path) -> Result<(), SearchArtifactError> {
315 OpenOptions::new()
316 .read(true)
317 .write(true)
318 .open(path)
319 .and_then(|file| file.sync_all())
320 .map_err(|source| io("sync vector snapshot", path, source))
321}
322
323pub fn read_vector_snapshot<C>(
330 artifact_dir: &Path,
331 dimension: usize,
332 limits: VectorStoreLimits,
333 mut checkpoint: C,
334) -> Result<Vec<StoredVector>, SearchArtifactError>
335where
336 C: FnMut() -> Result<(), SearchArtifactError>,
337{
338 checkpoint()?;
339 let path = artifact_dir.join(VECTOR_DATA_FILE);
340 let metadata = std::fs::metadata(&path).map_err(|source| corrupt_io(&path, &source))?;
341 if metadata.len() > limits.parquet_bytes {
342 return Err(SearchArtifactError::ResourceExhausted {
343 resource: "vector_parquet_bytes",
344 limit: limits.parquet_bytes,
345 });
346 }
347 let expected = vector_schema(dimension, limits)?;
348 let file = File::open(&path).map_err(|source| corrupt_io(&path, &source))?;
349 let reader = ParquetRecordBatchReaderBuilder::try_new(file)
350 .map_err(|error| corrupt(&path, format!("open Parquet: {error}")))?
351 .build()
352 .map_err(|error| corrupt(&path, format!("build Parquet reader: {error}")))?;
353 let mut rows = Vec::new();
354 for batch in reader {
355 checkpoint()?;
356 let batch = batch.map_err(|error| corrupt(&path, format!("read Parquet: {error}")))?;
357 if batch.schema().as_ref() != expected.as_ref() {
358 return Err(corrupt(
359 &path,
360 format!(
361 "schema mismatch: expected {expected:?}, found {:?}",
362 batch.schema()
363 ),
364 ));
365 }
366 if rows.len().saturating_add(batch.num_rows()) > limits.stored_vectors {
367 return Err(exhausted("stored_vectors", limits.stored_vectors));
368 }
369 checked_cells(rows.len() + batch.num_rows(), dimension, limits)?;
370 decode_batch(&path, &batch, dimension, &mut rows, &mut checkpoint)?;
371 }
372 validate_rows(&rows, Some(dimension), limits, &path)?;
373 Ok(rows)
374}
375
376pub fn validate_published_vectors<C>(
382 artifact: &PublishedSearchArtifact,
383 limits: VectorStoreLimits,
384 checkpoint: C,
385) -> Result<(), SearchArtifactError>
386where
387 C: FnMut() -> Result<(), SearchArtifactError>,
388{
389 let dimension = published_vector_dimension(artifact)?;
390 read_vector_snapshot(&artifact.path, dimension, limits, checkpoint).map(|_| ())
391}
392
393pub fn exact_cosine_search<C>(
402 rows: &[StoredVector],
403 query: &[f32],
404 eligible_nodes: &BTreeSet<[u8; 16]>,
405 limit: usize,
406 limits: VectorStoreLimits,
407 mut checkpoint: C,
408) -> Result<Vec<VectorSearchHit>, SearchArtifactError>
409where
410 C: FnMut() -> Result<(), SearchArtifactError>,
411{
412 let query_norm = validate_vector(query, limits)?;
413 validate_rows(rows, Some(query.len()), limits, Path::new("<memory>"))?;
414 if eligible_nodes.len() > limits.eligible_nodes {
415 return Err(exhausted("eligible_nodes", limits.eligible_nodes));
416 }
417 if limit == 0 {
418 return Err(invalid("limit", "must be greater than zero"));
419 }
420 if limit > limits.results {
421 return Err(exhausted("search_results", limits.results));
422 }
423 checked_cells(rows.len(), query.len(), limits)?;
424
425 let mut hits = Vec::with_capacity(rows.len().min(limit));
426 for row in rows {
427 checkpoint()?;
428 if !eligible_nodes.contains(&row.node_uuid) {
429 continue;
430 }
431 let (dot, stored_norm) = dot_and_norm(query, &row.vector)?;
432 let denominator = (query_norm * stored_norm).sqrt();
433 let score = (dot / denominator).clamp(-1.0, 1.0);
434 if !score.is_finite() {
435 return Err(corrupt(
436 Path::new("<memory>"),
437 "cosine score is not finite".to_owned(),
438 ));
439 }
440 hits.push(VectorSearchHit {
441 node_uuid: row.node_uuid,
442 score,
443 });
444 }
445 hits.sort_unstable_by(|left, right| {
446 right
447 .score
448 .total_cmp(&left.score)
449 .then_with(|| left.node_uuid.cmp(&right.node_uuid))
450 });
451 hits.truncate(limit);
452 Ok(hits)
453}
454
455pub fn search_published_vectors<C>(
461 artifact: &PublishedSearchArtifact,
462 query: &[f32],
463 eligible_nodes: &BTreeSet<[u8; 16]>,
464 limit: usize,
465 limits: VectorStoreLimits,
466 mut checkpoint: C,
467) -> Result<Vec<VectorSearchHit>, SearchArtifactError>
468where
469 C: FnMut() -> Result<(), SearchArtifactError>,
470{
471 let dimension = published_vector_dimension(artifact)?;
472 if query.len() != dimension {
473 return Err(invalid(
474 "vector",
475 format!(
476 "dimension {} does not match stored dimension {dimension}",
477 query.len()
478 ),
479 ));
480 }
481 let rows = read_vector_snapshot(&artifact.path, dimension, limits, &mut checkpoint)?;
482 exact_cosine_search(&rows, query, eligible_nodes, limit, limits, checkpoint)
483}
484
485fn rows_to_batch<C>(
486 rows: &[StoredVector],
487 schema: SchemaRef,
488 dimension: usize,
489 checkpoint: &mut C,
490) -> Result<RecordBatch, SearchArtifactError>
491where
492 C: FnMut() -> Result<(), SearchArtifactError>,
493{
494 let mut uuids = FixedSizeBinaryBuilder::with_capacity(rows.len(), 16);
495 let values = Float32Builder::with_capacity(rows.len().saturating_mul(dimension));
496 let mut vectors = FixedSizeListBuilder::with_capacity(
497 values,
498 i32::try_from(dimension).map_err(|_| build("vector dimension exceeds Arrow i32"))?,
499 rows.len(),
500 )
501 .with_field(Arc::new(Field::new("item", DataType::Float32, false)));
502 let timestamp_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()));
503 let mut updated =
504 TimestampMicrosecondBuilder::with_capacity(rows.len()).with_data_type(timestamp_type);
505 for row in rows {
506 checkpoint()?;
507 uuids
508 .append_value(row.node_uuid)
509 .map_err(|error| build(format!("append node_uuid: {error}")))?;
510 for &value in &row.vector {
511 vectors.values().append_value(value);
512 }
513 vectors.append(true);
514 updated.append_value(row.updated_at_micros);
515 }
516 RecordBatch::try_new(
517 schema,
518 vec![
519 Arc::new(uuids.finish()) as ArrayRef,
520 Arc::new(vectors.finish()) as ArrayRef,
521 Arc::new(updated.finish()) as ArrayRef,
522 ],
523 )
524 .map_err(|error| build(format!("construct vector record batch: {error}")))
525}
526
527fn validate_vector_manifest(
528 artifact: &PublishedSearchArtifact,
529 key: &SearchArtifactKey,
530 dimension: u32,
531) -> Result<(), SearchArtifactError> {
532 let manifest = &artifact.manifest;
533 if manifest.index_kind != key.kind()
534 || manifest.label != key.label()
535 || manifest.space.as_deref() != key.space()
536 || manifest.properties.is_some()
537 || manifest.backend_version != VECTOR_BACKEND_VERSION
538 || manifest.contract_version != VECTOR_CONTRACT_VERSION
539 || !manifest.completed
540 {
541 return Err(corrupt(
542 &artifact.path,
543 "vector manifest does not match the requested backend key",
544 ));
545 }
546 if manifest.dimension != Some(dimension) {
547 return Err(invalid(
548 "vector",
549 format!(
550 "dimension {dimension} does not match stored dimension {}",
551 manifest.dimension.unwrap_or_default()
552 ),
553 ));
554 }
555 Ok(())
556}
557
558fn published_vector_dimension(
559 artifact: &PublishedSearchArtifact,
560) -> Result<usize, SearchArtifactError> {
561 let manifest = &artifact.manifest;
562 let space = manifest
563 .space
564 .as_deref()
565 .ok_or_else(|| corrupt(&artifact.path, "vector manifest omits its space"))?;
566 let key = SearchArtifactKey::vector(&manifest.label, space)
567 .map_err(|error| corrupt(&artifact.path, error.to_string()))?;
568 let dimension = manifest
569 .dimension
570 .ok_or_else(|| corrupt(&artifact.path, "vector manifest omits its fixed dimension"))?;
571 validate_vector_manifest(artifact, &key, dimension)
572 .map_err(|error| corrupt(&artifact.path, error.to_string()))?;
573 usize::try_from(dimension)
574 .map_err(|_| corrupt(&artifact.path, "vector dimension exceeds usize"))
575}
576
577fn decode_batch<C>(
578 path: &Path,
579 batch: &RecordBatch,
580 dimension: usize,
581 rows: &mut Vec<StoredVector>,
582 checkpoint: &mut C,
583) -> Result<(), SearchArtifactError>
584where
585 C: FnMut() -> Result<(), SearchArtifactError>,
586{
587 let uuids = downcast::<FixedSizeBinaryArray>(path, batch, 0, "node_uuid")?;
588 let vectors = downcast::<FixedSizeListArray>(path, batch, 1, "vector")?;
589 let values = vectors
590 .values()
591 .as_any()
592 .downcast_ref::<Float32Array>()
593 .ok_or_else(|| corrupt(path, "vector child is not Float32".to_owned()))?;
594 let updated = downcast::<TimestampMicrosecondArray>(path, batch, 2, "updated_at")?;
595 if uuids.null_count() != 0
596 || vectors.null_count() != 0
597 || values.null_count() != 0
598 || updated.null_count() != 0
599 {
600 return Err(corrupt(
601 path,
602 "vector snapshot contains NULL values".to_owned(),
603 ));
604 }
605 for row in 0..batch.num_rows() {
606 checkpoint()?;
607 let uuid: [u8; 16] = uuids
608 .value(row)
609 .try_into()
610 .map_err(|_| corrupt(path, "node_uuid is not 16 bytes".to_owned()))?;
611 let offset = row
612 .checked_mul(dimension)
613 .ok_or_else(|| corrupt(path, "vector offset overflow".to_owned()))?;
614 let end = offset
615 .checked_add(dimension)
616 .ok_or_else(|| corrupt(path, "vector offset overflow".to_owned()))?;
617 if end > values.len() {
618 return Err(corrupt(path, "vector child length is truncated".to_owned()));
619 }
620 rows.push(StoredVector {
621 node_uuid: uuid,
622 vector: (offset..end).map(|index| values.value(index)).collect(),
623 updated_at_micros: updated.value(row),
624 });
625 }
626 Ok(())
627}
628
629fn downcast<'a, T: 'static>(
630 path: &Path,
631 batch: &'a RecordBatch,
632 index: usize,
633 name: &str,
634) -> Result<&'a T, SearchArtifactError> {
635 batch
636 .column(index)
637 .as_any()
638 .downcast_ref::<T>()
639 .ok_or_else(|| corrupt(path, format!("column {name:?} has the wrong Arrow type")))
640}
641
642fn validate_rows(
643 rows: &[StoredVector],
644 expected_dimension: Option<usize>,
645 limits: VectorStoreLimits,
646 path: &Path,
647) -> Result<(), SearchArtifactError> {
648 if rows.len() > limits.stored_vectors {
649 return Err(exhausted("stored_vectors", limits.stored_vectors));
650 }
651 let dimension = expected_dimension.or_else(|| rows.first().map(|row| row.vector.len()));
652 if let Some(dimension) = dimension {
653 validate_dimension(dimension, limits)?;
654 checked_cells(rows.len(), dimension, limits)?;
655 for (index, row) in rows.iter().enumerate() {
656 if row.vector.len() != dimension {
657 return Err(corrupt(
658 path,
659 format!(
660 "row {index} dimension {} does not match {dimension}",
661 row.vector.len()
662 ),
663 ));
664 }
665 validate_vector(&row.vector, limits)
666 .map_err(|error| corrupt(path, error.to_string()))?;
667 }
668 }
669 for pair in rows.windows(2) {
670 if pair[0].node_uuid >= pair[1].node_uuid {
671 let reason = if pair[0].node_uuid == pair[1].node_uuid {
672 "duplicate node_uuid"
673 } else {
674 "rows are not sorted by ascending node_uuid"
675 };
676 return Err(corrupt(path, reason.to_owned()));
677 }
678 }
679 Ok(())
680}
681
682fn dot_and_norm(left: &[f32], right: &[f32]) -> Result<(f64, f64), SearchArtifactError> {
683 let mut dot = 0.0_f64;
684 let mut norm = 0.0_f64;
685 for (&left, &right) in left.iter().zip(right) {
686 let left = f64::from(left);
687 let right = f64::from(right);
688 dot = left.mul_add(right, dot);
689 norm = right.mul_add(right, norm);
690 if !dot.is_finite() || !norm.is_finite() {
691 return Err(corrupt(
692 Path::new("<memory>"),
693 "cosine accumulation exceeds supported range".to_owned(),
694 ));
695 }
696 }
697 if norm == 0.0 {
698 return Err(corrupt(
699 Path::new("<memory>"),
700 "stored vector has zero norm".to_owned(),
701 ));
702 }
703 Ok((dot, norm))
704}
705
706fn validate_dimension(
707 dimension: usize,
708 limits: VectorStoreLimits,
709) -> Result<(), SearchArtifactError> {
710 if dimension == 0 {
711 return Err(invalid("vector", "dimension must be greater than zero"));
712 }
713 if dimension > limits.dimensions || i32::try_from(dimension).is_err() {
714 return Err(exhausted("vector_dimensions", limits.dimensions));
715 }
716 Ok(())
717}
718
719fn checked_cells(
720 rows: usize,
721 dimension: usize,
722 limits: VectorStoreLimits,
723) -> Result<usize, SearchArtifactError> {
724 let cells = rows
725 .checked_mul(dimension)
726 .ok_or_else(|| exhausted("vector_cells", limits.vector_cells))?;
727 if cells > limits.vector_cells {
728 return Err(exhausted("vector_cells", limits.vector_cells));
729 }
730 Ok(cells)
731}
732
733fn invalid(field: &'static str, reason: impl Into<String>) -> SearchArtifactError {
734 SearchArtifactError::InvalidSelector {
735 field,
736 reason: reason.into(),
737 }
738}
739
740fn exhausted(resource: &'static str, limit: usize) -> SearchArtifactError {
741 SearchArtifactError::ResourceExhausted {
742 resource,
743 limit: u64::try_from(limit).unwrap_or(u64::MAX),
744 }
745}
746
747fn build(reason: impl Into<String>) -> SearchArtifactError {
748 SearchArtifactError::Build(reason.into())
749}
750
751fn corrupt(path: &Path, reason: impl Into<String>) -> SearchArtifactError {
752 SearchArtifactError::CorruptPrimaryVectors {
753 path: path.to_path_buf(),
754 reason: reason.into(),
755 }
756}
757
758fn corrupt_io(path: &Path, source: &std::io::Error) -> SearchArtifactError {
759 corrupt(path, source.to_string())
760}
761
762fn io(operation: &'static str, path: &Path, source: std::io::Error) -> SearchArtifactError {
763 SearchArtifactError::Io {
764 operation,
765 path: path.to_path_buf(),
766 source,
767 }
768}
769
770#[cfg(test)]
771mod tests {
772 use std::cell::Cell;
773
774 use tempfile::TempDir;
775
776 use super::*;
777 use crate::{SearchManifest, current_search_artifact};
778
779 fn uuid(value: u8) -> [u8; 16] {
780 let mut uuid = [0_u8; 16];
781 uuid[15] = value;
782 uuid
783 }
784
785 fn row(value: u8, vector: &[f32], updated_at_micros: i64) -> StoredVector {
786 StoredVector {
787 node_uuid: uuid(value),
788 vector: vector.to_vec(),
789 updated_at_micros,
790 }
791 }
792
793 fn source_snapshot() -> Result<SearchSourceSnapshot, SearchArtifactError> {
794 Ok(SearchSourceSnapshot {
795 generation: 7,
796 fingerprint: format!("gf-fnv1a256:{:064x}", 7),
797 })
798 }
799
800 #[test]
801 fn schema_is_uuid_fixed_vector_and_utc_timestamp() {
802 let schema = vector_schema(3, VectorStoreLimits::default()).unwrap();
803 assert_eq!(schema.field(0).name(), "node_uuid");
804 assert_eq!(schema.field(0).data_type(), &DataType::FixedSizeBinary(16));
805 assert_eq!(
806 schema.field(1).data_type(),
807 &DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3,)
808 );
809 assert_eq!(
810 schema.field(2).data_type(),
811 &DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()))
812 );
813 assert!(schema.fields().iter().all(|field| !field.is_nullable()));
814 }
815
816 #[test]
817 fn upsert_is_sorted_idempotent_and_replaces_atomically() {
818 let limits = VectorStoreLimits::default();
819 let mut rows = vec![row(2, &[0.0, 1.0], 20)];
820 assert_eq!(
821 apply_vector_upsert(&mut rows, uuid(1), &[1.0, 0.0], 10, limits).unwrap(),
822 VectorUpsertChange::Changed
823 );
824 assert_eq!(
825 rows.iter().map(|row| row.node_uuid).collect::<Vec<_>>(),
826 vec![uuid(1), uuid(2)]
827 );
828 assert_eq!(
829 apply_vector_upsert(&mut rows, uuid(1), &[1.0, 0.0], 99, limits).unwrap(),
830 VectorUpsertChange::Unchanged
831 );
832 assert_eq!(rows[0].updated_at_micros, 10);
833 assert_eq!(
834 apply_vector_upsert(&mut rows, uuid(1), &[-1.0, 0.0], 30, limits).unwrap(),
835 VectorUpsertChange::Changed
836 );
837 assert_eq!(rows[0], row(1, &[-1.0, 0.0], 30));
838 }
839
840 #[test]
841 fn vector_validation_rejects_dimension_nonfinite_and_zero_norm() {
842 let limits = VectorStoreLimits {
843 dimensions: 2,
844 ..VectorStoreLimits::default()
845 };
846 for vector in [
847 &[][..],
848 &[0.0, 0.0],
849 &[f32::NAN],
850 &[f32::INFINITY],
851 &[1.0, 2.0, 3.0],
852 ] {
853 assert!(validate_vector(vector, limits).is_err(), "{vector:?}");
854 }
855 assert!((validate_vector(&[3.0, 4.0], limits).unwrap() - 25.0).abs() < f64::EPSILON);
856 }
857
858 #[test]
859 fn parquet_round_trip_preserves_fixed_dimension_and_order() {
860 let dir = TempDir::new().unwrap();
861 let rows = vec![row(1, &[1.0, 2.0], 11), row(3, &[-3.0, 4.0], 33)];
862 let writes = Cell::new(0);
863 let path =
864 write_vector_snapshot(dir.path(), &rows, 2, VectorStoreLimits::default(), || {
865 writes.set(writes.get() + 1);
866 Ok(())
867 })
868 .unwrap();
869 assert_eq!(path, dir.path().join(VECTOR_DATA_FILE));
870 assert!(writes.get() >= 3);
871 let decoded =
872 read_vector_snapshot(dir.path(), 2, VectorStoreLimits::default(), || Ok(())).unwrap();
873 assert_eq!(decoded, rows);
874 }
875
876 #[test]
877 fn vector_snapshot_sync_preserves_file_contents() {
878 let dir = TempDir::new().unwrap();
879 let path = dir.path().join(VECTOR_DATA_FILE);
880 std::fs::write(&path, b"complete snapshot").unwrap();
881
882 sync_vector_snapshot(&path).unwrap();
883
884 assert_eq!(std::fs::read(path).unwrap(), b"complete snapshot");
885 }
886
887 #[test]
888 fn write_enforces_the_parquet_byte_limit_before_publication() {
889 let dir = TempDir::new().unwrap();
890 let limits = VectorStoreLimits {
891 parquet_bytes: 1,
892 ..VectorStoreLimits::default()
893 };
894 assert!(matches!(
895 write_vector_snapshot(dir.path(), &[row(1, &[1.0], 1)], 1, limits, || Ok(())),
896 Err(SearchArtifactError::ResourceExhausted {
897 resource: "vector_parquet_bytes",
898 limit: 1,
899 })
900 ));
901 }
902
903 #[test]
904 fn published_artifact_validation_and_search_use_pinned_contracts() {
905 let dir = TempDir::new().unwrap();
906 let limits = VectorStoreLimits::default();
907 let rows = vec![row(1, &[1.0, 0.0], 11), row(2, &[0.0, 1.0], 22)];
908 write_vector_snapshot(dir.path(), &rows, 2, limits, || Ok(())).unwrap();
909 let key = SearchArtifactKey::vector("Person", "semantic").unwrap();
910 let manifest = SearchManifest::for_key(
911 &key,
912 VECTOR_BACKEND_VERSION,
913 VECTOR_CONTRACT_VERSION,
914 Some(2),
915 &SearchSourceSnapshot {
916 generation: 1,
917 fingerprint: format!("gf-fnv1a256:{:064x}", 1),
918 },
919 true,
920 )
921 .unwrap();
922 let artifact = PublishedSearchArtifact {
923 path: dir.path().to_path_buf(),
924 manifest,
925 };
926 validate_published_vectors(&artifact, limits, || Ok(())).unwrap();
927 let hits = search_published_vectors(
928 &artifact,
929 &[1.0, 0.0],
930 &BTreeSet::from([uuid(1), uuid(2)]),
931 1,
932 limits,
933 || Ok(()),
934 )
935 .unwrap();
936 assert_eq!(
937 hits,
938 vec![VectorSearchHit {
939 node_uuid: uuid(1),
940 score: 1.0,
941 }]
942 );
943 }
944
945 #[test]
946 fn published_reads_reject_a_mismatched_backend_contract() {
947 let dir = TempDir::new().unwrap();
948 let limits = VectorStoreLimits::default();
949 write_vector_snapshot(dir.path(), &[row(1, &[1.0], 1)], 1, limits, || Ok(())).unwrap();
950 let key = SearchArtifactKey::vector("Person", "semantic").unwrap();
951 let artifact = PublishedSearchArtifact {
952 path: dir.path().to_path_buf(),
953 manifest: SearchManifest::for_key(
954 &key,
955 "other-backend-1",
956 VECTOR_CONTRACT_VERSION,
957 Some(1),
958 &source_snapshot().unwrap(),
959 true,
960 )
961 .unwrap(),
962 };
963 assert!(matches!(
964 validate_published_vectors(&artifact, limits, || Ok(())),
965 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
966 ));
967 }
968
969 #[test]
970 fn atomic_upsert_publishes_reuses_replaces_and_preserves_dimension() {
971 let dir = TempDir::new().unwrap();
972 let limits = VectorStoreLimits::default();
973 let first = upsert_published_vector(
974 dir.path(),
975 "Person",
976 "semantic",
977 uuid(1),
978 &[1.0, 0.0],
979 11,
980 limits,
981 SearchCoordinationLimits::default(),
982 source_snapshot,
983 |_| Ok(true),
984 || Ok(()),
985 )
986 .unwrap();
987 let first_path = match first {
988 SearchPublicationOutcome::Published {
989 artifact,
990 attempts: 1,
991 } => artifact.path,
992 other => panic!("unexpected first outcome: {other:?}"),
993 };
994
995 let repeated = upsert_published_vector(
996 dir.path(),
997 "Person",
998 "semantic",
999 uuid(1),
1000 &[1.0, 0.0],
1001 99,
1002 limits,
1003 SearchCoordinationLimits::default(),
1004 source_snapshot,
1005 |_| Ok(true),
1006 || Ok(()),
1007 )
1008 .unwrap();
1009 assert!(matches!(
1010 repeated,
1011 SearchPublicationOutcome::Reused(ref artifact) if artifact.path == first_path
1012 ));
1013 assert_eq!(
1014 read_vector_snapshot(&first_path, 2, limits, || Ok(())).unwrap()[0].updated_at_micros,
1015 11
1016 );
1017
1018 let replaced = upsert_published_vector(
1019 dir.path(),
1020 "Person",
1021 "semantic",
1022 uuid(1),
1023 &[-1.0, 0.0],
1024 22,
1025 limits,
1026 SearchCoordinationLimits::default(),
1027 source_snapshot,
1028 |_| Ok(true),
1029 || Ok(()),
1030 )
1031 .unwrap();
1032 let replaced_path = match replaced {
1033 SearchPublicationOutcome::Published { artifact, .. } => artifact.path,
1034 other @ SearchPublicationOutcome::Reused(_) => {
1035 panic!("unexpected replacement outcome: {other:?}")
1036 }
1037 };
1038 assert_ne!(replaced_path, first_path);
1039 assert!(first_path.exists());
1040
1041 let key = SearchArtifactKey::vector("Person", "semantic").unwrap();
1042 let before_error = current_search_artifact(dir.path(), &key)
1043 .unwrap()
1044 .unwrap()
1045 .path;
1046 let mismatch = upsert_published_vector(
1047 dir.path(),
1048 "Person",
1049 "semantic",
1050 uuid(2),
1051 &[1.0, 0.0, 0.0],
1052 33,
1053 limits,
1054 SearchCoordinationLimits::default(),
1055 source_snapshot,
1056 |_| Ok(true),
1057 || Ok(()),
1058 );
1059 assert!(matches!(
1060 mismatch,
1061 Err(SearchArtifactError::InvalidSelector {
1062 field: "vector",
1063 ..
1064 })
1065 ));
1066 assert_eq!(
1067 current_search_artifact(dir.path(), &key)
1068 .unwrap()
1069 .unwrap()
1070 .path,
1071 before_error
1072 );
1073 }
1074
1075 #[test]
1076 fn atomic_upsert_requires_current_label_membership() {
1077 let dir = TempDir::new().unwrap();
1078 let result = upsert_published_vector(
1079 dir.path(),
1080 "Person",
1081 "semantic",
1082 uuid(1),
1083 &[1.0],
1084 1,
1085 VectorStoreLimits::default(),
1086 SearchCoordinationLimits::default(),
1087 || {
1088 Ok(SearchSourceSnapshot {
1089 generation: 1,
1090 fingerprint: format!("gf-fnv1a256:{:064x}", 1),
1091 })
1092 },
1093 |_| Ok(false),
1094 || Ok(()),
1095 );
1096 assert!(matches!(
1097 result,
1098 Err(SearchArtifactError::InvalidSelector { field: "node", .. })
1099 ));
1100 assert!(
1101 current_search_artifact(
1102 dir.path(),
1103 &SearchArtifactKey::vector("Person", "semantic").unwrap()
1104 )
1105 .unwrap()
1106 .is_none()
1107 );
1108 }
1109
1110 #[test]
1111 fn cosine_keeps_negative_scores_filters_orphans_and_breaks_ties_by_uuid() {
1112 let rows = vec![
1113 row(1, &[1.0, 0.0], 1),
1114 row(2, &[1.0, 0.0], 2),
1115 row(3, &[0.0, 1.0], 3),
1116 row(4, &[-1.0, 0.0], 4),
1117 ];
1118 let eligible = BTreeSet::from([uuid(1), uuid(2), uuid(4)]);
1119 let hits = exact_cosine_search(
1120 &rows,
1121 &[1.0, 0.0],
1122 &eligible,
1123 10,
1124 VectorStoreLimits::default(),
1125 || Ok(()),
1126 )
1127 .unwrap();
1128 assert_eq!(
1129 hits,
1130 vec![
1131 VectorSearchHit {
1132 node_uuid: uuid(1),
1133 score: 1.0
1134 },
1135 VectorSearchHit {
1136 node_uuid: uuid(2),
1137 score: 1.0
1138 },
1139 VectorSearchHit {
1140 node_uuid: uuid(4),
1141 score: -1.0
1142 },
1143 ]
1144 );
1145 }
1146
1147 #[test]
1148 fn insertion_order_and_result_limit_do_not_change_ranking() {
1149 let canonical = vec![
1150 row(1, &[1.0, 1.0], 1),
1151 row(2, &[2.0, 2.0], 2),
1152 row(3, &[0.0, 1.0], 3),
1153 ];
1154 let eligible = canonical.iter().map(|row| row.node_uuid).collect();
1155 let hits = exact_cosine_search(
1156 &canonical,
1157 &[1.0, 1.0],
1158 &eligible,
1159 2,
1160 VectorStoreLimits::default(),
1161 || Ok(()),
1162 )
1163 .unwrap();
1164 assert_eq!(
1165 hits.iter().map(|hit| hit.node_uuid).collect::<Vec<_>>(),
1166 vec![uuid(1), uuid(2)]
1167 );
1168 }
1169
1170 #[test]
1171 fn malformed_primary_data_is_never_repaired_or_silently_sorted() {
1172 let limits = VectorStoreLimits::default();
1173 let duplicate = vec![row(1, &[1.0], 1), row(1, &[2.0], 2)];
1174 assert!(matches!(
1175 exact_cosine_search(
1176 &duplicate,
1177 &[1.0],
1178 &BTreeSet::from([uuid(1)]),
1179 1,
1180 limits,
1181 || Ok(())
1182 ),
1183 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
1184 ));
1185
1186 let dir = TempDir::new().unwrap();
1187 std::fs::write(dir.path().join(VECTOR_DATA_FILE), b"not parquet").unwrap();
1188 assert!(matches!(
1189 read_vector_snapshot(dir.path(), 1, limits, || Ok(())),
1190 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
1191 ));
1192 }
1193
1194 #[test]
1195 fn limits_and_cancellation_return_no_partial_hits() {
1196 let rows = vec![row(1, &[1.0, 0.0], 1), row(2, &[0.0, 1.0], 2)];
1197 let eligible = BTreeSet::from([uuid(1), uuid(2)]);
1198 let limited = VectorStoreLimits {
1199 vector_cells: 3,
1200 ..VectorStoreLimits::default()
1201 };
1202 assert!(matches!(
1203 exact_cosine_search(&rows, &[1.0, 0.0], &eligible, 2, limited, || Ok(())),
1204 Err(SearchArtifactError::ResourceExhausted {
1205 resource: "vector_cells",
1206 ..
1207 })
1208 ));
1209
1210 let checks = Cell::new(0);
1211 let result = exact_cosine_search(
1212 &rows,
1213 &[1.0, 0.0],
1214 &eligible,
1215 2,
1216 VectorStoreLimits::default(),
1217 || {
1218 checks.set(checks.get() + 1);
1219 if checks.get() == 2 {
1220 Err(SearchArtifactError::Cancelled)
1221 } else {
1222 Ok(())
1223 }
1224 },
1225 );
1226 assert!(matches!(result, Err(SearchArtifactError::Cancelled)));
1227 }
1228
1229 #[test]
1230 fn every_vector_row_search_and_persistence_limit_is_named_and_atomic() {
1231 let base = VectorStoreLimits::default();
1232 let rows = vec![row(1, &[1.0, 0.0], 1), row(2, &[0.0, 1.0], 2)];
1233 let eligible = BTreeSet::from([uuid(1), uuid(2)]);
1234
1235 for (limits, resource) in [
1236 (
1237 VectorStoreLimits {
1238 stored_vectors: 1,
1239 ..base
1240 },
1241 "stored_vectors",
1242 ),
1243 (
1244 VectorStoreLimits {
1245 vector_cells: 3,
1246 ..base
1247 },
1248 "vector_cells",
1249 ),
1250 ] {
1251 assert!(matches!(
1252 validate_rows(&rows, Some(2), limits, Path::new("<memory>")),
1253 Err(SearchArtifactError::ResourceExhausted { resource: actual, .. }) if actual == resource
1254 ));
1255 }
1256 assert!(matches!(
1257 exact_cosine_search(&rows, &[1.0, 0.0], &eligible, 0, base, || Ok(())),
1258 Err(SearchArtifactError::InvalidSelector { field: "limit", .. })
1259 ));
1260 assert!(matches!(
1261 exact_cosine_search(
1262 &rows,
1263 &[1.0, 0.0],
1264 &eligible,
1265 2,
1266 VectorStoreLimits {
1267 eligible_nodes: 1,
1268 ..base
1269 },
1270 || Ok(())
1271 ),
1272 Err(SearchArtifactError::ResourceExhausted {
1273 resource: "eligible_nodes",
1274 ..
1275 })
1276 ));
1277 assert!(matches!(
1278 exact_cosine_search(
1279 &rows,
1280 &[1.0, 0.0],
1281 &eligible,
1282 2,
1283 VectorStoreLimits { results: 1, ..base },
1284 || Ok(())
1285 ),
1286 Err(SearchArtifactError::ResourceExhausted {
1287 resource: "search_results",
1288 ..
1289 })
1290 ));
1291
1292 let mut at_limit = vec![row(1, &[1.0], 1)];
1293 let before = at_limit.clone();
1294 assert!(matches!(
1295 apply_vector_upsert(
1296 &mut at_limit,
1297 uuid(2),
1298 &[1.0],
1299 2,
1300 VectorStoreLimits {
1301 stored_vectors: 1,
1302 ..base
1303 }
1304 ),
1305 Err(SearchArtifactError::ResourceExhausted {
1306 resource: "stored_vectors",
1307 ..
1308 })
1309 ));
1310 assert_eq!(at_limit, before);
1311 }
1312
1313 #[test]
1314 fn vector_row_validation_rejects_dimension_order_duplicate_and_zero_norm() {
1315 let limits = VectorStoreLimits::default();
1316 for (rows, fragment) in [
1317 (vec![row(1, &[1.0, 0.0], 1), row(2, &[1.0], 2)], "dimension"),
1318 (vec![row(2, &[1.0], 1), row(1, &[1.0], 2)], "not sorted"),
1319 (
1320 vec![row(1, &[1.0], 1), row(1, &[1.0], 2)],
1321 "duplicate node_uuid",
1322 ),
1323 (vec![row(1, &[0.0], 1)], "zero"),
1324 ] {
1325 let error = validate_rows(&rows, None, limits, Path::new("fixture")).unwrap_err();
1326 assert!(error.to_string().contains(fragment), "{error}");
1327 }
1328 assert!(validate_rows(&[], None, limits, Path::new("fixture")).is_ok());
1329 assert!(matches!(
1330 vector_schema(0, limits),
1331 Err(SearchArtifactError::InvalidSelector {
1332 field: "vector",
1333 ..
1334 })
1335 ));
1336 assert!(matches!(
1337 vector_schema(
1338 2,
1339 VectorStoreLimits {
1340 dimensions: 1,
1341 ..limits
1342 }
1343 ),
1344 Err(SearchArtifactError::ResourceExhausted {
1345 resource: "vector_dimensions",
1346 ..
1347 })
1348 ));
1349 }
1350
1351 #[test]
1352 fn snapshot_read_write_failures_preserve_structured_error_kinds() {
1353 let missing = TempDir::new().unwrap();
1354 assert!(matches!(
1355 read_vector_snapshot(missing.path(), 1, VectorStoreLimits::default(), || Ok(())),
1356 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
1357 ));
1358
1359 let cancelled_write = TempDir::new().unwrap();
1360 assert!(matches!(
1361 write_vector_snapshot(
1362 cancelled_write.path(),
1363 &[row(1, &[1.0], 1)],
1364 1,
1365 VectorStoreLimits::default(),
1366 || Err(SearchArtifactError::Cancelled)
1367 ),
1368 Err(SearchArtifactError::Cancelled)
1369 ));
1370 assert!(!cancelled_write.path().join(VECTOR_DATA_FILE).exists());
1371
1372 let dir = TempDir::new().unwrap();
1373 write_vector_snapshot(
1374 dir.path(),
1375 &[row(1, &[1.0], 1)],
1376 1,
1377 VectorStoreLimits::default(),
1378 || Ok(()),
1379 )
1380 .unwrap();
1381 assert!(matches!(
1382 read_vector_snapshot(
1383 dir.path(),
1384 1,
1385 VectorStoreLimits {
1386 parquet_bytes: 1,
1387 ..VectorStoreLimits::default()
1388 },
1389 || Ok(())
1390 ),
1391 Err(SearchArtifactError::ResourceExhausted {
1392 resource: "vector_parquet_bytes",
1393 ..
1394 })
1395 ));
1396 }
1397
1398 #[test]
1399 fn wave10_private_vector_error_contracts_are_structured() {
1400 let limits = VectorStoreLimits::default();
1401 assert!(matches!(
1402 dot_and_norm(&[1.0], &[0.0]),
1403 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
1404 ));
1405 assert!(matches!(
1406 validate_rows(
1407 &[row(1, &[1.0, 0.0], 1), row(2, &[1.0], 2)],
1408 Some(2),
1409 limits,
1410 Path::new("fixture")
1411 ),
1412 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
1413 ));
1414 assert!(matches!(
1415 checked_cells(
1416 2,
1417 2,
1418 VectorStoreLimits {
1419 vector_cells: 3,
1420 ..limits
1421 }
1422 ),
1423 Err(SearchArtifactError::ResourceExhausted {
1424 resource: "vector_cells",
1425 ..
1426 })
1427 ));
1428 assert!(matches!(build("failed"), SearchArtifactError::Build(_)));
1429 assert!(matches!(
1430 io(
1431 "read",
1432 Path::new("fixture"),
1433 std::io::Error::other("failed")
1434 ),
1435 SearchArtifactError::Io { .. }
1436 ));
1437 }
1438
1439 #[test]
1440 fn wave12_snapshot_reader_rejects_schema_row_and_query_dimension_mismatches() {
1441 let dir = TempDir::new().unwrap();
1442 let limits = VectorStoreLimits::default();
1443 write_vector_snapshot(dir.path(), &[row(1, &[1.0, 0.0], 1)], 2, limits, || Ok(())).unwrap();
1444
1445 assert!(matches!(
1446 read_vector_snapshot(dir.path(), 1, limits, || Ok(())),
1447 Err(SearchArtifactError::CorruptPrimaryVectors { .. })
1448 ));
1449 assert!(matches!(
1450 read_vector_snapshot(
1451 dir.path(),
1452 2,
1453 VectorStoreLimits {
1454 stored_vectors: 0,
1455 ..limits
1456 },
1457 || Ok(())
1458 ),
1459 Err(SearchArtifactError::ResourceExhausted {
1460 resource: "stored_vectors",
1461 ..
1462 })
1463 ));
1464
1465 let key = SearchArtifactKey::vector("Person", "semantic").unwrap();
1466 let artifact = PublishedSearchArtifact {
1467 path: dir.path().to_path_buf(),
1468 manifest: SearchManifest::for_key(
1469 &key,
1470 VECTOR_BACKEND_VERSION,
1471 VECTOR_CONTRACT_VERSION,
1472 Some(2),
1473 &source_snapshot().unwrap(),
1474 true,
1475 )
1476 .unwrap(),
1477 };
1478 assert!(matches!(
1479 search_published_vectors(&artifact, &[1.0], &BTreeSet::new(), 1, limits, || Ok(())),
1480 Err(SearchArtifactError::InvalidSelector {
1481 field: "vector",
1482 ..
1483 })
1484 ));
1485 }
1486}