1use crate::{Vector, VectorData, VectorError};
4use half::f16;
5
6#[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#[derive(Debug, Clone)]
32pub enum AdaptiveQuality {
33 Fast, Balanced, BestRatio, }
37
38pub 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#[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#[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}