summa_core/index/

primary_key.rs

//! Primary key deduplication index.
//!
//! Uses a bloom filter + a map of latest staged rows to reject duplicate adds
//! at `add_document()` time. Committed keys are checked via fast-field
//! `TextDictReader::ordinal()` (binary search, O(log n)).
//!
//! The bloom filter is persisted to `pk_bloom.bin` so that restarts don't need
//! to re-iterate every committed key. On load, only keys from segments that
//! appeared since the last persist are iterated.

#[cfg(feature = "native")]
use std::collections::HashSet;

use super::staged_row::StagedRow;
use crate::dsl::{Document, FieldValue, Schema};
use byteorder::{LittleEndian, WriteBytesExt};
use rustc_hash::{FxHashMap, FxHashSet};
use std::sync::Arc;

use crate::dsl::Field;
use crate::error::{Error, Result};
#[cfg(feature = "native")]
use crate::segment::SegmentSnapshot;
use crate::structures::BloomFilter;

/// Bloom filter sizing: 10 bits/key ≈ 1% false positive rate.
const BLOOM_BITS_PER_KEY: usize = 10;

/// Extra capacity added to bloom filter beyond known keys.
const BLOOM_HEADROOM: usize = 100_000;

/// File name for the persisted primary-key bloom filter.
#[cfg(feature = "native")]
pub const PK_BLOOM_FILE: &str = "pk_bloom.bin";

/// Magic bytes for the persisted bloom file.
const PK_BLOOM_MAGIC: u32 = 0x504B424C; // "PKBL"

/// The reservation, deletion target and persisted single-value column must
/// identify the same key. Validate before mutating any writer state.
pub(super) fn document_key(doc: &crate::dsl::Document, field: crate::Field) -> Result<&str> {
    let mut values = doc.get_all(field);
    let key = values
        .next()
        .ok_or_else(|| Error::Document("Missing primary key field".into()))?
        .as_text()
        .ok_or_else(|| Error::Document("Primary key must be text".into()))?;
    if values.next().is_some() {
        return Err(Error::Document(
            "primary key requires exactly one text value".into(),
        ));
    }
    if key.is_empty() {
        return Err(Error::Document("Primary key must not be empty".into()));
    }
    if key.len() > 64 * 1024 {
        return Err(Error::Document(
            "primary key must contain 1..=65536 bytes".into(),
        ));
    }
    Ok(key)
}

/// Lightweight per-segment data for primary key lookups.
///
/// Only holds fast-field readers (text dictionaries), not full `SegmentReader`s.
/// This avoids loading DimensionTables, SSTable FSTs, bloom filters, etc.
pub struct PkSegmentData {
    pub segment_id: String,
    pub(super) content_hash: Option<super::content_hash::ContentHashLookup>,
    pub deletion_meta: Option<crate::segment::DeletionMeta>,
    pub alive_docs: Option<std::sync::Arc<crate::query::DocBitset>>,
    pub live_key_ordinals: Option<crate::query::DocBitset>,
    pub fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
}

impl PkSegmentData {
    // One scan at visibility refresh; duplicate checks stay dictionary lookup
    // plus O(1) membership, even for a heavily deleted segment.
    pub(crate) fn prepare_live_keys(&mut self, field: Field) {
        if self.live_key_ordinals.is_some() {
            return;
        }
        let Some(alive) = &self.alive_docs else {
            return;
        };
        let Some(ff) = self.fast_fields.get(&field.0) else {
            return;
        };
        let Some(dict) = ff.text_dict() else {
            return;
        };
        let mut keys = crate::query::DocBitset::new(dict.len());
        ff.scan_single_values(|doc, ordinal| {
            if alive.contains(doc) && ordinal < u64::from(dict.len()) {
                keys.set(ordinal as u32);
            }
        });
        self.live_key_ordinals = Some(keys);
    }
}

/// Thread-safe primary key deduplication index.
///
/// Sync dedup in the hot path: `BloomFilter::may_contain()`,
/// `FxHashMap::contains_key()`, and `TextDictReader::ordinal()` are all sync.
///
/// Interior mutability for the mutable state (bloom + latest staged map) is
/// behind `parking_lot::Mutex`. The committed data is only mutated via
/// `&mut self` methods (commit/abort path), so no lock is needed for it.
pub struct PrimaryKeyIndex {
    field: Field,
    state: parking_lot::Mutex<PrimaryKeyState>,
    /// Lightweight per-segment fast-field data for checking committed keys.
    /// Only mutated by `&mut self` methods (refresh/clear) — no lock needed.
    committed_data: Vec<PkSegmentData>,
    /// Holds ref counts so segments aren't deleted while we hold readers.
    #[cfg(feature = "native")]
    _snapshot: Option<std::sync::Arc<SegmentSnapshot>>,
}

