diskann_disk/storage/quant/compressor.rs
1/*
2 * Copyright (c) Microsoft Corporation.
3 * Licensed under the MIT license.
4 */
5
6use diskann::{utils::VectorRepr, ANNResult};
7use diskann_utils::views::{MatrixView, MutMatrixView};
8
9/// [`CompressionStage`] defines the stage of compression
10/// used by the compressor to determine how to instantiate the quantizer
11/// (e.g., whether to initialize codebooks for a product quantizer).
12/// Passed to [`QuantCompressor::new_at_stage`] when creating quantizer instances.
13#[derive(Clone, Copy, Debug, PartialEq)]
14pub enum CompressionStage {
15 Start,
16 Resume,
17}
18
19/// [`QuantCompressor`] defines the interface for quantizer with [`QuantDataGenerator`]
20///
21/// This trait serves as a general wrapper for different quantizers, allowing them to be
22/// used interchangeably with QuantDataGenerator. Any type implementing this trait
23/// can be used to compress vector data during the data generation process.
24///
25/// # Type Parameters
26/// - `T`: The data type of the input vectors. Must impl Copy + Into<f32> + Pod + Sync
27/// so that the [`QuantDataGenerator`] can parallelize computation, call compress_into and read from data file.
28///
29/// # Associated Types
30/// - [`CompressorContext`]: An overloadable type that provides initialization parameters for the compressor
31///
32/// # Methods
33/// - `new_at_stage`: Constructs a new compressor instance with the provided context for the stage of compression.
34/// - `compress`: Compresses a batch of vectors into the output buffer.
35/// - `compressed_bytes`: Returns the size in bytes of each compressed vector
36pub trait QuantCompressor<T>: Sized + Sync
37where
38 T: VectorRepr,
39{
40 type CompressorContext;
41
42 fn new_at_stage(stage: CompressionStage, context: &Self::CompressorContext) -> ANNResult<Self>;
43 fn compress(&self, vector: MatrixView<f32>, output: MutMatrixView<u8>) -> ANNResult<()>;
44 fn compressed_bytes(&self) -> usize;
45}