luci/vector/

global.rs

//! Single global HNSW per index per dense_vector field.
//!
//! Per [[global-vector-indices]] Alternative B: vector indexes
//! decouple from the segment model. Every dense_vector field gets one
//! HNSW graph for the whole index, not one per segment.
//!
//! - **Write**: [`GlobalHnsw::add_vector`] appends a vector to the
//!   field's HNSW builder and records the (segment_id, local_doc_id) the
//!   vector was attached to, so kNN hits can be returned in the
//!   segment-local doc-id space the rest of the engine speaks.
//! - **Read**: [`GlobalHnsw::search`] runs HNSW once and returns hits
//!   as (segment_id, doc_id, distance). The bound query buckets by
//!   segment so each `SegmentReader`-scoped scorer sees only the docs
//!   in its own segment. A lazily-built [`HnswIndex`] cache services
//!   each search; the cache is invalidated on every mutation.
//! - **Persistence**: [`GlobalHnsw::field_to_bytes`] / [`load_field`]
//!   serialize one field's graph + resolver at a time. The bytes are
//!   stored in their own file extent via the storage layer's
//!   `write_vector_index` API rather than in `user_metadata`.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use crate::core::{DocId, FieldId, LuciError, SegmentId};
use crate::mapping::Mapping;

use super::DistanceMetric;
use super::hnsw::{
    BuildThreads, HnswBuilder, HnswIndex, HnswParams, checked_len, read_u32, read_u64, take_bytes,
};

/// HNSW M parameter used for every dense_vector field today. Hardcoded
/// in [`SegmentBuilder::new`]; mirrored here.
const HNSW_M: usize = 16;

/// HNSW ef_construction parameter, also hardcoded. Same provenance as
/// [`HNSW_M`].
const HNSW_EF_CONSTRUCTION: usize = 100;

/// Per-field blob magic — `b"VIDX"` (Vector InDeX).
const VECTOR_INDEX_MAGIC: [u8; 4] = *b"VIDX";

/// Current per-field blob format version.
const VECTOR_INDEX_VERSION: u8 = 1;

/// One field's slice of the global graph: the HNSW builder, the
/// `global_ord → (SegmentId, local_doc_id)` resolver, and a lazily-built
/// search-side [`HnswIndex`] cache.
struct FieldGlobalHnsw {
    builder: HnswBuilder,
    /// Indexed by global_ord. Each entry is the segment + per-segment
    /// doc id where the vector's owning document lives. Search hits
    /// flow back through this table.
    resolver: Vec<(SegmentId, u32)>,
    /// Cached search-side index built from `builder`. Reset to `None`
    /// on every mutation (add_vector, rewrite_after_merge); rebuilt
    /// from the builder on the next search/persist.
    cached: Option<Arc<HnswIndex>>,
}

impl FieldGlobalHnsw {
    fn new(params: HnswParams) -> Self {
        Self {
            builder: HnswBuilder::new(params),
            resolver: Vec::new(),
            cached: None,
        }
    }

    fn invalidate_cache(&mut self) {
        self.cached = None;
    }

    /// Build the search-side index from the current builder if not
    /// already cached, then return a cloned `Arc<HnswIndex>` handle.
    ///
    /// Building involves cloning the builder (so it stays writable)
    /// and running quantization for `Int8` fields. The cost is paid
    /// once per cache-bust; subsequent searches are zero-cost lookups.
    fn get_or_build_index(&mut self) -> Arc<HnswIndex> {
        // Every consumer (commit persist, merge re-persist, read-side
        // search) must see a fully-linked graph. A non-empty pending tail
        // here means `connect_pending` did not run before persist — the
        // disconnected-graph bug. See [[optimization-concurrent-hnsw-insert]].
        debug_assert!(
            !self.builder.has_pending_tail(),
            "get_or_build_index called with an unlinked pending tail; \
             connect_pending was not run before persist",
        );
        if self.cached.is_none() {
            self.cached = Some(Arc::new(self.builder.clone().build()));
        }
        Arc::clone(self.cached.as_ref().unwrap())
    }
}

/// One global HNSW graph per dense_vector field in the schema.
///
/// Owned by [`IndexWriter`] on the write side and held by the
/// [`SegmentStore`] (behind an Arc) on the read side. Concurrency:
/// access is serialized through a single `Mutex`. Writes today are
/// already single-threaded (one `IndexWriter`, sequential `add()` /
/// `bulk()`), so this lock doesn't serialize anything that wasn't
/// already serialized. If concurrent vector inserts are ever needed,
/// the lock has to be replaced (per-node atomics or a sharded map
/// are the two patterns prior art uses) — that work is not currently
/// scheduled.
pub struct GlobalHnsw {
    per_field: Mutex<HashMap<FieldId, FieldGlobalHnsw>>,
}

