velesdb-core 7.0.0

//! Backend adapter for NativeHnsw to replace hnsw_rs dependency.
//!
//! This module provides:
//! - `NativeNeighbour`: Drop-in replacement for `hnsw_rs::prelude::Neighbour`
//! - `NativeHnswBackend`: Trait for HNSW operations without hnsw_rs dependency
//! - Additional methods for `NativeHnsw` to match backend trait
//! - Parallel insertion using rayon
//!
//! Graph persistence (dump/load) is in [`super::graph_io`].
//! `BatchEfSchedule` is in [`super::batch_schedule`].

use super::batch_schedule::compute_batch_ef_schedule;
use super::distance::DistanceEngine;
use super::graph::{NativeHnsw, NO_ENTRY_POINT};
use super::layer::NodeId;
use crate::distance::DistanceMetric;
use rayon::prelude::*;
use std::path::Path;

// ============================================================================
// NativeHnswBackend Trait - Independent of hnsw_rs
// ============================================================================

/// Trait for HNSW backend operations - independent of `hnsw_rs`.
///
/// This trait mirrors `HnswBackend` but uses our own `NativeNeighbour` type,
/// allowing complete independence from the `hnsw_rs` crate.
///
/// # Thread Safety
///
/// All implementations must be `Send + Sync` to support concurrent access.
pub trait NativeHnswBackend: Send + Sync {
    /// Searches the HNSW graph and returns neighbors with distances.
    ///
    /// # Arguments
    ///
    /// * `query` - The query vector
    /// * `k` - Number of nearest neighbors to return
    /// * `ef_search` - Search expansion factor (higher = more accurate, slower)
    fn search(&self, query: &[f32], k: usize, ef_search: usize) -> Vec<NativeNeighbour>;

    /// Inserts a single vector into the HNSW graph.
    ///
    /// # Arguments
    ///
    /// * `data` - Tuple of (vector slice, internal index). The index is not
    ///   used: the graph allocates the slot itself (#2246). A caller that needs
    ///   the slot calls `NativeHnsw::insert`, which returns it.
    ///
    /// # Errors
    ///
    /// Returns an error if allocation or insertion fails.
    fn insert(&self, data: (&[f32], usize)) -> crate::error::Result<()>;

    /// Batch parallel insert into the HNSW graph.
    ///
    /// Returns the slot the graph gave each vector, one per input vector, in
    /// the same order as `data`. The index in each tuple is not used: the
    /// graph allocates the slots itself, and the returned ones are those to map
    /// ids to (#2246).
    ///
    /// # Errors
    ///
    /// Returns an error if any insertion fails.
    fn parallel_insert(&self, data: &[(&[f32], usize)]) -> crate::error::Result<Vec<usize>>;

    /// Sets the index to searching mode after bulk insertions.
    fn set_searching_mode(&mut self, mode: bool);

    /// Dumps the HNSW graph to files for persistence.
    ///
    /// # Errors
    ///
    /// Returns `io::Error` if file operations fail.
    fn file_dump(&self, path: &Path, basename: &str) -> std::io::Result<()>;

    /// Transforms raw distance to appropriate score based on metric type.
    ///
    /// For Euclidean metric, assumes the input is **squared L2** as produced
    /// by `CachedSimdDistance`. Other distance engines that already return
    /// actual Euclidean distance
    /// should **not** have their results passed through this function, as
    /// it would incorrectly apply `sqrt()` to an already-sqrt'd value.
    fn transform_score(&self, raw_distance: f32) -> f32;

    /// Returns the number of elements in the index.
    fn len(&self) -> usize;

