Skip to main content

hermes_core/structures/postings/sparse/
config.rs

1//! Configuration types for sparse vector posting lists
2
3use serde::{Deserialize, Serialize};
4
5/// Sparse vector index format
6///
7/// Determines the on-disk layout and query execution strategy:
8/// - **MaxScore**: Per-dimension variable-size blocks (DAAT — document-at-a-time).
9///   Default, optimal for general sparse retrieval with block-max pruning.
10/// - **Bmp**: Fixed doc_id range blocks (BAAT — block-at-a-time).
11///   Based on Mallia, Suel & Tonellotto (SIGIR 2024). Divides the document
12///   space into fixed-size blocks and processes them in decreasing upper-bound
13///   order, enabling aggressive early termination.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
15pub enum SparseFormat {
16    /// Per-dimension variable-size blocks (existing format, DAAT MaxScore)
17    #[default]
18    MaxScore,
19    /// Fixed doc_id range blocks (BMP, BAAT block-at-a-time)
20    Bmp,
21}
22
23impl SparseFormat {
24    fn is_default(&self) -> bool {
25        *self == Self::MaxScore
26    }
27}
28
29/// Size of the index (term/dimension ID) in sparse vectors
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
31#[repr(u8)]
32pub enum IndexSize {
33    /// 16-bit index (0-65535), ideal for SPLADE vocabularies
34    U16 = 0,
35    /// 32-bit index (0-4B), for large vocabularies
36    #[default]
37    U32 = 1,
38}
39
40impl IndexSize {
41    /// Bytes per index
42    pub fn bytes(&self) -> usize {
43        match self {
44            IndexSize::U16 => 2,
45            IndexSize::U32 => 4,
46        }
47    }
48
49    /// Maximum value representable
50    pub fn max_value(&self) -> u32 {
51        match self {
52            IndexSize::U16 => u16::MAX as u32,
53            IndexSize::U32 => u32::MAX,
54        }
55    }
56
57    pub(crate) fn from_u8(v: u8) -> Option<Self> {
58        match v {
59            0 => Some(IndexSize::U16),
60            1 => Some(IndexSize::U32),
61            _ => None,
62        }
63    }
64}
65
66/// Quantization format for sparse vector weights
67///
68/// Research-validated compression/effectiveness trade-offs (Pati, 2025):
69/// - **UInt8**: 4x compression, ~1-2% nDCG@10 loss (RECOMMENDED for production)
70/// - **Float16**: 2x compression, <1% nDCG@10 loss
71/// - **Float32**: No compression, baseline effectiveness
72/// - **UInt4**: 8x compression, ~3-5% nDCG@10 loss (experimental)
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
74#[repr(u8)]
75pub enum WeightQuantization {
76    /// Full 32-bit float precision
77    #[default]
78    Float32 = 0,
79    /// 16-bit float (half precision) - 2x compression, <1% effectiveness loss
80    Float16 = 1,
81    /// 8-bit unsigned integer with scale factor - 4x compression, ~1-2% effectiveness loss (RECOMMENDED)
82    UInt8 = 2,
83    /// 4-bit unsigned integer with scale factor (packed, 2 per byte) - 8x compression, ~3-5% effectiveness loss
84    UInt4 = 3,
85}
86
87impl WeightQuantization {
88    /// Bytes per weight (approximate for UInt4)
89    pub fn bytes_per_weight(&self) -> f32 {
90        match self {
91            WeightQuantization::Float32 => 4.0,
92            WeightQuantization::Float16 => 2.0,
93            WeightQuantization::UInt8 => 1.0,
94            WeightQuantization::UInt4 => 0.5,
95        }
96    }
97
98    pub(crate) fn from_u8(v: u8) -> Option<Self> {
99        match v {
100            0 => Some(WeightQuantization::Float32),
101            1 => Some(WeightQuantization::Float16),
102            2 => Some(WeightQuantization::UInt8),
103            3 => Some(WeightQuantization::UInt4),
104            _ => None,
105        }
106    }
107}
108
109/// Query-time weighting strategy for sparse vector queries
110#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
111#[serde(rename_all = "snake_case")]
112pub enum QueryWeighting {
113    /// All terms get weight 1.0
114    #[default]
115    One,
116    /// Terms weighted by IDF (inverse document frequency) from global index statistics
117    /// Uses ln(N/df) where N = total docs, df = docs containing dimension
118    Idf,
119    /// Terms weighted by pre-computed IDF from model's idf.json file
120    /// Loaded from HuggingFace model repo. No fallback to global stats.
121    IdfFile,
122}
123
124/// Query-time configuration for sparse vectors
125///
126/// Research-validated query optimization strategies:
127/// - **weight_threshold (0.01-0.05)**: Drop query dimensions with weight below threshold
128///   - Filters low-IDF tokens that add latency without improving relevance
129/// - **max_query_dims (10-20)**: Process only top-k dimensions by weight
130///   - 30-50% latency reduction with <2% nDCG loss (Qiao et al., 2023)
131/// - **heap_factor (0.8)**: Skip blocks with low max score contribution
132///   - ~20% speedup with minor recall loss (SEISMIC-style)
133#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
134pub struct SparseQueryConfig {
135    /// HuggingFace tokenizer path/name for query-time tokenization
136    /// Example: "Alibaba-NLP/gte-Qwen2-1.5B-instruct"
137    #[serde(default, skip_serializing_if = "Option::is_none")]
138    pub tokenizer: Option<String>,
139    /// Weighting strategy for tokenized query terms
140    #[serde(default)]
141    pub weighting: QueryWeighting,
142    /// Heap factor for approximate search (SEISMIC-style optimization)
143    /// A block is skipped if its max possible score < heap_factor * threshold
144    ///
145    /// Research recommendation:
146    /// - 1.0 = exact search (default)
147    /// - 0.8 = approximate, ~20% faster with minor recall loss (RECOMMENDED for production)
148    /// - 0.5 = very approximate, much faster but higher recall loss
149    #[serde(default = "default_heap_factor")]
150    pub heap_factor: f32,
151    /// Minimum weight for query dimensions (query-time pruning)
152    /// Dimensions with abs(weight) below this threshold are dropped before search.
153    /// Useful for filtering low-IDF tokens that add latency without improving relevance.
154    ///
155    /// - 0.0 = no filtering (default)
156    /// - 0.01-0.05 = recommended for SPLADE/learned sparse models
157    #[serde(default)]
158    pub weight_threshold: f32,
159    /// Maximum number of query dimensions to process (query pruning)
160    /// Processes only the top-k dimensions by weight
161    ///
162    /// Research recommendation (Multiple papers 2022-2024):
163    /// - None = process all dimensions (default, exact)
164    /// - Some(10-20) = process top 10-20 dimensions only (RECOMMENDED for SPLADE)
165    ///   - 30-50% latency reduction
166    ///   - <2% nDCG@10 loss
167    #[serde(default, skip_serializing_if = "Option::is_none")]
168    pub max_query_dims: Option<usize>,
169    /// Fraction of query dimensions to keep (0.0-1.0), same semantics as
170    /// indexing-time `pruning`: sort by abs(weight) descending,
171    /// keep top fraction. None or 1.0 = no pruning.
172    #[serde(default, skip_serializing_if = "Option::is_none")]
173    pub pruning: Option<f32>,
174    /// Minimum number of query dimensions before pruning and weight_threshold
175    /// filtering are applied. Protects short queries from losing most signal.
176    ///
177    /// Default: 4. Set to 0 to always apply pruning/filtering.
178    #[serde(default = "default_min_terms")]
179    pub min_query_dims: usize,
180    /// Maximum number of superblocks to visit (LSP/0 gamma cap).
181    /// 0 = unlimited (default). Only applies to BMP format.
182    #[serde(default)]
183    pub max_superblocks: usize,
184}
185
186fn default_heap_factor() -> f32 {
187    1.0
188}
189
190impl Default for SparseQueryConfig {
191    fn default() -> Self {
192        Self {
193            tokenizer: None,
194            weighting: QueryWeighting::One,
195            heap_factor: 1.0,
196            weight_threshold: 0.0,
197            max_query_dims: None,
198            pruning: None,
199            min_query_dims: 4,
200            max_superblocks: 0,
201        }
202    }
203}
204
205/// Configuration for sparse vector storage
206///
207/// Research-validated optimizations for learned sparse retrieval (SPLADE, uniCOIL, etc.):
208/// - **Weight threshold (0.01-0.05)**: Removes ~30-50% of postings with minimal nDCG impact
209/// - **Posting list pruning (0.1)**: Keeps top 10% per dimension, 50-70% index reduction, <1% nDCG loss
210/// - **Query pruning (top 10-20 dims)**: 30-50% latency reduction, <2% nDCG loss
211/// - **UInt8 quantization**: 4x compression, 1-2% nDCG loss (optimal trade-off)
212#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
213pub struct SparseVectorConfig {
214    /// Index format: MaxScore (DAAT) or BMP (BAAT)
215    #[serde(default, skip_serializing_if = "SparseFormat::is_default")]
216    pub format: SparseFormat,
217    /// Size of dimension/term indices
218    pub index_size: IndexSize,
219    /// Quantization for weights (see WeightQuantization docs for trade-offs)
220    pub weight_quantization: WeightQuantization,
221    /// Minimum weight threshold - weights below this value are not indexed
222    ///
223    /// Research recommendation (Guo et al., 2022; SPLADE v2):
224    /// - 0.01-0.05 for SPLADE models removes ~30-50% of postings
225    /// - Minimal impact on nDCG@10 (<1% loss)
226    /// - Major reduction in index size and query latency
227    #[serde(default)]
228    pub weight_threshold: f32,
229    /// Document-side mass cropping: keep the top-|weight| entries covering
230    /// this fraction of a sparse vector's total |weight| mass; the excessive
231    /// tail is dropped at indexing time.
232    ///
233    /// SPLADE-style vectors concentrate importance in a few head terms; the
234    /// long tail inflates the index and query cost with little relevance
235    /// signal. 0.9-0.95 typically drops 20-40% of postings with <1% nDCG loss.
236    ///
237    /// - None or >= 1.0 = keep all entries (default)
238    /// - Applied after `weight_threshold`; vectors with <= `min_terms`
239    ///   entries are never cropped.
240    #[serde(default, skip_serializing_if = "Option::is_none")]
241    pub doc_mass: Option<f32>,
242    /// Block size for posting lists (must be power of 2, default 128 for SIMD)
243    /// Larger blocks = better compression, smaller blocks = faster seeks.
244    /// Used by MaxScore format only.
245    #[serde(default = "default_block_size")]
246    pub block_size: usize,
247    /// BMP block size: number of consecutive doc_ids per block (must be power of 2).
248    /// Default 64. Only used when format = Bmp.
249    /// Smaller = better pruning granularity, larger = less overhead.
250    #[serde(default = "default_bmp_block_size")]
251    pub bmp_block_size: u32,
252    /// Maximum BMP grid memory in bytes. If the grid (num_dims × num_blocks)
253    /// would exceed this, bmp_block_size is automatically increased to cap memory.
254    /// Default: 256MB. Set to 0 to disable the cap.
255    #[serde(default = "default_max_bmp_grid_bytes")]
256    pub max_bmp_grid_bytes: u64,
257    /// BMP superblock size: number of consecutive blocks grouped for hierarchical
258    /// pruning (Carlson et al., SIGIR 2025). Must be power of 2.
259    /// Default 64. Set to 0 to disable superblock pruning (flat BMP scoring).
260    /// Only used when format = Bmp.
261    #[serde(default = "default_bmp_superblock_size")]
262    pub bmp_superblock_size: u32,
263    /// Static pruning: fraction of postings to keep per inverted list (SEISMIC-style)
264    /// Lists are sorted by weight descending and truncated to top fraction.
265    ///
266    /// Research recommendation (SPLADE v2, Formal et al., 2021):
267    /// - None = keep all postings (default, exact)
268    /// - Some(0.1) = keep top 10% of postings per dimension
269    ///   - 50-70% index size reduction
270    ///   - <1% nDCG@10 loss
271    ///   - Exploits "concentration of importance" in learned representations
272    ///
273    /// Applied only during initial segment build, not during merge.
274    #[serde(default, skip_serializing_if = "Option::is_none")]
275    pub pruning: Option<f32>,
276    /// Query-time configuration (tokenizer, weighting)
277    #[serde(default, skip_serializing_if = "Option::is_none")]
278    pub query_config: Option<SparseQueryConfig>,
279    /// Fixed vocabulary size (number of dimensions) for BMP format.
280    ///
281    /// When set, all BMP segments use the same grid dimensions (rows = dims),
282    /// enabling zero-copy block-copy merge. The grid is indexed by dim_id directly
283    /// (no dim_ids Section C needed).
284    ///
285    /// Required for BMP format. Typical values:
286    /// - SPLADE/BERT: 30522 or 105879 (WordPiece / Unigram vocabulary)
287    /// - uniCOIL: 30522
288    /// - Custom models: set to vocabulary size
289    ///
290    /// If None, BMP builder derives dims from observed data (V10 behavior).
291    #[serde(default, skip_serializing_if = "Option::is_none")]
292    pub dims: Option<u32>,
293    /// Fixed max weight scale for BMP format.
294    ///
295    /// When set, all BMP segments use the same quantization scale
296    /// (`max_weight_scale = max_weight`), eliminating rescaling during merge.
297    ///
298    /// For SPLADE models: 5.0 (covers typical weight range 0-5).
299    /// If None, BMP builder derives scale from data (V10 behavior).
300    #[serde(default, skip_serializing_if = "Option::is_none")]
301    pub max_weight: Option<f32>,
302    /// Minimum number of postings in a dimension before pruning and
303    /// weight_threshold filtering are applied. Protects dimensions with
304    /// very few postings from losing most of their signal.
305    ///
306    /// Default: 4. Set to 0 to always apply pruning/filtering.
307    #[serde(default = "default_min_terms")]
308    pub min_terms: usize,
309}
310
311fn default_block_size() -> usize {
312    128
313}
314
315fn default_bmp_block_size() -> u32 {
316    64
317}
318
319fn default_max_bmp_grid_bytes() -> u64 {
320    0 // disabled by default — masks eliminate DRAM stalls during scoring
321}
322
323fn default_bmp_superblock_size() -> u32 {
324    64
325}
326
327fn default_min_terms() -> usize {
328    4
329}
330
331impl Default for SparseVectorConfig {
332    fn default() -> Self {
333        Self {
334            format: SparseFormat::MaxScore,
335            index_size: IndexSize::U32,
336            weight_quantization: WeightQuantization::Float32,
337            weight_threshold: 0.0,
338            doc_mass: None,
339            block_size: 128,
340            bmp_block_size: 64,
341            max_bmp_grid_bytes: 0,
342            bmp_superblock_size: 64,
343            pruning: None,
344            query_config: None,
345
346            dims: None,
347            max_weight: None,
348            min_terms: 4,
349        }
350    }
351}
352
353impl SparseVectorConfig {
354    /// SPLADE-optimized config with research-validated defaults
355    ///
356    /// Optimized for SPLADE, uniCOIL, and similar learned sparse retrieval models.
357    /// Based on research findings from:
358    /// - Pati (2025): UInt8 quantization = 4x compression, 1-2% nDCG loss
359    /// - Formal et al. (2021): SPLADE v2 posting list pruning
360    /// - Qiao et al. (2023): Query dimension pruning and approximate search
361    /// - Guo et al. (2022): Weight thresholding for efficiency
362    ///
363    /// Expected performance vs. full precision baseline:
364    /// - Index size: ~15-25% of original (combined effect of all optimizations)
365    /// - Query latency: 40-60% faster
366    /// - Effectiveness: 2-4% nDCG@10 loss (typically acceptable for production)
367    ///
368    /// Vocabulary: ~30K dimensions (fits in u16)
369    pub fn splade() -> Self {
370        Self {
371            format: SparseFormat::MaxScore,
372            index_size: IndexSize::U16,
373            weight_quantization: WeightQuantization::UInt8,
374            weight_threshold: 0.01, // Remove ~30-50% of low-weight postings
375            doc_mass: None,
376            block_size: 128,
377            bmp_block_size: 64,
378            max_bmp_grid_bytes: 0,
379            bmp_superblock_size: 64,
380            pruning: Some(0.1), // Keep top 10% per dimension
381            query_config: Some(SparseQueryConfig {
382                tokenizer: None,
383                weighting: QueryWeighting::One,
384                heap_factor: 0.8,         // 20% faster approximate search
385                weight_threshold: 0.01,   // Drop low-IDF query tokens
386                max_query_dims: Some(20), // Process top 20 query dimensions
387                pruning: Some(0.1),       // Keep top 10% of query dims
388                min_query_dims: 4,
389                max_superblocks: 0,
390            }),
391
392            dims: None,
393            max_weight: None,
394            min_terms: 4,
395        }
396    }
397
398    /// SPLADE-optimized config with BMP (Block-Max Pruning) format
399    ///
400    /// Same optimization settings as `splade()` but uses the BMP block-at-a-time
401    /// format (Mallia, Suel & Tonellotto, SIGIR 2024) instead of MaxScore.
402    /// BMP divides the document space into fixed-size blocks and processes them
403    /// in decreasing upper-bound order, enabling aggressive early termination.
404    pub fn splade_bmp() -> Self {
405        Self {
406            format: SparseFormat::Bmp,
407            index_size: IndexSize::U16,
408            weight_quantization: WeightQuantization::UInt8,
409            weight_threshold: 0.01,
410            doc_mass: None,
411            block_size: 128,
412            bmp_block_size: 64,
413            max_bmp_grid_bytes: 0,
414            bmp_superblock_size: 64,
415            pruning: Some(0.1),
416            query_config: Some(SparseQueryConfig {
417                tokenizer: None,
418                weighting: QueryWeighting::One,
419                heap_factor: 0.8,
420                weight_threshold: 0.01,
421                max_query_dims: Some(20),
422                pruning: Some(0.1),
423                min_query_dims: 4,
424                max_superblocks: 0,
425            }),
426
427            dims: Some(105879),
428            max_weight: Some(5.0),
429            min_terms: 4,
430        }
431    }
432
433    /// Compact config: Maximum compression (experimental)
434    ///
435    /// Uses aggressive UInt4 quantization for smallest possible index size.
436    /// Expected trade-offs:
437    /// - Index size: ~10-15% of Float32 baseline
438    /// - Effectiveness: ~3-5% nDCG@10 loss
439    ///
440    /// Recommended for: Memory-constrained environments, cache-heavy workloads
441    pub fn compact() -> Self {
442        Self {
443            format: SparseFormat::MaxScore,
444            index_size: IndexSize::U16,
445            weight_quantization: WeightQuantization::UInt4,
446            weight_threshold: 0.02, // Slightly higher threshold for UInt4
447            doc_mass: None,
448            block_size: 128,
449            bmp_block_size: 64,
450            max_bmp_grid_bytes: 0,
451            bmp_superblock_size: 64,
452            pruning: Some(0.15), // Keep top 15% per dimension
453            query_config: Some(SparseQueryConfig {
454                tokenizer: None,
455                weighting: QueryWeighting::One,
456                heap_factor: 0.7,         // More aggressive approximate search
457                weight_threshold: 0.02,   // Drop low-IDF query tokens
458                max_query_dims: Some(15), // Fewer query dimensions
459                pruning: Some(0.15),      // Keep top 15% of query dims
460                min_query_dims: 4,
461                max_superblocks: 0,
462            }),
463
464            dims: None,
465            max_weight: None,
466            min_terms: 4,
467        }
468    }
469
470    /// Full precision config: No compression, baseline effectiveness
471    ///
472    /// Use for: Research baselines, when effectiveness is critical
473    pub fn full_precision() -> Self {
474        Self {
475            format: SparseFormat::MaxScore,
476            index_size: IndexSize::U32,
477            weight_quantization: WeightQuantization::Float32,
478            weight_threshold: 0.0,
479            doc_mass: None,
480            block_size: 128,
481            bmp_block_size: 64,
482            max_bmp_grid_bytes: 0,
483            bmp_superblock_size: 64,
484            pruning: None,
485            query_config: None,
486
487            dims: None,
488            max_weight: None,
489            min_terms: 4,
490        }
491    }
492
493    /// Conservative config: Mild optimizations, minimal effectiveness loss
494    ///
495    /// Balances compression and effectiveness with conservative defaults.
496    /// Expected trade-offs:
497    /// - Index size: ~40-50% of Float32 baseline
498    /// - Query latency: ~20-30% faster
499    /// - Effectiveness: <1% nDCG@10 loss
500    ///
501    /// Recommended for: Production deployments prioritizing effectiveness
502    pub fn conservative() -> Self {
503        Self {
504            format: SparseFormat::MaxScore,
505            index_size: IndexSize::U32,
506            weight_quantization: WeightQuantization::Float16,
507            weight_threshold: 0.005, // Minimal pruning
508            doc_mass: None,
509            block_size: 128,
510            bmp_block_size: 64,
511            max_bmp_grid_bytes: 0,
512            bmp_superblock_size: 64,
513            pruning: None, // No posting list pruning
514            query_config: Some(SparseQueryConfig {
515                tokenizer: None,
516                weighting: QueryWeighting::One,
517                heap_factor: 0.9,         // Nearly exact search
518                weight_threshold: 0.005,  // Minimal query pruning
519                max_query_dims: Some(50), // Process more dimensions
520                pruning: None,            // No fraction-based pruning
521                min_query_dims: 4,
522                max_superblocks: 0,
523            }),
524
525            dims: None,
526            max_weight: None,
527            min_terms: 4,
528        }
529    }
530
531    /// Set weight threshold (builder pattern)
532    pub fn with_weight_threshold(mut self, threshold: f32) -> Self {
533        self.weight_threshold = threshold;
534        self
535    }
536
537    /// Set document-side mass cropping fraction (builder pattern)
538    /// e.g., 0.9 = keep top-weight entries covering 90% of each vector's mass
539    pub fn with_doc_mass(mut self, fraction: f32) -> Self {
540        self.doc_mass = Some(fraction.clamp(0.0, 1.0));
541        self
542    }
543
544    /// Set posting list pruning fraction (builder pattern)
545    /// e.g., 0.1 = keep top 10% of postings per dimension
546    pub fn with_pruning(mut self, fraction: f32) -> Self {
547        self.pruning = Some(fraction.clamp(0.0, 1.0));
548        self
549    }
550
551    /// Bytes per entry (index + weight)
552    pub fn bytes_per_entry(&self) -> f32 {
553        self.index_size.bytes() as f32 + self.weight_quantization.bytes_per_weight()
554    }
555
556    /// Serialize config to a single byte.
557    ///
558    /// Layout: bits 7-4 = IndexSize, bit 3 = format (0=MaxScore, 1=BMP), bits 2-0 = WeightQuantization
559    pub fn to_byte(&self) -> u8 {
560        let format_bit = if self.format == SparseFormat::Bmp {
561            0x08
562        } else {
563            0
564        };
565        ((self.index_size as u8) << 4) | format_bit | (self.weight_quantization as u8)
566    }
567
568    /// Deserialize config from a single byte.
569    ///
570    /// Note: weight_threshold, block_size, bmp_block_size, and query_config are not
571    /// serialized in the byte — they come from the schema.
572    pub fn from_byte(b: u8) -> Option<Self> {
573        let index_size = IndexSize::from_u8((b >> 4) & 0x03)?;
574        let format = if b & 0x08 != 0 {
575            SparseFormat::Bmp
576        } else {
577            SparseFormat::MaxScore
578        };
579        let weight_quantization = WeightQuantization::from_u8(b & 0x07)?;
580        Some(Self {
581            format,
582            index_size,
583            weight_quantization,
584            weight_threshold: 0.0,
585            doc_mass: None,
586            block_size: 128,
587            bmp_block_size: 64,
588            max_bmp_grid_bytes: 0,
589            bmp_superblock_size: 64,
590            pruning: None,
591            query_config: None,
592
593            dims: None,
594            max_weight: None,
595            min_terms: 4,
596        })
597    }
598
599    /// Set block size (builder pattern)
600    /// Must be power of 2, recommended: 64, 128, 256
601    pub fn with_block_size(mut self, size: usize) -> Self {
602        self.block_size = size.next_power_of_two();
603        self
604    }
605
606    /// Set query configuration (builder pattern)
607    pub fn with_query_config(mut self, config: SparseQueryConfig) -> Self {
608        self.query_config = Some(config);
609        self
610    }
611}
612
613/// A sparse vector entry: (dimension_id, weight)
614#[derive(Debug, Clone, Copy, PartialEq)]
615pub struct SparseEntry {
616    pub dim_id: u32,
617    pub weight: f32,
618}
619
620/// Sparse vector representation
621#[derive(Debug, Clone, Default)]
622pub struct SparseVector {
623    pub(super) entries: Vec<SparseEntry>,
624}
625
626impl SparseVector {
627    /// Create a new sparse vector
628    pub fn new() -> Self {
629        Self {
630            entries: Vec::new(),
631        }
632    }
633
634    /// Create with pre-allocated capacity
635    pub fn with_capacity(capacity: usize) -> Self {
636        Self {
637            entries: Vec::with_capacity(capacity),
638        }
639    }
640
641    /// Create from dimension IDs and weights
642    pub fn from_entries(dim_ids: &[u32], weights: &[f32]) -> Self {
643        assert_eq!(dim_ids.len(), weights.len());
644        let mut entries: Vec<SparseEntry> = dim_ids
645            .iter()
646            .zip(weights.iter())
647            .map(|(&dim_id, &weight)| SparseEntry { dim_id, weight })
648            .collect();
649        // Sort by dimension ID for efficient intersection
650        entries.sort_by_key(|e| e.dim_id);
651        Self { entries }
652    }
653
654    /// Add an entry (must maintain sorted order by dim_id)
655    pub fn push(&mut self, dim_id: u32, weight: f32) {
656        debug_assert!(
657            self.entries.is_empty() || self.entries.last().unwrap().dim_id < dim_id,
658            "Entries must be added in sorted order by dim_id"
659        );
660        self.entries.push(SparseEntry { dim_id, weight });
661    }
662
663    /// Number of non-zero entries
664    pub fn len(&self) -> usize {
665        self.entries.len()
666    }
667
668    /// Check if empty
669    pub fn is_empty(&self) -> bool {
670        self.entries.is_empty()
671    }
672
673    /// Iterate over entries
674    pub fn iter(&self) -> impl Iterator<Item = &SparseEntry> {
675        self.entries.iter()
676    }
677
678    /// Sort by dimension ID (required for posting list encoding)
679    pub fn sort_by_dim(&mut self) {
680        self.entries.sort_by_key(|e| e.dim_id);
681    }
682
683    /// Sort by weight descending
684    pub fn sort_by_weight_desc(&mut self) {
685        self.entries.sort_by(|a, b| {
686            b.weight
687                .partial_cmp(&a.weight)
688                .unwrap_or(std::cmp::Ordering::Equal)
689        });
690    }
691
692    /// Get top-k entries by weight
693    pub fn top_k(&self, k: usize) -> Vec<SparseEntry> {
694        let mut sorted = self.entries.clone();
695        sorted.sort_by(|a, b| {
696            b.weight
697                .partial_cmp(&a.weight)
698                .unwrap_or(std::cmp::Ordering::Equal)
699        });
700        sorted.truncate(k);
701        sorted
702    }
703
704    /// Compute dot product with another sparse vector
705    pub fn dot(&self, other: &SparseVector) -> f32 {
706        let mut result = 0.0f32;
707        let mut i = 0;
708        let mut j = 0;
709
710        while i < self.entries.len() && j < other.entries.len() {
711            let a = &self.entries[i];
712            let b = &other.entries[j];
713
714            match a.dim_id.cmp(&b.dim_id) {
715                std::cmp::Ordering::Less => i += 1,
716                std::cmp::Ordering::Greater => j += 1,
717                std::cmp::Ordering::Equal => {
718                    result += a.weight * b.weight;
719                    i += 1;
720                    j += 1;
721                }
722            }
723        }
724
725        result
726    }
727
728    /// L2 norm squared
729    pub fn norm_squared(&self) -> f32 {
730        self.entries.iter().map(|e| e.weight * e.weight).sum()
731    }
732
733    /// L2 norm
734    pub fn norm(&self) -> f32 {
735        self.norm_squared().sqrt()
736    }
737
738    /// Prune dimensions below a weight threshold
739    pub fn filter_by_weight(&self, min_weight: f32) -> Self {
740        let entries: Vec<SparseEntry> = self
741            .entries
742            .iter()
743            .filter(|e| e.weight.abs() >= min_weight)
744            .cloned()
745            .collect();
746        Self { entries }
747    }
748}
749
750impl From<Vec<(u32, f32)>> for SparseVector {
751    fn from(pairs: Vec<(u32, f32)>) -> Self {
752        Self {
753            entries: pairs
754                .into_iter()
755                .map(|(dim_id, weight)| SparseEntry { dim_id, weight })
756                .collect(),
757        }
758    }
759}
760
761impl From<SparseVector> for Vec<(u32, f32)> {
762    fn from(vec: SparseVector) -> Self {
763        vec.entries
764            .into_iter()
765            .map(|e| (e.dim_id, e.weight))
766            .collect()
767    }
768}