/// Result of a global kNN search: which segment owns the hit, the
/// per-segment doc_id, and the raw HNSW distance.
#[derive(Clone, Copy, Debug)]
pub struct GlobalHit {
    pub segment_id: SegmentId,
    pub doc_id: DocId,
    pub distance: f32,
}

impl GlobalHnsw {
    /// Initialize the per-field builders from the schema.
    ///
    /// Walks every field; for each dense_vector field, constructs an
    /// `HnswBuilder` with identical params to what
    /// [`SegmentBuilder::new`] used in the per-segment-HNSW era. Schema
    /// fields that aren't dense_vector are ignored.
    ///
    /// Panics: if `schema.field_id(name)` returns `None` for a name we
    /// just read out of `schema.fields()`. That's an internal-consistency
    /// invariant; reaching the panic indicates a mapping bug, not user
    /// input.
    pub fn new(schema: &Mapping) -> Self {
        let mut per_field = HashMap::new();
        for mapping in schema.fields() {
            let Some(dims) = mapping.field_type.vector_dims() else {
                continue;
            };
            if dims == 0 {
                continue;
            }
            let field_id = schema.field_id(&mapping.name).unwrap_or_else(|| {
                panic!(
                    "schema.fields() returned mapping for {:?} but \
                     field_id() couldn't find it; schema is internally \
                     inconsistent",
                    mapping.name
                );
            });
            let quantization = mapping
                .field_type
                .vector_quantization()
                .expect("dense_vector mapping must carry quantization");
            per_field.insert(
                field_id,
                FieldGlobalHnsw::new(HnswParams {
                    dims,
                    m: HNSW_M,
                    ef_construction: HNSW_EF_CONSTRUCTION,
                    metric: DistanceMetric::Cosine,
                    quantization,
                }),
            );
        }
        Self {
            per_field: Mutex::new(per_field),
        }
    }

    /// Append a vector to the global graph for `field_id` and record
    /// the (segment_id, local_doc_id) it belongs to. Returns the
    /// assigned global ordinal.
    ///
    /// Cosine vectors are normalized in place by `HnswBuilder` before
    /// insertion under the [[optimize-cosine-norm-precompute]]
    /// invariant. Zero / non-finite vectors return
    /// `LuciError::InvalidQuery`.
    ///
    /// Returns `LuciError::InvalidQuery` if `field_id` is not a
    /// dense_vector field in the schema — that's a writer bug, not
    /// user input, but bubbling it as an error rather than panicking
    /// lets `add()` propagate it.
    pub fn add_vector(
        &self,
        field_id: FieldId,
        segment_id: SegmentId,
        local_doc_id: u32,
        vector: Vec<f32>,
    ) -> Result<u32, LuciError> {
        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        let field = guard.get_mut(&field_id).ok_or_else(|| {
            LuciError::InvalidQuery(format!(
                "GlobalHnsw::add_vector called for field {field_id:?} which is \
                 not a dense_vector field in the schema; this is an internal \
                 wiring bug",
            ))
        })?;
        let ord = field.builder.len() as u32;
        field.builder.add_vector(vector)?;
        field.resolver.push((segment_id, local_doc_id));
        field.invalidate_cache();
        debug_assert_eq!(
            field.resolver.len(),
            field.builder.len(),
            "resolver and builder lengths must agree after add_vector",
        );
        Ok(ord)
    }

    /// Store a vector for `field_id` **without linking it into the
    /// graph** — the deferred half of the write model
    /// ([[optimization-concurrent-hnsw-insert]]). The node joins the
    /// builder's pending tail and is linked by [`Self::connect_pending`]
    /// at commit. Mirrors [`Self::add_vector`]'s resolver/length
    /// invariant. Cheap and sequential; called per-doc from
    /// `IndexWriter::add`.
    pub fn store_vector(
        &self,
        field_id: FieldId,
        segment_id: SegmentId,
        local_doc_id: u32,
        vector: Vec<f32>,
    ) -> Result<u32, LuciError> {
        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        let field = guard.get_mut(&field_id).ok_or_else(|| {
            LuciError::InvalidQuery(format!(
                "GlobalHnsw::store_vector called for field {field_id:?} which is \
                 not a dense_vector field in the schema; this is an internal \
                 wiring bug",
            ))
        })?;
        let ord = field.builder.len() as u32;
        field.builder.store_vector(vector)?;
        field.resolver.push((segment_id, local_doc_id));
        field.invalidate_cache();
        debug_assert_eq!(
            field.resolver.len(),
            field.builder.len(),
            "resolver and builder lengths must agree after store_vector",
        );
        Ok(ord)
    }