    /// Returns true if the index is empty.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

/// Native neighbour type - drop-in replacement for `hnsw_rs::prelude::Neighbour`.
///
/// This allows `NativeHnsw` to implement `HnswBackend` without depending on `hnsw_rs`.
#[derive(Debug, Clone, PartialEq)]
pub struct NativeNeighbour {
    /// Data index (maps to external ID via `HnswIndex` mappings)
    pub d_id: usize,
    /// Distance from query vector
    pub distance: f32,
}

impl NativeNeighbour {
    /// Creates a new neighbour result.
    #[must_use]
    pub fn new(d_id: usize, distance: f32) -> Self {
        Self { d_id, distance }
    }
}

// ============================================================================
// Extended NativeHnsw methods for HnswBackend compatibility
// ============================================================================

/// Below this many nodes a batch is linked one node at a time: connecting in
/// parallel costs more than it saves.
///
/// `pub(crate)` because it decides two things that must agree: whether
/// `place_batch` enters rayon at all, and therefore whether
/// `HnswIndex::insert_batch_parallel` needs the dedicated pool. Reading the
/// same constant in both places keeps a batch from being dispatched to a pool
/// it will not use — a flat ~32 us of thread hand-off, measured, which costs
/// a small batch several times what placing it does.
/// `HnswIndex::insert_batch_parallel` carries the measured table; it is not
/// repeated here, because the figure this comment used to carry was a stale
/// copy of it and said the opposite of the measurement.
pub(crate) const PARALLEL_BATCH_MIN: usize = 100;

/// The query each node of a batch connects with, in assignment order.
#[derive(Clone, Copy)]
enum BatchQueries<'a> {
    /// The vectors the batch was placed from, each prepared as an insert
    /// prepares its vector ([`NativeHnsw::with_prepared_query`]).
    Raw(&'a [&'a [f32]]),
    /// Vectors read back from the arena, which stores each one in its
    /// prepared form.
    Prepared(&'a [&'a [f32]]),
}

impl<D: DistanceEngine + Send + Sync> NativeHnsw<D> {
    /// Parallel batch insert using rayon.
    ///
    /// Inserts multiple vectors in parallel for better throughput on multi-core systems.
    /// Returns the slot the graph gave each vector, one per input vector in order.
    ///
    /// # Arguments
    ///
    /// * `data` - Slice of (vector reference, internal index) pairs. The index
    ///   is not used: the graph allocates each slot itself, and the returned
    ///   slots are the ones to map ids to (#2246).
    ///
    /// # Errors
    ///
    /// Returns an error if any insertion fails.
    ///
    /// # Note
    ///
    /// Graph structure may differ from sequential insertion due to concurrent
    /// neighbor selection; every node stays reachable from the entry point
    /// either way (#2259).
    pub fn parallel_insert(&self, data: &[(&[f32], usize)]) -> crate::error::Result<Vec<usize>> {
        let vectors: Vec<&[f32]> = data.iter().map(|&(vector, _)| vector).collect();
        self.place_batch(&vectors)
    }

    /// Places a batch of vectors in the graph and returns the slot each one
    /// got, in input order (#2246).
    ///
    /// Batches under [`PARALLEL_BATCH_MIN`] vectors go one by one; larger ones
    /// are pushed into the arena in one allocation and connected in parallel.
    ///
    /// # Errors
    ///
    /// Returns an error if any insertion fails.
    pub(crate) fn place_batch(&self, vectors: &[&[f32]]) -> crate::error::Result<Vec<usize>> {
        // For small batches, sequential is faster due to parallelization overhead
        if vectors.len() < PARALLEL_BATCH_MIN {
            let mut assigned_ids = Vec::with_capacity(vectors.len());
            for vec in vectors {
                assigned_ids.push(self.insert(vec)?);
            }
            return Ok(assigned_ids);
        }

        // Phase A: Batch allocate — stores vectors, assigns layers (single lock scopes)
        let assignments = self.allocate_batch(vectors)?;
        if assignments.is_empty() {
            return Ok(Vec::new());
        }

        self.connect_batch(&assignments, BatchQueries::Raw(vectors));

        // Invalidate GPU caches — topology and vectors both changed.
        // Single `insert()` does this per-call; batch path must do it once
        // after all nodes are connected to avoid stale CSR/vector snapshots.
        // `connect_batch` has already released every `Vectors` write lock
        // taken during the insert loop, so the helper's mutex acquisition
        // is a flat acquire (rank 5) with nothing on the lock stack —
        // consistent with the declared global order.
        #[cfg(feature = "gpu")]
        self.invalidate_gpu_caches();

        // Return the graph-assigned node IDs in input order
        let assigned_ids: Vec<usize> = assignments.iter().map(|(node_id, _)| *node_id).collect();
        Ok(assigned_ids)
    }

