qdrant-edge 0.8.0

use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;

use crate::common::counter::hardware_counter::HardwareCounterCell;
use crate::common::types::{PointOffsetType, ScoredPointOffset, TelemetryDetail};
use crate::common::universal_io::MmapFile;
use half::f16;
use crate::sparse::common::types::{DimId, QuantizedU8};
use crate::sparse::index::inverted_index::InvertedIndex;
use crate::sparse::index::inverted_index::inverted_index_compressed_immutable_ram::InvertedIndexCompressedImmutableRam;
use crate::sparse::index::inverted_index::inverted_index_compressed_mmap::InvertedIndexCompressedMmap;
use crate::sparse::index::inverted_index::inverted_index_ram::InvertedIndexRam;

use super::hnsw_index::hnsw::HNSWIndex;
use super::plain_vector_index::PlainVectorIndex;
use super::sparse_index::sparse_vector_index::SparseVectorIndex;
use crate::segment::common::operation_error::OperationResult;
use crate::segment::data_types::query_context::VectorQueryContext;
use crate::segment::data_types::vectors::{QueryVector, VectorRef};
use crate::segment::telemetry::VectorIndexSearchesTelemetry;
use crate::segment::types::{Filter, SearchParams};

/// Read-only trait for vector index.
///
/// Defines all read operations on a vector index. Search and retrieval logic
/// only requires this trait, which makes it possible to implement read-only
/// segments without duplicating index code.
pub trait VectorIndexRead {
    /// Return list of Ids with fitting
    fn search(
        &self,
        vectors: &[&QueryVector],
        filter: Option<&Filter>,
        top: usize,
        params: Option<&SearchParams>,
        query_context: &VectorQueryContext,
    ) -> OperationResult<Vec<Vec<ScoredPointOffset>>>;

    fn get_telemetry_data(&self, detail: TelemetryDetail) -> VectorIndexSearchesTelemetry;

    /// The number of indexed vectors, currently accessible
    fn indexed_vector_count(&self) -> usize;

    /// Total size of all searchable vectors in bytes.
    fn size_of_searchable_vectors_in_bytes(&self) -> usize;

    /// Augment the IDF stats for the given dimensions over the given corpus
    /// and return the number of documents contributing to them: indexed
    /// vectors matching the corpus filter, or all indexed vectors when
    /// `corpus` is `None` (global statistics).
    ///
    /// Most indexes don't track IDF and should contribute no df counts.
    /// Sparse-vector indexes are the only ones that contribute. No default is
    /// provided on purpose so a new index implementation cannot silently skip
    /// this.
    fn fill_idf_statistics(
        &self,
        idf: &mut HashMap<DimId, usize>,
        corpus: Option<&Filter>,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<usize>;

    /// Whether this is a "real" index rather than a plain (full-scan) one.
    ///
    /// Used by reporting code to decide whether to count vectors as indexed.
    fn is_index(&self) -> bool;
}

/// Trait for vector index with mutating operations.
pub trait VectorIndex: VectorIndexRead {
    fn files(&self) -> Vec<PathBuf>;

    fn immutable_files(&self) -> Vec<PathBuf> {
        Vec::new()
    }