    /// Link every stored-but-unlinked vector for `field_id` — the pending
    /// tail — into the graph (in parallel per `threads`), then advance
    /// the cursor. A no-op when the tail is empty or `field_id` is not a
    /// vector field. Called as the **first** step of
    /// [`crate::writer::IndexWriter::commit`], before the persist loop,
    /// so `field_to_bytes` always serializes a fully-linked graph. Holds
    /// the `per_field` mutex across the whole parallel phase — fine, since
    /// Luci is single-writer and readers use a separate disk-reloaded
    /// graph. See [[optimization-concurrent-hnsw-insert]] §Write model.
    pub fn connect_pending(&self, field_id: FieldId, threads: BuildThreads) {
        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        if let Some(field) = guard.get_mut(&field_id) {
            field.builder.connect_pending(threads);
            field.invalidate_cache();
        }
    }

    /// Run a global kNN search and return up to `k` hits in
    /// segment-local doc-id space, sorted by ascending distance.
    ///
    /// `ef` is the HNSW beam width. Returns `Ok(None)` if `field_id`
    /// is not a dense_vector field in the schema (lets the caller
    /// short-circuit cleanly). An empty graph returns
    /// `Ok(Some((vec![], metric)))`.
    pub fn search(
        &self,
        field_id: FieldId,
        query: &[f32],
        k: usize,
        ef: usize,
    ) -> Result<Option<(Vec<GlobalHit>, DistanceMetric)>, LuciError> {
        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        let field = match guard.get_mut(&field_id) {
            Some(f) => f,
            None => return Ok(None),
        };
        let metric = field.builder.params().metric;
        let resolver = field.resolver.clone();
        let index = field.get_or_build_index();
        // Drop the lock before running search — the index is now a
        // standalone Arc and search is the bulk of the work.
        drop(guard);

        let raw = index.search(query, k, ef)?;
        let hits = raw
            .into_iter()
            .map(|(global_ord, dist)| {
                let (seg, doc) = resolver[global_ord as usize];
                GlobalHit {
                    segment_id: seg,
                    doc_id: DocId::new(doc),
                    distance: dist,
                }
            })
            .collect();
        Ok(Some((hits, metric)))
    }

    /// Number of vectors stored in the global graph for `field_id`.
    /// Returns `None` if `field_id` isn't a dense_vector field.
    pub fn len(&self, field_id: FieldId) -> Option<usize> {
        let guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        guard.get(&field_id).map(|f| f.builder.len())
    }

    /// True if every field in the global graph is empty.
    pub fn is_empty(&self) -> bool {
        let guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        guard.values().all(|f| f.builder.is_empty())
    }

    /// Apply a merge ord-map to the resolver: for every entry that
    /// referenced one of the source segments, rewrite it to point at
    /// the merged segment with its new local doc id.
    ///
    /// `merge_map` is keyed by `(old_segment_id, old_local_doc_id)`
    /// and yields `(new_segment_id, new_local_doc_id)`. Entries whose
    /// key is absent from the map are left untouched — they belong to
    /// segments not involved in this merge. Mutating the resolver
    /// invalidates the cached search index for every field.
    pub fn rewrite_after_merge(&self, merge_map: &HashMap<(SegmentId, u32), (SegmentId, u32)>) {
        if merge_map.is_empty() {
            return;
        }
        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        for field in guard.values_mut() {
            let mut changed = false;
            for entry in &mut field.resolver {
                if let Some(&(new_seg, new_doc)) = merge_map.get(entry) {
                    *entry = (new_seg, new_doc);
                    changed = true;
                }
            }
            if changed {
                field.invalidate_cache();
            }
        }
    }

    /// Field ids that have a global HNSW (including empty ones).
    pub fn field_ids(&self) -> Vec<FieldId> {
        let guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        guard.keys().copied().collect()
    }

