velesdb-core 5.0.0

//! Label, sparse-vector, and deferred-indexing helpers for CRUD operations.
//!
//! Extracted from `crud.rs` to reduce NLOC.

use crate::collection::types::Collection;
use crate::error::Result;
use crate::point::Point;
use crate::quantization::StorageMode;
use crate::storage::VectorStorage;
use std::collections::{BTreeMap, HashMap};

use super::crud_helpers::QuantizationGuards;

impl Collection {
    /// Checks whether label index updates are needed for this batch.
    pub(super) fn needs_label_updates(
        points: &[Point],
        old_payloads: &[Option<serde_json::Value>],
    ) -> bool {
        Self::any_point_has_labels(points)
            || old_payloads
                .iter()
                .any(|opt| opt.as_ref().is_some_and(|v| v.get("_labels").is_some()))
    }

    /// Pre-allocates the label update buffer when needed.
    pub(super) fn alloc_label_buffer(
        needed: bool,
        capacity: usize,
    ) -> Vec<(u64, Option<serde_json::Value>, Option<serde_json::Value>)> {
        if needed {
            Vec::with_capacity(capacity)
        } else {
            Vec::new()
        }
    }

    /// Returns `true` if any point carries `_labels` in its payload.
    pub(super) fn any_point_has_labels(points: &[Point]) -> bool {
        points.iter().any(|p| {
            p.payload
                .as_ref()
                .is_some_and(|v| v.get("_labels").is_some())
        })
    }