struct PrimaryKeyState {
    bloom: BloomFilter,
    uncommitted: FxHashMap<Vec<u8>, PendingKey>,
    pending_bytes: usize,
    cancelled_bytes: usize,
    deletes: FxHashSet<String>,
    delete_bytes: usize,
}

const MAX_PENDING_KEY_BYTES: usize = 64 * 1024 * 1024;
// Vec<u32> initially allocates four slots; later growth uses at most two per row.
const CANCELLED_ROW_BYTES: usize = 4 * std::mem::size_of::<u32>();

struct PendingKey {
    row: Arc<StagedRow>,
    hash: Option<FieldValue>,
    bytes: usize,
}

// Twice entry size covers table occupancy, control bytes, and allocation overhead.
const KEY_SLOT_BYTES: usize =
    2 * (std::mem::size_of::<PendingKey>() + std::mem::size_of::<Vec<u8>>());

fn pending_bytes(key: &str, hash: Option<&FieldValue>) -> usize {
    KEY_SLOT_BYTES
        + std::mem::size_of::<StagedRow>()
        + 2 * std::mem::size_of::<usize>()
        + key.len()
        + match hash {
            Some(FieldValue::Text(value)) => value.len(),
            Some(FieldValue::Bytes(value)) => value.len(),
            _ => 0,
        }
}

fn check_pending_budget(
    state: &PrimaryKeyState,
    old: usize,
    new: usize,
    cancelled: usize,
) -> Result<()> {
    let next_len = state.uncommitted.len() + usize::from(new > 0) - usize::from(old > 0);
    // HashMap may grow on the accepted insertion. Reserve a conservative next
    // capacity before queue admission; retained capacity still counts after clear.
    let capacity = if next_len > state.uncommitted.capacity() {
        (next_len * 2).max(3)
    } else {
        state.uncommitted.capacity()
    };
    let payload = state.pending_bytes - old + new - next_len * KEY_SLOT_BYTES;
    let used = payload + capacity * KEY_SLOT_BYTES + state.cancelled_bytes + cancelled;
    if used > MAX_PENDING_KEY_BYTES {
        return Err(Error::Document(
            "pending primary-key metadata exceeds 64 MiB; commit before continuing".into(),
        ));
    }
    Ok(())
}

fn validate_mutation_key(key: &str) -> Result<()> {
    if key.is_empty() || key.len() > 64 * 1024 {
        return Err(Error::Document(
            "deletion key must contain 1..=65536 bytes".into(),
        ));
    }
    Ok(())
}

fn check_delete_budget(state: &PrimaryKeyState, key: &str) -> Result<()> {
    if !state.deletes.contains(key)
        && (state.deletes.len() >= 100_000 || state.delete_bytes + key.len() > 8 * 1024 * 1024)
    {
        return Err(Error::Document(
            "pending deletions exceed 100000 keys or 8 MiB; commit before continuing".into(),
        ));
    }
    Ok(())
}

fn stage_delete(state: &mut PrimaryKeyState, key: &str) -> bool {
    if state.deletes.contains(key) {
        return false;
    }
    state.delete_bytes += key.len();
    state.deletes.insert(key.to_owned());
    true
}

impl PrimaryKeyIndex {
    /// Create a new PrimaryKeyIndex by scanning committed segments.
    ///
    /// Iterates each segment's fast-field text dictionary to populate the bloom
    /// filter with all existing primary key values. The snapshot keeps ref counts
    /// alive so segments aren't deleted while we hold data.
    ///
    /// **CPU-intensive** — call from `spawn_blocking`, not the async runtime.
    #[cfg(feature = "native")]
    pub fn new(field: Field, pk_data: Vec<PkSegmentData>, snapshot: SegmentSnapshot) -> Self {
        let mut index = Self::build(field, pk_data);
        index._snapshot = Some(std::sync::Arc::new(snapshot));
        index
    }

