laurus 0.10.0

Unified search library for lexical, vector, and semantic retrieval
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

//! Integration tests for searcher warmup (Issue #677).
//!
//! `VectorStore::warmup` eagerly builds the cached searcher and invokes the
//! searcher's `warmup`, which for HNSW pre-faults on-disk (`Mmap` / `OnDemand`)
//! vector data into the OS page cache. Warmup must be transparent (results
//! unchanged), idempotent, and must exercise the HNSW `OnDemand` path without
//! error. Cold-start *timing* is environment-dependent and intentionally not
//! asserted.

use std::collections::BTreeSet;
use std::sync::Arc;

use laurus::PrecomputedEmbedder;
use laurus::storage::file::FileStorageConfig;
use laurus::storage::memory::{MemoryStorage, MemoryStorageConfig};
use laurus::storage::{StorageConfig, StorageFactory};
use laurus::vector::core::field::HnswOption;
use laurus::vector::store::config::VectorFieldConfig;
use laurus::vector::{
    DistanceMetric, FieldOption, VectorIndexConfig, VectorSearchRequestBuilder, VectorStore,
};
use laurus::{DataValue, Document};

fn hnsw_field(dimension: usize) -> VectorFieldConfig {
    VectorFieldConfig {
        vector: Some(FieldOption::Hnsw(HnswOption {
            dimension,
            distance: DistanceMetric::Cosine,
            m: 16,
            ef_construction: 100,
            default_ef_search: None,
            base_weight: 1.0,
            quantizer: Default::default(),
            rerank_storage: None,
            embedder: None,
            pq_codebook_path: None,
        })),
        lexical: None,
    }
}

fn config(dimension: usize) -> VectorIndexConfig {
    VectorIndexConfig::builder()
        .embedder(Arc::new(PrecomputedEmbedder::new()))
        .field("v", hnsw_field(dimension))
        .build()
        .unwrap()
}

async fn add_docs(store: &VectorStore, vectors: Vec<Vec<f32>>) {
    for (i, vec) in vectors.into_iter().enumerate() {
        let doc = Document::builder()
            .add_field("v", DataValue::Vector(vec))
            .build();
        store
            .upsert_document_by_internal_id((i + 1) as u64, doc)
            .await
            .unwrap();
    }
    store.commit().await.unwrap();
}

fn search_ids(store: &VectorStore, query: Vec<f32>) -> BTreeSet<u64> {
    let request = VectorSearchRequestBuilder::new()
        .add_vector("v", query)
        .limit(3)
        .build();
    store
        .search(request)
        .unwrap()
        .hits
        .into_iter()
        .map(|h| h.doc_id)
        .collect()
}

/// Warmup is transparent (results unchanged) and idempotent on an in-memory
/// (Eager) HNSW store — exercises `VectorStore::warmup` and the `Owned*`
/// early-return in `HnswSearcher::warmup`.
#[tokio::test(flavor = "multi_thread")]
async fn warmup_is_transparent_and_idempotent() {
    let storage = Arc::new(MemoryStorage::new(MemoryStorageConfig::default()));
    let store = VectorStore::new(storage, config(3)).unwrap();
    add_docs(
        &store,
        vec![
            vec![1.0, 0.0, 0.0],
            vec![0.0, 1.0, 0.0],
            vec![0.0, 0.0, 1.0],
            vec![1.0, 1.0, 0.0],
        ],
    )
    .await;

    let query = vec![1.0, 0.1, 0.0];
    let before = search_ids(&store, query.clone());
    assert!(!before.is_empty(), "baseline search must return hits");

    store.warmup().unwrap();
    store.warmup().unwrap(); // idempotent

    let after = search_ids(&store, query);
    assert_eq!(after, before, "warmup must not change search results");
}

/// Warmup exercises the HNSW `OnDemand` (Mmap) page-fault loop without error
/// and leaves search working. `use_mmap` is forced on so the storage is `Lazy`
/// on every platform (the default is platform-specific).
#[tokio::test(flavor = "multi_thread")]
async fn warmup_on_demand_hnsw_pre_faults() {
    let dir = tempfile::tempdir().unwrap();
    let mut file_config = FileStorageConfig::new(dir.path());
    file_config.use_mmap = true; // force Lazy / OnDemand on all platforms
    let storage = StorageFactory::create(StorageConfig::File(file_config)).unwrap();
    let store = VectorStore::new(storage, config(3)).unwrap();
    add_docs(
        &store,
        vec![
            vec![1.0, 0.0, 0.0],
            vec![0.0, 1.0, 0.0],
            vec![0.0, 0.0, 1.0],
        ],
    )
    .await;

    // Warmup before any query: builds the cached searcher and runs the
    // OnDemand page-fault pass.
    store.warmup().unwrap();

    let hits = search_ids(&store, vec![1.0, 0.1, 0.0]);
    assert!(
        !hits.is_empty(),
        "search after warmup on an OnDemand store must return hits"
    );
}