[TQDT] TurboVectorStorage Implementation (#9331)

* Partial function impls for TurboVectorStorage

* create/load functionality + tests

# Conflicts:
#	lib/segment/src/vector_storage/turbo/mod.rs
#	lib/segment/src/vector_storage/turbo/turbo_encoded_vectors.rs

* Fix vectors are not rotated back

* update_from impl

* Rebase fixes

* Proper update_from for ChunkedMmap + better tests

* Improve test coverage

* Move update_from to new DenseTQVectorStorage trait

* Review remark quantization::DistanceType

* Add populate and clear_cache to TurboVectorStorage and TurboEncodedVectorStorage

lib/quantization/src/encoded_storage.rs

@@ -20,6 +20,8 @@ use fs_err::File;
 pub trait EncodedStorage {
     fn get_vector_data(&self, index: PointOffsetType) -> Cow<'_, [u8]>;
 
+    fn get_vector_data_opt(&self, index: PointOffsetType) -> Option<Cow<'_, [u8]>>;
+
     fn iter_batch(
         &self,
         offsets: &[PointOffsetType],
@@ -136,6 +138,11 @@ impl TestEncodedStorage {
 #[cfg(feature = "testing")]
 impl EncodedStorage for TestEncodedStorage {
     fn get_vector_data(&self, index: PointOffsetType) -> Cow<'_, [u8]> {
+        self.get_vector_data_opt(index)
+            .unwrap_or(Cow::Borrowed(&[]))
+    }
+
+    fn get_vector_data_opt(&self, index: PointOffsetType) -> Option<Cow<'_, [u8]>> {
         let start = self
             .quantized_vector_size
             .get()
@@ -145,7 +152,7 @@ impl EncodedStorage for TestEncodedStorage {
             .get()
             .saturating_mul(index as usize + 1);
 
-        Cow::Borrowed(self.data.get(start..end).unwrap_or(&[]))
+        Some(Cow::Borrowed(self.data.get(start..end)?))
     }
 
     fn upsert_vector(

lib/quantization/src/turboquant/quantization.rs

@@ -1,3 +1,5 @@
+use num_traits::AsPrimitive;
+
 use crate::DistanceType;
 use crate::turboquant::encoding::TqVectorExtras;
 use crate::turboquant::rotation::HadamardRotation;
@@ -10,7 +12,7 @@ use crate::turboquant::{EncodedQueryTQ, EncodedQueryTQData, TQBits, TQMode};
 
 /// Quantize vectors using TurboQuant.
 pub struct TurboQuantizer {
-    pub(super) rotation: HadamardRotation,
+    pub rotation: HadamardRotation,
     pub(super) bits: TQBits,
     pub(super) mode: TQMode,
     pub(super) distance: DistanceType,
@@ -284,7 +286,11 @@ impl TurboQuantizer {
         norm
     }
 
-    pub fn dequantize(&self, quantized: &[u8]) -> Vec<f64> {
+    pub fn dequantize<T>(&self, quantized: &[u8]) -> Vec<T>
+    where
+        T: Copy + 'static,
+        f64: AsPrimitive<T>,
+    {
         let (unpacked_iter, extras) = self.unpack_vector(quantized);
         let scaling_factor = f64::from(extras.scaling_factor());
         // Materialize the unpacked centroids once. `unpack_vector` returns a
@@ -329,13 +335,17 @@ impl TurboQuantizer {
                 .enumerate()
                 .map(|(i, x)| {
                     let rescaled = x / f64::from(ec.scale[i]) - f64::from(ec.shift[i]);
-                    rescaled * scale
+                    (rescaled * scale).as_()
                 })
                 .collect(),
-            None => unpacked.into_iter().map(|x| x * scale).collect(),
+            None => unpacked.into_iter().map(|x| (x * scale).as_()).collect(),
         }
     }
 
+    pub fn get_padded_dim(&self) -> usize {
+        self.padded_dim
+    }
+
     /// Similarity score between two vectors that were both encoded with this
     /// quantizer. Returns an approximate `<v1, v2>` for Dot and `cos(θ)` for
     /// Cosine.

lib/segment/src/types.rs

@@ -364,6 +364,22 @@ impl Distance {
     }
 }
 
+/// Map a segment [`Distance`] to the TurboQuant [`DistanceType`].
+///
+/// Uses the true Cosine mapping (`Cosine → Cosine`); the legacy quantizers fold
+/// Cosine into Dot for backwards-compat, but do so with an explicit match rather
+/// than this conversion.
+impl From<Distance> for quantization::DistanceType {
+    fn from(distance: Distance) -> Self {
+        match distance {
+            Distance::Cosine => quantization::DistanceType::Cosine,
+            Distance::Euclid => quantization::DistanceType::L2,
+            Distance::Dot => quantization::DistanceType::Dot,
+            Distance::Manhattan => quantization::DistanceType::L1,
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Clone, Copy)]
 pub enum Order {
     LargeBetter,

lib/segment/src/vector_storage/quantized/quantized_chunked_mmap_storage/read_only.rs

@@ -44,9 +44,11 @@ impl<S: UniversalRead> QuantizedChunkedStorageRead<S> {
 
 impl<S: UniversalRead> quantization::EncodedStorage for QuantizedChunkedStorageRead<S> {
     fn get_vector_data(&self, index: PointOffsetType) -> Cow<'_, [u8]> {
-        self.data
-            .get::<Random>(index as VectorOffsetType)
-            .unwrap_or_default()
+        self.get_vector_data_opt(index).unwrap_or_default()
+    }
+
+    fn get_vector_data_opt(&self, index: PointOffsetType) -> Option<Cow<'_, [u8]>> {
+        self.data.get::<Random>(index as VectorOffsetType)
     }
 
     fn iter_batch(

lib/segment/src/vector_storage/quantized/quantized_chunked_mmap_storage/read_write.rs

@@ -55,9 +55,11 @@ impl<S: UniversalWrite + Send + 'static> quantization::EncodedStorage
     for QuantizedChunkedStorage<S>
 {
     fn get_vector_data(&self, index: PointOffsetType) -> Cow<'_, [u8]> {
-        self.data
-            .get::<Random>(index as VectorOffsetType)
-            .unwrap_or_default()
+        self.get_vector_data_opt(index).unwrap_or_default()
+    }
+
+    fn get_vector_data_opt(&self, index: PointOffsetType) -> Option<Cow<'_, [u8]>> {
+        self.data.get::<Random>(index as VectorOffsetType)
     }
 
     fn iter_batch(

lib/segment/src/vector_storage/quantized/quantized_ram_storage.rs

@@ -71,6 +71,12 @@ impl quantization::EncodedStorage for QuantizedRamStorage {
         Cow::Borrowed(self.vectors.get(index as VectorOffsetType))
     }
 
+    fn get_vector_data_opt(&self, index: PointOffsetType) -> Option<Cow<'_, [u8]>> {
+        Some(Cow::Borrowed(
+            self.vectors.get_opt(index as VectorOffsetType)?,
+        ))
+    }
+
     fn upsert_vector(
         &mut self,
         id: PointOffsetType,

lib/segment/src/vector_storage/quantized/quantized_storage.rs

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::io::BufWriter;
 use std::marker::PhantomData;
 use std::num::NonZeroUsize;
 use std::path::{Path, PathBuf};
@@ -7,7 +8,9 @@ use common::counter::hardware_counter::HardwareCounterCell;
 use common::generic_consts::Random;
 use common::mmap::{AdviceSetting, MmapFlusher, advice};
 use common::types::PointOffsetType;
-use common::universal_io::{OpenOptions, Populate, ReadOnly, ReadRange, UniversalRead};
+use common::universal_io::{
+    MmapFile, MmapFs, OpenOptions, Populate, ReadOnly, ReadRange, UniversalRead,
+};
 use fs_err as fs;
 use memmap2::MmapMut;
 
@@ -39,6 +42,28 @@ impl<S: UniversalRead> QuantizedStorage<S> {
     }
 }
 
+impl QuantizedStorage<MmapFile> {
+    /// Open the backing file for build-time bulk appends, bypassing the read-only mmap.
+    pub(crate) fn open_appender(&self) -> std::io::Result<BufWriter<fs::File>> {
+        Ok(BufWriter::new(open_append(&self.path)?))
+    }
+
+    /// Re-mmap after the file grew so reads observe appended vectors. Build-time only.
+    pub(crate) fn reload(&mut self) -> OperationResult<()> {
+        *self = Self::from_file(
+            &MmapFs,
+            &self.path.clone(),
+            self.quantized_vector_size.get(),
+        )?;
+        Ok(())
+    }
+}
+
+/// Open a file shortly for appending.
+fn open_append(path: &Path) -> std::io::Result<fs::File> {
+    fs::OpenOptions::new().append(true).open(path)
+}
+
 pub struct QuantizedStorageBuilder<S> {
     mmap: MmapMut,
     cursor_pos: usize,
@@ -82,10 +107,42 @@ impl<S: UniversalRead> QuantizedStorage<S> {
             path: path.to_path_buf(),
         })
     }
+
+    /// Open the encoded vectors at `path`, creating an empty storage if the file does not yet exist.
+    pub fn open(
+        fs: &S::Fs,
+        path: &Path,
+        quantized_vector_size: usize,
+        prefault: bool,
+    ) -> OperationResult<QuantizedStorage<S>> {
+        if let Some(parent) = path.parent() {
+            fs::create_dir_all(parent)?;
+        }
+
+        // Ensure the backing file exists without clobbering existing data:
+        // `from_file` mmaps the file read-only and fails if it is missing.
+        fs_err::OpenOptions::new()
+            .write(true)
+            .create(true)
+            .truncate(false)
+            .open(path)?;
+
+        let storage = Self::from_file(fs, path, quantized_vector_size)?;
+
+        if prefault {
+            storage.populate();
+        }
+
+        Ok(storage)
+    }
 }
 
 impl<S: UniversalRead> quantization::EncodedStorage for QuantizedStorage<S> {
     fn get_vector_data(&self, index: PointOffsetType) -> Cow<'_, [u8]> {
+        self.get_vector_data_opt(index).expect("vector exists")
+    }
+
+    fn get_vector_data_opt(&self, index: PointOffsetType) -> Option<Cow<'_, [u8]>> {
         let start = (self.quantized_vector_size.get() * index as usize) as u64;
         let length = self.quantized_vector_size.get() as u64;
         self.storage
@@ -93,7 +150,7 @@ impl<S: UniversalRead> quantization::EncodedStorage for QuantizedStorage<S> {
                 byte_offset: start,
                 length,
             })
-            .expect("vector exists")
+            .ok()
     }
 
     fn upsert_vector(

lib/segment/src/vector_storage/turbo/turbo_encoded_vectors.rs

@@ -1,52 +1,91 @@
 use std::borrow::Cow;
-use std::path::PathBuf;
+use std::io::{self, Write};
+use std::ops::Range;
+use std::path::{Path, PathBuf};
+use std::sync::atomic::AtomicBool;
 
+use common::counter::hardware_counter::HardwareCounterCell;
 use common::mmap::MmapFlusher;
 use common::types::PointOffsetType;
-use common::universal_io::MmapFile;
+use common::universal_io::{MmapFile, MmapFs};
 use quantization::EncodedStorage;
 
+use crate::common::operation_error::{OperationResult, check_process_stopped};
 use crate::vector_storage::quantized::quantized_chunked_mmap_storage::QuantizedChunkedStorage;
-use crate::vector_storage::quantized::quantized_ram_storage::QuantizedRamStorage;
 use crate::vector_storage::quantized::quantized_storage::QuantizedStorage;
 
 /// Raw quantized storage backend for the TurboQuant-encoded bytes.
 pub(super) enum TurboEncodedVectorStorage {
-    /// In-memory encoded vectors.
-    Ram(QuantizedRamStorage),
     /// Single mem-mapped file of encoded vectors.
     Mmap(QuantizedStorage<MmapFile>),
+
     /// Chunked mem-mapped encoded vectors (appendable).
     ChunkedMmap(QuantizedChunkedStorage<MmapFile>),
 }
 
 impl TurboEncodedVectorStorage {
+    /// Open (create-or-load) the single mem-mapped file backend (non-appendable).
+    pub(super) fn open_mmap(
+        path: &Path,
+        quantized_vector_size: usize,
+        populate: bool,
+    ) -> OperationResult<Self> {
+        Ok(Self::Mmap(QuantizedStorage::<MmapFile>::open(
+            &MmapFs,
+            path,
+            quantized_vector_size,
+            populate,
+        )?))
+    }
+
+    /// Open (create-or-load) the appendable chunked mem-mapped backend.
+    pub(super) fn open_chunked_mmap(
+        path: &Path,
+        quantized_vector_size: usize,
+        in_ram: bool,
+    ) -> OperationResult<Self> {
+        Ok(Self::ChunkedMmap(QuantizedChunkedStorage::new(
+            MmapFs,
+            path,
+            quantized_vector_size,
+            in_ram,
+        )?))
+    }
+
     /// Raw encoded blob for one vector (no dequantization).
     pub(super) fn get_quantized_vector(&self, key: PointOffsetType) -> Cow<'_, [u8]> {
         match self {
-            Self::Ram(s) => s.get_vector_data(key),
             Self::Mmap(s) => s.get_vector_data(key),
             Self::ChunkedMmap(s) => s.get_vector_data(key),
         }
     }
 
+    /// Raw encoded blob for one vector (no dequantization).
+    pub(super) fn get_quantized_vector_opt(&self, key: PointOffsetType) -> Option<Cow<'_, [u8]>> {
+        match self {
+            Self::Mmap(s) => s.get_vector_data_opt(key),
+            Self::ChunkedMmap(s) => s.get_vector_data_opt(key),
+        }
+    }
+
     /// Number of encoded vectors (including soft-deleted).
     pub(super) fn vectors_count(&self) -> usize {
         match self {
-            Self::Ram(s) => s.vectors_count(),
             Self::Mmap(s) => s.vectors_count(),
             Self::ChunkedMmap(s) => s.vectors_count(),
         }
     }
 
     pub(super) fn is_on_disk(&self) -> bool {
-        unimplemented!("TurboEncodedVectorStorage::is_on_disk")
+        match self {
+            Self::Mmap(s) => s.is_on_disk(),
+            Self::ChunkedMmap(s) => s.is_on_disk(),
+        }
     }
 
     /// All on-disk files backing the encoded vectors.
     pub(super) fn files(&self) -> Vec<PathBuf> {
         match self {
-            Self::Ram(s) => s.files(),
             Self::Mmap(s) => s.files(),
             Self::ChunkedMmap(s) => s.files(),
         }
@@ -54,7 +93,6 @@ impl TurboEncodedVectorStorage {
 
     pub(super) fn immutable_files(&self) -> Vec<PathBuf> {
         match self {
-            Self::Ram(s) => s.immutable_files(),
             Self::Mmap(s) => s.immutable_files(),
             Self::ChunkedMmap(s) => s.immutable_files(),
         }
@@ -62,9 +100,103 @@ impl TurboEncodedVectorStorage {
 
     pub(super) fn flusher(&self) -> MmapFlusher {
         match self {
-            Self::Ram(s) => s.flusher(),
             Self::Mmap(s) => s.flusher(),
             Self::ChunkedMmap(s) => s.flusher(),
         }
     }
+
+    /// Populate all pages of the encoded vectors into the page cache.
+    pub(super) fn populate(&self) -> OperationResult<()> {
+        match self {
+            Self::Mmap(s) => s.populate(),
+            Self::ChunkedMmap(s) => s.populate()?,
+        }
+        Ok(())
+    }
+
+    /// Drop the disk cache for the encoded vectors.
+    pub(super) fn clear_cache(&self) -> OperationResult<()> {
+        match self {
+            Self::Mmap(s) => s.clear_cache(),
+            Self::ChunkedMmap(s) => s.clear_cache()?,
+        }
+        Ok(())
+    }
+
+    pub(super) fn upsert_vector(
+        &mut self,
+        id: PointOffsetType,
+        vector: &[u8],
+        hw_counter: &HardwareCounterCell,
+    ) -> std::io::Result<()> {
+        match self {
+            TurboEncodedVectorStorage::Mmap(storage) => {
+                // We let the underlying storage decide the error instead of doing it here.
+                // Therefore, we don't assume it's read-only here and pretend to write.
+                storage.upsert_vector(id, vector, hw_counter)
+            }
+            TurboEncodedVectorStorage::ChunkedMmap(storage) => {
+                storage.upsert_vector(id, vector, hw_counter)
+            }
+        }
+    }
+
+    /// Bulk-ingest already-encoded vectors, dispatching to the backend implementation.
+    pub(super) fn update_from<'a>(
+        &mut self,
+        vectors: impl Iterator<Item = Cow<'a, [u8]>>,
+        stopped: &AtomicBool,
+    ) -> OperationResult<Range<PointOffsetType>> {
+        match self {
+            Self::Mmap(storage) => Self::update_from_mmap(storage, vectors, stopped),
+            Self::ChunkedMmap(storage) => Self::update_from_chunked_mmap(storage, vectors, stopped),
+        }
+    }
+
+    /// Single-file backend: bulk-append encoded bytes to the file, then re-mmap once
+    /// (mirrors `DenseVectorStorageImpl::update_from`).
+    fn update_from_mmap<'a>(
+        storage: &mut QuantizedStorage<MmapFile>,
+        vectors: impl Iterator<Item = Cow<'a, [u8]>>,
+        stopped: &AtomicBool,
+    ) -> OperationResult<Range<PointOffsetType>> {
+        let start_index = storage.vectors_count() as PointOffsetType;
+        let mut end_index = start_index;
+
+        let mut writer = storage.open_appender()?;
+        for vector in vectors {
+            check_process_stopped(stopped)?;
+            writer.write_all(&vector)?;
+            end_index += 1;
+        }
+
+        // Persist + re-mmap so reads observe the appended vectors.
+        writer.flush()?;
+        let file = writer
+            .into_inner()
+            .map_err(io::IntoInnerError::into_error)?;
+        file.sync_data()?;
+        storage.reload()?;
+
+        Ok(start_index..end_index)
+    }
+
+    /// Chunked backend: append each encoded vector through the chunked structure.
+    fn update_from_chunked_mmap<'a>(
+        storage: &mut QuantizedChunkedStorage<MmapFile>,
+        vectors: impl Iterator<Item = Cow<'a, [u8]>>,
+        stopped: &AtomicBool,
+    ) -> OperationResult<Range<PointOffsetType>> {
+        let disposed_hw = HardwareCounterCell::disposable();
+        let start_index = storage.vectors_count() as PointOffsetType;
+        let mut key = start_index;
+
+        for vector in vectors {
+            check_process_stopped(stopped)?;
+            storage.upsert_vector(key, &vector, &disposed_hw)?;
+            key += 1;
+        }
+
+        Ok(start_index..key)
+    }
 }