    pub(crate) fn build(field: Field, mut pk_data: Vec<PkSegmentData>) -> Self {
        for data in &mut pk_data {
            data.prepare_live_keys(field);
        }
        // Count total unique keys across all segments for bloom sizing.
        let mut total_keys: usize = 0;
        for data in &pk_data {
            if let Some(ff) = data.fast_fields.get(&field.0)
                && let Some(dict) = ff.text_dict()
            {
                total_keys += dict.len() as usize;
            }
        }

        let mut bloom = BloomFilter::new(total_keys + BLOOM_HEADROOM, BLOOM_BITS_PER_KEY);

        // Insert all committed keys into the bloom filter.
        for data in &pk_data {
            if let Some(ff) = data.fast_fields.get(&field.0)
                && let Some(dict) = ff.text_dict()
            {
                for key in dict.iter() {
                    bloom.insert(key.as_bytes());
                }
            }
        }

        let bloom_bytes = bloom.size_bytes();
        log::info!(
            "[primary_key] bloom filter: {} keys, {}",
            total_keys,
            crate::format_bytes(bloom_bytes as u64),
        );

        Self {
            field,
            state: parking_lot::Mutex::new(PrimaryKeyState {
                bloom,
                uncommitted: FxHashMap::default(),
                pending_bytes: 0,
                cancelled_bytes: 0,
                deletes: FxHashSet::default(),
                delete_bytes: 0,
            }),
            committed_data: pk_data,
            #[cfg(feature = "native")]
            _snapshot: None,
        }
    }

    /// Create from a pre-loaded bloom filter (loaded from `pk_bloom.bin`).
    ///
    /// Skips dictionary iteration because the caller has already extended the
    /// persisted bloom with any segments it did not cover. `pk_data` contains
    /// data for all current segments.
    #[cfg(feature = "native")]
    pub fn from_persisted(
        field: Field,
        bloom: BloomFilter,
        mut pk_data: Vec<PkSegmentData>,
        snapshot: SegmentSnapshot,
    ) -> Self {
        for data in &mut pk_data {
            data.prepare_live_keys(field);
        }
        log::info!(
            "[primary_key] bloom filter loaded from cache: {}",
            crate::format_bytes(bloom.size_bytes() as u64),
        );

        Self {
            field,
            state: parking_lot::Mutex::new(PrimaryKeyState {
                bloom,
                uncommitted: FxHashMap::default(),
                pending_bytes: 0,
                cancelled_bytes: 0,
                deletes: FxHashSet::default(),
                delete_bytes: 0,
            }),
            committed_data: pk_data,
            _snapshot: Some(std::sync::Arc::new(snapshot)),
        }
    }

    /// Stream the complete primary-key cache without a corpus-sized
    /// intermediate allocation.
    pub fn write_bloom_cache(
        &self,
        segment_ids: &[String],
        writer: &mut (impl std::io::Write + ?Sized),
    ) -> std::io::Result<()> {
        let state = self.state.lock();
        write_pk_bloom(writer, segment_ids, &state.bloom)
    }

    /// Memory used by the bloom filter and latest staged map.
    pub fn memory_bytes(&self) -> usize {
        let state = self.state.lock();
        state.bloom.size_bytes()
            + state.pending_bytes
            + (state.uncommitted.capacity() - state.uncommitted.len()) * KEY_SLOT_BYTES
            + state.cancelled_bytes
            + state.delete_bytes
            + state.deletes.capacity() * std::mem::size_of::<String>()
            + self
                .committed_data
                .iter()
                .map(|data| {
                    data.content_hash
                        .as_ref()
                        .map_or(0, |lookup| lookup.memory_bytes())
                        + data
                            .alive_docs
                            .as_ref()
                            .map_or(0, |bits| bits.bits.len() * 8)
                        + data
                            .live_key_ordinals
                            .as_ref()
                            .map_or(0, |bits| bits.bits.len() * 8)
                })
                .sum::<usize>()
    }

