Skip to main content

oxirs_vec/
compression_types.rs

1//! Codec types, compression stats, config enums, error types for vector compression.
2
3use crate::{Vector, VectorData, VectorError};
4use half::f16;
5
6/// Compression method selection
7#[derive(Debug, Clone, Default)]
8pub enum CompressionMethod {
9    #[default]
10    None,
11    Zstd {
12        level: i32,
13    },
14    Quantization {
15        bits: u8,
16    },
17    ProductQuantization {
18        subvectors: usize,
19        codebook_size: usize,
20    },
21    Pca {
22        components: usize,
23    },
24    Adaptive {
25        quality_level: AdaptiveQuality,
26        analysis_samples: usize,
27    },
28}
29
30/// Adaptive quality level for adaptive compression
31#[derive(Debug, Clone)]
32pub enum AdaptiveQuality {
33    Fast,      // Prioritize speed, moderate compression
34    Balanced,  // Balance speed and compression ratio
35    BestRatio, // Prioritize compression ratio over speed
36}
37
38/// Trait implemented by all vector compressors
39pub trait VectorCompressor: Send + Sync {
40    fn compress(&self, vector: &Vector) -> Result<Vec<u8>, VectorError>;
41    fn decompress(&self, data: &[u8], dimensions: usize) -> Result<Vector, VectorError>;
42    fn compression_ratio(&self) -> f32;
43}
44
45/// Performance metrics for a compressor instance
46#[derive(Debug, Clone)]
47pub struct CompressionMetrics {
48    pub vectors_compressed: usize,
49    pub total_original_size: usize,
50    pub total_compressed_size: usize,
51    pub compression_time_ms: f64,
52    pub decompression_time_ms: f64,
53    pub current_ratio: f32,
54    pub method_switches: usize,
55}
56
57impl Default for CompressionMetrics {
58    fn default() -> Self {
59        Self {
60            vectors_compressed: 0,
61            total_original_size: 0,
62            total_compressed_size: 0,
63            compression_time_ms: 0.0,
64            decompression_time_ms: 0.0,
65            current_ratio: 1.0,
66            method_switches: 0,
67        }
68    }
69}
70
71/// Vector characteristics analysis for adaptive compression selection
72#[derive(Debug, Clone)]
73pub struct VectorAnalysis {
74    pub sparsity: f32,
75    pub range: f32,
76    pub mean: f32,
77    pub std_dev: f32,
78    pub entropy: f32,
79    pub dominant_patterns: Vec<f32>,
80    pub recommended_method: CompressionMethod,
81    pub expected_ratio: f32,
82}
83
84impl VectorAnalysis {
85    pub fn analyze(vectors: &[Vector], quality: &AdaptiveQuality) -> Result<Self, VectorError> {
86        if vectors.is_empty() {
87            return Err(VectorError::InvalidDimensions(
88                "No vectors to analyze".to_string(),
89            ));
90        }
91
92        let mut all_values = Vec::new();
93        let mut dimensions = 0;
94
95        for vector in vectors {
96            let values = match &vector.values {
97                VectorData::F32(v) => v.clone(),
98                VectorData::F64(v) => v.iter().map(|&x| x as f32).collect(),
99                VectorData::F16(v) => v.iter().map(|&x| f16::from_bits(x).to_f32()).collect(),
100                VectorData::I8(v) => v.iter().map(|&x| x as f32).collect(),
101                VectorData::Binary(_) => {
102                    return Ok(Self::binary_analysis(vectors.len()));
103                }
104            };
105            if dimensions == 0 {
106                dimensions = values.len();
107            }
108            all_values.extend(values);
109        }
110
111        if all_values.is_empty() {
112            return Err(VectorError::InvalidDimensions(
113                "No values to analyze".to_string(),
114            ));
115        }
116
117        let min_val = all_values.iter().copied().fold(f32::INFINITY, f32::min);
118        let max_val = all_values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
119        let range = max_val - min_val;
120        let mean = all_values.iter().sum::<f32>() / all_values.len() as f32;
121
122        let variance =
123            all_values.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / all_values.len() as f32;
124        let std_dev = variance.sqrt();
125
126        let epsilon = std_dev * 0.01;
127        let near_zero_count = all_values.iter().filter(|&&x| x.abs() < epsilon).count();
128        let sparsity = near_zero_count as f32 / all_values.len() as f32;
129
130        let entropy = Self::calculate_entropy(&all_values);
131        let dominant_patterns = Self::find_dominant_patterns(&all_values);
132
133        let (recommended_method, expected_ratio) =
134            Self::select_optimal_method(sparsity, range, std_dev, entropy, dimensions, quality);
135
136        Ok(Self {
137            sparsity,
138            range,
139            mean,
140            std_dev,
141            entropy,
142            dominant_patterns,
143            recommended_method,
144            expected_ratio,
145        })
146    }
147
148    fn binary_analysis(_vector_count: usize) -> Self {
149        Self {
150            sparsity: 0.0,
151            range: 1.0,
152            mean: 0.5,
153            std_dev: 0.5,
154            entropy: 1.0,
155            dominant_patterns: vec![0.0, 1.0],
156            recommended_method: CompressionMethod::Zstd { level: 1 },
157            expected_ratio: 0.125,
158        }
159    }
160
161    fn calculate_entropy(values: &[f32]) -> f32 {
162        let mut histogram = std::collections::HashMap::new();
163        let bins = 64;
164
165        if values.is_empty() {
166            return 0.0;
167        }
168
169        let min_val = values.iter().copied().fold(f32::INFINITY, f32::min);
170        let max_val = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
171        let range = max_val - min_val;
172
173        if range == 0.0 {
174            return 0.0;
175        }
176
177        for &value in values {
178            let bin = ((value - min_val) / range * (bins - 1) as f32) as usize;
179            let bin = bin.min(bins - 1);
180            *histogram.entry(bin).or_insert(0) += 1;
181        }
182
183        let total = values.len() as f32;
184        let mut entropy = 0.0;
185
186        for count in histogram.values() {
187            let probability = *count as f32 / total;
188            if probability > 0.0 {
189                entropy -= probability * probability.log2();
190            }
191        }
192
193        entropy
194    }
195
196    fn find_dominant_patterns(values: &[f32]) -> Vec<f32> {
197        let mut value_counts = std::collections::HashMap::new();
198
199        for &value in values {
200            let quantized = (value * 1000.0).round() / 1000.0;
201            *value_counts.entry(quantized.to_bits()).or_insert(0) += 1;
202        }
203
204        let mut patterns: Vec<_> = value_counts.into_iter().collect();
205        patterns.sort_by_key(|b| std::cmp::Reverse(b.1));
206
207        patterns
208            .into_iter()
209            .take(5)
210            .map(|(bits, _)| f32::from_bits(bits))
211            .collect()
212    }
213
214    pub(crate) fn select_optimal_method(
215        sparsity: f32,
216        range: f32,
217        std_dev: f32,
218        entropy: f32,
219        dimensions: usize,
220        quality: &AdaptiveQuality,
221    ) -> (CompressionMethod, f32) {
222        if sparsity > 0.7 {
223            return match quality {
224                AdaptiveQuality::Fast => (CompressionMethod::Zstd { level: 1 }, 0.3),
225                AdaptiveQuality::Balanced => (CompressionMethod::Zstd { level: 6 }, 0.2),
226                AdaptiveQuality::BestRatio => (CompressionMethod::Zstd { level: 19 }, 0.15),
227            };
228        }
229
230        if entropy < 2.0 {
231            return match quality {
232                AdaptiveQuality::Fast => (CompressionMethod::Zstd { level: 3 }, 0.4),
233                AdaptiveQuality::Balanced => (CompressionMethod::Zstd { level: 9 }, 0.3),
234                AdaptiveQuality::BestRatio => (CompressionMethod::Zstd { level: 22 }, 0.2),
235            };
236        }
237
238        if range < 2.0 && std_dev < 0.5 {
239            return match quality {
240                AdaptiveQuality::Fast => (CompressionMethod::Quantization { bits: 8 }, 0.25),
241                AdaptiveQuality::Balanced => (CompressionMethod::Quantization { bits: 6 }, 0.1875),
242                AdaptiveQuality::BestRatio => (CompressionMethod::Quantization { bits: 4 }, 0.125),
243            };
244        }
245
246        if dimensions > 128 {
247            let components = match quality {
248                AdaptiveQuality::Fast => dimensions * 7 / 10,
249                AdaptiveQuality::Balanced => dimensions / 2,
250                AdaptiveQuality::BestRatio => dimensions / 3,
251            };
252            return (
253                CompressionMethod::Pca { components },
254                components as f32 / dimensions as f32,
255            );
256        }
257
258        match quality {
259            AdaptiveQuality::Fast => (CompressionMethod::Zstd { level: 3 }, 0.6),
260            AdaptiveQuality::Balanced => (CompressionMethod::Zstd { level: 6 }, 0.5),
261            AdaptiveQuality::BestRatio => (CompressionMethod::Zstd { level: 12 }, 0.4),
262        }
263    }
264}