    /// Links into the graph `nodes` the direct writer placed with
    /// [`Self::push_unlinked`], where they are: nothing is pushed (#2264).
    ///
    /// Each node connects with the vector its slot holds. The arena stores a
    /// vector in the form an insert prepares it in ([`Self::stores_unit_norm`]
    /// decides both), so the stored vector is read back and not prepared a
    /// second time. The nodes are linked the way a batch of their number is
    /// placed: one at a time below [`PARALLEL_BATCH_MIN`], in parallel chunks
    /// from there.
    ///
    /// The caller passes each node once, none of them linked yet, and holds
    /// the index read guard, so no renumbering moves a slot in between.
    ///
    /// # Errors
    ///
    /// Returns an error, and links nothing, if the arena does not hold one of
    /// `nodes`.
    // Reason: its one caller, the async builder's drain, is persistence-gated.
    #[cfg_attr(not(feature = "persistence"), allow(dead_code))]
    pub(in crate::index::hnsw) fn link_placed(&self, nodes: &[NodeId]) -> crate::error::Result<()> {
        if nodes.is_empty() {
            return Ok(());
        }
        let (stored, dimension) = self
            .with_vectors_read(|arena| {
                let mut stored = Vec::with_capacity(nodes.len() * arena.dimension());
                for &node in nodes {
                    stored.extend_from_slice(arena.get(node)?);
                }
                Some((stored, arena.dimension()))
            })
            .ok_or_else(|| {
                crate::error::Error::Internal("a node to link is not in the arena".to_string())
            })?;
        let queries: Vec<&[f32]> = stored.chunks_exact(dimension).collect();
        if nodes.len() < PARALLEL_BATCH_MIN {
            for (&node, query) in nodes.iter().zip(&queries) {
                self.link_node(node, query);
            }
        } else {
            let assignments = self.assign_layers(nodes);
            self.connect_batch(&assignments, BatchQueries::Prepared(&queries));
        }
        // The topology changed, and the vectors with it: `push_unlinked`
        // placed them without invalidating.
        #[cfg(feature = "gpu")]
        self.invalidate_gpu_caches();
        Ok(())
    }

    /// Phases B and C of a batch insert, for nodes the arena already holds:
    /// the first node claims an empty graph's entry point, the others connect
    /// in chunks, and the batch's highest node is promoted.
    fn connect_batch(&self, assignments: &[(NodeId, usize)], queries: BatchQueries<'_>) {
        let connect_start = self.bootstrap_entry_point(assignments);
        self.connect_batch_chunked(assignments, queries, connect_start);
        self.finalize_batch(assignments, connect_start);
    }

    /// Establishes the first node as entry point if the index is empty.
    ///
    /// Returns the number of nodes consumed by bootstrapping (0 or 1).
    /// Consumed nodes are excluded from the parallel connect phase because
    /// they have no valid entry point to search from.
    ///
    /// One CAS decides: the first node is consumed only when it claimed an
    /// empty graph's entry point. A batch that loses that claim to a
    /// concurrent insert connects its first node like the others, where it
    /// used to leave it out of the connect phase with no edge (#2259).
    pub(in crate::index::hnsw::native) fn bootstrap_entry_point(
        &self,
        assignments: &[(NodeId, usize)],
    ) -> usize {
        let (node_id, layer) = assignments[0];
        usize::from(self.claim_first_entry_point(node_id, layer).is_none())
    }

    /// Final promotion of the highest-layer node and bootstrap count update.
    ///
    /// Called after `connect_batch_chunked` completes. Ensures the global
    /// entry point reflects the best candidate across the entire batch, and
    /// accounts for any bootstrapped node that was not counted by the
    /// chunked phase.
    fn finalize_batch(&self, assignments: &[(NodeId, usize)], connect_start: usize) {
        if let Some(best) = assignments.iter().max_by_key(|(_, layer)| *layer) {
            self.promote_entry_point(best.0, best.1);
        }
        if connect_start > 0 {
            self.count
                .fetch_add(connect_start, std::sync::atomic::Ordering::Relaxed);
        }
    }

    /// Computes the chunk size for batched Phase B insertion.
    ///
    /// Balances parallelism (larger chunks) against entry-point staleness
    /// (smaller chunks refresh the EP more often). The formula scales
    /// linearly with batch size, clamped to `[1000, 5000]`.
    #[must_use]
    pub(in crate::index::hnsw::native) fn compute_chunk_size(batch_len: usize) -> usize {
        const DEFAULT_CHUNK: usize = 1000;
        const MAX_CHUNK: usize = 5000;
        (batch_len / 50).clamp(DEFAULT_CHUNK, MAX_CHUNK)
    }