    /// Resolve committed content only when there is no latest staged row.
    pub(super) fn content_hash_target(
        &self,
        key: &str,
    ) -> Result<Option<super::content_hash::ContentHashTarget>> {
        let state = self.state.lock();
        if !state.bloom.may_contain(key.as_bytes()) {
            return Ok(None);
        }
        if state.uncommitted.contains_key(key.as_bytes()) {
            return Ok(None);
        }
        if state.deletes.contains(key) {
            return Ok(None);
        }
        drop(state);
        let mut target = None;
        for data in &self.committed_data {
            let column = data
                .fast_fields
                .get(&self.field.0)
                .ok_or_else(|| Error::Corruption("primary-key column is missing".into()))?;
            if let Some(ordinal) = column.text_ordinal(key)
                && data
                    .live_key_ordinals
                    .as_ref()
                    .is_none_or(|keys| keys.contains(ordinal as u32))
            {
                if target.is_some() {
                    return Err(Error::Corruption(
                        "multiple live segments share a primary key".into(),
                    ));
                }
                let lookup = data.content_hash.as_ref().ok_or_else(|| {
                    Error::Internal("content hash lookup was not initialized".into())
                })?;
                target = Some(super::content_hash::ContentHashTarget {
                    store: std::sync::Arc::clone(&lookup.store),
                    row: lookup.row(ordinal, column, data)?,
                    #[cfg(feature = "native")]
                    snapshot: self._snapshot.clone(),
                });
            }
        }
        Ok(target)
    }

    /// Stage deletion of the latest row, including an unpublished insertion.
    pub(crate) fn delete(&self, key: &str) -> Result<bool> {
        validate_mutation_key(key)?;
        let mut state = self.state.lock();
        check_delete_budget(&state, key)?;
        if let Some(old) = state.uncommitted.get(key.as_bytes()) {
            check_pending_budget(&state, old.bytes, 0, CANCELLED_ROW_BYTES)?;
        }
        let staged = stage_delete(&mut state, key);
        if let Some(old) = state.uncommitted.remove(key.as_bytes()) {
            state.pending_bytes -= old.bytes;
            state.cancelled_bytes += CANCELLED_ROW_BYTES;
            old.row.cancel();
        }
        Ok(staged)
    }

    /// Some(false) means a staged version exists and committed content must
    /// not be consulted. Missing hashes are never equality assertions.
    pub(super) fn staged_hash_matches(&self, key: &str, hash: &FieldValue) -> Option<bool> {
        self.state
            .lock()
            .uncommitted
            .get(key.as_bytes())
            .map(|pending| pending.hash.as_ref() == Some(hash))
    }

    pub(super) fn admit_document(
        &self,
        doc: Document,
        schema: &Schema,
        replace: bool,
        accept: impl FnOnce(Document, Arc<StagedRow>) -> Result<()>,
    ) -> Result<()> {
        let key = document_key(&doc, self.field)?.to_owned();
        let hash = super::content_hash::document_hash(&doc, schema)?;
        // Check the budget before cloning a caller-controlled hash.
        let bytes = pending_bytes(&key, hash);
        let mut state = self.state.lock();
        self.check_admission(&state, &key, replace, bytes)?;
        let hash = hash.cloned();
        let row = Arc::new(StagedRow::default());
        accept(doc, Arc::clone(&row))?;
        self.finish_admission(&mut state, key, PendingKey { row, hash, bytes }, replace);
        Ok(())
    }

    fn check_admission(
        &self,
        state: &PrimaryKeyState,
        key: &str,
        replace: bool,
        bytes: usize,
    ) -> Result<()> {
        if replace {
            check_delete_budget(state, key)?;
        } else if state.uncommitted.contains_key(key.as_bytes()) {
            return Err(Error::DuplicatePrimaryKey(key.to_owned()));
        } else if !state.deletes.contains(key) && state.bloom.may_contain(key.as_bytes()) {
            for data in &self.committed_data {
                if let Some(ff) = data.fast_fields.get(&self.field.0)
                    && let Some(ordinal) = ff.text_ordinal(key)
                    && data
                        .live_key_ordinals
                        .as_ref()
                        .is_none_or(|keys| keys.contains(ordinal as u32))
                {
                    return Err(Error::DuplicatePrimaryKey(key.to_owned()));
                }
            }
        }
        let old_bytes = state
            .uncommitted
            .get(key.as_bytes())
            .map_or(0, |old| old.bytes);
        check_pending_budget(
            state,
            old_bytes,
            bytes,
            if old_bytes > 0 {
                CANCELLED_ROW_BYTES
            } else {
                0
            },
        )?;
        Ok(())
    }

