From 8e13c84e95e84db29f97e866ab360a37b606d0ec Mon Sep 17 00:00:00 2001 From: Minoru OSUKA Date: Sun, 2 Aug 2026 21:24:23 +0900 Subject: [PATCH] feat(vector/search): Stage-2 rerank for Flat/IVF via the shared pipeline (closes #932) --- .../src/concepts/indexing/vector_indexing.md | 4 +- docs/ja/src/laurus-cli/schema_format.md | 10 +- docs/ja/src/laurus-server/grpc_api.md | 2 +- docs/src/concepts/indexing/vector_indexing.md | 4 +- docs/src/laurus-cli/schema_format.md | 14 +- docs/src/laurus-server/grpc_api.md | 2 +- laurus/src/vector/index/flat/reader.rs | 26 +++ laurus/src/vector/index/flat/searcher.rs | 77 ++++++- laurus/src/vector/index/flat/writer.rs | 25 +++ laurus/src/vector/index/ivf/reader.rs | 26 +++ laurus/src/vector/index/ivf/searcher.rs | 85 ++++++- laurus/src/vector/index/ivf/writer.rs | 29 ++- laurus/src/vector/index/rerank_sidecar.rs | 74 ++++++ laurus/src/vector/store.rs | 8 + laurus/tests/vector_rerank_flat_ivf_test.rs | 210 ++++++++++++++++++ 15 files changed, 561 insertions(+), 35 deletions(-) create mode 100644 laurus/tests/vector_rerank_flat_ivf_test.rs diff --git a/docs/ja/src/concepts/indexing/vector_indexing.md b/docs/ja/src/concepts/indexing/vector_indexing.md index f381342d..29dc6c15 100644 --- a/docs/ja/src/concepts/indexing/vector_indexing.md +++ b/docs/ja/src/concepts/indexing/vector_indexing.md @@ -262,7 +262,9 @@ Stage 1 ではベクトルを int8 のみで保持します。グラフ検索は ベクトルで再スコアする。 3. 新しいランキングを `top_k` に切り詰めて返す。 -Stage 2 はフィールド単位で +Issue #932 以降、同じ sidecar 機構は共有 `RerankPipeline`（#650）を +通じて Flat / IVF でも機能します（以下の説明は Stage 2 の元ホストである +HNSW を例にしています）。Stage 2 はフィールド単位で [`HnswOption.rerank_storage`](../../laurus-cli/schema_format.md#rerank-storage) で opt-in します: diff --git a/docs/ja/src/laurus-cli/schema_format.md b/docs/ja/src/laurus-cli/schema_format.md index 8ebf1d67..abf73903 100644 --- a/docs/ja/src/laurus-cli/schema_format.md +++ b/docs/ja/src/laurus-cli/schema_format.md @@ -209,7 +209,7 @@ base_weight = 1.0 | `distance` | `string` | `"Cosine"` | 距離メトリクス（[距離メトリクス](#距離メトリクス)を参照） | | `base_weight` | `float` | `1.0` | ハイブリッド検索のスコア融合における重み | | `quantizer` | `object` | `"Scalar8Bit"` | 量子化方式（[量子化](#量子化)を参照）。必須。デフォルトは Issue #481 Stage 1 で導入された int8 形式を保つ。 | -| `rerank_storage` | `string` | *（省略）* | [Rerank Storage](#rerank-storage) 用に予約。現状 sidecar を書き出すのは HNSW writer のみで、Flat / IVF はスキーマの対称性のためにフィールドを受け付けるが sidecar の書き出し・読み込みは行わない。 | +| `rerank_storage` | `string` | *（省略）* | Stage 2 rerank sidecar（[Rerank Storage](#rerank-storage)）。#932 以降、3 つのベクトルインデックスタイプすべてでサポート。`"F32"` でフィールド単位の f32 sidecar を有効化し、検索時に int8 候補を元のベクトルで再スコアできる。 | #### Ivf @@ -232,7 +232,7 @@ base_weight = 1.0 | `n_probe` | `integer` | `1` | クエリ時に検索するクラスタ数。大きいほど再現率が向上するが遅くなる | | `base_weight` | `float` | `1.0` | ハイブリッド検索のスコア融合における重み | | `quantizer` | `object` | `"Scalar8Bit"` | 量子化方式（[量子化](#量子化)を参照）。必須。デフォルトは Issue #481 Stage 1 で導入された int8 形式を保つ。 | -| `rerank_storage` | `string` | *（省略）* | [Rerank Storage](#rerank-storage) 用に予約。現状 sidecar を書き出すのは HNSW writer のみで、Flat / IVF はスキーマの対称性のためにフィールドを受け付けるが sidecar の書き出し・読み込みは行わない。 | +| `rerank_storage` | `string` | *（省略）* | Stage 2 rerank sidecar（[Rerank Storage](#rerank-storage)）。#932 以降、3 つのベクトルインデックスタイプすべてでサポート。`"F32"` でフィールド単位の f32 sidecar を有効化し、検索時に int8 候補を元のベクトルで再スコアできる。 | > **注意:** Hnsw および Flat とは異なり、Ivf の `dimension` フィールドは**必須**であり、デフォルト値はありません。 @@ -318,9 +318,11 @@ commit は大幅に高速化し、小さな per-commit segment も PQ を維持 ## Rerank Storage 任意の Stage 2 sidecar（Issue #481）。元の完全精度ベクトルを int8 -セグメントの隣に保持し、HNSW searcher が int8 で広めに候補を取得 +セグメントの隣に保持し、searcher が int8 で広めに候補を取得 （高速）してから上位 `top_k * rerank_factor` 件を完全な f32 値で -再スコア（高精度）できるようにします。 +再スコア（高精度）できるようにします。#932 以降、HNSW / Flat / +IVF の 3 タイプすべてでサポートされます（Flat / IVF の再スコアは +フィールド指定クエリに適用）。 sidecar はフィールド単位で `rerank_storage` で設定します: diff --git a/docs/ja/src/laurus-server/grpc_api.md b/docs/ja/src/laurus-server/grpc_api.md index 5c19a34b..90271f76 100644 --- a/docs/ja/src/laurus-server/grpc_api.md +++ b/docs/ja/src/laurus-server/grpc_api.md @@ -117,7 +117,7 @@ message AnalyzerDefinition { `NONE`（量子化なし）は Issue #481 Stage 1 で廃止されました。proto enum 値 0（`QUANTIZATION_METHOD_NONE`）は wire 互換のため予約されていますが、サーバ側で受信すると `Default::default()`（`SCALAR_8BIT`）にフォールバックします。 -**Rerank storage:** オプションの `rerank_storage` フィールド（enum `RerankStorageKind`: `UNSPECIFIED` = サイドカーなし、`F32`）は Stage-2 rerank サイドカー（Issue #481 / #793）を有効化します。HNSW フィールドで `F32` を設定すると、commit 時に完全精度の `.hnsw.f32` サイドカーを追加で書き出し、`rerank_factor` を指定した検索が int8 候補を元のベクトルで再スコアします。フィールドを省略（または `UNSPECIFIED`）すると Stage-1 の int8 のみのランキングになります。スキーマの round-trip 整合のため `FlatOption` / `IvfOption` にも保持されますが、これらのインデックスはまだサイドカーを出力しません。 +**Rerank storage:** オプションの `rerank_storage` フィールド（enum `RerankStorageKind`: `UNSPECIFIED` = サイドカーなし、`F32`）は Stage-2 rerank サイドカー（Issue #481 / #793）を有効化します。HNSW フィールドで `F32` を設定すると、commit 時に完全精度の `.hnsw.f32` サイドカーを追加で書き出し、`rerank_factor` を指定した検索が int8 候補を元のベクトルで再スコアします。フィールドを省略（または `UNSPECIFIED`）すると Stage-1 の int8 のみのランキングになります。#932 以降、サイドカーは 3 つのベクトルインデックスタイプ（HNSW / Flat / IVF）すべてで出力・利用されます（Flat / IVF の再スコアはフィールド指定クエリに適用）。 **共有 PQ codebook:** `HnswOption` のオプションフィールド `pq_codebook_path`（Issue #631）は、`laurus train pq-codebook` CLI コマンドで一度だけ学習するストレージ相対の共有 PQ codebook ファイルを指定します。設定すると segment は commit / merge のたびに k-means を再学習する代わりに、学習済み codebook で encode されます。`PRODUCT_QUANTIZATION` quantizer との組み合わせでのみ意味を持ち、設定済みで未学習の場合、commit は学習コマンドを示すエラーで失敗します（per-segment 学習への無言のフォールバック無し）。未設定なら per-segment 学習のままです。 diff --git a/docs/src/concepts/indexing/vector_indexing.md b/docs/src/concepts/indexing/vector_indexing.md index 87232c17..a444df3f 100644 --- a/docs/src/concepts/indexing/vector_indexing.md +++ b/docs/src/concepts/indexing/vector_indexing.md @@ -261,7 +261,9 @@ original full-precision vectors: [LRS1 sidecar](#lrs1-rerank-sidecar) (`*.hnsw.f32`). 3. The new ranking is truncated to `top_k` and returned. -Stage 2 is opt-in per field via +Since #932 the same sidecar mechanism serves Flat and IVF too (the +shared `RerankPipeline`, #650); the description below uses HNSW, the +original Stage-2 host. Stage 2 is opt-in per field via [`HnswOption.rerank_storage`](../../laurus-cli/schema_format.md#rerank-storage): ```rust diff --git a/docs/src/laurus-cli/schema_format.md b/docs/src/laurus-cli/schema_format.md index e4431da4..6fb3862f 100644 --- a/docs/src/laurus-cli/schema_format.md +++ b/docs/src/laurus-cli/schema_format.md @@ -209,7 +209,7 @@ base_weight = 1.0 | `distance` | `string` | `"Cosine"` | Distance metric (see [Distance Metrics](#distance-metrics)) | | `base_weight` | `float` | `1.0` | Scoring weight in hybrid search fusion | | `quantizer` | `object` | `"Scalar8Bit"` | Quantization method (see [Quantization](#quantization)). Mandatory; default keeps the int8 format introduced in Issue #481 Stage 1. | -| `rerank_storage` | `string` | *(omit)* | Reserved for [Rerank Storage](#rerank-storage). Currently emitted only by the HNSW writer; Flat / IVF accept the field for schema symmetry but do not yet write or consume the sidecar. | +| `rerank_storage` | `string` | *(omit)* | Optional Stage 2 rerank sidecar (see [Rerank Storage](#rerank-storage)); supported by all three vector index types since #932. `"F32"` enables the per-field f32 sidecar so search can rescore int8 candidates against the original vectors. | #### Ivf @@ -232,7 +232,7 @@ base_weight = 1.0 | `n_probe` | `integer` | `1` | Number of clusters to search at query time. Higher = better recall, slower | | `base_weight` | `float` | `1.0` | Scoring weight in hybrid search fusion | | `quantizer` | `object` | `"Scalar8Bit"` | Quantization method (see [Quantization](#quantization)). Mandatory; default keeps the int8 format introduced in Issue #481 Stage 1. | -| `rerank_storage` | `string` | *(omit)* | Reserved for [Rerank Storage](#rerank-storage). Currently emitted only by the HNSW writer; Flat / IVF accept the field for schema symmetry but do not yet write or consume the sidecar. | +| `rerank_storage` | `string` | *(omit)* | Optional Stage 2 rerank sidecar (see [Rerank Storage](#rerank-storage)); supported by all three vector index types since #932. `"F32"` enables the per-field f32 sidecar so search can rescore int8 candidates against the original vectors. | > **Note:** Unlike Hnsw and Flat, the `dimension` field in Ivf is **required** and has no default value. @@ -317,10 +317,12 @@ with an error naming the `laurus train pq-codebook` command to run ## Rerank Storage Optional Stage 2 sidecar (Issue #481) that keeps the original -full-precision vectors alongside the int8 segment so the HNSW -searcher can do a wide candidate fetch over int8 (cheap) and then -rescore the top `top_k * rerank_factor` candidates against the -exact f32 values (accurate). +full-precision vectors alongside the int8 segment so the searcher +can do a wide candidate fetch over int8 (cheap) and then rescore +the top `top_k * rerank_factor` candidates against the exact f32 +values (accurate). Supported by all three vector index types — +HNSW, Flat, and IVF (#932); on Flat/IVF the rescoring applies to +field-routed queries. The sidecar is configured per field with `rerank_storage`: diff --git a/docs/src/laurus-server/grpc_api.md b/docs/src/laurus-server/grpc_api.md index 0175c5ee..4978a80b 100644 --- a/docs/src/laurus-server/grpc_api.md +++ b/docs/src/laurus-server/grpc_api.md @@ -117,7 +117,7 @@ The `embedder` field in vector options specifies the name of an embedder defined `NONE` (no quantization) was removed in Issue #481 Stage 1. The proto enum value `0` (`QUANTIZATION_METHOD_NONE`) is kept as a wire-compat reservation; if the server receives it, it falls back to `SCALAR_8BIT` via `Default::default()`. -**Rerank storage:** the optional `rerank_storage` field (enum `RerankStorageKind`: `UNSPECIFIED` = no sidecar, `F32`) enables the Stage-2 rerank sidecar (Issue #481 / #793). When set to `F32` on an HNSW field, commit writes an extra full-precision `.hnsw.f32` sidecar so searches that set `rerank_factor` rescore int8 candidates against the original vectors. Omitting the field (or `UNSPECIFIED`) keeps Stage-1 int8-only ranking. The field is also carried on `FlatOption` / `IvfOption` for schema round-tripping, but those indexes do not emit a sidecar yet. +**Rerank storage:** the optional `rerank_storage` field (enum `RerankStorageKind`: `UNSPECIFIED` = no sidecar, `F32`) enables the Stage-2 rerank sidecar (Issue #481 / #793). When set to `F32` on an HNSW field, commit writes an extra full-precision `.hnsw.f32` sidecar so searches that set `rerank_factor` rescore int8 candidates against the original vectors. Omitting the field (or `UNSPECIFIED`) keeps Stage-1 int8-only ranking. Since #932 the sidecar is emitted and consumed by all three vector index types (HNSW / Flat / IVF); on Flat/IVF the rescoring applies to field-routed queries. **Shared PQ codebook:** the optional `pq_codebook_path` field on `HnswOption` (Issue #631) names a storage-relative shared PQ codebook file, trained once via the `laurus train pq-codebook` CLI command. Segments are then encoded against the pre-trained codebook instead of re-training k-means on every commit and merge. Only meaningful with a `PRODUCT_QUANTIZATION` quantizer; when set but not yet trained, commits fail with an error naming the training command (no silent fallback to per-segment training). Unset keeps per-segment training. diff --git a/laurus/src/vector/index/flat/reader.rs b/laurus/src/vector/index/flat/reader.rs index a51ed221..6bfd0cde 100644 --- a/laurus/src/vector/index/flat/reader.rs +++ b/laurus/src/vector/index/flat/reader.rs @@ -37,6 +37,10 @@ pub struct FlatVectorIndexReader { /// once at load so `doc_ids_for_field` returns a refcount-shared /// slice without re-cloning `vector_ids`. #405. vector_ids_by_field: std::collections::HashMap>, + /// Stage 2 rerank sidecar pool (Issue #481, extended to Flat by + /// #650 PR-2 / #932). `Some` only when the `.f32` sidecar exists and + /// the loading mode is Eager; absence keeps Stage 1 behavior. + rerank_storage: Option>, } /// Group `vector_ids` by field name into refcount-shared slices. @@ -263,6 +267,20 @@ impl FlatVectorIndexReader { }; let vector_ids_by_field = build_vector_ids_by_field(&vector_ids, &field_dict); + + // Stage 2 rerank sidecar (Issue #481, extended to Flat by #650 + // PR-2 / #932): loaded eagerly when present, mirroring HNSW. The + // pool's positions pair with `vector_ids` (the record order the + // writer also used for the sidecar payload) — an identity mapping. + // Lazy mode skips the sidecar to honor its memory-savings promise. + let rerank_storage = crate::vector::index::rerank_sidecar::load_rerank_sidecar( + storage.as_ref(), + &file_name, + dimension, + &vector_ids, + &field_dict, + )?; + Ok(Self { vectors, vector_ids, @@ -271,9 +289,17 @@ impl FlatVectorIndexReader { distance_metric, deletion_bitmap: None, vector_ids_by_field, + rerank_storage, }) } + /// Borrow the optional Stage 2 rerank storage pool (#932). + pub fn rerank_storage( + &self, + ) -> Option<&Arc> { + self.rerank_storage.as_ref() + } + pub fn set_deletion_bitmap(&mut self, bitmap: Arc) { self.deletion_bitmap = Some(bitmap); } diff --git a/laurus/src/vector/index/flat/searcher.rs b/laurus/src/vector/index/flat/searcher.rs index cec3ad68..e0407a6a 100644 --- a/laurus/src/vector/index/flat/searcher.rs +++ b/laurus/src/vector/index/flat/searcher.rs @@ -33,14 +33,13 @@ impl VectorIndexSearcher for FlatVectorSearcher { // filter. The store's post-filter still runs but becomes a no-op for // the already-filtered results, so recall is unchanged. - // Issue #481 Stage 2 (rerank) -- API surface only in Stage 1. - if request.params.rerank_factor.is_some() { - return Err(crate::error::LaurusError::NotImplemented( - "Two-stage rerank (Issue #481 Stage 2) is not yet implemented. \ - Pass rerank_factor = None for the Stage 1 quantized search." - .to_string(), - )); - } + // Issue #481 Stage 2 rerank (extended to Flat by #650 PR-2 / + // #932): honored on the field-filtered path below when the reader + // has the `.f32` sidecar loaded. On the unfiltered path (or when + // the sidecar is absent) `rerank_factor` silently falls back to + // Stage 1 ranking — the same convention as HNSW, where a missing + // prerequisite (here: a single field to key the sidecar position + // index) cannot be recovered at query time. let start = Timer::now(); let mut results = VectorIndexQueryResults::new(); @@ -137,6 +136,50 @@ impl VectorIndexSearcher for FlatVectorSearcher { // would make top-k membership arbitrary; distance stays precise. candidates.sort_unstable_by(|a, b| a.2.total_cmp(&b.2).then(a.0.cmp(&b.0))); + // Stage 2 (Issue #481 / #932): run the shared rerank pipeline + // over the sorted quantized candidates. When it applies, the + // survivors carry exact f32 distances — stamped below as the + // fan-out's `score_basis` (#927). + let mut rerank_applied = false; + if let (Some(factor), Some(pool)) = ( + request.params.rerank_factor, + flat_reader.and_then(|r| r.rerank_storage()), + ) { + use crate::vector::search::rerank::{ + F32SidecarStage, RerankCandidates, RerankPipeline, + }; + let pipeline = RerankPipeline::new( + vec![Box::new(F32SidecarStage::new( + std::sync::Arc::clone(pool), + field_name, + metric, + ))], + vec![factor], + ); + let mut rc = RerankCandidates::with_capacity(candidates.len()); + for (doc_id, _, distance, _) in &candidates { + rc.push(*doc_id, *distance); + } + rerank_applied = pipeline.run(&request.query, &mut rc, request.params.top_k)?; + if rerank_applied { + // Rebuild on the exact basis; vectors are re-fetched in + // the result loop only when `include_vectors` asks. + candidates = rc + .doc_ids + .iter() + .zip(&rc.distances) + .map(|(&doc_id, &distance)| { + ( + doc_id, + metric.distance_to_similarity(distance), + distance, + Vector::new(Vec::new()), + ) + }) + .collect(); + } + } + let top_k = request.params.top_k.min(candidates.len()); for (doc_id, similarity, distance, vector) in candidates.into_iter().take(top_k) { if similarity < request.params.min_similarity { @@ -144,7 +187,13 @@ impl VectorIndexSearcher for FlatVectorSearcher { } let vector_output = if request.params.include_vectors { - Some(vector) + if rerank_applied { + // The rescored tuple carries a placeholder; fetch + // the real vector only for the final results. + self.index_reader.get_vector(doc_id, field_name)? + } else { + Some(vector) + } } else { None }; @@ -159,6 +208,16 @@ impl VectorIndexSearcher for FlatVectorSearcher { vector: vector_output, }); } + + // Issue #927: exact-f32 scores must be flagged so the + // multi-segment fan-out keeps them instead of overwriting + // with its dequantized rescore. + if rerank_applied { + results.query_metadata.insert( + crate::vector::search::searcher::SCORE_BASIS_METADATA_KEY.to_string(), + crate::vector::search::searcher::SCORE_BASIS_F32_RERANK.to_string(), + ); + } } else { // Unfiltered path: each doc may belong to a different field, so the // field name must travel with each candidate. diff --git a/laurus/src/vector/index/flat/writer.rs b/laurus/src/vector/index/flat/writer.rs index 29f478dc..78aee619 100644 --- a/laurus/src/vector/index/flat/writer.rs +++ b/laurus/src/vector/index/flat/writer.rs @@ -485,6 +485,31 @@ impl VectorIndexWriter for FlatIndexWriter { // power loss could surface a published-but-hollow segment file. output.close()?; storage.rename_file(&tmp_name, &file_name)?; + + // Stage 2 (Issue #481, extended to Flat by #650 PR-2 / #932): emit + // the optional LRS1 rerank sidecar alongside the main int8 segment. + // The payload follows `self.vectors` — the exact record emission + // order above — so the reader's (sidecar position) -> (record + // position) mapping is the identity, mirroring HNSW. + if let Some(rerank_kind) = self.index_config.rerank_storage { + let sidecar_name = format!("{}.f32", file_name); + let sidecar_tmp = format!("{}.f32.tmp", file_name); + let mut sidecar_out = storage.create_output(&sidecar_tmp)?; + let mut payload: Vec = + Vec::with_capacity(self.vectors.len() * self.index_config.dimension); + for (_, _, v) in &self.vectors { + payload.extend_from_slice(&v.data); + } + crate::vector::index::rerank_sidecar::write_sidecar( + &mut sidecar_out, + rerank_kind, + self.index_config.dimension as u32, + &payload, + )?; + sidecar_out.flush()?; + drop(sidecar_out); + storage.rename_file(&sidecar_tmp, &sidecar_name)?; + } Ok(()) } diff --git a/laurus/src/vector/index/ivf/reader.rs b/laurus/src/vector/index/ivf/reader.rs index 771f92f5..d6e119d1 100644 --- a/laurus/src/vector/index/ivf/reader.rs +++ b/laurus/src/vector/index/ivf/reader.rs @@ -46,6 +46,10 @@ pub struct IvfIndexReader { /// once at load so `doc_ids_for_field` returns a refcount-shared /// slice without re-cloning `vector_ids`. #405. vector_ids_by_field: HashMap>, + /// Stage 2 rerank sidecar pool (Issue #481, extended to IVF by + /// #650 PR-2 / #932). `Some` only when the `.f32` sidecar exists and + /// the loading mode is Eager; absence keeps Stage 1 behavior. + rerank_storage: Option>, } /// Group `vector_ids` by field name into refcount-shared slices. @@ -312,6 +316,20 @@ impl IvfIndexReader { }; let vector_ids_by_field = build_vector_ids_by_field(&vector_ids, &field_dict); + + // Stage 2 rerank sidecar (Issue #481, extended to IVF by #650 + // PR-2 / #932): loaded eagerly when present, mirroring HNSW. The + // pool's positions pair with `vector_ids` (the cluster-grouped + // record order the writer also used for the sidecar payload) — + // an identity mapping. + let rerank_storage = crate::vector::index::rerank_sidecar::load_rerank_sidecar( + storage.as_ref(), + &file_name, + dimension, + &vector_ids, + &field_dict, + )?; + Ok(Self { vectors, vector_ids, @@ -324,9 +342,17 @@ impl IvfIndexReader { field_dict, deletion_bitmap: None, vector_ids_by_field, + rerank_storage, }) } + /// Borrow the optional Stage 2 rerank storage pool (#932). + pub fn rerank_storage( + &self, + ) -> Option<&Arc> { + self.rerank_storage.as_ref() + } + pub fn set_deletion_bitmap(&mut self, bitmap: Arc) { self.deletion_bitmap = Some(bitmap); } diff --git a/laurus/src/vector/index/ivf/searcher.rs b/laurus/src/vector/index/ivf/searcher.rs index 08c7ac48..b53d1f4f 100644 --- a/laurus/src/vector/index/ivf/searcher.rs +++ b/laurus/src/vector/index/ivf/searcher.rs @@ -197,14 +197,12 @@ impl VectorIndexSearcher for IvfSearcher { // filter. The store's post-filter still runs but becomes a no-op for // the already-filtered results, so recall is unchanged. - // Issue #481 Stage 2 (rerank) -- API surface only in Stage 1. - if request.params.rerank_factor.is_some() { - return Err(crate::error::LaurusError::NotImplemented( - "Two-stage rerank (Issue #481 Stage 2) is not yet implemented. \ - Pass rerank_factor = None for the Stage 1 quantized search." - .to_string(), - )); - } + // Issue #481 Stage 2 rerank (extended to IVF by #650 PR-2 / + // #932): honored below when the query is field-routed and the + // reader has the `.f32` sidecar loaded. Otherwise `rerank_factor` + // silently falls back to Stage 1 ranking — the same convention as + // HNSW, where a missing prerequisite (here: a single field to key + // the sidecar position index) cannot be recovered at query time. let start = Timer::now(); let mut results = VectorIndexQueryResults::new(); @@ -294,6 +292,59 @@ impl VectorIndexSearcher for IvfSearcher { // membership arbitrary; distance stays precise at any range. candidates.sort_unstable_by(|a, b| a.3.total_cmp(&b.3).then(a.0.cmp(&b.0))); + // Stage 2 (Issue #481 / #932): run the shared rerank pipeline over + // the sorted quantized candidates when the query is field-routed + // (every candidate then shares one sidecar position index). The + // survivors carry exact f32 distances — stamped below as the + // fan-out's `score_basis` (#927). + let mut rerank_applied = false; + if let (Some(factor), Some(pool), Some(field_name)) = ( + request.params.rerank_factor, + ivf_reader.and_then(|r| r.rerank_storage()), + request.field_name.as_deref(), + ) { + use crate::vector::search::rerank::{ + F32SidecarStage, RerankCandidates, RerankPipeline, + }; + let pipeline = RerankPipeline::new( + vec![Box::new(F32SidecarStage::new( + std::sync::Arc::clone(pool), + field_name, + metric, + ))], + vec![factor], + ); + let mut rc = RerankCandidates::with_capacity(candidates.len()); + for (doc_id, _, _, distance, _) in &candidates { + rc.push(*doc_id, *distance); + } + rerank_applied = pipeline.run(&request.query, &mut rc, request.params.top_k)?; + if rerank_applied { + // Every candidate shares the routed field; recover its + // interned id for the rebuild (guaranteed present — the + // probe already filtered on it). + let fid = field_dict + .iter() + .position(|f| &**f == field_name) + .map(|i| i as u16) + .unwrap_or_default(); + candidates = rc + .doc_ids + .iter() + .zip(&rc.distances) + .map(|(&doc_id, &distance)| { + ( + doc_id, + fid, + metric.distance_to_similarity(distance), + distance, + Vector::new(Vec::new()), + ) + }) + .collect(); + } + } + // Take top_k results let candidates_len = candidates.len(); let top_k = request.params.top_k.min(candidates_len); @@ -306,7 +357,13 @@ impl VectorIndexSearcher for IvfSearcher { } let vector_output = if request.params.include_vectors { - Some(vector) + if rerank_applied { + // The rescored tuple carries a placeholder; fetch the + // real vector only for the final results. + self.index_reader.get_vector(doc_id, &field_name)? + } else { + Some(vector) + } } else { None }; @@ -322,6 +379,16 @@ impl VectorIndexSearcher for IvfSearcher { }); } + // Issue #927: exact-f32 scores must be flagged so the multi-segment + // fan-out keeps them instead of overwriting with its dequantized + // rescore. + if rerank_applied { + results.query_metadata.insert( + crate::vector::search::searcher::SCORE_BASIS_METADATA_KEY.to_string(), + crate::vector::search::searcher::SCORE_BASIS_F32_RERANK.to_string(), + ); + } + results.search_time_ms = start.elapsed().as_secs_f64() * 1000.0; results.candidates_examined = candidates_len; Ok(results) diff --git a/laurus/src/vector/index/ivf/writer.rs b/laurus/src/vector/index/ivf/writer.rs index 3719f3c2..fef9c915 100644 --- a/laurus/src/vector/index/ivf/writer.rs +++ b/laurus/src/vector/index/ivf/writer.rs @@ -1234,14 +1234,37 @@ impl VectorIndexWriter for IvfIndexWriter { } } - // Close with an fsync BEFORE the rename (mirrors HNSW's #882 review - // fix): a flush alone leaves the content in the page cache, so a - // power loss could surface a published-but-hollow segment file. // Close with an fsync BEFORE the rename (mirrors HNSW's #882 review // fix): a flush alone leaves the content in the page cache, so a // power loss could surface a published-but-hollow segment file. output.close()?; storage.rename_file(&tmp_name, &file_name)?; + + // Stage 2 (Issue #481, extended to IVF by #650 PR-2 / #932): emit + // the optional LRS1 rerank sidecar alongside the main int8 segment. + // The payload reuses `all_vectors` — materialized above in the same + // cluster-grouped flatten order the records were emitted in — so + // the reader's (sidecar position) -> (record position) mapping is + // the identity, mirroring HNSW. + if let Some(rerank_kind) = self.index_config.rerank_storage { + let sidecar_name = format!("{}.f32", file_name); + let sidecar_tmp = format!("{}.f32.tmp", file_name); + let mut sidecar_out = storage.create_output(&sidecar_tmp)?; + let mut payload: Vec = + Vec::with_capacity(all_vectors.len() * self.index_config.dimension); + for v in &all_vectors { + payload.extend_from_slice(&v.data); + } + crate::vector::index::rerank_sidecar::write_sidecar( + &mut sidecar_out, + rerank_kind, + self.index_config.dimension as u32, + &payload, + )?; + sidecar_out.flush()?; + drop(sidecar_out); + storage.rename_file(&sidecar_tmp, &sidecar_name)?; + } Ok(()) } diff --git a/laurus/src/vector/index/rerank_sidecar.rs b/laurus/src/vector/index/rerank_sidecar.rs index ec81e99d..fa055340 100644 --- a/laurus/src/vector/index/rerank_sidecar.rs +++ b/laurus/src/vector/index/rerank_sidecar.rs @@ -471,6 +471,80 @@ pub fn read_sidecar( } } +/// Load a segment's optional `.f32` rerank sidecar into a +/// [`crate::vector::index::rerank_storage::RerankStoragePool`] +/// (Issue #481; shared across HNSW/Flat/IVF readers by #650 PR-2 / #932). +/// +/// Lenient-if-absent: returns `Ok(None)` when no sidecar file exists or +/// when `storage` is in Lazy loading mode (the sidecar is skipped to honor +/// Lazy's memory-savings promise — Stage 2 segments opened Lazy silently +/// degrade to Stage 1). A present sidecar whose `dim`/`vector_count` +/// disagree with the segment fails loudly. +/// +/// The pool's positions pair with `vector_ids` (the segment's record +/// order, which the writer also used for the sidecar payload), giving an +/// identity (sidecar position) -> (record position) mapping. +/// +/// # Arguments +/// +/// * `storage` - The segment's storage backend. +/// * `file_name` - The main segment file name (`"{id}.hnsw"` / `.flat` / +/// `.ivf`); the sidecar is `"{file_name}.f32"`. +/// * `dimension` - The segment's vector dimension. +/// * `vector_ids` - Interned `(doc_id, field_id)` records in segment +/// order. +/// * `field_dict` - The segment's field-name dictionary. +/// +/// # Errors +/// +/// Forwards [`read_sidecar`] errors and fails on `dim` / `vector_count` +/// mismatches. +pub(crate) fn load_rerank_sidecar( + storage: &dyn crate::storage::Storage, + file_name: &str, + dimension: usize, + vector_ids: &[(u64, u16)], + field_dict: &[std::sync::Arc], +) -> Result>> { + if !matches!(storage.loading_mode(), crate::storage::LoadingMode::Eager) { + return Ok(None); + } + let sidecar_name = format!("{file_name}.f32"); + if !storage.file_exists(&sidecar_name) { + return Ok(None); + } + let mut sidecar_in = storage.open_input(&sidecar_name)?; + let sidecar_size = sidecar_in.size()?; + let (header, payload) = read_sidecar(&mut sidecar_in, sidecar_size)?; + if header.dim as usize != dimension { + return Err(LaurusError::InvalidOperation(format!( + "rerank sidecar dim mismatch: segment uses {dimension}, sidecar uses {}", + header.dim + ))); + } + if header.vector_count as usize != vector_ids.len() { + return Err(LaurusError::InvalidOperation(format!( + "rerank sidecar vector_count mismatch: segment has {} vectors, sidecar has {}", + vector_ids.len(), + header.vector_count + ))); + } + // Transient rehydration for the pool's String-shaped assignment input + // (eager-only path; nothing retained). + let assignment: Vec<(u64, String)> = vector_ids + .iter() + .map(|&(id, fid)| (id, field_dict[fid as usize].to_string())) + .collect(); + let pool = crate::vector::index::rerank_storage::RerankStoragePool::from_sidecar_payload( + header.storage_kind, + dimension, + header.vector_count as usize, + payload, + &assignment, + )?; + Ok(Some(std::sync::Arc::new(pool))) +} + #[cfg(test)] mod tests { use super::*; diff --git a/laurus/src/vector/store.rs b/laurus/src/vector/store.rs index 6eb9805f..79e82eed 100644 --- a/laurus/src/vector/store.rs +++ b/laurus/src/vector/store.rs @@ -130,6 +130,10 @@ impl VectorStore { FieldOption::Flat(opt) => VectorIndexTypeConfig::Flat(FlatIndexConfig { dimension: opt.dimension, distance_metric: opt.distance, + // Stage-2 rerank sidecar (#481, wired for Flat by + // #932 — previously carried on the option but + // dropped here, so the writer never emitted it). + rerank_storage: opt.rerank_storage, // L2-normalize only for the magnitude-invariant // Cosine metric (Issue #794); normalizing a // Euclidean/DotProduct/Manhattan field changes its @@ -161,6 +165,10 @@ impl VectorStore { distance_metric: opt.distance, n_clusters: opt.n_clusters, n_probe: opt.n_probe, + // Stage-2 rerank sidecar (#481, wired for IVF by + // #932 — previously carried on the option but + // dropped here, so the writer never emitted it). + rerank_storage: opt.rerank_storage, // L2-normalize only for the magnitude-invariant // Cosine metric (Issue #794). normalize_vectors: opt.distance diff --git a/laurus/tests/vector_rerank_flat_ivf_test.rs b/laurus/tests/vector_rerank_flat_ivf_test.rs new file mode 100644 index 00000000..f1b828c2 --- /dev/null +++ b/laurus/tests/vector_rerank_flat_ivf_test.rs @@ -0,0 +1,210 @@ +//! End-to-end engine wiring tests for Stage-2 rerank on Flat and IVF +//! (Issue #650 PR-2 / #932), mirroring `vector_rerank_engine_test.rs`. +//! +//! Pins the full loop per index type: `rerank_storage: Some(F32)` on the +//! schema field makes commit emit the `.f32` sidecar next to each sealed +//! segment, the reader loads it, and a search with `rerank_factor` +//! rescores against the exact f32 vectors (measurably different scores +//! from the int8 path on non-grid data). A field without the sidecar +//! keeps the silent Stage-1 fallback — including that `rerank_factor` no +//! longer returns `NotImplemented` (the pre-#932 behavior). + +use tempfile::TempDir; + +use laurus::DistanceMetric; +use laurus::Engine; +use laurus::SearchRequestBuilder; +use laurus::storage::file::FileStorageConfig; +use laurus::storage::{StorageConfig, StorageFactory}; +use laurus::vector::core::rerank::RerankStorageKind; +use laurus::vector::{FlatOption, IvfOption, Vector}; +use laurus::{DataValue, Document}; +use laurus::{FieldOption, QueryVector, Schema, VectorSearchQuery}; + +/// Non-grid components (not exactly representable on the int8 affine +/// grid) so the exact-f32 rerank score measurably differs from the int8 +/// score; a query between the docs keeps similarities off any clamp. +const DOCS: [(&str, [f32; 4]); 4] = [ + ("doc1", [0.92, 0.31, 0.17, 0.05]), + ("doc2", [0.13, 0.83, 0.41, 0.27]), + ("doc3", [0.05, 0.19, 0.77, 0.61]), + ("doc4", [0.33, 0.47, 0.29, 0.71]), +]; +const QUERY: [f32; 4] = [0.87, 0.36, 0.21, 0.09]; + +fn vector_request(rerank_factor: Option) -> laurus::SearchRequest { + let mut builder = SearchRequestBuilder::new() + .vector_query(VectorSearchQuery::Vectors(vec![QueryVector { + vector: Vector::new(QUERY.to_vec()), + weight: 1.0, + fields: Some(vec!["embedding".to_string()]), + }])) + .limit(1); + if let Some(factor) = rerank_factor { + builder = builder.vector_rerank_factor(factor); + } + builder.build() +} + +async fn ingest(engine: &Engine) -> laurus::Result<()> { + for (id, vec) in &DOCS { + let doc = Document::builder() + .add_field("embedding", DataValue::Vector(vec.to_vec())) + .build(); + engine.put_document(id, doc).await?; + } + engine.commit().await +} + +/// Shared body: build the engine over `field_option`, ingest + commit, +/// assert the sidecar exists at `sidecar_name`, and assert rerank both +/// works (doc1 wins) and *changes* the score vs the int8 path. +async fn assert_rerank_end_to_end( + field_option: FieldOption, + sidecar_name: &str, +) -> laurus::Result<()> { + let temp_dir = TempDir::new().unwrap(); + let storage = + StorageFactory::create(StorageConfig::File(FileStorageConfig::new(temp_dir.path())))?; + let schema = Schema::builder() + .add_field("embedding", field_option) + .build(); + let engine = Engine::new(storage.clone(), schema).await?; + ingest(&engine).await?; + + assert!( + storage.file_exists(sidecar_name), + "commit with rerank_storage: Some(F32) must emit {sidecar_name}" + ); + + let with_rerank = engine.search(vector_request(Some(4))).await?; + assert_eq!(with_rerank.len(), 1, "expected exactly 1 hit"); + assert_eq!( + with_rerank[0].id, "doc1", + "doc1 is the closest match under exact f32 distances" + ); + + let without_rerank = engine.search(vector_request(None)).await?; + assert_eq!(without_rerank[0].id, "doc1"); + assert_ne!( + with_rerank[0].score.to_bits(), + without_rerank[0].score.to_bits(), + "rerank_factor must change the score via the f32 sidecar pool; \ + identical scores mean the silent Stage-1 fallback was taken" + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn flat_rerank_succeeds_on_stage2_field() -> laurus::Result<()> { + assert_rerank_end_to_end( + FieldOption::Flat(FlatOption { + dimension: 4, + distance: DistanceMetric::Cosine, + rerank_storage: Some(RerankStorageKind::F32), + ..FlatOption::default() + }), + // Segment-per-commit is the Flat default since #907. + "vector/segment_000000.flat.f32", + ) + .await +} + +#[tokio::test(flavor = "multi_thread")] +async fn ivf_rerank_succeeds_on_stage2_field() -> laurus::Result<()> { + assert_rerank_end_to_end( + FieldOption::Ivf(IvfOption { + dimension: 4, + distance: DistanceMetric::Cosine, + n_clusters: 2, + n_probe: 2, + rerank_storage: Some(RerankStorageKind::F32), + ..IvfOption::default() + }), + "vector/segment_000000.ivf.f32", + ) + .await +} + +/// A Stage-1 field (no `rerank_storage`) with `rerank_factor` set must +/// silently fall back — succeeding with int8-identical scores instead of +/// the pre-#932 `NotImplemented` error. +#[tokio::test(flavor = "multi_thread")] +async fn flat_and_ivf_rerank_factor_falls_back_silently_on_stage1_field() -> laurus::Result<()> { + for field_option in [ + FieldOption::Flat(FlatOption { + dimension: 4, + distance: DistanceMetric::Cosine, + ..FlatOption::default() + }), + FieldOption::Ivf(IvfOption { + dimension: 4, + distance: DistanceMetric::Cosine, + n_clusters: 2, + n_probe: 2, + ..IvfOption::default() + }), + ] { + let temp_dir = TempDir::new().unwrap(); + let storage = + StorageFactory::create(StorageConfig::File(FileStorageConfig::new(temp_dir.path())))?; + let schema = Schema::builder() + .add_field("embedding", field_option) + .build(); + let engine = Engine::new(storage, schema).await?; + ingest(&engine).await?; + + let with_factor = engine.search(vector_request(Some(4))).await?; + let without = engine.search(vector_request(None)).await?; + assert_eq!(with_factor[0].id, "doc1"); + assert_eq!( + with_factor[0].score.to_bits(), + without[0].score.to_bits(), + "a sidecar-less field must silently keep Stage-1 scores" + ); + } + Ok(()) +} + +/// Multi-commit (multi-segment) Flat index under rerank: the per-segment +/// exact-f32 scores carry the `score_basis` stamp, so the fan-out keeps +/// them (#927) and the cross-segment order matches the exact computation. +#[tokio::test(flavor = "multi_thread")] +async fn segmented_flat_rerank_orders_across_segments() -> laurus::Result<()> { + let temp_dir = TempDir::new().unwrap(); + let storage = + StorageFactory::create(StorageConfig::File(FileStorageConfig::new(temp_dir.path())))?; + let schema = Schema::builder() + .add_field( + "embedding", + FieldOption::Flat(FlatOption { + dimension: 4, + distance: DistanceMetric::Cosine, + rerank_storage: Some(RerankStorageKind::F32), + ..FlatOption::default() + }), + ) + .build(); + let engine = Engine::new(storage.clone(), schema).await?; + + // Two commits = two sealed segments, best match in the second. + for (id, vec) in &DOCS[..2] { + let doc = Document::builder() + .add_field("embedding", DataValue::Vector(vec.to_vec())) + .build(); + engine.put_document(id, doc).await?; + } + engine.commit().await?; + for (id, vec) in &DOCS[2..] { + let doc = Document::builder() + .add_field("embedding", DataValue::Vector(vec.to_vec())) + .build(); + engine.put_document(id, doc).await?; + } + engine.commit().await?; + assert!(storage.file_exists("vector/segment_000001.flat.f32")); + + let hits = engine.search(vector_request(Some(4))).await?; + assert_eq!(hits[0].id, "doc1", "cross-segment rerank must find doc1"); + Ok(()) +}