    /// Computes the effective `ef_construction` for a batch of the given size.
    ///
    /// For large batches, the full `ef_construction` search budget is wasteful
    /// because the graph scaffold built by earlier vectors already provides
    /// sufficient connectivity for neighbor discovery. Reducing the beam width
    /// proportionally to batch size matches the strategy used by Qdrant and
    /// hnswlib for bulk loading.
    ///
    /// The returned value is always >= `max_connections` to guarantee that
    /// each inserted node can discover enough neighbors for a well-connected
    /// graph.
    ///
    /// Returns `(effective_ef, stagnation_limit)`.
    #[allow(dead_code)] // Reason: Tested in backend_adapter_tests; wired in batch insert V2
    #[must_use]
    pub(in crate::index::hnsw::native) fn adaptive_ef_for_batch(
        &self,
        batch_size: usize,
    ) -> (usize, usize) {
        let base = self.ef_construction;

        // Conservative scaling: the original 0.25/0.50 reduction destroyed
        // graph quality at 100K+ (recall dropped from 97% to 64%).
        // Malkov & Yashunin 2018 recommends ef_construction >= 2*M;
        // these floors keep ef well above that while still accelerating
        // bulk loads vs single-insert.
        let scale = if batch_size > 50_000 {
            0.60
        } else if batch_size > 10_000 {
            0.75
        } else if batch_size > 1_000 {
            0.85
        } else {
            return (base, 0);
        };

        // Reason: f64 product of two small positive values fits in usize.
        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
        let scaled = (base as f64 * scale) as usize;

        // Floor at 4*M to guarantee adequate neighbor diversity budget.
        let effective_ef = scaled.max(self.max_connections * 4);

        // Stagnation-based early termination: ef/2 gives the beam search
        // enough runway to escape local clusters at scale (was ef/3, which
        // caused premature termination at 100K+).
        let stagnation = effective_ef / 2;

        (effective_ef, stagnation)
    }

    /// Connects nodes in chunks, refreshing the entry point between chunks.
    ///
    /// Each chunk runs `par_iter` over its assignments, then promotes the
    /// highest-layer node and increments the count. This keeps the entry
    /// point fresher than a single monolithic `par_iter` over the entire
    /// batch, improving recall for large insertions.
    ///
    /// For batches > 1K vectors, uses adaptive `ef_construction` reduction
    /// to lower the search budget proportionally, matching the bulk-loading
    /// strategies of Qdrant and hnswlib. Single-vector insert is unaffected.
    ///
    /// # Errors
    ///
    /// Returns an error if any node connection fails.
    /// Connects nodes in chunks with graduated ef\_construction.
    ///
    /// Uses a 3-phase schedule (VAMANA/DiskANN pattern):
    /// - **Phase 1** (first 10%): full ef — builds a quality scaffold
    /// - **Phase 2** (10%-90%): reduced ef (0.5x) — graph is dense enough
    /// - **Phase 3** (last 10%): moderate ef (0.75x) — finalizes connections
    ///
    /// Skips the first `connect_start` assignments, which bootstrapping
    /// consumed. Each node connects with its entry in `queries`: a raw input
    /// vector is prepared on demand via `with_prepared_query`, which reuses
    /// one thread-local scratch buffer per rayon worker instead of holding one
    /// owned normalized buffer per batch element alive (PERF2); a vector read
    /// back from the arena is used as it is.
    fn connect_batch_chunked(
        &self,
        assignments: &[(NodeId, usize)],
        queries: BatchQueries<'_>,
        connect_start: usize,
    ) {
        let first_node = assignments[0].0;
        let to_connect = &assignments[connect_start..];
        let chunk_size = Self::compute_chunk_size(to_connect.len());
        let schedule =
            compute_batch_ef_schedule(self.ef_construction, to_connect.len(), self.max_connections);
        let mut nodes_connected: usize = 0;

        for chunk in to_connect.chunks(chunk_size) {
            let loaded = self.entry_point.load(std::sync::atomic::Ordering::Acquire);
            let ep_id = if loaded == NO_ENTRY_POINT {
                first_node
            } else {
                loaded
            };

            let chunk_offset = nodes_connected;

            // Each node is anchored as soon as one of its neighbours is; batch
            // mates still connecting are anchored by the pass after the chunk
            // (#2259).
            let pending: Vec<NodeId> = chunk
                .par_iter()
                .enumerate()
                .filter_map(|(i, (node_id, layer))| {
                    let position = connect_start + chunk_offset + i;
                    let anchored = self.with_batch_query(queries, position, |query| {
                        let current_ep = self.greedy_descent_upper_layers(query, *layer, ep_id);
                        let ef = schedule.ef_for_position(chunk_offset + i);
                        let stagnation = ef / 2;
                        self.connect_node_with_ef(
                            *node_id, query, *layer, current_ep, ef, stagnation,
                        );
                        self.with_vectors_and_layers_read(|arena, layers| {
                            self.anchor_node(*node_id, arena, layers)
                        })
                    });
                    (!anchored).then_some(*node_id)
                })
                .collect();
            self.anchor_pending(pending);

            if let Some(best) = chunk.iter().max_by_key(|(_, layer)| *layer) {
                self.promote_entry_point(best.0, best.1);
            }
            self.count
                .fetch_add(chunk.len(), std::sync::atomic::Ordering::Relaxed);
            nodes_connected += chunk.len();
        }
    }