    fn finish_admission(
        &self,
        state: &mut PrimaryKeyState,
        key: String,
        pending: PendingKey,
        replace: bool,
    ) {
        if replace {
            stage_delete(state, &key);
        }
        state.bloom.insert(key.as_bytes());
        state.pending_bytes += pending.bytes;
        if let Some(old) = state.uncommitted.insert(key.into_bytes(), pending) {
            state.pending_bytes -= old.bytes;
            state.cancelled_bytes += CANCELLED_ROW_BYTES;
            old.row.cancel();
        }
    }

    // Clear at the publication boundary, independently of the fallible PK
    // cache refresh. Retrying a refresh must never delete the replacement rows.
    pub(crate) fn mark_deletes_published(&self) {
        let mut state = self.state.lock();
        state.deletes.clear();
        state.delete_bytes = 0;
    }

    pub(crate) fn pending_deletes(&self) -> Vec<String> {
        self.state.lock().deletes.iter().cloned().collect()
    }

    /// Check whether a document's primary key is unique, and if so, register it.
    ///
    /// Returns `Ok(())` if the key is new (inserted into bloom + latest staged map).
    /// Returns `Err(DuplicatePrimaryKey)` if the key already exists.
    /// Returns `Err(Document)` if the primary key field is missing or empty.
    pub fn check_and_insert(&self, doc: &Document) -> Result<()> {
        let key = document_key(doc, self.field)?;
        let bytes = pending_bytes(key, None);
        let mut state = self.state.lock();
        self.check_admission(&state, key, false, bytes)?;
        self.finish_admission(
            &mut state,
            key.to_owned(),
            PendingKey {
                row: Arc::new(StagedRow::default()),
                hash: None,
                bytes,
            },
            false,
        );
        Ok(())
    }

    /// Refresh after commit: merge new segment data, prune removed segments,
    /// insert new keys into bloom, and clear latest staged map.
    ///
    /// Only `new_data` (segments not already held) need to be loaded by the
    /// caller. Existing data for segments still in `snapshot` is retained.
    /// The snapshot keeps ref counts alive so segments aren't deleted.
    #[cfg(feature = "native")]
    pub fn refresh_incremental(&mut self, new_data: Vec<PkSegmentData>, snapshot: SegmentSnapshot) {
        self.refresh_data(new_data, snapshot.segment_ids());
        self._snapshot = Some(std::sync::Arc::new(snapshot));
    }

    pub(crate) fn refresh_data(&mut self, new_data: Vec<PkSegmentData>, segment_ids: &[String]) {
        // Insert new segments' keys into bloom (these were uncommitted before).
        // get_mut() bypasses the mutex — safe because we have &mut self.
        let state = self.state.get_mut();
        for data in &new_data {
            if let Some(ff) = data.fast_fields.get(&self.field.0)
                && let Some(dict) = ff.text_dict()
            {
                for key in dict.iter() {
                    state.bloom.insert(key.as_bytes());
                }
            }
        }
        state.uncommitted.clear();
        state.pending_bytes = 0;
        state.cancelled_bytes = 0;
        state.deletes.clear();
        state.delete_bytes = 0;
        self.replace_committed_data(new_data, segment_ids);
    }

    /// Refresh segment readers after a topology-only replacement.
    ///
    /// Merge/reorder outputs contain only keys from their sources (compaction
    /// can remove deleted keys). Their keys are already represented in the
    /// monotonic bloom filter, and any live ingestion reservations must remain
    /// registered while only the committed segment topology changes.
    #[cfg(feature = "native")]
    pub fn refresh_replacement(&mut self, new_data: Vec<PkSegmentData>, snapshot: SegmentSnapshot) {
        self.replace_committed_data(new_data, snapshot.segment_ids());
        self._snapshot = Some(std::sync::Arc::new(snapshot));
    }

    fn replace_committed_data(&mut self, new_data: Vec<PkSegmentData>, segment_ids: &[String]) {
        let new_seg_ids: FxHashSet<&str> = segment_ids.iter().map(|s| s.as_str()).collect();
        let replaced: FxHashSet<&str> = new_data
            .iter()
            .map(|data| data.segment_id.as_str())
            .collect();
        let mut kept: Vec<PkSegmentData> = self
            .committed_data
            .drain(..)
            .filter(|d| {
                new_seg_ids.contains(d.segment_id.as_str())
                    && !replaced.contains(d.segment_id.as_str())
            })
            .collect();
        kept.extend(new_data);
        self.committed_data = kept;
    }

