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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

//! End-to-end engine tests for the FastScan shared PQ codebook (Issue
//! #920 sub-item 2).
//!
//! FastScan mirror of `engine_pq_codebook_test.rs`: with a
//! `ProductQuantizationFastScan` quantizer and a `pq_codebook_path` on
//! the schema, `Engine::train_pq_codebook` must train a k=16 codebook,
//! a commit before training must hard-error (before this fix the path
//! was silently ignored and every segment retrained k-means), and a
//! reopened engine must encode against the shared codebook.

#![cfg(feature = "pq-fastscan")]

use laurus::storage::file::FileStorageConfig;
use laurus::storage::{StorageConfig, StorageFactory};
use laurus::vector::HnswOption;
use laurus::vector::Vector;
use laurus::vector::core::distance::DistanceMetric;
use laurus::vector::core::quantization::QuantizationMethod;
use laurus::vector::index::pq_codebook::read_pq_codebook;
use laurus::{DataValue, Document, Engine, FieldOption, Schema};
use tempfile::TempDir;

const DIM: usize = 32;
const M: usize = 4;

/// Deterministic pseudo-random vectors (same LCG as
/// `engine_pq_codebook_test.rs`); distinct seeds give distinct corpora.
fn make_vectors(count: usize, seed: u64) -> Vec<Vector> {
    let mut state = seed;
    (0..count)
        .map(|_| {
            let data: Vec<f32> = (0..DIM)
                .map(|_| {
                    state = state
                        .wrapping_mul(6_364_136_223_846_793_005)
                        .wrapping_add(1_442_695_040_888_963_407);
                    ((state >> 33) as f32 / u32::MAX as f32) * 2.0 - 1.0
                })
                .collect();
            Vector::new(data)
        })
        .collect()
}

/// Schema with one HNSW + FastScan field pointing at a shared codebook.
fn fastscan_schema(pq_codebook_path: Option<&str>) -> Schema {
    Schema::builder()
        .add_field(
            "embedding",
            FieldOption::Hnsw(HnswOption {
                dimension: DIM,
                distance: DistanceMetric::Euclidean,
                m: 8,
                ef_construction: 32,
                quantizer: QuantizationMethod::ProductQuantizationFastScan { subvector_count: M },
                pq_codebook_path: pq_codebook_path.map(str::to_string),
                ..HnswOption::default()
            }),
        )
        .build()
}

/// The full schema-level loop for FastScan: (1) a commit before training
/// hard-errors naming the training step, (2) `train_pq_codebook` trains
/// a k=16 codebook into the vector namespace, (3) a reopened engine
/// commits a tiny batch successfully — far below the FastScan min-train
/// threshold, so a silent fallback (the pre-#920 behavior of ignoring
/// the path, or the demotion guard firing) would produce a different
/// outcome.
#[tokio::test(flavor = "multi_thread")]
async fn fastscan_train_then_reopen_commits_with_the_shared_codebook() -> laurus::Result<()> {
    let dir = TempDir::new().unwrap();
    let storage = StorageFactory::create(StorageConfig::File(FileStorageConfig::new(dir.path())))?;

    // (1) Before training: the configured-but-missing codebook must fail
    // the commit loudly (pre-#920 this was silently ignored for FastScan).
    {
        let engine = Engine::new(storage.clone(), fastscan_schema(Some("embedding.pqcb"))).await?;
        let doc = Document::builder()
            .add_field(
                "embedding",
                DataValue::Vector(make_vectors(1, 7)[0].data.to_vec()),
            )
            .build();
        engine.put_document("doc-early", doc).await?;
        let err = engine.commit().await.unwrap_err();
        assert!(
            err.to_string().contains("train"),
            "an untrained configured codebook must point at the training step: {err}"
        );
    }

    // (2) Train — the FastScan quantizer variant must yield k=16.
    let fresh_dir = TempDir::new().unwrap();
    let fresh_storage = StorageFactory::create(StorageConfig::File(FileStorageConfig::new(
        fresh_dir.path(),
    )))?;
    let trainer = Engine::new(
        fresh_storage.clone(),
        fastscan_schema(Some("embedding.pqcb")),
    )
    .await?;
    let info = trainer.train_pq_codebook("embedding", &make_vectors(400, 0x0F0F), None)?;
    assert_eq!(info.path, "embedding.pqcb");
    assert_eq!(info.subvector_count, M);
    assert_eq!(
        info.centroids, 16,
        "a FastScan field must train a k=16 codebook, not the standard 256"
    );
    assert_eq!(info.sub_dimension, DIM / M);
    assert!(fresh_storage.file_exists("vector/embedding.pqcb"));
    drop(trainer);

    // The persisted file itself carries k=16.
    let vector_ns =
        laurus::storage::prefixed::PrefixedStorage::new("vector", fresh_storage.clone());
    let loaded = read_pq_codebook(&vector_ns, "embedding.pqcb")?;
    assert_eq!(loaded.params.k, 16);

    // (3) A reopened engine picks the codebook up at open and commits a
    // tiny batch — 5 docs, far below PQ_FASTSCAN_MIN_TRAIN_VECTORS (16),
    // so this only succeeds because the shared codebook exempts the
    // segment from both the demotion guard and per-segment training.
    let engine = Engine::new(
        fresh_storage.clone(),
        fastscan_schema(Some("embedding.pqcb")),
    )
    .await?;
    for (i, v) in make_vectors(5, 0xBEEF).into_iter().enumerate() {
        let doc = Document::builder()
            .add_field("embedding", DataValue::Vector(v.data.to_vec()))
            .build();
        engine.put_document(&format!("d{i}"), doc).await?;
    }
    engine.commit().await?;

    Ok(())
}

/// A k=256 (standard PQ) codebook configured on a FastScan field must
/// fail the commit loudly with the k mismatch named — not encode
/// garbage or silently retrain.
#[tokio::test(flavor = "multi_thread")]
async fn k256_codebook_on_fastscan_field_fails_the_commit() -> laurus::Result<()> {
    let dir = TempDir::new().unwrap();
    let storage = StorageFactory::create(StorageConfig::File(FileStorageConfig::new(dir.path())))?;

    // Train a k=256 codebook by declaring the field as standard PQ first.
    let pq_schema = Schema::builder()
        .add_field(
            "embedding",
            FieldOption::Hnsw(HnswOption {
                dimension: DIM,
                distance: DistanceMetric::Euclidean,
                m: 8,
                ef_construction: 32,
                quantizer: QuantizationMethod::ProductQuantization { subvector_count: M },
                pq_codebook_path: Some("embedding.pqcb".to_string()),
                ..HnswOption::default()
            }),
        )
        .build();
    let trainer = Engine::new(storage.clone(), pq_schema).await?;
    let info = trainer.train_pq_codebook("embedding", &make_vectors(400, 0x0F0F), None)?;
    assert_eq!(info.centroids, 256);
    drop(trainer);

    // Reopen with the SAME codebook file but a FastScan quantizer.
    let engine = Engine::new(storage.clone(), fastscan_schema(Some("embedding.pqcb"))).await?;
    for (i, v) in make_vectors(20, 0xBEEF).into_iter().enumerate() {
        let doc = Document::builder()
            .add_field("embedding", DataValue::Vector(v.data.to_vec()))
            .build();
        engine.put_document(&format!("d{i}"), doc).await?;
    }
    let err = engine.commit().await.unwrap_err();
    let message = err.to_string();
    assert!(
        message.contains("k = 256") && message.contains("k = 16"),
        "the k mismatch must be named: {message}"
    );

    Ok(())
}