1use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
26use std::io::{self, Write};
27
28use super::posting_common::{read_vint, write_vint};
29use crate::DocId;
30
31pub const POSITION_BLOCK_SIZE: usize = 128;
33
34pub const MAX_TOKEN_POSITION: u32 = (1 << 20) - 1;
36
37pub const MAX_ELEMENT_ORDINAL: u32 = (1 << 12) - 1;
39
40#[inline]
42pub fn encode_position(element_ordinal: u32, token_position: u32) -> u32 {
43 debug_assert!(
44 element_ordinal <= MAX_ELEMENT_ORDINAL,
45 "Element ordinal {} exceeds maximum {}",
46 element_ordinal,
47 MAX_ELEMENT_ORDINAL
48 );
49 debug_assert!(
50 token_position <= MAX_TOKEN_POSITION,
51 "Token position {} exceeds maximum {}",
52 token_position,
53 MAX_TOKEN_POSITION
54 );
55 (element_ordinal << 20) | (token_position & MAX_TOKEN_POSITION)
56}
57
58#[inline]
60pub fn decode_element_ordinal(position: u32) -> u32 {
61 position >> 20
62}
63
64#[inline]
66pub fn decode_token_position(position: u32) -> u32 {
67 position & MAX_TOKEN_POSITION
68}
69
70#[derive(Debug, Clone, PartialEq, Eq)]
72pub struct PostingWithPositions {
73 pub doc_id: DocId,
74 pub term_freq: u32,
75 pub positions: Vec<u32>,
77}
78
79#[derive(Debug, Clone)]
84pub struct PositionPostingList {
85 skip_list: Vec<(DocId, DocId, u64)>,
88 data: Vec<u8>,
90 doc_count: u32,
92}
93
94impl Default for PositionPostingList {
95 fn default() -> Self {
96 Self::new()
97 }
98}
99
100impl PositionPostingList {
101 pub fn new() -> Self {
102 Self {
103 skip_list: Vec::new(),
104 data: Vec::new(),
105 doc_count: 0,
106 }
107 }
108
109 pub fn with_capacity(capacity: usize) -> Self {
110 Self {
111 skip_list: Vec::with_capacity(capacity / POSITION_BLOCK_SIZE + 1),
112 data: Vec::with_capacity(capacity * 8), doc_count: 0,
114 }
115 }
116
117 pub fn from_postings(postings: &[PostingWithPositions]) -> io::Result<Self> {
119 if postings.is_empty() {
120 return Ok(Self::new());
121 }
122
123 let mut skip_list = Vec::new();
124 let mut data = Vec::new();
125 let mut i = 0;
126
127 while i < postings.len() {
128 let block_start = data.len() as u64;
129 let block_end = (i + POSITION_BLOCK_SIZE).min(postings.len());
130 let block = &postings[i..block_end];
131
132 let base_doc_id = block.first().unwrap().doc_id;
134 let last_doc_id = block.last().unwrap().doc_id;
135 skip_list.push((base_doc_id, last_doc_id, block_start));
136
137 data.write_u32::<LittleEndian>(block.len() as u32)?;
139 data.write_u32::<LittleEndian>(base_doc_id)?;
140
141 let mut prev_doc_id = base_doc_id;
142 for (j, posting) in block.iter().enumerate() {
143 if j > 0 {
144 let delta = posting.doc_id - prev_doc_id;
145 write_vint(&mut data, delta as u64)?;
146 }
147 prev_doc_id = posting.doc_id;
148
149 write_vint(&mut data, posting.positions.len() as u64)?;
151 for &pos in &posting.positions {
152 write_vint(&mut data, pos as u64)?;
153 }
154 }
155
156 i = block_end;
157 }
158
159 Ok(Self {
160 skip_list,
161 data,
162 doc_count: postings.len() as u32,
163 })
164 }
165
166 pub fn push(&mut self, doc_id: DocId, positions: Vec<u32>) {
168 let posting = PostingWithPositions {
171 doc_id,
172 term_freq: positions.len() as u32,
173 positions,
174 };
175
176 let block_start = self.data.len() as u64;
178
179 let need_new_block =
181 self.skip_list.is_empty() || self.doc_count.is_multiple_of(POSITION_BLOCK_SIZE as u32);
182
183 if need_new_block {
184 self.skip_list.push((doc_id, doc_id, block_start));
186 self.data.write_u32::<LittleEndian>(1u32).unwrap();
187 self.data.write_u32::<LittleEndian>(doc_id).unwrap();
188 } else {
189 let last_block = self.skip_list.last_mut().unwrap();
191 let prev_doc_id = last_block.1;
192 last_block.1 = doc_id;
193
194 let count_offset = last_block.2 as usize;
196 let old_count = u32::from_le_bytes(
197 self.data[count_offset..count_offset + 4]
198 .try_into()
199 .unwrap(),
200 );
201 self.data[count_offset..count_offset + 4]
202 .copy_from_slice(&(old_count + 1).to_le_bytes());
203
204 let delta = doc_id - prev_doc_id;
205 write_vint(&mut self.data, delta as u64).unwrap();
206 }
207
208 write_vint(&mut self.data, posting.positions.len() as u64).unwrap();
210 for &pos in &posting.positions {
211 write_vint(&mut self.data, pos as u64).unwrap();
212 }
213
214 self.doc_count += 1;
215 }
216
217 pub fn doc_count(&self) -> u32 {
218 self.doc_count
219 }
220
221 pub fn len(&self) -> usize {
222 self.doc_count as usize
223 }
224
225 pub fn is_empty(&self) -> bool {
226 self.doc_count == 0
227 }
228
229 pub fn get_positions_into(&self, target_doc_id: DocId, out: &mut Vec<u32>) -> bool {
232 out.clear();
233 self.get_positions_impl(target_doc_id, Some(out)).is_some()
234 }
235
236 pub fn get_positions(&self, target_doc_id: DocId) -> Option<Vec<u32>> {
238 self.get_positions_impl(target_doc_id, None)
239 }
240
241 fn get_positions_impl(
242 &self,
243 target_doc_id: DocId,
244 mut out: Option<&mut Vec<u32>>,
245 ) -> Option<Vec<u32>> {
246 if self.skip_list.is_empty() {
247 return None;
248 }
249
250 let block_idx = match self.skip_list.binary_search_by(|&(base, last, _)| {
252 if target_doc_id < base {
253 std::cmp::Ordering::Greater
254 } else if target_doc_id > last {
255 std::cmp::Ordering::Less
256 } else {
257 std::cmp::Ordering::Equal
258 }
259 }) {
260 Ok(idx) => idx,
261 Err(_) => return None, };
263
264 let offset = self.skip_list[block_idx].2 as usize;
266 let mut reader = &self.data[offset..];
267
268 let count = reader.read_u32::<LittleEndian>().ok()? as usize;
270 let first_doc = reader.read_u32::<LittleEndian>().ok()?;
271 let mut prev_doc_id = first_doc;
272
273 for i in 0..count {
274 let doc_id = if i == 0 {
275 first_doc
276 } else {
277 let delta = read_vint(&mut reader).ok()? as u32;
278 prev_doc_id + delta
279 };
280 prev_doc_id = doc_id;
281
282 let num_positions = read_vint(&mut reader).ok()? as usize;
283
284 if doc_id == target_doc_id {
285 if let Some(buf) = &mut out {
287 buf.reserve(num_positions);
289 for _ in 0..num_positions {
290 let pos = read_vint(&mut reader).ok()? as u32;
291 buf.push(pos);
292 }
293 return Some(Vec::new()); }
295 let mut positions = Vec::with_capacity(num_positions);
296 for _ in 0..num_positions {
297 let pos = read_vint(&mut reader).ok()? as u32;
298 positions.push(pos);
299 }
300 return Some(positions);
301 } else {
302 for _ in 0..num_positions {
304 if read_vint(&mut reader).is_err() {
305 return None;
306 }
307 }
308 }
309 }
310
311 None
312 }
313
314 const SKIP_ENTRY_SIZE: usize = 20;
317
318 pub fn serialize<W: Write>(&self, writer: &mut W) -> io::Result<()> {
327 writer.write_all(&self.data)?;
329
330 for (i, (base_doc_id, last_doc_id, offset)) in self.skip_list.iter().enumerate() {
332 let next_offset = if i + 1 < self.skip_list.len() {
333 self.skip_list[i + 1].2
334 } else {
335 self.data.len() as u64
336 };
337 let length = (next_offset - offset) as u32;
338 writer.write_u32::<LittleEndian>(*base_doc_id)?;
339 writer.write_u32::<LittleEndian>(*last_doc_id)?;
340 writer.write_u64::<LittleEndian>(*offset)?;
341 writer.write_u32::<LittleEndian>(length)?;
342 }
343
344 writer.write_u64::<LittleEndian>(self.data.len() as u64)?;
346 writer.write_u32::<LittleEndian>(self.skip_list.len() as u32)?;
347 writer.write_u32::<LittleEndian>(self.doc_count)?;
348
349 Ok(())
350 }
351
352 pub fn deserialize(raw: &[u8]) -> io::Result<Self> {
354 if raw.len() < 16 {
355 return Err(io::Error::new(
356 io::ErrorKind::InvalidData,
357 "position data too short",
358 ));
359 }
360
361 let f = raw.len() - 16;
363 let data_len = u64::from_le_bytes(raw[f..f + 8].try_into().unwrap()) as usize;
364 let skip_count = u32::from_le_bytes(raw[f + 8..f + 12].try_into().unwrap()) as usize;
365 let doc_count = u32::from_le_bytes(raw[f + 12..f + 16].try_into().unwrap());
366
367 let mut skip_list = Vec::with_capacity(skip_count);
369 let mut pos = data_len;
370 for _ in 0..skip_count {
371 let base = u32::from_le_bytes(raw[pos..pos + 4].try_into().unwrap());
372 let last = u32::from_le_bytes(raw[pos + 4..pos + 8].try_into().unwrap());
373 let offset = u64::from_le_bytes(raw[pos + 8..pos + 16].try_into().unwrap());
374 skip_list.push((base, last, offset));
376 pos += Self::SKIP_ENTRY_SIZE;
377 }
378
379 let data = raw[..data_len].to_vec();
380
381 Ok(Self {
382 skip_list,
383 data,
384 doc_count,
385 })
386 }
387
388 pub fn concatenate_blocks(sources: &[(PositionPostingList, u32)]) -> io::Result<Self> {
390 let mut skip_list = Vec::new();
391 let mut data = Vec::new();
392 let mut total_docs = 0u32;
393
394 for (source, doc_offset) in sources {
395 for block_idx in 0..source.skip_list.len() {
396 let (base, last, src_offset) = source.skip_list[block_idx];
397 let next_offset = if block_idx + 1 < source.skip_list.len() {
398 source.skip_list[block_idx + 1].2 as usize
399 } else {
400 source.data.len()
401 };
402
403 let new_base = base + doc_offset;
404 let new_last = last + doc_offset;
405 let new_offset = data.len() as u64;
406
407 let block_bytes = &source.data[src_offset as usize..next_offset];
409
410 let count = u32::from_le_bytes(block_bytes[0..4].try_into().unwrap());
412 let first_doc = u32::from_le_bytes(block_bytes[4..8].try_into().unwrap());
413
414 data.write_u32::<LittleEndian>(count)?;
416 data.write_u32::<LittleEndian>(first_doc + doc_offset)?;
417 data.extend_from_slice(&block_bytes[8..]);
418
419 skip_list.push((new_base, new_last, new_offset));
420 total_docs += count;
421 }
422 }
423
424 Ok(Self {
425 skip_list,
426 data,
427 doc_count: total_docs,
428 })
429 }
430
431 pub fn concatenate_streaming<W: Write>(
446 sources: &[(&[u8], u32)],
447 writer: &mut W,
448 ) -> crate::Result<(u32, usize)> {
449 struct SourceMeta {
451 data_len: usize,
452 skip_count: usize,
453 }
454
455 let mut metas: Vec<SourceMeta> = Vec::with_capacity(sources.len());
456 let mut total_docs = 0u32;
457
458 for (source_index, (raw, _)) in sources.iter().enumerate() {
459 if raw.len() < 16 {
460 return Err(crate::Error::Corruption(format!(
461 "position posting source {} is shorter than its footer: {} bytes < 16",
462 source_index,
463 raw.len(),
464 )));
465 }
466 let f = raw.len() - 16;
467 let data_len = u64::from_le_bytes(raw[f..f + 8].try_into().unwrap()) as usize;
468 let skip_count = u32::from_le_bytes(raw[f + 8..f + 12].try_into().unwrap()) as usize;
469 let doc_count = u32::from_le_bytes(raw[f + 12..f + 16].try_into().unwrap());
470 total_docs += doc_count;
471 metas.push(SourceMeta {
472 data_len,
473 skip_count,
474 });
475 }
476
477 let mut out_skip: Vec<u8> = Vec::new();
480 let mut out_skip_count = 0u32;
481 let mut data_written = 0u64;
482 let mut patch_buf = [0u8; 8];
483 let es = Self::SKIP_ENTRY_SIZE;
484
485 for (src_idx, meta) in metas.iter().enumerate() {
486 let (raw, doc_offset) = &sources[src_idx];
487 let skip_base = meta.data_len;
488 let data = &raw[..meta.data_len];
489
490 for i in 0..meta.skip_count {
491 let p = skip_base + i * es;
493 let base = u32::from_le_bytes(raw[p..p + 4].try_into().unwrap());
494 let last = u32::from_le_bytes(raw[p + 4..p + 8].try_into().unwrap());
495 let offset = u64::from_le_bytes(raw[p + 8..p + 16].try_into().unwrap());
496 let length = u32::from_le_bytes(raw[p + 16..p + 20].try_into().unwrap());
497
498 let block = &data[offset as usize..(offset as usize + length as usize)];
499
500 out_skip.extend_from_slice(&(base + doc_offset).to_le_bytes());
502 out_skip.extend_from_slice(&(last + doc_offset).to_le_bytes());
503 out_skip.extend_from_slice(&data_written.to_le_bytes());
504 out_skip.extend_from_slice(&length.to_le_bytes());
505 out_skip_count += 1;
506
507 patch_buf[0..4].copy_from_slice(&block[0..4]);
509 let first_doc = u32::from_le_bytes(block[4..8].try_into().unwrap());
510 patch_buf[4..8].copy_from_slice(&(first_doc + doc_offset).to_le_bytes());
511 writer.write_all(&patch_buf)?;
512 writer.write_all(&block[8..])?;
513
514 data_written += block.len() as u64;
515 }
516 }
517
518 writer.write_all(&out_skip)?;
520
521 writer.write_u64::<LittleEndian>(data_written)?;
522 writer.write_u32::<LittleEndian>(out_skip_count)?;
523 writer.write_u32::<LittleEndian>(total_docs)?;
524
525 let total_bytes = data_written as usize + out_skip.len() + 16;
526 Ok((total_docs, total_bytes))
527 }
528
529 pub fn iter(&self) -> PositionPostingIterator<'_> {
531 PositionPostingIterator::new(self)
532 }
533}
534
535pub struct PositionPostingIterator<'a> {
541 list: &'a PositionPostingList,
542 current_block: usize,
543 position_in_block: usize,
544 block_count: usize,
546 block_doc_ids: Vec<DocId>,
548 block_term_freqs: Vec<u32>,
550 block_positions: Vec<u32>,
552 block_pos_offsets: Vec<usize>,
554 exhausted: bool,
555}
556
557impl<'a> PositionPostingIterator<'a> {
558 pub fn new(list: &'a PositionPostingList) -> Self {
559 let exhausted = list.skip_list.is_empty();
560 let mut iter = Self {
561 list,
562 current_block: 0,
563 position_in_block: 0,
564 block_count: 0,
565 block_doc_ids: Vec::with_capacity(POSITION_BLOCK_SIZE),
566 block_term_freqs: Vec::with_capacity(POSITION_BLOCK_SIZE),
567 block_positions: Vec::new(),
568 block_pos_offsets: Vec::with_capacity(POSITION_BLOCK_SIZE + 1),
569 exhausted,
570 };
571 if !iter.exhausted {
572 iter.load_block(0);
573 }
574 iter
575 }
576
577 fn load_block(&mut self, block_idx: usize) {
578 if block_idx >= self.list.skip_list.len() {
579 self.exhausted = true;
580 return;
581 }
582
583 self.current_block = block_idx;
584 self.position_in_block = 0;
585
586 let offset = self.list.skip_list[block_idx].2 as usize;
587 let mut reader = &self.list.data[offset..];
588
589 let count = reader.read_u32::<LittleEndian>().unwrap_or(0) as usize;
591 let first_doc = reader.read_u32::<LittleEndian>().unwrap_or(0);
592
593 self.block_count = count;
594 self.block_doc_ids.clear();
595 self.block_term_freqs.clear();
596 self.block_positions.clear();
597 self.block_pos_offsets.clear();
598
599 let mut prev_doc_id = first_doc;
600
601 for i in 0..count {
602 let doc_id = if i == 0 {
603 first_doc
604 } else {
605 let delta = read_vint(&mut reader).unwrap_or(0) as u32;
606 prev_doc_id + delta
607 };
608 prev_doc_id = doc_id;
609
610 let num_positions = read_vint(&mut reader).unwrap_or(0) as usize;
611 self.block_doc_ids.push(doc_id);
612 self.block_term_freqs.push(num_positions as u32);
613 self.block_pos_offsets.push(self.block_positions.len());
614 for _ in 0..num_positions {
615 let pos = read_vint(&mut reader).unwrap_or(0) as u32;
616 self.block_positions.push(pos);
617 }
618 }
619 self.block_pos_offsets.push(self.block_positions.len());
621 }
622
623 pub fn doc(&self) -> DocId {
624 if self.exhausted || self.position_in_block >= self.block_count {
625 u32::MAX
626 } else {
627 self.block_doc_ids[self.position_in_block]
628 }
629 }
630
631 pub fn term_freq(&self) -> u32 {
632 if self.exhausted || self.position_in_block >= self.block_count {
633 0
634 } else {
635 self.block_term_freqs[self.position_in_block]
636 }
637 }
638
639 pub fn positions(&self) -> &[u32] {
640 if self.exhausted || self.position_in_block >= self.block_count {
641 &[]
642 } else {
643 let start = self.block_pos_offsets[self.position_in_block];
644 let end = self.block_pos_offsets[self.position_in_block + 1];
645 &self.block_positions[start..end]
646 }
647 }
648
649 pub fn advance(&mut self) {
650 if self.exhausted {
651 return;
652 }
653
654 self.position_in_block += 1;
655 if self.position_in_block >= self.block_count {
656 self.load_block(self.current_block + 1);
657 }
658 }
659
660 pub fn seek(&mut self, target: DocId) {
661 if self.exhausted {
662 return;
663 }
664
665 if let Some((_, last, _)) = self.list.skip_list.get(self.current_block)
667 && target <= *last
668 {
669 let remaining = &self.block_doc_ids[self.position_in_block..self.block_count];
671 let offset = crate::structures::simd::find_first_ge_u32(remaining, target);
672 self.position_in_block += offset;
673 if self.position_in_block >= self.block_count {
674 self.load_block(self.current_block + 1);
675 self.seek(target); }
677 return;
678 }
679
680 let block_idx = match self.list.skip_list.binary_search_by(|&(base, last, _)| {
682 if target < base {
683 std::cmp::Ordering::Greater
684 } else if target > last {
685 std::cmp::Ordering::Less
686 } else {
687 std::cmp::Ordering::Equal
688 }
689 }) {
690 Ok(idx) => idx,
691 Err(idx) => idx, };
693
694 if block_idx >= self.list.skip_list.len() {
695 self.exhausted = true;
696 return;
697 }
698
699 self.load_block(block_idx);
700
701 let remaining = &self.block_doc_ids[self.position_in_block..self.block_count];
703 let offset = crate::structures::simd::find_first_ge_u32(remaining, target);
704 self.position_in_block += offset;
705
706 if self.position_in_block >= self.block_count {
707 self.load_block(self.current_block + 1);
708 }
709 }
710}
711
712#[cfg(test)]
713mod tests {
714 use super::*;
715
716 #[test]
717 fn test_position_encoding() {
718 let pos = encode_position(0, 5);
720 assert_eq!(decode_element_ordinal(pos), 0);
721 assert_eq!(decode_token_position(pos), 5);
722
723 let pos = encode_position(3, 100);
725 assert_eq!(decode_element_ordinal(pos), 3);
726 assert_eq!(decode_token_position(pos), 100);
727
728 let pos = encode_position(MAX_ELEMENT_ORDINAL, MAX_TOKEN_POSITION);
730 assert_eq!(decode_element_ordinal(pos), MAX_ELEMENT_ORDINAL);
731 assert_eq!(decode_token_position(pos), MAX_TOKEN_POSITION);
732 }
733
734 #[test]
735 fn test_position_posting_list_build() {
736 let postings = vec![
738 PostingWithPositions {
739 doc_id: 1,
740 term_freq: 2,
741 positions: vec![encode_position(0, 0), encode_position(0, 2)],
742 },
743 PostingWithPositions {
744 doc_id: 3,
745 term_freq: 1,
746 positions: vec![encode_position(1, 0)],
747 },
748 ];
749
750 let list = PositionPostingList::from_postings(&postings).unwrap();
751 assert_eq!(list.doc_count(), 2);
752
753 let pos = list.get_positions(1).unwrap();
755 assert_eq!(pos.len(), 2);
756
757 let pos = list.get_positions(3).unwrap();
758 assert_eq!(pos.len(), 1);
759
760 assert!(list.get_positions(2).is_none());
762 assert!(list.get_positions(99).is_none());
763 }
764
765 #[test]
766 fn test_serialization_roundtrip() {
767 let postings = vec![
768 PostingWithPositions {
769 doc_id: 1,
770 term_freq: 2,
771 positions: vec![encode_position(0, 0), encode_position(0, 5)],
772 },
773 PostingWithPositions {
774 doc_id: 3,
775 term_freq: 1,
776 positions: vec![encode_position(1, 0)],
777 },
778 PostingWithPositions {
779 doc_id: 5,
780 term_freq: 1,
781 positions: vec![encode_position(0, 10)],
782 },
783 ];
784
785 let list = PositionPostingList::from_postings(&postings).unwrap();
786
787 let mut bytes = Vec::new();
788 list.serialize(&mut bytes).unwrap();
789
790 let deserialized = PositionPostingList::deserialize(&bytes).unwrap();
791
792 assert_eq!(list.doc_count(), deserialized.doc_count());
793
794 let pos = deserialized.get_positions(1).unwrap();
796 assert_eq!(pos, vec![encode_position(0, 0), encode_position(0, 5)]);
797
798 let pos = deserialized.get_positions(3).unwrap();
799 assert_eq!(pos, vec![encode_position(1, 0)]);
800 }
801
802 #[test]
803 fn test_binary_search_many_blocks() {
804 let mut postings = Vec::new();
806 for i in 0..300 {
807 postings.push(PostingWithPositions {
808 doc_id: i * 2, term_freq: 1,
810 positions: vec![encode_position(0, i)],
811 });
812 }
813
814 let list = PositionPostingList::from_postings(&postings).unwrap();
815 assert_eq!(list.doc_count(), 300);
816
817 assert_eq!(list.skip_list.len(), 3);
819
820 let pos = list.get_positions(0).unwrap();
822 assert_eq!(pos, vec![encode_position(0, 0)]);
823
824 let pos = list.get_positions(256).unwrap(); assert_eq!(pos, vec![encode_position(0, 128)]);
826
827 let pos = list.get_positions(598).unwrap(); assert_eq!(pos, vec![encode_position(0, 299)]);
829
830 assert!(list.get_positions(1).is_none());
832 assert!(list.get_positions(257).is_none());
833 }
834
835 #[test]
836 fn test_concatenate_blocks_merge() {
837 let postings1 = vec![
839 PostingWithPositions {
840 doc_id: 0,
841 term_freq: 1,
842 positions: vec![0],
843 },
844 PostingWithPositions {
845 doc_id: 1,
846 term_freq: 1,
847 positions: vec![5],
848 },
849 PostingWithPositions {
850 doc_id: 2,
851 term_freq: 1,
852 positions: vec![10],
853 },
854 ];
855 let list1 = PositionPostingList::from_postings(&postings1).unwrap();
856
857 let postings2 = vec![
858 PostingWithPositions {
859 doc_id: 0,
860 term_freq: 1,
861 positions: vec![100],
862 },
863 PostingWithPositions {
864 doc_id: 1,
865 term_freq: 1,
866 positions: vec![105],
867 },
868 ];
869 let list2 = PositionPostingList::from_postings(&postings2).unwrap();
870
871 let combined = PositionPostingList::concatenate_blocks(&[
873 (list1, 0), (list2, 3), ])
876 .unwrap();
877
878 assert_eq!(combined.doc_count(), 5);
879
880 assert!(combined.get_positions(0).is_some());
882 assert!(combined.get_positions(1).is_some());
883 assert!(combined.get_positions(2).is_some());
884 assert!(combined.get_positions(3).is_some()); assert!(combined.get_positions(4).is_some()); }
887
888 #[test]
889 fn test_iterator() {
890 let postings = vec![
891 PostingWithPositions {
892 doc_id: 1,
893 term_freq: 2,
894 positions: vec![0, 5],
895 },
896 PostingWithPositions {
897 doc_id: 3,
898 term_freq: 1,
899 positions: vec![10],
900 },
901 PostingWithPositions {
902 doc_id: 5,
903 term_freq: 1,
904 positions: vec![15],
905 },
906 ];
907
908 let list = PositionPostingList::from_postings(&postings).unwrap();
909 let mut iter = list.iter();
910
911 assert_eq!(iter.doc(), 1);
912 assert_eq!(iter.positions(), &[0, 5]);
913
914 iter.advance();
915 assert_eq!(iter.doc(), 3);
916
917 iter.seek(5);
918 assert_eq!(iter.doc(), 5);
919 assert_eq!(iter.positions(), &[15]);
920
921 iter.advance();
922 assert_eq!(iter.doc(), u32::MAX); }
924
925 fn build_ppl(entries: &[(u32, Vec<u32>)]) -> PositionPostingList {
927 let postings: Vec<PostingWithPositions> = entries
928 .iter()
929 .map(|(doc_id, positions)| PostingWithPositions {
930 doc_id: *doc_id,
931 term_freq: positions.len() as u32,
932 positions: positions.clone(),
933 })
934 .collect();
935 PositionPostingList::from_postings(&postings).unwrap()
936 }
937
938 fn serialize_ppl(ppl: &PositionPostingList) -> Vec<u8> {
940 let mut buf = Vec::new();
941 ppl.serialize(&mut buf).unwrap();
942 buf
943 }
944
945 fn collect_positions(ppl: &PositionPostingList) -> Vec<(u32, Vec<u32>)> {
947 let mut result = Vec::new();
948 let mut it = ppl.iter();
949 while it.doc() != u32::MAX {
950 result.push((it.doc(), it.positions().to_vec()));
951 it.advance();
952 }
953 result
954 }
955
956 #[test]
957 fn test_concatenate_streaming_matches_blocks() {
958 let seg_a: Vec<(u32, Vec<u32>)> = (0..150)
960 .map(|i| (i * 2, vec![i * 10, i * 10 + 3]))
961 .collect();
962 let seg_b: Vec<(u32, Vec<u32>)> = (0..100).map(|i| (i * 5, vec![i * 7])).collect();
963 let seg_c: Vec<(u32, Vec<u32>)> = (0..80).map(|i| (i * 3, vec![i, i + 1, i + 2])).collect();
964
965 let ppl_a = build_ppl(&seg_a);
966 let ppl_b = build_ppl(&seg_b);
967 let ppl_c = build_ppl(&seg_c);
968
969 let offset_b = 500u32;
970 let offset_c = 1000u32;
971
972 let ref_merged = PositionPostingList::concatenate_blocks(&[
974 (ppl_a.clone(), 0),
975 (ppl_b.clone(), offset_b),
976 (ppl_c.clone(), offset_c),
977 ])
978 .unwrap();
979 let mut ref_buf = Vec::new();
980 ref_merged.serialize(&mut ref_buf).unwrap();
981
982 let bytes_a = serialize_ppl(&ppl_a);
984 let bytes_b = serialize_ppl(&ppl_b);
985 let bytes_c = serialize_ppl(&ppl_c);
986
987 let sources: Vec<(&[u8], u32)> =
988 vec![(&bytes_a, 0), (&bytes_b, offset_b), (&bytes_c, offset_c)];
989 let mut stream_buf = Vec::new();
990 let (doc_count, bytes_written) =
991 PositionPostingList::concatenate_streaming(&sources, &mut stream_buf).unwrap();
992
993 assert_eq!(doc_count, 330);
994 assert_eq!(bytes_written, stream_buf.len());
995
996 let ref_posts = collect_positions(&PositionPostingList::deserialize(&ref_buf).unwrap());
998 let stream_posts =
999 collect_positions(&PositionPostingList::deserialize(&stream_buf).unwrap());
1000
1001 assert_eq!(ref_posts.len(), stream_posts.len());
1002 for (i, (r, s)) in ref_posts.iter().zip(stream_posts.iter()).enumerate() {
1003 assert_eq!(r.0, s.0, "doc_id mismatch at {}", i);
1004 assert_eq!(r.1, s.1, "positions mismatch at doc {}", r.0);
1005 }
1006 }
1007
1008 #[test]
1009 fn test_position_concatenate_streaming_short_source_returns_corruption() {
1010 let ppl_a = build_ppl(&[(0, vec![1, 5]), (3, vec![2])]);
1014 let ppl_c = build_ppl(&[(1, vec![4]), (2, vec![7, 9])]);
1015 let bytes_a = serialize_ppl(&ppl_a);
1016 let bytes_c = serialize_ppl(&ppl_c);
1017 let short = vec![0u8; 15]; let sources: Vec<(&[u8], u32)> = vec![(&bytes_a, 0), (&short, 100), (&bytes_c, 200)];
1020 let mut out = Vec::new();
1021 let result = PositionPostingList::concatenate_streaming(&sources, &mut out);
1022 assert!(
1023 matches!(result, Err(crate::Error::Corruption(_))),
1024 "short/corrupt source must be a Corruption error, not silently skipped: {:?}",
1025 result.map(|r| r.0)
1026 );
1027 }
1028
1029 #[test]
1030 fn test_positions_multi_round_merge() {
1031 let segments: Vec<Vec<(u32, Vec<u32>)>> = (0..4)
1033 .map(|seg| {
1034 (0..200)
1035 .map(|i| {
1036 let pos_count = (i % 3) + 1;
1037 let positions: Vec<u32> = (0..pos_count)
1038 .map(|p| (seg * 1000 + i * 10 + p) as u32)
1039 .collect();
1040 (i as u32 * 3, positions)
1041 })
1042 .collect()
1043 })
1044 .collect();
1045
1046 let ppls: Vec<PositionPostingList> = segments.iter().map(|s| build_ppl(s)).collect();
1047 let serialized: Vec<Vec<u8>> = ppls.iter().map(serialize_ppl).collect();
1048
1049 let mut merged_01 = Vec::new();
1051 let sources_01: Vec<(&[u8], u32)> = vec![(&serialized[0], 0), (&serialized[1], 600)];
1052 let (dc_01, _) =
1053 PositionPostingList::concatenate_streaming(&sources_01, &mut merged_01).unwrap();
1054 assert_eq!(dc_01, 400);
1055
1056 let mut merged_23 = Vec::new();
1057 let sources_23: Vec<(&[u8], u32)> = vec![(&serialized[2], 0), (&serialized[3], 600)];
1058 let (dc_23, _) =
1059 PositionPostingList::concatenate_streaming(&sources_23, &mut merged_23).unwrap();
1060 assert_eq!(dc_23, 400);
1061
1062 let mut final_merged = Vec::new();
1064 let sources_final: Vec<(&[u8], u32)> = vec![(&merged_01, 0), (&merged_23, 1200)];
1065 let (dc_final, _) =
1066 PositionPostingList::concatenate_streaming(&sources_final, &mut final_merged).unwrap();
1067 assert_eq!(dc_final, 800);
1068
1069 let final_ppl = PositionPostingList::deserialize(&final_merged).unwrap();
1071 let all = collect_positions(&final_ppl);
1072 assert_eq!(all.len(), 800);
1073
1074 assert_eq!(all[0].0, 0);
1077 assert_eq!(all[0].1, vec![0]); assert_eq!(all[200].0, 600);
1081 assert_eq!(all[200].1, vec![1000]); assert_eq!(all[400].0, 1200);
1085 assert_eq!(all[400].1, vec![2000]); let mut it = final_ppl.iter();
1089 it.seek(1200);
1090 assert_eq!(it.doc(), 1200);
1091 assert_eq!(it.positions(), &[2000]);
1092
1093 let pos = final_ppl.get_positions(600).unwrap();
1095 assert_eq!(pos, vec![1000]);
1096 }
1097
1098 #[test]
1099 fn test_positions_large_scale_merge() {
1100 let num_segments = 5usize;
1102 let docs_per_segment = 500usize;
1103
1104 let segments: Vec<Vec<(u32, Vec<u32>)>> = (0..num_segments)
1105 .map(|seg| {
1106 (0..docs_per_segment)
1107 .map(|i| {
1108 let n_pos = (i % 4) + 1;
1109 let positions: Vec<u32> =
1110 (0..n_pos).map(|p| (p * 5 + seg) as u32).collect();
1111 (i as u32 * 2, positions)
1112 })
1113 .collect()
1114 })
1115 .collect();
1116
1117 let ppls: Vec<PositionPostingList> = segments.iter().map(|s| build_ppl(s)).collect();
1118 let serialized: Vec<Vec<u8>> = ppls.iter().map(serialize_ppl).collect();
1119
1120 let max_doc = (docs_per_segment as u32 - 1) * 2;
1121 let offsets: Vec<u32> = (0..num_segments)
1122 .map(|i| i as u32 * (max_doc + 1))
1123 .collect();
1124
1125 let sources: Vec<(&[u8], u32)> = serialized
1126 .iter()
1127 .zip(offsets.iter())
1128 .map(|(b, o)| (b.as_slice(), *o))
1129 .collect();
1130
1131 let mut merged = Vec::new();
1132 let (doc_count, _) =
1133 PositionPostingList::concatenate_streaming(&sources, &mut merged).unwrap();
1134 assert_eq!(doc_count, (num_segments * docs_per_segment) as u32);
1135
1136 let merged_ppl = PositionPostingList::deserialize(&merged).unwrap();
1138 let all = collect_positions(&merged_ppl);
1139 assert_eq!(all.len(), num_segments * docs_per_segment);
1140
1141 for (seg, &offset) in offsets.iter().enumerate() {
1143 for i in 0..docs_per_segment {
1144 let idx = seg * docs_per_segment + i;
1145 let expected_doc = i as u32 * 2 + offset;
1146 assert_eq!(all[idx].0, expected_doc, "seg={} i={}", seg, i);
1147
1148 let n_pos = (i % 4) + 1;
1149 let expected_positions: Vec<u32> =
1150 (0..n_pos).map(|p| (p * 5 + seg) as u32).collect();
1151 assert_eq!(
1152 all[idx].1, expected_positions,
1153 "positions mismatch seg={} i={}",
1154 seg, i
1155 );
1156 }
1157 }
1158 }
1159
1160 #[test]
1161 fn test_positions_streaming_single_source() {
1162 let entries: Vec<(u32, Vec<u32>)> =
1163 (0..300).map(|i| (i * 4, vec![i * 2, i * 2 + 1])).collect();
1164 let ppl = build_ppl(&entries);
1165 let direct = serialize_ppl(&ppl);
1166
1167 let sources: Vec<(&[u8], u32)> = vec![(&direct, 0)];
1168 let mut streamed = Vec::new();
1169 PositionPostingList::concatenate_streaming(&sources, &mut streamed).unwrap();
1170
1171 let p1 = collect_positions(&PositionPostingList::deserialize(&direct).unwrap());
1172 let p2 = collect_positions(&PositionPostingList::deserialize(&streamed).unwrap());
1173 assert_eq!(p1, p2);
1174 }
1175
1176 #[test]
1177 fn test_positions_edge_single_doc() {
1178 let ppl_a = build_ppl(&[(0, vec![42, 43, 44])]);
1179 let ppl_b = build_ppl(&[(0, vec![100])]);
1180
1181 let merged = PositionPostingList::concatenate_blocks(&[(ppl_a, 0), (ppl_b, 1)]).unwrap();
1182
1183 assert_eq!(merged.doc_count(), 2);
1184 assert_eq!(merged.get_positions(0).unwrap(), vec![42, 43, 44]);
1185 assert_eq!(merged.get_positions(1).unwrap(), vec![100]);
1186 }
1187}