    #[cfg(all(feature = "wasm", not(feature = "native")))]
    pub(crate) fn segment_data(&self) -> &[PkSegmentData] {
        &self.committed_data
    }

    /// Iterator over segment IDs already held in this PK index.
    pub fn committed_segment_ids(&self) -> impl Iterator<Item = &str> {
        self.committed_data.iter().map(|d| d.segment_id.as_str())
    }

    pub(crate) fn committed_visibility(
        &self,
    ) -> impl Iterator<Item = (&str, Option<&crate::segment::DeletionMeta>)> {
        self.committed_data
            .iter()
            .map(|data| (data.segment_id.as_str(), data.deletion_meta.as_ref()))
    }

    /// Roll back an uncommitted key registration (e.g. when channel send fails
    /// after check_and_insert succeeded). Bloom may retain the key but that only
    /// causes harmless false positives, never missed duplicates.
    pub fn rollback_uncommitted_key(&self, doc: &crate::dsl::Document) {
        if let Some(value) = doc.get_first(self.field)
            && let Some(key) = value.as_text()
        {
            let mut state = self.state.lock();
            if let Some(old) = state.uncommitted.remove(key.as_bytes()) {
                state.pending_bytes -= old.bytes;
                state.cancelled_bytes += CANCELLED_ROW_BYTES;
                old.row.cancel();
            }
        }
    }

    /// Clear uncommitted keys (e.g. on abort). Bloom may retain stale entries
    /// but that only causes harmless false positives (extra committed-segment
    /// lookups), never missed duplicates.
    pub fn clear_uncommitted(&mut self) {
        let state = self.state.get_mut();
        state.uncommitted.clear();
        state.pending_bytes = 0;
        state.cancelled_bytes = 0;
        state.deletes.clear();
        state.delete_bytes = 0;
    }
}

/// Write a bloom filter with the segment IDs it covers in `pk_bloom.bin` format.
///
/// Layout: `[magic:u32][num_segs:u32][seg_id_hex × 32 bytes each...][bloom_bytes...]`
fn write_pk_bloom(
    writer: &mut (impl std::io::Write + ?Sized),
    segment_ids: &[String],
    bloom: &BloomFilter,
) -> std::io::Result<()> {
    writer.write_u32::<LittleEndian>(PK_BLOOM_MAGIC)?;
    writer.write_u32::<LittleEndian>(u32::try_from(segment_ids.len()).map_err(|_| {
        std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            "primary-key bloom segment count exceeds u32::MAX",
        )
    })?)?;
    for seg_id in segment_ids {
        let bytes = seg_id.as_bytes();
        if bytes.len() > 32 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "primary-key bloom segment ID exceeds 32 bytes",
            ));
        }
        writer.write_all(bytes)?;
        // Pad to 32 bytes (segment IDs are 32-char hex strings)
        writer.write_all(&[0u8; 32][..32 - bytes.len()])?;
    }
    bloom.write_to(writer)
}

/// Deserialize `pk_bloom.bin`. Returns the set of covered segment IDs and the bloom filter,
/// or `None` if the data is corrupt / wrong magic.
#[cfg(feature = "native")]
pub fn deserialize_pk_bloom(data: &[u8]) -> Option<(HashSet<String>, BloomFilter)> {
    if data.len() < 8 {
        return None;
    }
    let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
    if magic != PK_BLOOM_MAGIC {
        return None;
    }
    let num_segments = u32::from_le_bytes([data[4], data[5], data[6], data[7]]) as usize;
    let header_end = 8 + num_segments * 32;
    if data.len() < header_end + BloomFilter::SERIALIZED_HEADER_SIZE {
        return None;
    }
    let mut segment_ids = HashSet::with_capacity(num_segments);
    for i in 0..num_segments {
        let start = 8 + i * 32;
        let raw = &data[start..start + 32];
        let end = raw.iter().position(|&b| b == 0).unwrap_or(32);
        let hex = std::str::from_utf8(&raw[..end]).ok()?;
        segment_ids.insert(hex.to_string());
    }
    let bloom = BloomFilter::from_bytes_mutable(&data[header_end..]).ok()?;
    Some((segment_ids, bloom))
}

#[cfg(all(test, feature = "native"))]
mod tests {
    use std::sync::Arc;