    /// Field ids whose graph has at least one vector. Used by the
    /// writer to decide which fields to persist on commit.
    pub fn non_empty_field_ids(&self) -> Vec<FieldId> {
        let guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        let mut ids: Vec<FieldId> = guard
            .iter()
            .filter(|(_, f)| !f.builder.is_empty())
            .map(|(fid, _)| *fid)
            .collect();
        ids.sort();
        ids
    }

    /// Serialize one field's HNSW graph + resolver as the bytes the
    /// storage layer persists via `write_vector_index(field_id, ..)`.
    ///
    /// Layout:
    /// ```text
    /// [magic: 4 bytes "VIDX"][version: u8]
    /// [hnsw_len: u32][hnsw_bytes]
    /// [resolver_len: u32]
    /// per resolver entry { [segment_id: u64][local_doc_id: u32] }
    /// ```
    ///
    /// Returns `None` if `field_id` isn't a dense_vector field — the
    /// caller should skip it rather than persisting an empty blob.
    pub fn field_to_bytes(&self, field_id: FieldId) -> Option<Vec<u8>> {
        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        let field = guard.get_mut(&field_id)?;

        let mut buf = Vec::new();
        buf.extend_from_slice(&VECTOR_INDEX_MAGIC);
        buf.push(VECTOR_INDEX_VERSION);

        let index = field.get_or_build_index();
        let hnsw_bytes = index.to_bytes();
        buf.extend_from_slice(&(hnsw_bytes.len() as u32).to_le_bytes());
        buf.extend_from_slice(&hnsw_bytes);

        buf.extend_from_slice(&(field.resolver.len() as u32).to_le_bytes());
        for (seg, doc) in &field.resolver {
            buf.extend_from_slice(&seg.as_u64().to_le_bytes());
            buf.extend_from_slice(&doc.to_le_bytes());
        }

        Some(buf)
    }

