1use byteorder::{ReadBytesExt, WriteBytesExt};
12use std::io::{self, Read, Write};
13
14use super::elias_fano::EliasFanoPostingList;
15use super::horizontal_bp128::HorizontalBP128PostingList;
16use super::opt_p4d::OptP4DPostingList;
17use super::partitioned_ef::PartitionedEFPostingList;
18use super::roaring::RoaringPostingList;
19use super::vertical_bp128::VerticalBP128PostingList;
20
21pub const INLINE_THRESHOLD: usize = 3;
23pub const ROARING_THRESHOLD_RATIO: f32 = 0.01; pub const PARTITIONED_EF_THRESHOLD: usize = 20_000;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
29pub enum IndexOptimization {
30 #[default]
32 Adaptive,
33 SizeOptimized,
35 PerformanceOptimized,
37}
38
39impl IndexOptimization {
40 pub fn zstd_level(&self) -> i32 {
42 match self {
43 IndexOptimization::Adaptive => 9,
44 IndexOptimization::SizeOptimized => 22,
45 IndexOptimization::PerformanceOptimized => 1,
46 }
47 }
48
49 pub fn default_posting_codec(&self) -> super::posting::PostingCodec {
53 match self {
54 IndexOptimization::SizeOptimized => super::posting::PostingCodec::Pfor,
55 IndexOptimization::Adaptive | IndexOptimization::PerformanceOptimized => {
56 super::posting::PostingCodec::Rounded
57 }
58 }
59 }
60
61 pub fn parse(s: &str) -> Option<Self> {
63 match s.to_lowercase().as_str() {
64 "adaptive" | "balanced" | "default" => Some(IndexOptimization::Adaptive),
65 "size" | "size-optimized" | "small" | "compact" => {
66 Some(IndexOptimization::SizeOptimized)
67 }
68 "performance" | "perf" | "fast" | "speed" => {
69 Some(IndexOptimization::PerformanceOptimized)
70 }
71 _ => None,
72 }
73 }
74}
75
76const FORMAT_HORIZONTAL_BP128: u8 = 0;
78const FORMAT_ELIAS_FANO: u8 = 1;
79const FORMAT_ROARING: u8 = 2;
80const FORMAT_VERTICAL_BP128: u8 = 3;
81const FORMAT_PARTITIONED_EF: u8 = 4;
82const FORMAT_OPT_P4D: u8 = 5;
83
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub enum PostingFormat {
87 HorizontalBP128,
89 VerticalBP128,
91 EliasFano,
93 PartitionedEF,
95 Roaring,
97 OptP4D,
99}
100
101impl PostingFormat {
102 pub fn select(doc_count: usize, total_docs: usize) -> Self {
109 Self::select_with_optimization(doc_count, total_docs, IndexOptimization::Adaptive)
110 }
111
112 pub fn select_with_optimization(
114 doc_count: usize,
115 total_docs: usize,
116 optimization: IndexOptimization,
117 ) -> Self {
118 let frequency_ratio = doc_count as f32 / total_docs.max(1) as f32;
119
120 match optimization {
121 IndexOptimization::Adaptive => {
122 if frequency_ratio >= ROARING_THRESHOLD_RATIO
124 && doc_count >= PARTITIONED_EF_THRESHOLD
125 {
126 PostingFormat::Roaring
127 } else if doc_count >= PARTITIONED_EF_THRESHOLD {
128 PostingFormat::PartitionedEF
129 } else {
130 PostingFormat::HorizontalBP128
131 }
132 }
133 IndexOptimization::SizeOptimized => {
134 if doc_count >= 128 {
136 PostingFormat::OptP4D
137 } else {
138 PostingFormat::HorizontalBP128
139 }
140 }
141 IndexOptimization::PerformanceOptimized => {
142 if doc_count >= 64 {
144 PostingFormat::Roaring
145 } else {
146 PostingFormat::HorizontalBP128
147 }
148 }
149 }
150 }
151}
152
153#[derive(Debug, Clone)]
155pub enum CompressedPostingList {
156 HorizontalBP128(HorizontalBP128PostingList),
157 VerticalBP128(VerticalBP128PostingList),
158 EliasFano(EliasFanoPostingList),
159 PartitionedEF(PartitionedEFPostingList),
160 Roaring(RoaringPostingList),
161 OptP4D(OptP4DPostingList),
162}
163
164impl CompressedPostingList {
165 pub fn from_postings(doc_ids: &[u32], term_freqs: &[u32], total_docs: usize, idf: f32) -> Self {
167 let format = PostingFormat::select(doc_ids.len(), total_docs);
168
169 match format {
170 PostingFormat::HorizontalBP128 => CompressedPostingList::HorizontalBP128(
171 HorizontalBP128PostingList::from_postings(doc_ids, term_freqs, idf),
172 ),
173 PostingFormat::VerticalBP128 => CompressedPostingList::VerticalBP128(
174 VerticalBP128PostingList::from_postings(doc_ids, term_freqs, idf),
175 ),
176 PostingFormat::EliasFano => CompressedPostingList::EliasFano(
177 EliasFanoPostingList::from_postings(doc_ids, term_freqs),
178 ),
179 PostingFormat::PartitionedEF => CompressedPostingList::PartitionedEF(
180 PartitionedEFPostingList::from_postings_with_idf(doc_ids, term_freqs, idf),
181 ),
182 PostingFormat::Roaring => CompressedPostingList::Roaring(
183 RoaringPostingList::from_postings(doc_ids, term_freqs),
184 ),
185 PostingFormat::OptP4D => CompressedPostingList::OptP4D(
186 OptP4DPostingList::from_postings(doc_ids, term_freqs, idf),
187 ),
188 }
189 }
190
191 pub fn from_postings_with_format(
193 doc_ids: &[u32],
194 term_freqs: &[u32],
195 format: PostingFormat,
196 idf: f32,
197 ) -> Self {
198 match format {
199 PostingFormat::HorizontalBP128 => CompressedPostingList::HorizontalBP128(
200 HorizontalBP128PostingList::from_postings(doc_ids, term_freqs, idf),
201 ),
202 PostingFormat::VerticalBP128 => CompressedPostingList::VerticalBP128(
203 VerticalBP128PostingList::from_postings(doc_ids, term_freqs, idf),
204 ),
205 PostingFormat::EliasFano => CompressedPostingList::EliasFano(
206 EliasFanoPostingList::from_postings(doc_ids, term_freqs),
207 ),
208 PostingFormat::PartitionedEF => CompressedPostingList::PartitionedEF(
209 PartitionedEFPostingList::from_postings_with_idf(doc_ids, term_freqs, idf),
210 ),
211 PostingFormat::Roaring => CompressedPostingList::Roaring(
212 RoaringPostingList::from_postings(doc_ids, term_freqs),
213 ),
214 PostingFormat::OptP4D => CompressedPostingList::OptP4D(
215 OptP4DPostingList::from_postings(doc_ids, term_freqs, idf),
216 ),
217 }
218 }
219
220 pub fn doc_count(&self) -> u32 {
222 match self {
223 CompressedPostingList::HorizontalBP128(p) => p.doc_count,
224 CompressedPostingList::VerticalBP128(p) => p.doc_count,
225 CompressedPostingList::EliasFano(p) => p.len(),
226 CompressedPostingList::PartitionedEF(p) => p.len(),
227 CompressedPostingList::Roaring(p) => p.len(),
228 CompressedPostingList::OptP4D(p) => p.len(),
229 }
230 }
231
232 pub fn max_tf(&self) -> u32 {
234 match self {
235 CompressedPostingList::HorizontalBP128(p) => p.max_score as u32, CompressedPostingList::VerticalBP128(p) => {
237 p.blocks.iter().map(|b| b.max_tf).max().unwrap_or(0)
238 }
239 CompressedPostingList::EliasFano(p) => p.max_tf,
240 CompressedPostingList::PartitionedEF(p) => p.max_tf,
241 CompressedPostingList::Roaring(p) => p.max_tf,
242 CompressedPostingList::OptP4D(p) => {
243 p.blocks.iter().map(|b| b.max_tf).max().unwrap_or(0)
244 }
245 }
246 }
247
248 pub fn format(&self) -> PostingFormat {
250 match self {
251 CompressedPostingList::HorizontalBP128(_) => PostingFormat::HorizontalBP128,
252 CompressedPostingList::VerticalBP128(_) => PostingFormat::VerticalBP128,
253 CompressedPostingList::EliasFano(_) => PostingFormat::EliasFano,
254 CompressedPostingList::PartitionedEF(_) => PostingFormat::PartitionedEF,
255 CompressedPostingList::Roaring(_) => PostingFormat::Roaring,
256 CompressedPostingList::OptP4D(_) => PostingFormat::OptP4D,
257 }
258 }
259
260 pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
262 match self {
263 CompressedPostingList::HorizontalBP128(p) => {
264 writer.write_u8(FORMAT_HORIZONTAL_BP128)?;
265 p.serialize(writer)
266 }
267 CompressedPostingList::VerticalBP128(p) => {
268 writer.write_u8(FORMAT_VERTICAL_BP128)?;
269 p.serialize(writer)
270 }
271 CompressedPostingList::EliasFano(p) => {
272 writer.write_u8(FORMAT_ELIAS_FANO)?;
273 p.serialize(writer)
274 }
275 CompressedPostingList::PartitionedEF(p) => {
276 writer.write_u8(FORMAT_PARTITIONED_EF)?;
277 p.serialize(writer)
278 }
279 CompressedPostingList::Roaring(p) => {
280 writer.write_u8(FORMAT_ROARING)?;
281 p.serialize(writer)
282 }
283 CompressedPostingList::OptP4D(p) => {
284 writer.write_u8(FORMAT_OPT_P4D)?;
285 p.serialize(writer)
286 }
287 }
288 }
289
290 pub fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
292 let format = reader.read_u8()?;
293 match format {
294 FORMAT_HORIZONTAL_BP128 => Ok(CompressedPostingList::HorizontalBP128(
295 HorizontalBP128PostingList::deserialize(reader)?,
296 )),
297 FORMAT_VERTICAL_BP128 => Ok(CompressedPostingList::VerticalBP128(
298 VerticalBP128PostingList::deserialize(reader)?,
299 )),
300 FORMAT_ELIAS_FANO => Ok(CompressedPostingList::EliasFano(
301 EliasFanoPostingList::deserialize(reader)?,
302 )),
303 FORMAT_PARTITIONED_EF => Ok(CompressedPostingList::PartitionedEF(
304 PartitionedEFPostingList::deserialize(reader)?,
305 )),
306 FORMAT_ROARING => Ok(CompressedPostingList::Roaring(
307 RoaringPostingList::deserialize(reader)?,
308 )),
309 FORMAT_OPT_P4D => Ok(CompressedPostingList::OptP4D(
310 OptP4DPostingList::deserialize(reader)?,
311 )),
312 _ => Err(io::Error::new(
313 io::ErrorKind::InvalidData,
314 format!("Unknown posting list format: {}", format),
315 )),
316 }
317 }
318
319 pub fn iterator(&self) -> CompressedPostingIterator<'_> {
321 match self {
322 CompressedPostingList::HorizontalBP128(p) => {
323 CompressedPostingIterator::HorizontalBP128(p.iterator())
324 }
325 CompressedPostingList::VerticalBP128(p) => {
326 CompressedPostingIterator::VerticalBP128(p.iterator())
327 }
328 CompressedPostingList::EliasFano(p) => {
329 CompressedPostingIterator::EliasFano(p.iterator())
330 }
331 CompressedPostingList::PartitionedEF(p) => {
332 CompressedPostingIterator::PartitionedEF(p.iterator())
333 }
334 CompressedPostingList::Roaring(p) => {
335 let mut iter = p.iterator();
336 iter.init();
337 CompressedPostingIterator::Roaring(iter)
338 }
339 CompressedPostingList::OptP4D(p) => CompressedPostingIterator::OptP4D(p.iterator()),
340 }
341 }
342}
343
344pub enum CompressedPostingIterator<'a> {
346 HorizontalBP128(super::horizontal_bp128::HorizontalBP128Iterator<'a>),
347 VerticalBP128(super::vertical_bp128::VerticalBP128Iterator<'a>),
348 EliasFano(super::elias_fano::EliasFanoPostingIterator<'a>),
349 PartitionedEF(super::partitioned_ef::PartitionedEFPostingIterator<'a>),
350 Roaring(super::roaring::RoaringPostingIterator<'a>),
351 OptP4D(super::opt_p4d::OptP4DIterator<'a>),
352}
353
354impl<'a> CompressedPostingIterator<'a> {
355 pub fn doc(&self) -> u32 {
357 match self {
358 CompressedPostingIterator::HorizontalBP128(i) => i.doc(),
359 CompressedPostingIterator::VerticalBP128(i) => i.doc(),
360 CompressedPostingIterator::EliasFano(i) => i.doc(),
361 CompressedPostingIterator::PartitionedEF(i) => i.doc(),
362 CompressedPostingIterator::Roaring(i) => i.doc(),
363 CompressedPostingIterator::OptP4D(i) => i.doc(),
364 }
365 }
366
367 pub fn term_freq(&self) -> u32 {
369 match self {
370 CompressedPostingIterator::HorizontalBP128(i) => i.term_freq(),
371 CompressedPostingIterator::VerticalBP128(i) => i.term_freq(),
372 CompressedPostingIterator::EliasFano(i) => i.term_freq(),
373 CompressedPostingIterator::PartitionedEF(i) => i.term_freq(),
374 CompressedPostingIterator::Roaring(i) => i.term_freq(),
375 CompressedPostingIterator::OptP4D(i) => i.term_freq(),
376 }
377 }
378
379 pub fn advance(&mut self) -> u32 {
381 match self {
382 CompressedPostingIterator::HorizontalBP128(i) => i.advance(),
383 CompressedPostingIterator::VerticalBP128(i) => i.advance(),
384 CompressedPostingIterator::EliasFano(i) => i.advance(),
385 CompressedPostingIterator::PartitionedEF(i) => i.advance(),
386 CompressedPostingIterator::Roaring(i) => i.advance(),
387 CompressedPostingIterator::OptP4D(i) => i.advance(),
388 }
389 }
390
391 pub fn seek(&mut self, target: u32) -> u32 {
393 match self {
394 CompressedPostingIterator::HorizontalBP128(i) => i.seek(target),
395 CompressedPostingIterator::VerticalBP128(i) => i.seek(target),
396 CompressedPostingIterator::EliasFano(i) => i.seek(target),
397 CompressedPostingIterator::PartitionedEF(i) => i.seek(target),
398 CompressedPostingIterator::Roaring(i) => i.seek(target),
399 CompressedPostingIterator::OptP4D(i) => i.seek(target),
400 }
401 }
402
403 pub fn is_exhausted(&self) -> bool {
405 self.doc() == u32::MAX
406 }
407}
408
409#[derive(Debug, Default, Clone)]
411pub struct CompressionStats {
412 pub bitpacked_count: u32,
413 pub bitpacked_docs: u64,
414 pub simd_bp128_count: u32,
415 pub simd_bp128_docs: u64,
416 pub elias_fano_count: u32,
417 pub elias_fano_docs: u64,
418 pub partitioned_ef_count: u32,
419 pub partitioned_ef_docs: u64,
420 pub roaring_count: u32,
421 pub roaring_docs: u64,
422 pub inline_count: u32,
423 pub inline_docs: u64,
424}
425
426impl CompressionStats {
427 pub fn record(&mut self, format: PostingFormat, doc_count: u32) {
428 match format {
429 PostingFormat::HorizontalBP128 => {
430 self.bitpacked_count += 1;
431 self.bitpacked_docs += doc_count as u64;
432 }
433 PostingFormat::VerticalBP128 => {
434 self.simd_bp128_count += 1;
435 self.simd_bp128_docs += doc_count as u64;
436 }
437 PostingFormat::EliasFano => {
438 self.elias_fano_count += 1;
439 self.elias_fano_docs += doc_count as u64;
440 }
441 PostingFormat::PartitionedEF => {
442 self.partitioned_ef_count += 1;
443 self.partitioned_ef_docs += doc_count as u64;
444 }
445 PostingFormat::Roaring => {
446 self.roaring_count += 1;
447 self.roaring_docs += doc_count as u64;
448 }
449 PostingFormat::OptP4D => {
450 self.bitpacked_count += 1;
452 self.bitpacked_docs += doc_count as u64;
453 }
454 }
455 }
456
457 pub fn record_inline(&mut self, doc_count: u32) {
458 self.inline_count += 1;
459 self.inline_docs += doc_count as u64;
460 }
461
462 pub fn total_terms(&self) -> u32 {
463 self.bitpacked_count
464 + self.simd_bp128_count
465 + self.elias_fano_count
466 + self.partitioned_ef_count
467 + self.roaring_count
468 + self.inline_count
469 }
470
471 pub fn total_postings(&self) -> u64 {
472 self.bitpacked_docs
473 + self.simd_bp128_docs
474 + self.elias_fano_docs
475 + self.partitioned_ef_docs
476 + self.roaring_docs
477 + self.inline_docs
478 }
479}
480
481#[cfg(test)]
482mod tests {
483 use super::*;
484
485 #[test]
486 fn test_format_selection() {
487 assert_eq!(
489 PostingFormat::select(100, 1_000_000),
490 PostingFormat::HorizontalBP128
491 );
492
493 assert_eq!(
495 PostingFormat::select(500, 1_000_000),
496 PostingFormat::HorizontalBP128
497 );
498
499 assert_eq!(
500 PostingFormat::select(5_000, 1_000_000),
501 PostingFormat::HorizontalBP128
502 );
503
504 assert_eq!(
505 PostingFormat::select(15_000, 10_000_000),
506 PostingFormat::HorizontalBP128
507 );
508
509 assert_eq!(
511 PostingFormat::select(25_000, 10_000_000),
512 PostingFormat::PartitionedEF
513 );
514
515 assert_eq!(
517 PostingFormat::select(50_000, 1_000_000),
518 PostingFormat::Roaring
519 );
520 }
521
522 #[test]
523 fn test_compressed_posting_list_small() {
524 let doc_ids: Vec<u32> = (0..100).map(|i| i * 2).collect();
525 let term_freqs: Vec<u32> = vec![1; 100];
526
527 let list = CompressedPostingList::from_postings(&doc_ids, &term_freqs, 1_000_000, 1.0);
528
529 assert_eq!(list.format(), PostingFormat::HorizontalBP128);
531 assert_eq!(list.doc_count(), 100);
532
533 let mut iter = list.iterator();
534 for (i, &expected) in doc_ids.iter().enumerate() {
535 assert_eq!(iter.doc(), expected, "Mismatch at {}", i);
536 iter.advance();
537 }
538 }
539
540 #[test]
541 fn test_compressed_posting_list_bitpacked() {
542 let doc_ids: Vec<u32> = (0..15_000).map(|i| i * 2).collect();
543 let term_freqs: Vec<u32> = vec![1; 15_000];
544
545 let list = CompressedPostingList::from_postings(&doc_ids, &term_freqs, 10_000_000, 1.0);
547
548 assert_eq!(list.format(), PostingFormat::HorizontalBP128);
549 assert_eq!(list.doc_count(), 15_000);
550 }
551
552 #[test]
553 fn test_compressed_posting_list_serialization() {
554 let doc_ids: Vec<u32> = (0..500).map(|i| i * 3).collect();
555 let term_freqs: Vec<u32> = (0..500).map(|i| (i % 5) + 1).collect();
556
557 let list = CompressedPostingList::from_postings(&doc_ids, &term_freqs, 1_000_000, 1.0);
558
559 let mut buffer = Vec::new();
560 list.serialize(&mut buffer).unwrap();
561
562 let restored = CompressedPostingList::deserialize(&mut &buffer[..]).unwrap();
563
564 assert_eq!(restored.format(), list.format());
565 assert_eq!(restored.doc_count(), list.doc_count());
566
567 let mut iter1 = list.iterator();
569 let mut iter2 = restored.iterator();
570
571 while iter1.doc() != u32::MAX {
572 assert_eq!(iter1.doc(), iter2.doc());
573 assert_eq!(iter1.term_freq(), iter2.term_freq());
574 iter1.advance();
575 iter2.advance();
576 }
577 }
578
579 #[test]
580 fn test_iterator_seek() {
581 let doc_ids: Vec<u32> = vec![10, 20, 30, 100, 200, 300, 1000, 2000];
582 let term_freqs: Vec<u32> = vec![1; 8];
583
584 let list = CompressedPostingList::from_postings(&doc_ids, &term_freqs, 1_000_000, 1.0);
585 let mut iter = list.iterator();
586
587 assert_eq!(iter.seek(25), 30);
588 assert_eq!(iter.seek(100), 100);
589 assert_eq!(iter.seek(500), 1000);
590 assert_eq!(iter.seek(3000), u32::MAX);
591 }
592
593 #[test]
594 fn test_opt_p4d_via_unified_interface() {
595 let doc_ids: Vec<u32> = (0..500).map(|i| i * 3).collect();
596 let term_freqs: Vec<u32> = (0..500).map(|i| (i % 5) + 1).collect();
597
598 let list = CompressedPostingList::from_postings_with_format(
600 &doc_ids,
601 &term_freqs,
602 PostingFormat::OptP4D,
603 1.0,
604 );
605
606 assert_eq!(list.format(), PostingFormat::OptP4D);
607 assert_eq!(list.doc_count(), 500);
608
609 let mut buffer = Vec::new();
611 list.serialize(&mut buffer).unwrap();
612 let restored = CompressedPostingList::deserialize(&mut &buffer[..]).unwrap();
613
614 assert_eq!(restored.format(), PostingFormat::OptP4D);
615 assert_eq!(restored.doc_count(), 500);
616
617 let mut iter1 = list.iterator();
619 let mut iter2 = restored.iterator();
620
621 while iter1.doc() != u32::MAX {
622 assert_eq!(iter1.doc(), iter2.doc());
623 assert_eq!(iter1.term_freq(), iter2.term_freq());
624 iter1.advance();
625 iter2.advance();
626 }
627
628 let mut iter = list.iterator();
630 assert_eq!(iter.seek(100), 102); assert_eq!(iter.seek(500), 501); }
633}