    use super::*;
    use crate::dsl::{Document, Field};
    use crate::segment::SegmentTracker;

    fn make_doc(field: Field, key: &str) -> Document {
        let mut doc = Document::new();
        doc.add_text(field, key);
        doc
    }

    fn empty_snapshot() -> SegmentSnapshot {
        SegmentSnapshot::new(Arc::new(SegmentTracker::new()), vec![])
    }

    #[test]
    fn test_new_empty_readers() {
        let field = Field(0);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());
        // Should construct without panicking
        let doc = make_doc(field, "key1");
        assert!(pk.check_and_insert(&doc).is_ok());
    }

    #[test]
    fn test_unique_keys_accepted() {
        let field = Field(0);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        assert!(pk.check_and_insert(&make_doc(field, "a")).is_ok());
        assert!(pk.check_and_insert(&make_doc(field, "b")).is_ok());
        assert!(pk.check_and_insert(&make_doc(field, "c")).is_ok());
    }

    #[test]
    fn test_duplicate_uncommitted_rejected() {
        let field = Field(0);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_ok());
        let result = pk.check_and_insert(&make_doc(field, "key1"));
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::DuplicatePrimaryKey(k) => assert_eq!(k, "key1"),
            other => panic!("Expected DuplicatePrimaryKey, got {:?}", other),
        }
    }

    #[test]
    fn test_missing_field_rejected() {
        let field = Field(0);
        let other_field = Field(1);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        // Document has a different field, not the primary key field
        let doc = make_doc(other_field, "value");
        let result = pk.check_and_insert(&doc);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::Document(msg) => assert!(msg.contains("Missing"), "{}", msg),
            other => panic!("Expected Document error, got {:?}", other),
        }
    }

    #[test]
    fn test_empty_key_rejected() {
        let field = Field(0);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        let result = pk.check_and_insert(&make_doc(field, ""));
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::Document(msg) => assert!(msg.contains("empty"), "{}", msg),
            other => panic!("Expected Document error, got {:?}", other),
        }
    }

    #[test]
    fn test_clear_uncommitted() {
        let field = Field(0);
        let mut pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        // Insert key1
        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_ok());
        // Duplicate should fail
        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_err());

        // Clear uncommitted
        pk.clear_uncommitted();

        // After clear, bloom still has key1 but uncommitted doesn't.
        // With no committed readers, the key should be allowed again
        // (bloom positive → check uncommitted (not found) → check committed (empty) → accept)
        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_ok());
    }

    #[test]
    fn test_many_unique_keys() {
        let field = Field(0);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        for i in 0..1000 {
            let key = format!("key_{}", i);
            assert!(pk.check_and_insert(&make_doc(field, &key)).is_ok());
        }

        // All should be duplicates now
        for i in 0..1000 {
            let key = format!("key_{}", i);
            assert!(pk.check_and_insert(&make_doc(field, &key)).is_err());
        }
    }

    #[test]
    fn test_refresh_clears_uncommitted() {
        let field = Field(0);
        let mut pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_ok());
        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_err());

        // Refresh with empty data (simulates commit where segments
        // don't have fast fields — edge case)
        pk.refresh_incremental(vec![], empty_snapshot());

        // After refresh, uncommitted is cleared and no committed data has
        // the key, so it should be accepted again
        assert!(pk.check_and_insert(&make_doc(field, "key1")).is_ok());
    }

    #[test]
    fn replacement_refresh_preserves_uncommitted_reservations() {
        let field = Field(0);
        let mut pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());

        assert!(pk.check_and_insert(&make_doc(field, "queued")).is_ok());
        pk.refresh_replacement(vec![], empty_snapshot());

        assert!(
            pk.check_and_insert(&make_doc(field, "queued")).is_err(),
            "topology-only BP refresh must not erase a queued key reservation"
        );
    }

    #[test]
    fn test_pk_bloom_serialize_roundtrip() {
        let field = Field(0);
        let pk = PrimaryKeyIndex::new(field, vec![], empty_snapshot());
        for i in 0..100 {
            pk.check_and_insert(&make_doc(field, &format!("key_{}", i)))
                .unwrap();
        }

        let seg_ids = vec![
            "00000000000000000000000000000001".to_string(),
            "00000000000000000000000000000002".to_string(),
        ];
        let mut data = Vec::new();
        pk.write_bloom_cache(&seg_ids, &mut data).unwrap();
        let (got_ids, got_bloom) = deserialize_pk_bloom(&data).expect("deserialize failed");

        assert_eq!(got_ids.len(), 2);
        assert!(got_ids.contains(&seg_ids[0]));
        assert!(got_ids.contains(&seg_ids[1]));

        // Verify the loaded bloom recognizes previously inserted keys.
        for i in 0..100 {
            let key = format!("key_{}", i);
            assert!(
                got_bloom.may_contain(key.as_bytes()),
                "bloom miss for {}",
                key
            );
        }
    }

    #[test]
    fn test_pk_bloom_deserialize_bad_data() {
        assert!(deserialize_pk_bloom(&[]).is_none());
        assert!(deserialize_pk_bloom(&[0; 7]).is_none());
        assert!(deserialize_pk_bloom(&[0; 8]).is_none()); // wrong magic
    }

    #[test]
    fn test_concurrent_access() {
        use std::sync::Arc;

        let field = Field(0);
        let pk = Arc::new(PrimaryKeyIndex::new(field, vec![], empty_snapshot()));

        // Spawn multiple threads trying to insert the same key
        let mut handles = vec![];
        for _ in 0..10 {
            let pk = Arc::clone(&pk);
            handles.push(std::thread::spawn(move || {
                pk.check_and_insert(&make_doc(field, "contested_key"))
            }));
        }

        let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
        let successes = results.iter().filter(|r| r.is_ok()).count();
        let failures = results.iter().filter(|r| r.is_err()).count();

        // Exactly one thread should succeed, rest should get DuplicatePrimaryKey
        assert_eq!(successes, 1, "Exactly one insert should succeed");
        assert_eq!(failures, 9, "Rest should fail with duplicate");
    }
}