    /// Runs `f` on the query the node at `position` in a batch connects with.
    fn with_batch_query<R>(
        &self,
        queries: BatchQueries<'_>,
        position: usize,
        f: impl FnOnce(&[f32]) -> R,
    ) -> R {
        match queries {
            BatchQueries::Raw(vectors) => self.with_prepared_query(vectors[position], f),
            BatchQueries::Prepared(queries) => f(queries[position]),
        }
    }

    /// Sets the index to searching mode after bulk insertions.
    ///
    /// For `NativeHnsw`, this is currently a no-op as our implementation
    /// doesn't require mode switching. Kept for API compatibility.
    ///
    /// # Arguments
    ///
    /// * `_mode` - `true` to enable searching mode, `false` to disable
    pub fn set_searching_mode(&mut self, _mode: bool) {
        // No-op for NativeHnsw - our implementation doesn't need mode switching
        // hnsw_rs uses this to optimize internal data structures after bulk insert
    }

    /// Searches and returns results in `NativeNeighbour` format.
    ///
    /// This is the HnswBackend-compatible search method.
    #[must_use]
    pub fn search_neighbours(
        &self,
        query: &[f32],
        k: usize,
        ef_search: usize,
    ) -> Vec<NativeNeighbour> {
        self.search(query, k, ef_search)
            .into_iter()
            .map(|(id, dist)| NativeNeighbour::new(id, dist))
            .collect()
    }

    /// Transforms raw distance to appropriate score based on metric type.
    ///
    /// - **Cosine** and **Jaccard**: `1.0 - distance`, clamped through
    ///   [`DistanceMetric::clamp_score`]. The engine traverses on
    ///   `1 - similarity` for both, and `DistanceMetric::higher_is_better()`
    ///   declares both similarities, so every downstream sort/top-k/merge
    ///   expects the similarity back — but the two do **not** share a range,
    ///   and taking it from the metric is what keeps this path agreeing with
    ///   brute force and the exact rerank. Flooring cosine at zero, as a
    ///   shared arm with Jaccard once did, tied every anti-correlated match at
    ///   `0.0` and made the score depend on which path `k` happened to select.
    /// - **Euclidean**: `sqrt(raw_distance)` — the search loop stores squared L2
    ///   to skip redundant sqrt during traversal; this restores the actual
    ///   Euclidean distance for user-visible scores.
    /// - **Hamming**: raw distance (lower is better)
    /// - **DotProduct**: `-distance` (negated for consistency)
    #[must_use]
    pub fn transform_score(&self, raw_distance: f32) -> f32 {
        let metric = self.distance.metric();
        match metric {
            DistanceMetric::Cosine | DistanceMetric::Jaccard => {
                metric.clamp_score(1.0 - raw_distance)
            }
            // Reason: CachedSimdDistance stores squared L2 during HNSW traversal
            // to avoid per-comparison sqrt. Apply sqrt here on the final k results.
            DistanceMetric::Euclidean => raw_distance.sqrt(),
            DistanceMetric::Hamming => raw_distance,
            DistanceMetric::DotProduct => -raw_distance,
        }
    }
}

// ============================================================================
// NativeHnswBackend implementation for NativeHnsw
// ============================================================================

impl<D: DistanceEngine + Send + Sync> NativeHnswBackend for NativeHnsw<D> {
    fn search(&self, query: &[f32], k: usize, ef_search: usize) -> Vec<NativeNeighbour> {
        self.search_neighbours(query, k, ef_search)
    }

    fn insert(&self, data: (&[f32], usize)) -> crate::error::Result<()> {
        NativeHnsw::insert(self, data.0)?;
        Ok(())
    }

    fn parallel_insert(&self, data: &[(&[f32], usize)]) -> crate::error::Result<Vec<usize>> {
        NativeHnsw::parallel_insert(self, data)
    }

    fn set_searching_mode(&mut self, mode: bool) {
        NativeHnsw::set_searching_mode(self, mode);
    }

    fn file_dump(&self, path: &Path, basename: &str) -> std::io::Result<()> {
        NativeHnsw::file_dump(self, path, basename)
    }

    fn transform_score(&self, raw_distance: f32) -> f32 {
        NativeHnsw::transform_score(self, raw_distance)
    }

    fn len(&self) -> usize {
        NativeHnsw::len(self)
    }
}