    /// Update index for a single vector
    ///
    /// # Arguments
    /// - `id` - sequential vector id, offset in the vector storage
    /// - `vector` - new vector value,
    ///   if None - vector will be removed from the index marked as deleted in storage.
    ///   Note: inserting None vector is not equal to removing vector from the storage.
    ///   Unlike removing, it will always result in storage growth.
    ///   Proper removing should be performed by the optimizer.
    fn update_vector(
        &mut self,
        id: PointOffsetType,
        vector: Option<VectorRef>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()>;

    /// Byte-blob analogue of [`VectorIndex::update_vector`]: `vector` is the
    /// storage-native serialized form (the [`retrieve_raw`] format), letting
    /// requantized storages ingest bytes verbatim instead of a lossy
    /// decode/re-encode round-trip. `None` behaves exactly like
    /// `update_vector(id, None, _)`.
    ///
    /// [`retrieve_raw`]: crate::entry::entry_point::ReadSegmentEntry::retrieve_raw
    fn update_vector_raw(
        &mut self,
        id: PointOffsetType,
        vector: Option<&[u8]>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()>;
}

#[derive(Debug)]
pub enum VectorIndexEnum {
    Plain(PlainVectorIndex),
    Hnsw(HNSWIndex),
    SparseRam(SparseVectorIndex<InvertedIndexRam>),
    SparseCompressedImmutableRamF32(SparseVectorIndex<InvertedIndexCompressedImmutableRam<f32>>),
    SparseCompressedImmutableRamF16(SparseVectorIndex<InvertedIndexCompressedImmutableRam<f16>>),
    SparseCompressedImmutableRamU8(
        SparseVectorIndex<InvertedIndexCompressedImmutableRam<QuantizedU8>>,
    ),
    SparseCompressedMmapF32(SparseVectorIndex<InvertedIndexCompressedMmap<f32, MmapFile>>),
    SparseCompressedMmapF16(SparseVectorIndex<InvertedIndexCompressedMmap<f16, MmapFile>>),
    SparseCompressedMmapU8(SparseVectorIndex<InvertedIndexCompressedMmap<QuantizedU8, MmapFile>>),
}

impl VectorIndexEnum {
    /// Returns true if underlying storage is configured to be stored on disk without
    /// actively holding data in RAM
    pub fn is_on_disk(&self) -> bool {
        match self {
            Self::Plain(_) => false,
            Self::Hnsw(index) => index.is_on_disk(),
            Self::SparseRam(index) => index.inverted_index().is_on_disk(),
            Self::SparseCompressedImmutableRamF32(index) => index.inverted_index().is_on_disk(),
            Self::SparseCompressedImmutableRamF16(index) => index.inverted_index().is_on_disk(),
            Self::SparseCompressedImmutableRamU8(index) => index.inverted_index().is_on_disk(),
            Self::SparseCompressedMmapF32(index) => index.inverted_index().is_on_disk(),
            Self::SparseCompressedMmapF16(index) => index.inverted_index().is_on_disk(),
            Self::SparseCompressedMmapU8(index) => index.inverted_index().is_on_disk(),
        }
    }

    pub fn populate(&self) -> OperationResult<()> {
        match self {
            Self::Plain(_) => {}
            Self::Hnsw(index) => index.populate()?,
            Self::SparseRam(_) => {}
            Self::SparseCompressedImmutableRamF32(_) => {}
            Self::SparseCompressedImmutableRamF16(_) => {}
            Self::SparseCompressedImmutableRamU8(_) => {}
            Self::SparseCompressedMmapF32(index) => index.inverted_index().populate()?,
            Self::SparseCompressedMmapF16(index) => index.inverted_index().populate()?,
            Self::SparseCompressedMmapU8(index) => index.inverted_index().populate()?,
        };
        Ok(())
    }

    pub fn clear_cache(&self) -> OperationResult<()> {
        match self {
            Self::Plain(_) => {}
            Self::Hnsw(index) => index.clear_cache()?,
            Self::SparseRam(_) => {}
            Self::SparseCompressedImmutableRamF32(_) => {}
            Self::SparseCompressedImmutableRamF16(_) => {}
            Self::SparseCompressedImmutableRamU8(_) => {}
            Self::SparseCompressedMmapF32(index) => index.inverted_index().clear_cache()?,
            Self::SparseCompressedMmapF16(index) => index.inverted_index().clear_cache()?,
            Self::SparseCompressedMmapU8(index) => index.inverted_index().clear_cache()?,
        };
        Ok(())
    }

    pub fn as_hnsw(&self) -> Option<&HNSWIndex> {
        match self {
            VectorIndexEnum::Plain(_) => None,
            VectorIndexEnum::Hnsw(index) => Some(index),
            VectorIndexEnum::SparseRam(_) => None,
            VectorIndexEnum::SparseCompressedImmutableRamF32(_) => None,
            VectorIndexEnum::SparseCompressedImmutableRamF16(_) => None,
            VectorIndexEnum::SparseCompressedImmutableRamU8(_) => None,
            VectorIndexEnum::SparseCompressedMmapF32(_) => None,
            VectorIndexEnum::SparseCompressedMmapF16(_) => None,
            VectorIndexEnum::SparseCompressedMmapU8(_) => None,
        }
    }
}

impl VectorIndexRead for VectorIndexEnum {
    fn search(
        &self,
        vectors: &[&QueryVector],
        filter: Option<&Filter>,
        top: usize,
        params: Option<&SearchParams>,
        query_context: &VectorQueryContext,
    ) -> OperationResult<Vec<Vec<ScoredPointOffset>>> {
        match self {
            VectorIndexEnum::Plain(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::Hnsw(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseRam(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseCompressedImmutableRamF32(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseCompressedImmutableRamF16(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseCompressedImmutableRamU8(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseCompressedMmapF32(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseCompressedMmapF16(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
            VectorIndexEnum::SparseCompressedMmapU8(index) => {
                index.search(vectors, filter, top, params, query_context)
            }
        }
    }

    fn get_telemetry_data(&self, detail: TelemetryDetail) -> VectorIndexSearchesTelemetry {
        match self {
            VectorIndexEnum::Plain(index) => index.get_telemetry_data(detail),
            VectorIndexEnum::Hnsw(index) => index.get_telemetry_data(detail),
            VectorIndexEnum::SparseRam(index) => index.get_telemetry_data(detail),
            VectorIndexEnum::SparseCompressedImmutableRamF32(index) => {
                index.get_telemetry_data(detail)
            }
            VectorIndexEnum::SparseCompressedImmutableRamF16(index) => {
                index.get_telemetry_data(detail)
            }
            VectorIndexEnum::SparseCompressedImmutableRamU8(index) => {
                index.get_telemetry_data(detail)
            }
            VectorIndexEnum::SparseCompressedMmapF32(index) => index.get_telemetry_data(detail),
            VectorIndexEnum::SparseCompressedMmapF16(index) => index.get_telemetry_data(detail),
            VectorIndexEnum::SparseCompressedMmapU8(index) => index.get_telemetry_data(detail),
        }
    }

    fn indexed_vector_count(&self) -> usize {
        match self {
            Self::Plain(index) => index.indexed_vector_count(),
            Self::Hnsw(index) => index.indexed_vector_count(),
            Self::SparseRam(index) => index.indexed_vector_count(),
            Self::SparseCompressedImmutableRamF32(index) => index.indexed_vector_count(),
            Self::SparseCompressedImmutableRamF16(index) => index.indexed_vector_count(),
            Self::SparseCompressedImmutableRamU8(index) => index.indexed_vector_count(),
            Self::SparseCompressedMmapF32(index) => index.indexed_vector_count(),
            Self::SparseCompressedMmapF16(index) => index.indexed_vector_count(),
            Self::SparseCompressedMmapU8(index) => index.indexed_vector_count(),
        }
    }

    fn size_of_searchable_vectors_in_bytes(&self) -> usize {
        match self {
            Self::Plain(index) => index.size_of_searchable_vectors_in_bytes(),
            Self::Hnsw(index) => index.size_of_searchable_vectors_in_bytes(),
            Self::SparseRam(index) => index.size_of_searchable_vectors_in_bytes(),
            Self::SparseCompressedImmutableRamF32(index) => {
                index.size_of_searchable_vectors_in_bytes()
            }
            Self::SparseCompressedImmutableRamF16(index) => {
                index.size_of_searchable_vectors_in_bytes()
            }
            Self::SparseCompressedImmutableRamU8(index) => {
                index.size_of_searchable_vectors_in_bytes()
            }
            Self::SparseCompressedMmapF32(index) => index.size_of_searchable_vectors_in_bytes(),
            Self::SparseCompressedMmapF16(index) => index.size_of_searchable_vectors_in_bytes(),
            Self::SparseCompressedMmapU8(index) => index.size_of_searchable_vectors_in_bytes(),
        }
    }

    fn is_index(&self) -> bool {
        match self {
            Self::Plain(_) => false,
            Self::Hnsw(_) => true,
            Self::SparseRam(_) => true,
            Self::SparseCompressedImmutableRamF32(_) => true,
            Self::SparseCompressedImmutableRamF16(_) => true,
            Self::SparseCompressedImmutableRamU8(_) => true,
            Self::SparseCompressedMmapF32(_) => true,
            Self::SparseCompressedMmapF16(_) => true,
            Self::SparseCompressedMmapU8(_) => true,
        }
    }

    fn fill_idf_statistics(
        &self,
        idf: &mut HashMap<DimId, usize>,
        corpus: Option<&Filter>,
        is_stopped: &AtomicBool,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<usize> {
        match self {
            // Dense indexes contribute no df counts; document count matches
            // the previous segment-level `indexed_vector_count` accounting.
            Self::Plain(index) => Ok(match corpus {
                None => index.indexed_vector_count(),
                Some(_) => 0,
            }),
            Self::Hnsw(index) => Ok(match corpus {
                None => index.indexed_vector_count(),
                Some(_) => 0,
            }),
            Self::SparseRam(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
            Self::SparseCompressedImmutableRamF32(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
            Self::SparseCompressedImmutableRamF16(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
            Self::SparseCompressedImmutableRamU8(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
            Self::SparseCompressedMmapF32(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
            Self::SparseCompressedMmapF16(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
            Self::SparseCompressedMmapU8(index) => {
                index.fill_idf_statistics(idf, corpus, is_stopped, hw_counter)
            }
        }
    }
}

impl VectorIndex for VectorIndexEnum {
    fn files(&self) -> Vec<PathBuf> {
        match self {
            VectorIndexEnum::Plain(index) => index.files(),
            VectorIndexEnum::Hnsw(index) => index.files(),
            VectorIndexEnum::SparseRam(index) => index.files(),
            VectorIndexEnum::SparseCompressedImmutableRamF32(index) => index.files(),
            VectorIndexEnum::SparseCompressedImmutableRamF16(index) => index.files(),
            VectorIndexEnum::SparseCompressedImmutableRamU8(index) => index.files(),
            VectorIndexEnum::SparseCompressedMmapF32(index) => index.files(),
            VectorIndexEnum::SparseCompressedMmapF16(index) => index.files(),
            VectorIndexEnum::SparseCompressedMmapU8(index) => index.files(),
        }
    }

    fn immutable_files(&self) -> Vec<PathBuf> {
        match self {
            VectorIndexEnum::Plain(index) => index.immutable_files(),
            VectorIndexEnum::Hnsw(index) => index.immutable_files(),
            VectorIndexEnum::SparseRam(index) => index.immutable_files(),
            VectorIndexEnum::SparseCompressedImmutableRamF32(index) => index.immutable_files(),
            VectorIndexEnum::SparseCompressedImmutableRamF16(index) => index.immutable_files(),
            VectorIndexEnum::SparseCompressedImmutableRamU8(index) => index.immutable_files(),
            VectorIndexEnum::SparseCompressedMmapF32(index) => index.immutable_files(),
            VectorIndexEnum::SparseCompressedMmapF16(index) => index.immutable_files(),
            VectorIndexEnum::SparseCompressedMmapU8(index) => index.immutable_files(),
        }
    }

    fn update_vector(
        &mut self,
        id: PointOffsetType,
        vector: Option<VectorRef>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        match self {
            Self::Plain(index) => index.update_vector(id, vector, hw_counter),
            Self::Hnsw(index) => index.update_vector(id, vector, hw_counter),
            Self::SparseRam(index) => index.update_vector(id, vector, hw_counter),
            Self::SparseCompressedImmutableRamF32(index) => {
                index.update_vector(id, vector, hw_counter)
            }
            Self::SparseCompressedImmutableRamF16(index) => {
                index.update_vector(id, vector, hw_counter)
            }
            Self::SparseCompressedImmutableRamU8(index) => {
                index.update_vector(id, vector, hw_counter)
            }
            Self::SparseCompressedMmapF32(index) => index.update_vector(id, vector, hw_counter),
            Self::SparseCompressedMmapF16(index) => index.update_vector(id, vector, hw_counter),
            Self::SparseCompressedMmapU8(index) => index.update_vector(id, vector, hw_counter),
        }
    }

    fn update_vector_raw(
        &mut self,
        id: PointOffsetType,
        vector: Option<&[u8]>,
        hw_counter: &HardwareCounterCell,
    ) -> OperationResult<()> {
        match self {
            Self::Plain(index) => index.update_vector_raw(id, vector, hw_counter),
            Self::Hnsw(index) => index.update_vector_raw(id, vector, hw_counter),
            Self::SparseRam(index) => index.update_vector_raw(id, vector, hw_counter),
            Self::SparseCompressedImmutableRamF32(index) => {
                index.update_vector_raw(id, vector, hw_counter)
            }
            Self::SparseCompressedImmutableRamF16(index) => {
                index.update_vector_raw(id, vector, hw_counter)
            }
            Self::SparseCompressedImmutableRamU8(index) => {
                index.update_vector_raw(id, vector, hw_counter)
            }
            Self::SparseCompressedMmapF32(index) => index.update_vector_raw(id, vector, hw_counter),
            Self::SparseCompressedMmapF16(index) => index.update_vector_raw(id, vector, hw_counter),
            Self::SparseCompressedMmapU8(index) => index.update_vector_raw(id, vector, hw_counter),
        }
    }
}