/// Load only fast-field data for a segment (lightweight alternative to full SegmentReader).
pub(crate) async fn load_pk_segment_data<D: crate::directories::Directory>(
    dir: &D,
    seg_id_str: &str,
    schema: &crate::dsl::Schema,
    deletion: Option<(u32, crate::segment::DeletionMeta)>,
) -> Result<PkSegmentData> {
    let seg_id = crate::segment::SegmentId::from_hex(seg_id_str)
        .ok_or_else(|| Error::Internal(format!("Invalid segment id: {}", seg_id_str)))?;
    let files = crate::segment::SegmentFiles::new(seg_id.0);
    let fast_fields =
        crate::segment::reader::loader::load_fast_fields_file(dir, &files, schema).await?;
    let (deletion_meta, alive_docs) = match deletion {
        Some((num_docs, meta)) => {
            let alive = meta.load(dir, num_docs).await?;
            (Some(meta), Some(alive))
        }
        None => (None, None),
    };
    let data = PkSegmentData {
        deletion_meta,
        alive_docs,
        live_key_ordinals: None,
        segment_id: seg_id_str.to_string(),
        content_hash: None,
        fast_fields,
    };
    if schema.content_hash_field().is_none() {
        return Ok(data);
    }
    #[cfg(feature = "native")]
    let cache = super::shared_store_cache(32 * 1024 * 1024);
    #[cfg(not(feature = "native"))]
    let cache = std::sync::Arc::new(crate::segment::SharedStoreCache::new(0));
    let store = std::sync::Arc::new(
        crate::segment::AsyncStoreReader::open(
            dir.open_lazy(&files.store).await?,
            dir as *const D as usize,
            seg_id.0,
            cache,
        )
        .await?,
    );
    let field = schema
        .primary_field()
        .ok_or_else(|| Error::Schema("content_hash requires a primary key".into()))?;
    let prepare = move || {
        let mut data = data;
        let column = data
            .fast_fields
            .get(&field.0)
            .ok_or_else(|| Error::Corruption("primary-key column is missing".into()))?;
        let lookup = super::content_hash::ContentHashLookup::new(store, column, &data)?;
        data.content_hash = Some(lookup);
        Ok(data)
    };
    #[cfg(feature = "native")]
    {
        tokio::task::spawn_blocking(prepare)
            .await
            .map_err(|error| {
                Error::Internal(format!("content hash lookup preparation failed: {error}"))
            })?
    }
    #[cfg(not(feature = "native"))]
    {
        prepare()
    }
}