    /// Replace the (empty) builder for `field_id` with a graph
    /// deserialized from bytes written by [`Self::field_to_bytes`].
    ///
    /// Errors if `field_id` isn't a dense_vector field in the current
    /// schema — that means the field was removed since the index was
    /// last written, and the caller should `remove_vector_index` for
    /// it instead of loading.
    pub fn load_field(&self, field_id: FieldId, data: &[u8]) -> Result<(), LuciError> {
        if data.len() < 5 {
            return Err(LuciError::IndexCorrupted(format!(
                "vector index blob for field {field_id:?} too short: {} bytes",
                data.len()
            )));
        }
        if data[0..4] != VECTOR_INDEX_MAGIC {
            return Err(LuciError::IndexCorrupted(format!(
                "vector index blob for field {field_id:?} missing magic prefix"
            )));
        }
        if data[4] != VECTOR_INDEX_VERSION {
            return Err(LuciError::SegmentFormatUnknown(format!(
                "unknown vector index blob version {} for field {field_id:?}",
                data[4]
            )));
        }
        let mut pos = 5;
        let hnsw_len = read_u32(data, &mut pos)? as usize;
        let hnsw_bytes = take_bytes(data, &mut pos, hnsw_len)?;
        let index = HnswIndex::from_bytes(hnsw_bytes)?;
        let builder = HnswBuilder::from_index(index);

        // Guard `resolver_len` before reserving: each entry is a (SegmentId
        // u64, doc u32) pair = 12 bytes, so a corrupt count cannot drive a giant
        // speculative allocation. Per [[cross-platform-portability]] a valid
        // file is never truncated, so this rejects only corrupt input.
        let resolver_len = read_u32(data, &mut pos)? as usize;
        let mut resolver = Vec::with_capacity(checked_len(resolver_len, 12, data, pos)?);
        for _ in 0..resolver_len {
            let seg = SegmentId::new(read_u64(data, &mut pos)?);
            let doc = read_u32(data, &mut pos)?;
            resolver.push((seg, doc));
        }
        if resolver.len() != builder.len() {
            return Err(LuciError::IndexCorrupted(format!(
                "vector index resolver/graph mismatch for field {field_id:?}: \
                 graph has {} vectors, resolver has {}",
                builder.len(),
                resolver.len()
            )));
        }

        let mut guard = self.per_field.lock().expect("GlobalHnsw mutex poisoned");
        let field = guard.get_mut(&field_id).ok_or_else(|| {
            LuciError::InvalidQuery(format!(
                "GlobalHnsw::load_field called for field {field_id:?} which is \
                 not a dense_vector field in the current schema"
            ))
        })?;
        field.builder = builder;
        field.resolver = resolver;
        field.cached = None;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::mapping::{FieldType, Mapping};

    use super::*;

    fn vector_schema(name: &str, dims: usize) -> Mapping {
        Mapping::builder()
            .field(name, FieldType::dense_vector(dims))
            .build()
    }

    #[test]
    fn new_finds_dense_vector_fields() {
        let schema = vector_schema("embedding", 4);
        let g = GlobalHnsw::new(&schema);
        let ids = g.field_ids();
        assert_eq!(ids.len(), 1);
        let field_id = schema.field_id("embedding").unwrap();
        assert_eq!(g.len(field_id), Some(0));
    }

    #[test]
    fn new_with_no_vector_fields_is_empty() {
        let schema = Mapping::builder().build();
        let g = GlobalHnsw::new(&schema);
        assert!(g.is_empty());
    }

    #[test]
    fn add_vector_returns_increasing_ordinals() {
        let schema = vector_schema("embedding", 3);
        let field_id = schema.field_id("embedding").unwrap();
        let g = GlobalHnsw::new(&schema);
        let seg = SegmentId::new(1);
        let ord0 = g.add_vector(field_id, seg, 0, vec![1.0, 0.0, 0.0]).unwrap();
        let ord1 = g.add_vector(field_id, seg, 1, vec![0.0, 1.0, 0.0]).unwrap();
        let ord2 = g.add_vector(field_id, seg, 2, vec![0.0, 0.0, 1.0]).unwrap();
        assert_eq!(ord0, 0);
        assert_eq!(ord1, 1);
        assert_eq!(ord2, 2);
        assert_eq!(g.len(field_id), Some(3));
    }

    #[test]
    fn add_vector_for_unknown_field_errors() {
        let schema = vector_schema("embedding", 3);
        let g = GlobalHnsw::new(&schema);
        let seg = SegmentId::new(1);
        let result = g.add_vector(FieldId(999), seg, 0, vec![1.0, 0.0, 0.0]);
        assert!(matches!(result, Err(LuciError::InvalidQuery(_))));
    }

    #[test]
    fn cosine_zero_vector_rejected() {
        let schema = vector_schema("embedding", 3);
        let field_id = schema.field_id("embedding").unwrap();
        let g = GlobalHnsw::new(&schema);
        let seg = SegmentId::new(1);
        let result = g.add_vector(field_id, seg, 0, vec![0.0, 0.0, 0.0]);
        assert!(matches!(result, Err(LuciError::InvalidQuery(_))));
    }

    #[test]
    fn search_returns_hits_in_segment_local_doc_space() {
        let schema = vector_schema("embedding", 3);
        let field_id = schema.field_id("embedding").unwrap();
        let g = GlobalHnsw::new(&schema);
        let seg1 = SegmentId::new(1);
        let seg2 = SegmentId::new(2);
        g.add_vector(field_id, seg1, 0, vec![1.0, 0.0, 0.0])
            .unwrap();
        g.add_vector(field_id, seg1, 1, vec![0.0, 1.0, 0.0])
            .unwrap();
        g.add_vector(field_id, seg2, 0, vec![0.9, 0.1, 0.0])
            .unwrap();

        let (hits, metric) = g
            .search(field_id, &[1.0, 0.0, 0.0], 3, 16)
            .unwrap()
            .unwrap();
        assert_eq!(metric, DistanceMetric::Cosine);
        assert_eq!(hits.len(), 3);
        // Best match is the [1,0,0] vector in segment 1, doc 0.
        assert_eq!(hits[0].segment_id, seg1);
        assert_eq!(hits[0].doc_id, DocId::new(0));
    }

    #[test]
    fn roundtrip_field_to_bytes_load_field() {
        let schema = vector_schema("embedding", 3);
        let field_id = schema.field_id("embedding").unwrap();
        let g = GlobalHnsw::new(&schema);
        let seg1 = SegmentId::new(1);
        g.add_vector(field_id, seg1, 0, vec![1.0, 0.0, 0.0])
            .unwrap();
        g.add_vector(field_id, seg1, 1, vec![0.0, 1.0, 0.0])
            .unwrap();
        g.add_vector(field_id, seg1, 2, vec![0.0, 0.0, 1.0])
            .unwrap();

        let bytes = g.field_to_bytes(field_id).unwrap();
        let g2 = GlobalHnsw::new(&schema);
        g2.load_field(field_id, &bytes).unwrap();
        assert_eq!(g2.len(field_id), Some(3));

        let (hits, _) = g2
            .search(field_id, &[1.0, 0.0, 0.0], 1, 16)
            .unwrap()
            .unwrap();
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].segment_id, seg1);
        assert_eq!(hits[0].doc_id, DocId::new(0));
    }

    /// `load_field` must reject a malformed vector-index blob with an explicit
    /// [`LuciError`] rather than panicking out of bounds — the resolver-wrapper
    /// companion to `HnswIndex::from_bytes`'s own guard. See
    /// [[cross-platform-portability]].
    #[test]
    fn load_field_rejects_corrupt_blob() {
        let schema = vector_schema("embedding", 3);
        let field_id = schema.field_id("embedding").unwrap();
        let g = GlobalHnsw::new(&schema);
        let seg = SegmentId::new(1);
        g.add_vector(field_id, seg, 0, vec![1.0, 0.0, 0.0]).unwrap();
        g.add_vector(field_id, seg, 1, vec![0.0, 1.0, 0.0]).unwrap();
        let valid = g.field_to_bytes(field_id).unwrap();
        assert!(
            GlobalHnsw::new(&schema)
                .load_field(field_id, &valid)
                .is_ok(),
            "valid blob must load"
        );

        // Truncation anywhere past the 5-byte header must error, not panic —
        // exercises the bounded readers for hnsw_len, the hnsw blob,
        // resolver_len, and the resolver entries.
        for cut in [5usize, 6, 9, valid.len() / 2, valid.len() - 1] {
            assert!(
                GlobalHnsw::new(&schema)
                    .load_field(field_id, &valid[..cut])
                    .is_err(),
                "truncated-to-{cut} blob must be rejected, not panic"
            );
        }

        // A corrupt resolver length prefix must be rejected by `checked_len`
        // before it drives a giant speculative allocation. It sits at
        // 5 (header) + 4 (hnsw_len) + hnsw_len.
        let hnsw_len = u32::from_le_bytes(valid[5..9].try_into().unwrap()) as usize;
        let resolver_len_off = 5 + 4 + hnsw_len;
        let mut bad_resolver = valid.clone();
        bad_resolver[resolver_len_off..resolver_len_off + 4]
            .copy_from_slice(&u32::MAX.to_le_bytes());
        assert!(
            matches!(
                GlobalHnsw::new(&schema).load_field(field_id, &bad_resolver),
                Err(LuciError::IndexCorrupted(_))
            ),
            "corrupt resolver length must be IndexCorrupted, not OOM/panic"
        );
    }

    #[test]
    fn non_empty_field_ids_omits_empty_fields() {
        let schema = Mapping::builder()
            .field("a", FieldType::dense_vector(2))
            .field("b", FieldType::dense_vector(2))
            .build();
        let a = schema.field_id("a").unwrap();
        let b = schema.field_id("b").unwrap();
        let g = GlobalHnsw::new(&schema);
        // populate only "a"
        g.add_vector(a, SegmentId::new(1), 0, vec![1.0, 0.0])
            .unwrap();
        let ids = g.non_empty_field_ids();
        assert_eq!(ids, vec![a]);
        assert_eq!(g.len(b), Some(0));
    }

    #[test]
    fn rewrite_after_merge_remaps_resolver() {
        let schema = vector_schema("embedding", 3);
        let field_id = schema.field_id("embedding").unwrap();
        let g = GlobalHnsw::new(&schema);
        let s1 = SegmentId::new(1);
        let s2 = SegmentId::new(2);
        let s3 = SegmentId::new(3);
        g.add_vector(field_id, s1, 0, vec![1.0, 0.0, 0.0]).unwrap();
        g.add_vector(field_id, s2, 0, vec![0.0, 1.0, 0.0]).unwrap();

        // Merge s1 + s2 into s3: s1/0 -> s3/0, s2/0 -> s3/1.
        let mut merge_map = HashMap::new();
        merge_map.insert((s1, 0), (s3, 0));
        merge_map.insert((s2, 0), (s3, 1));
        g.rewrite_after_merge(&merge_map);

        let (hits, _) = g
            .search(field_id, &[1.0, 0.0, 0.0], 2, 16)
            .unwrap()
            .unwrap();
        assert_eq!(hits.len(), 2);
        // Both hits now point at s3.
        for hit in &hits {
            assert_eq!(hit.segment_id, s3);
        }
    }
}