    /// Resolves the effective "old payload" for a point, accounting for
    /// within-batch duplicate IDs.
    pub(super) fn resolve_effective_old<'a>(
        seen: &HashMap<u64, Option<&'a serde_json::Value>>,
        id: u64,
        pre_batch_old: Option<&'a serde_json::Value>,
    ) -> Option<&'a serde_json::Value> {
        if let Some(&inner) = seen.get(&id) {
            inner
        } else {
            pre_batch_old
        }
    }

    /// Conditionally caches a quantized vector for a single point.
    pub(super) fn maybe_quantize(
        collection: &Collection,
        point: &Point,
        storage_mode: StorageMode,
        quant_guards: &mut QuantizationGuards<'_>,
        quant_done: bool,
    ) {
        if !quant_done {
            let (sq8, binary, pq) = (
                quant_guards.sq8.as_deref_mut(),
                quant_guards.binary.as_deref_mut(),
                quant_guards.pq.as_deref_mut(),
            );
            collection.cache_quantized_vector(point, storage_mode, sq8, binary, pq);
        } else if matches!(storage_mode, StorageMode::ProductQuantization) {
            let pq = quant_guards.pq.as_deref_mut();
            collection.cache_quantized_vector(point, storage_mode, None, None, pq);
        }
    }

    /// Applies buffered label index updates in a single write lock scope.
    pub(super) fn apply_label_updates(
        label_index: &parking_lot::RwLock<crate::collection::graph::LabelIndex>,
        label_updates: &[(u64, Option<serde_json::Value>, Option<serde_json::Value>)],
    ) {
        if label_updates.is_empty() {
            return;
        }
        let mut label_idx = label_index.write();
        for (id, old, new) in label_updates {
            if let Some(old_val) = old {
                label_idx.remove_from_payload(*id, old_val);
            }
            if let Some(new_val) = new {
                label_idx.index_from_payload(*id, new_val);
            }
        }
    }

    /// Attempts parallel quantization for SQ8/Binary modes.
    pub(super) fn try_parallel_quantize(
        &self,
        points: &[Point],
        storage_mode: StorageMode,
    ) -> bool {
        #[cfg(feature = "persistence")]
        match storage_mode {
            StorageMode::SQ8 => {
                self.batch_quantize_sq8_parallel(points);
                true
            }
            StorageMode::Binary => {
                self.batch_quantize_binary_parallel(points);
                true
            }
            _ => false,
        }
        #[cfg(not(feature = "persistence"))]
        {
            let _ = (points, storage_mode);
            false
        }
    }

    /// Collects sparse vectors from a point into the batch buffer.
    pub(super) fn collect_sparse_vectors(
        point: &Point,
        sparse_batch: &mut Vec<(u64, BTreeMap<String, crate::index::sparse::SparseVector>)>,
    ) {
        if let Some(sv_map) = &point.sparse_vectors {
            if !sv_map.is_empty() {
                sparse_batch.push((point.id, sv_map.clone()));
            }
        }
    }

    /// Updates the BM25 text index for a single point (WAL-then-apply).
    ///
    /// Issue #389: appends the mutation to `bm25.wal` BEFORE calling
    /// `add_document` / `remove_document` on the in-memory index so
    /// that a crash between the two replays the mutation on next
    /// open. WAL append errors propagate; the in-memory mutation is
    /// skipped so in-memory and WAL state never diverge.
    pub(super) fn update_text_index(&self, point: &Point) -> Result<()> {
        if let Some(payload) = &point.payload {
            let text = Self::extract_text_from_payload(payload);
            if !text.is_empty() {
                #[cfg(feature = "persistence")]
                self.append_bm25_wal_add(point.id, &text)?;
                self.storage.text_index.add_document(point.id, &text);
            }
        } else {
            #[cfg(feature = "persistence")]
            self.append_bm25_wal_remove(point.id)?;
            self.storage.text_index.remove_document(point.id);
        }
        Ok(())
    }

    /// Applies the BM25 side of a whole batch under ONE durability barrier.
    ///
    /// The batched counterpart of [`Self::update_text_index`]. Calling that per
    /// point is what cost one `open` and one `fsync` PER DOCUMENT (#1797); here
    /// every frame goes into a single `wal_append_batch` call.
    ///
    /// # The three cases a mixed batch contains
    ///
    /// Classified explicitly rather than folded together, because they are not
    /// the same operation and the single-point path distinguishes them:
    ///
    /// * payload carrying text → an `Add` frame;
    /// * payload with NO text  → nothing at all, no WAL entry (empty text is
    ///   not indexed, matching [`Self::update_text_index`]);
    /// * no payload at all     → a `Remove` frame, preserving delete semantics.
    ///
    /// # Ordering
    ///
    /// WAL-before-apply, at BATCH granularity: every frame is written and
    /// fsynced first, and the in-memory index is touched only once that
    /// succeeded. A WAL failure returns `Err` with the index untouched — never
    /// partially updated for a batch that was never acknowledged, which matters
    /// because a lost WAL entry is NOT rebuilt when a BM25 snapshot exists.
    ///
    /// # Errors
    ///
    /// Propagates any WAL open / write / flush / fsync failure.
    pub(super) fn bulk_update_text_index(&self, points: &[Point]) -> Result<()> {
        let (adds, removes) = Self::classify_text_mutations(points);
        if adds.is_empty() && removes.is_empty() {
            return Ok(());
        }
        self.append_text_batch_to_wal(&adds, &removes)?;

        // Only now is the batch durable, so only now does memory change.
        for (id, text) in &adds {
            self.storage.text_index.add_document(*id, text);
        }
        for id in &removes {
            self.storage.text_index.remove_document(*id);
        }
        Ok(())
    }

    /// Splits a batch into the BM25 mutations it implies.
    ///
    /// The three cases stay separate here rather than being decided inline, so
    /// that "a payload with no indexable string writes nothing" is a visible
    /// decision and not an accident of control flow.
    fn classify_text_mutations(points: &[Point]) -> (Vec<(u64, String)>, Vec<u64>) {
        let mut adds: Vec<(u64, String)> = Vec::new();
        let mut removes: Vec<u64> = Vec::new();
        for point in points {
            if let Some(payload) = &point.payload {
                let text = Self::extract_text_from_payload(payload);
                if !text.is_empty() {
                    adds.push((point.id, text));
                }
            } else {
                removes.push(point.id);
            }
        }
        (adds, removes)
    }

    /// Writes the whole batch to the BM25 WAL under one durability barrier.
    ///
    /// Non-persistence builds have no on-disk WAL, so this is a no-op there and
    /// the caller proceeds straight to the in-memory update.
    #[allow(clippy::unused_self)] // Reason: needs `self.storage.path` under `persistence`.
    fn append_text_batch_to_wal(&self, adds: &[(u64, String)], removes: &[u64]) -> Result<()> {
        #[cfg(feature = "persistence")]
        {
            use crate::index::bm25_persistence_wal::{wal_append_batch, wal_path_for_bm25, WalOp};
            let ops: Vec<WalOp<'_>> = adds
                .iter()
                .map(|(id, text)| WalOp::Add {
                    id: *id,
                    text: text.as_str(),
                })
                .chain(removes.iter().map(|id| WalOp::Remove { id: *id }))
                .collect();
            wal_append_batch(&wal_path_for_bm25(&self.storage.path), &ops)?;
        }
        #[cfg(not(feature = "persistence"))]
        {
            let _ = (adds, removes);
        }
        Ok(())
    }

    /// Appends an `add_document` mutation to the BM25 WAL.
    ///
    /// Feature-gated — non-persistence builds have no on-disk WAL.
    /// Callers already gate on `feature = "persistence"` when they
    /// need crash-safety ordering.
    #[cfg(feature = "persistence")]
    #[inline]
    pub(super) fn append_bm25_wal_add(&self, id: u64, text: &str) -> Result<()> {
        let wal_path = crate::index::bm25_persistence_wal::wal_path_for_bm25(&self.storage.path);
        crate::index::bm25_persistence_wal::wal_append_add_document(&wal_path, id, text)
    }

    /// Appends a `remove_document` mutation to the BM25 WAL.
    #[cfg(feature = "persistence")]
    #[inline]
    pub(super) fn append_bm25_wal_remove(&self, id: u64) -> Result<()> {
        let wal_path = crate::index::bm25_persistence_wal::wal_path_for_bm25(&self.storage.path);
        crate::index::bm25_persistence_wal::wal_append_remove_document(&wal_path, id)
    }

    /// Appends `remove_document` mutations for a whole batch to the BM25 WAL
    /// under ONE durability barrier.
    ///
    /// Batched counterpart of [`Self::append_bm25_wal_remove`]: calling that
    /// in a loop pays one `open` + one fsync PER ID (the #1797 failure mode,
    /// resurfacing on the delete path — finding C3). The frames are identical
    /// to N sequential appends; only the syscall count differs. Callers keep
    /// WAL-before-apply at batch granularity: let this return `Ok` first,
    /// then remove the documents from the in-memory index.
    ///
    /// # Errors
    ///
    /// Propagates any WAL open / write / flush / fsync failure; nothing was
    /// acknowledged in that case and the in-memory index must not be touched.
    #[cfg(feature = "persistence")]
    pub(super) fn append_bm25_wal_remove_batch(&self, ids: &[u64]) -> Result<()> {
        use crate::index::bm25_persistence_wal::{wal_append_batch, wal_path_for_bm25, WalOp};
        let ops: Vec<WalOp<'_>> = ids.iter().map(|&id| WalOp::Remove { id }).collect();
        wal_append_batch(&wal_path_for_bm25(&self.storage.path), &ops)
    }

    /// Appends `(name, point_id, sparse_vector)` triples to the per-index
    /// sparse WAL under WAL-before-apply semantics.
    ///
    /// Centralises the `wal_path_for_name` + `wal_append_upsert` loop that was
    /// duplicated between `apply_sparse_batch_upsert` (single-point path) and
    /// `apply_sparse_batch_bulk` (bulk path). Callers keep ownership of their
    /// input shape (`Vec<(u64, BTreeMap)>` vs `BTreeMap<String, Vec<(u64,
    /// SparseVector)>>`) and build the iterator of triples themselves, which
    /// keeps this helper allocation-free.
    ///
    /// Feature-gated on `persistence` — on targets without persistence the
    /// sparse WAL does not exist and the caller short-circuits.
    ///
    /// Issue #450 Phase 3.1.
    #[cfg(feature = "persistence")]
    pub(super) fn append_sparse_wal_entries<'a, I>(&self, entries: I) -> Result<()>
    where
        I: IntoIterator<Item = (&'a str, u64, &'a crate::index::sparse::SparseVector)>,
    {
        // Cache wal_path across consecutive entries sharing the same index name
        // so callers that yield entries grouped by name (e.g. apply_sparse_batch_bulk)
        // retain the O(N_NAMES) path-resolution cost of the pre-refactor code
        // rather than paying O(N_ENTRIES). Mixed-name callers degrade gracefully
        // to one resolution per entry, matching the original per-triple cost.
        let mut cached: Option<(&'a str, std::path::PathBuf)> = None;
        for (name, point_id, sv) in entries {
            if cached.as_ref().map(|(cached_name, _)| *cached_name) != Some(name) {
                let wal_path =
                    crate::index::sparse::persistence::wal_path_for_name(&self.storage.path, name);
                cached = Some((name, wal_path));
            }
            let Some((_, wal_path)) = cached.as_ref() else {
                continue;
            };
            crate::index::sparse::persistence::wal_append_upsert(wal_path, point_id, sv)?;
        }
        Ok(())
    }

    /// Applies buffered sparse vector upserts with WAL-before-apply semantics.
    pub(super) fn apply_sparse_batch_upsert(
        &self,
        sparse_batch: &[(u64, BTreeMap<String, crate::index::sparse::SparseVector>)],
    ) -> Result<()> {
        if sparse_batch.is_empty() {
            return Ok(());
        }
        #[cfg(feature = "persistence")]
        {
            self.append_sparse_wal_entries(sparse_batch.iter().flat_map(|(point_id, sv_map)| {
                sv_map
                    .iter()
                    .map(move |(name, sv)| (name.as_str(), *point_id, sv))
            }))?;
        }
        let mut indexes = self.query.sparse_indexes.write();
        for (point_id, sv_map) in sparse_batch {
            for (name, sv) in sv_map {
                let idx = indexes.entry(name.clone()).or_default();
                idx.insert(*point_id, sv);
            }
        }
        Ok(())
    }

    /// Invalidates stats cache and bumps write generation.
    ///
    /// Also drops the payload mirror: any mutation path that does not
    /// explicitly maintain the mirror must invalidate it so stale columnar
    /// data can never serve queries (it is rebuilt lazily on demand).
    pub(super) fn invalidate_caches_and_bump_generation(&self) {
        *self.query.cached_stats.lock() = None;
        self.storage.payload_mirror.invalidate();
        self.generations
            .write_generation
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Like [`Self::invalidate_caches_and_bump_generation`], but keeps the
    /// payload mirror warm by applying the upserted points incrementally.
    pub(super) fn bump_generation_with_mirror_upserts(&self, points: &[crate::point::Point]) {
        *self.query.cached_stats.lock() = None;
        self.storage.payload_mirror.apply_upserts(points);
        self.generations
            .write_generation
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Like [`Self::invalidate_caches_and_bump_generation`], but keeps the
    /// payload mirror warm by tombstoning the deleted ids incrementally.
    pub(super) fn bump_generation_with_mirror_deletes(&self, ids: &[u64]) {
        *self.query.cached_stats.lock() = None;
        self.storage.payload_mirror.apply_deletes(ids);
        self.generations
            .write_generation
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
    }

    /// Drains the deferred indexer and batch-inserts into HNSW.
    #[cfg(feature = "persistence")]
    pub(super) fn merge_deferred_batch(&self, di: &crate::collection::streaming::DeferredIndexer) {
        let drained = di.swap_and_drain();
        if drained.is_empty() {
            return;
        }
        let storage = self.storage.vector_storage.read();
        let valid: Vec<(u64, &[f32])> = drained
            .iter()
            .filter(|(id, _)| storage.retrieve(*id).ok().flatten().is_some())
            .map(|(id, v)| (*id, v.as_slice()))
            .collect();
        drop(storage);
        let expected = valid.len();
        if valid.is_empty() {
            return;
        }
        let inserted = self.storage.index.insert_batch_parallel(valid);
        if inserted < expected {
            tracing::warn!("merge_deferred_batch: inserted {inserted}/{expected} vectors");
        }
    }

    /// Batch-inserts into HNSW or defers into the deferred indexer.
    pub(super) fn bulk_index_or_defer(&self, vector_refs: &[(u64, &[f32])]) -> usize {
        let count = vector_refs.len();
        #[cfg(feature = "persistence")]
        if let Some(ref di) = self.streaming.deferred_indexer {
            // PERF3: the per-vector `to_vec()` copy is intrinsic to the
            // deferred contract — `vector_refs` borrows the caller's data and
            // does not outlive this call, while the deferred buffer must own
            // the vectors until the next merge. `DeltaBuffer` stores one
            // exact-sized `Vec<f32>` per entry (its per-entry ownership is
            // what makes upsert-replace/remove O(1) data-move); the copies
            // run OUTSIDE the buffer's write lock (see `DeltaBuffer::extend`)
            // and the buffer is bounded by `merge_threshold` entries.
            di.extend(vector_refs.iter().map(|(id, v)| (*id, v.to_vec())));
            if di.should_merge() {
                self.merge_deferred_batch(di);
            }
            #[allow(clippy::cast_possible_truncation)]
            self.generations
                .inserts_since_last_hnsw_save
                .fetch_add(count as u64, std::sync::atomic::Ordering::Relaxed);
            return count;
        }
        let inserted = self
            .storage
            .index
            .insert_batch_parallel(vector_refs.iter().copied());
        #[allow(clippy::cast_possible_truncation)]
        self.generations
            .inserts_since_last_hnsw_save
            .fetch_add(count as u64, std::sync::atomic::Ordering::Relaxed);
        inserted
    }
}