1use super::bm25::TermScorer;
58use super::index::{FieldId, SegmentIndex};
59use super::postings::{
60 BLOCK_META_LEN, BlockImpact, BlockMeta, DEFAULT_POSTINGS_PER_BLOCK, HEADER_LEN, PostingList,
61 PostingsError, PostingsReader, block_impacts, encode_v2,
62};
63use super::search::FieldWeights;
64use crate::kernels::postings::{prefix_sum, unpack};
65use crate::meta::DocBitmap;
66
67const REGION_MAGIC: [u8; 4] = *b"ZFTS";
68const REGION_VERSION: u16 = 1;
69const REGION_HEADER_LEN: usize = 40;
70const REGION_SPAN_LEN: usize = 48;
71const REGION_FIELD_HEADER_LEN: usize = 8;
72
73#[derive(Clone, Copy, Debug, Eq, PartialEq)]
75pub enum SealedSegmentError {
76 Geometry(&'static str),
78 Postings(PostingsError),
80}
81
82impl std::fmt::Display for SealedSegmentError {
83 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84 match self {
85 Self::Geometry(detail) => write!(formatter, "sealed lexical geometry: {detail}"),
86 Self::Postings(error) => write!(formatter, "sealed lexical postings: {error}"),
87 }
88 }
89}
90
91impl std::error::Error for SealedSegmentError {}
92
93impl From<PostingsError> for SealedSegmentError {
94 fn from(error: PostingsError) -> Self {
95 Self::Postings(error)
96 }
97}
98
99#[derive(Clone, Debug)]
101struct FieldLengths {
102 field: FieldId,
103 lengths: Vec<u32>,
104}
105
106#[derive(Clone, Copy, Debug)]
113struct ListSpan {
114 field: FieldId,
115 term_start: u32,
116 term_len: u32,
117 meta_start: u32,
118 docids_start: u32,
119 tfs_start: u32,
120 block_count: u32,
121 doc_freq: u32,
122 overall_max_tf: u32,
128 overall_min_len: u16,
130 union_doc_freq: u32,
132 field_count: u32,
134}
135
136#[derive(Clone, Debug, Default)]
138pub struct SealedSegment {
139 terms: Vec<u8>,
141 spans: Vec<ListSpan>,
143 blob: Vec<u8>,
145 lengths: Vec<FieldLengths>,
146 total_lengths: Vec<u32>,
152 row_count: u32,
153 postings_per_block: u16,
154}
155
156impl SealedSegment {
157 pub(crate) fn resident_bytes(&self) -> Result<usize, SealedSegmentError> {
158 self.resident_bytes_controlled(&mut super::control::WorkCheck::new(|| {
159 Ok::<(), SealedSegmentError>(())
160 }))
161 }
162
163 pub(crate) fn resident_bytes_controlled<E: From<SealedSegmentError>>(
164 &self,
165 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
166 ) -> Result<usize, E> {
167 work.check_now()?;
168 let direct = self
169 .terms
170 .capacity()
171 .checked_add(self.blob.capacity())
172 .and_then(|bytes| {
173 self.spans
174 .capacity()
175 .checked_mul(std::mem::size_of::<ListSpan>())
176 .and_then(|spans| bytes.checked_add(spans))
177 })
178 .and_then(|bytes| {
179 self.lengths
180 .capacity()
181 .checked_mul(std::mem::size_of::<FieldLengths>())
182 .and_then(|lengths| bytes.checked_add(lengths))
183 })
184 .and_then(|bytes| {
185 self.total_lengths
186 .capacity()
187 .checked_mul(std::mem::size_of::<u32>())
188 .and_then(|lengths| bytes.checked_add(lengths))
189 })
190 .ok_or(SealedSegmentError::Geometry(
191 "sealed lexical resident bytes overflow",
192 ))?;
193 self.lengths.iter().try_fold(direct, |total, entry| {
194 work.step()?;
195 entry
196 .lengths
197 .capacity()
198 .checked_mul(std::mem::size_of::<u32>())
199 .and_then(|bytes| total.checked_add(bytes))
200 .ok_or(SealedSegmentError::Geometry(
201 "sealed lexical resident bytes overflow",
202 ))
203 .map_err(E::from)
204 })
205 }
206
207 pub fn terms(&self) -> impl Iterator<Item = &[u8]> {
213 (0..self.spans.len()).filter_map(move |index| {
214 let term = self.term_of(index);
215 (index == 0 || self.term_of(index.saturating_sub(1)) != term).then_some(term)
216 })
217 }
218
219 pub(crate) fn vocabulary_terms(&self) -> impl Iterator<Item = (&[u8], FieldId)> + Clone {
220 self.spans
221 .iter()
222 .enumerate()
223 .map(move |(index, span)| (self.term_of(index), span.field))
224 }
225
226 pub fn seal(segment: &SegmentIndex) -> Result<Self, PostingsError> {
235 Self::seal_controlled(
236 segment,
237 &mut super::control::WorkCheck::new(|| Ok::<(), PostingsError>(())),
238 )
239 }
240
241 pub(crate) fn seal_controlled<E: From<PostingsError>>(
242 segment: &SegmentIndex,
243 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
244 ) -> Result<Self, E> {
245 work.check_now()?;
246 let per_block = DEFAULT_POSTINGS_PER_BLOCK;
247 let mut sealed = Self {
248 terms: Vec::new(),
249 spans: Vec::new(),
250 blob: Vec::new(),
251 lengths: Vec::new(),
252 total_lengths: Vec::new(),
253 row_count: segment.row_count(),
254 postings_per_block: per_block,
255 };
256 for field in segment.fields() {
257 work.step()?;
258 let source = segment.field_lengths(field).unwrap_or(&[]);
259 let mut lengths = Vec::with_capacity(source.len());
260 for length in source {
261 work.step()?;
262 #[cfg(test)]
263 SEAL_LENGTH_PROBES.with(|value| value.set(value.get() + 1));
264 lengths.push(*length);
265 }
266 sealed.lengths.push(FieldLengths { field, lengths });
267 }
268
269 let rows = usize::try_from(sealed.row_count).unwrap_or(0);
270 sealed.total_lengths = vec![0_u32; rows];
271 for entry in &sealed.lengths {
272 work.step()?;
273 for (slot, length) in entry.lengths.iter().enumerate() {
274 work.step()?;
275 #[cfg(test)]
276 SEAL_TOTAL_PROBES.with(|value| value.set(value.get() + 1));
277 if let Some(total) = sealed.total_lengths.get_mut(slot) {
278 *total = total.saturating_add(*length);
279 }
280 }
281 }
282
283 let mut group_term: Option<&[u8]> = None;
289 let mut group_start = 0_usize;
290 let mut group_lists: Vec<&PostingList> = Vec::new();
291 for (key, list) in segment.postings() {
292 work.step()?;
293 if list.is_empty() {
294 continue;
295 }
296 if group_term != Some(key.term.as_slice()) {
297 finish_union_group(&mut sealed.spans, group_start, &group_lists, work)?;
298 group_term = Some(key.term.as_slice());
299 group_start = sealed.spans.len();
300 group_lists.clear();
301 }
302 group_lists.push(list);
303 let mut field_lengths = &[][..];
304 for entry in &sealed.lengths {
305 work.step()?;
306 if entry.field == key.field {
307 field_lengths = entry.lengths.as_slice();
308 break;
309 }
310 }
311 let impacts =
312 super::postings::block_impacts_controlled(list, per_block, field_lengths, work)?;
313 let encoded =
314 super::postings::encode_v2_controlled(list, per_block, &[], &impacts, work)?;
315 let bytes = encoded.as_bytes();
316 let reader = PostingsReader::open_controlled(bytes, work)?;
319 let base = u32::try_from(sealed.blob.len()).unwrap_or(u32::MAX);
320 let block_count = u32::try_from(reader.blocks().len()).unwrap_or(u32::MAX);
321 let meta_start = base.saturating_add(HEADER_LEN_U32);
322 let streams = meta_start.saturating_add(block_count.saturating_mul(META_LEN_U32));
323 let mut docid_bytes = 0;
324 for meta in reader.blocks() {
325 work.step()?;
326 docid_bytes =
327 docid_bytes.max(stream_end(meta.docids_offset, meta.count, meta.docid_bits));
328 }
329 let mut overall_max_tf = 0_u32;
330 let mut overall_min_len = u16::MAX;
331 for meta in reader.blocks() {
332 work.step()?;
333 let impact = BlockImpact::from_meta(meta);
334 overall_max_tf = overall_max_tf.max(impact.max_tf);
335 overall_min_len = overall_min_len.min(impact.min_len);
336 }
337
338 let term_start = u32::try_from(sealed.terms.len()).unwrap_or(u32::MAX);
339 super::control::extend_bytes(&mut sealed.terms, &key.term, work)?;
340 sealed.spans.push(ListSpan {
341 field: key.field,
342 term_start,
343 term_len: u32::try_from(key.term.len()).unwrap_or(u32::MAX),
344 meta_start,
345 docids_start: streams,
346 tfs_start: streams.saturating_add(docid_bytes),
347 block_count,
348 doc_freq: reader.document_frequency(),
349 overall_max_tf,
350 overall_min_len,
351 union_doc_freq: 0,
353 field_count: 0,
354 });
355 super::control::extend_bytes(&mut sealed.blob, bytes, work)?;
356 }
357 finish_union_group(&mut sealed.spans, group_start, &group_lists, work)?;
358 work.check_now()?;
359 Ok(sealed)
360 }
361
362 pub fn encode_region(&self) -> Result<Vec<u8>, SealedSegmentError> {
364 let span_count = u32::try_from(self.spans.len())
365 .map_err(|_| SealedSegmentError::Geometry("span count exceeds u32"))?;
366 let field_count = u32::try_from(self.lengths.len())
367 .map_err(|_| SealedSegmentError::Geometry("field count exceeds u32"))?;
368 let terms_len = u64::try_from(self.terms.len())
369 .map_err(|_| SealedSegmentError::Geometry("term bytes exceed u64"))?;
370 let blob_len = u64::try_from(self.blob.len())
371 .map_err(|_| SealedSegmentError::Geometry("postings bytes exceed u64"))?;
372 let spans_len = self
373 .spans
374 .len()
375 .checked_mul(REGION_SPAN_LEN)
376 .ok_or(SealedSegmentError::Geometry("span bytes overflow"))?;
377 let fields_len = self.lengths.iter().try_fold(0_usize, |total, entry| {
378 if entry.lengths.len() != self.row_count as usize {
379 return Err(SealedSegmentError::Geometry(
380 "field lengths do not cover every row",
381 ));
382 }
383 let body = entry
384 .lengths
385 .len()
386 .checked_mul(4)
387 .and_then(|bytes| bytes.checked_add(REGION_FIELD_HEADER_LEN))
388 .ok_or(SealedSegmentError::Geometry("field bytes overflow"))?;
389 total
390 .checked_add(body)
391 .ok_or(SealedSegmentError::Geometry("field bytes overflow"))
392 })?;
393 let capacity = REGION_HEADER_LEN
394 .checked_add(spans_len)
395 .and_then(|value| value.checked_add(fields_len))
396 .and_then(|value| value.checked_add(self.terms.len()))
397 .and_then(|value| value.checked_add(self.blob.len()))
398 .ok_or(SealedSegmentError::Geometry(
399 "lexical region length overflow",
400 ))?;
401 let mut bytes = Vec::with_capacity(capacity);
402 bytes.extend_from_slice(®ION_MAGIC);
403 bytes.extend_from_slice(®ION_VERSION.to_le_bytes());
404 bytes.extend_from_slice(&self.postings_per_block.to_le_bytes());
405 bytes.extend_from_slice(&self.row_count.to_le_bytes());
406 bytes.extend_from_slice(&span_count.to_le_bytes());
407 bytes.extend_from_slice(&field_count.to_le_bytes());
408 bytes.extend_from_slice(&0_u32.to_le_bytes());
409 bytes.extend_from_slice(&terms_len.to_le_bytes());
410 bytes.extend_from_slice(&blob_len.to_le_bytes());
411 for span in &self.spans {
412 bytes.extend_from_slice(&span.field.0.to_le_bytes());
413 bytes.extend_from_slice(&0_u16.to_le_bytes());
414 bytes.extend_from_slice(&span.term_start.to_le_bytes());
415 bytes.extend_from_slice(&span.term_len.to_le_bytes());
416 bytes.extend_from_slice(&span.meta_start.to_le_bytes());
417 bytes.extend_from_slice(&span.docids_start.to_le_bytes());
418 bytes.extend_from_slice(&span.tfs_start.to_le_bytes());
419 bytes.extend_from_slice(&span.block_count.to_le_bytes());
420 bytes.extend_from_slice(&span.doc_freq.to_le_bytes());
421 bytes.extend_from_slice(&span.overall_max_tf.to_le_bytes());
422 bytes.extend_from_slice(&span.overall_min_len.to_le_bytes());
423 bytes.extend_from_slice(&0_u16.to_le_bytes());
424 bytes.extend_from_slice(&span.union_doc_freq.to_le_bytes());
425 bytes.extend_from_slice(&span.field_count.to_le_bytes());
426 }
427 for entry in &self.lengths {
428 bytes.extend_from_slice(&entry.field.0.to_le_bytes());
429 bytes.extend_from_slice(&0_u16.to_le_bytes());
430 bytes.extend_from_slice(&self.row_count.to_le_bytes());
431 for length in &entry.lengths {
432 bytes.extend_from_slice(&length.to_le_bytes());
433 }
434 }
435 bytes.extend_from_slice(&self.terms);
436 bytes.extend_from_slice(&self.blob);
437 Ok(bytes)
438 }
439
440 pub fn decode_region(bytes: &[u8]) -> Result<Self, SealedSegmentError> {
442 let mut cursor = RegionCursor::new(bytes);
443 if cursor.take_array::<4>()? != REGION_MAGIC {
444 return Err(SealedSegmentError::Geometry("magic did not match"));
445 }
446 if cursor.read_u16()? != REGION_VERSION {
447 return Err(SealedSegmentError::Geometry("version is not readable"));
448 }
449 let postings_per_block = cursor.read_u16()?;
450 if postings_per_block == 0 {
451 return Err(SealedSegmentError::Geometry("postings per block is zero"));
452 }
453 let row_count = cursor.read_u32()?;
454 let span_count = usize::try_from(cursor.read_u32()?)
455 .map_err(|_| SealedSegmentError::Geometry("span count exceeds usize"))?;
456 let field_count = usize::try_from(cursor.read_u32()?)
457 .map_err(|_| SealedSegmentError::Geometry("field count exceeds usize"))?;
458 if cursor.read_u32()? != 0 {
459 return Err(SealedSegmentError::Geometry(
460 "header reserved field is nonzero",
461 ));
462 }
463 let terms_len = usize::try_from(cursor.read_u64()?)
464 .map_err(|_| SealedSegmentError::Geometry("term bytes exceed usize"))?;
465 let blob_len = usize::try_from(cursor.read_u64()?)
466 .map_err(|_| SealedSegmentError::Geometry("postings bytes exceed usize"))?;
467 let mut spans = Vec::with_capacity(span_count);
468 for _ in 0..span_count {
469 let field = FieldId(cursor.read_u16()?);
470 if cursor.read_u16()? != 0 {
471 return Err(SealedSegmentError::Geometry(
472 "span reserved field is nonzero",
473 ));
474 }
475 let span = ListSpan {
476 field,
477 term_start: cursor.read_u32()?,
478 term_len: cursor.read_u32()?,
479 meta_start: cursor.read_u32()?,
480 docids_start: cursor.read_u32()?,
481 tfs_start: cursor.read_u32()?,
482 block_count: cursor.read_u32()?,
483 doc_freq: cursor.read_u32()?,
484 overall_max_tf: cursor.read_u32()?,
485 overall_min_len: cursor.read_u16()?,
486 union_doc_freq: {
487 if cursor.read_u16()? != 0 {
488 return Err(SealedSegmentError::Geometry(
489 "span trailing reserved field is nonzero",
490 ));
491 }
492 cursor.read_u32()?
493 },
494 field_count: cursor.read_u32()?,
495 };
496 spans.push(span);
497 }
498 let mut lengths = Vec::with_capacity(field_count);
499 let rows = usize::try_from(row_count)
500 .map_err(|_| SealedSegmentError::Geometry("row count exceeds usize"))?;
501 let mut previous_field: Option<FieldId> = None;
502 for _ in 0..field_count {
503 let field = FieldId(cursor.read_u16()?);
504 if previous_field.is_some_and(|previous| previous >= field) {
505 return Err(SealedSegmentError::Geometry(
506 "fields are not strictly ascending",
507 ));
508 }
509 previous_field = Some(field);
510 if cursor.read_u16()? != 0 || cursor.read_u32()? != row_count {
511 return Err(SealedSegmentError::Geometry("field header is invalid"));
512 }
513 let mut field_lengths = Vec::with_capacity(rows);
514 for _ in 0..rows {
515 field_lengths.push(cursor.read_u32()?);
516 }
517 lengths.push(FieldLengths {
518 field,
519 lengths: field_lengths,
520 });
521 }
522 let terms = cursor.take(terms_len)?.to_vec();
523 let blob = cursor.take(blob_len)?.to_vec();
524 cursor.finish()?;
525 validate_decoded_spans(&spans, &terms, &blob, postings_per_block, row_count)?;
526 let mut total_lengths = vec![0_u32; rows];
527 for entry in &lengths {
528 for (slot, length) in entry.lengths.iter().enumerate() {
529 if let Some(total) = total_lengths.get_mut(slot) {
530 *total = total.saturating_add(*length);
531 }
532 }
533 }
534 Ok(Self {
535 terms,
536 spans,
537 blob,
538 lengths,
539 total_lengths,
540 row_count,
541 postings_per_block,
542 })
543 }
544
545 pub fn retain_rows(&self, survivors: &[usize]) -> Result<Self, SealedSegmentError> {
550 let old_rows = usize::try_from(self.row_count)
551 .map_err(|_| SealedSegmentError::Geometry("row count exceeds usize"))?;
552 let mut mapping = vec![None; old_rows];
553 let mut previous = None;
554 for (new_row, old_row) in survivors.iter().copied().enumerate() {
555 if old_row >= old_rows || previous.is_some_and(|value| value >= old_row) {
556 return Err(SealedSegmentError::Geometry(
557 "survivor rows are not ascending and in bounds",
558 ));
559 }
560 previous = Some(old_row);
561 let new_row = u32::try_from(new_row)
562 .map_err(|_| SealedSegmentError::Geometry("survivor row exceeds u32"))?;
563 let target = mapping
564 .get_mut(old_row)
565 .ok_or(SealedSegmentError::Geometry(
566 "survivor row is out of bounds",
567 ))?;
568 *target = Some(new_row);
569 }
570 let row_count = u32::try_from(survivors.len())
571 .map_err(|_| SealedSegmentError::Geometry("survivor count exceeds u32"))?;
572 let mut retained = Self {
573 terms: Vec::new(),
574 spans: Vec::new(),
575 blob: Vec::new(),
576 lengths: Vec::with_capacity(self.lengths.len()),
577 total_lengths: vec![0; survivors.len()],
578 row_count,
579 postings_per_block: self.postings_per_block,
580 };
581 for entry in &self.lengths {
582 let mut lengths = Vec::with_capacity(survivors.len());
583 for row in survivors {
584 lengths.push(entry.lengths.get(*row).copied().ok_or(
585 SealedSegmentError::Geometry("survivor field length is missing"),
586 )?);
587 }
588 retained.lengths.push(FieldLengths {
589 field: entry.field,
590 lengths,
591 });
592 }
593 for entry in &retained.lengths {
594 for (row, length) in entry.lengths.iter().copied().enumerate() {
595 let total =
596 retained
597 .total_lengths
598 .get_mut(row)
599 .ok_or(SealedSegmentError::Geometry(
600 "retained total-length row is missing",
601 ))?;
602 *total = total.saturating_add(length);
603 }
604 }
605
606 let mut lists = Vec::new();
607 for (index, span) in self.spans.iter().enumerate() {
608 let encoded = self.list_bytes(index)?;
609 let decoded = PostingsReader::open(encoded)?.decode_all()?;
610 let mut list = PostingList::new();
611 for posting in decoded.postings() {
612 let old_row = usize::try_from(posting.docid)
613 .map_err(|_| SealedSegmentError::Geometry("posting row exceeds usize"))?;
614 let Some(new_row) = mapping.get(old_row).copied().flatten() else {
615 continue;
616 };
617 let mut posting = posting.clone();
618 posting.docid = new_row;
619 list.push(posting)?;
620 }
621 if list.is_empty() {
622 continue;
623 }
624 let field_lengths = retained
625 .lengths
626 .iter()
627 .find(|entry| entry.field == span.field)
628 .map_or(&[][..], |entry| entry.lengths.as_slice());
629 let impacts = block_impacts(&list, retained.postings_per_block, field_lengths);
630 let encoded = encode_v2(&list, retained.postings_per_block, &[], &impacts)?;
631 let reader = PostingsReader::open(encoded.as_bytes())?;
632 let base = u32::try_from(retained.blob.len())
633 .map_err(|_| SealedSegmentError::Geometry("postings blob exceeds u32"))?;
634 let block_count = u32::try_from(reader.blocks().len())
635 .map_err(|_| SealedSegmentError::Geometry("posting blocks exceed u32"))?;
636 let meta_start =
637 base.checked_add(HEADER_LEN_U32)
638 .ok_or(SealedSegmentError::Geometry(
639 "posting metadata offset overflow",
640 ))?;
641 let streams = meta_start
642 .checked_add(block_count.saturating_mul(META_LEN_U32))
643 .ok_or(SealedSegmentError::Geometry(
644 "posting stream offset overflow",
645 ))?;
646 let docid_bytes = reader
647 .blocks()
648 .iter()
649 .map(|meta| stream_end(meta.docids_offset, meta.count, meta.docid_bits))
650 .max()
651 .unwrap_or(0);
652 let mut overall_max_tf = 0_u32;
653 let mut overall_min_len = u16::MAX;
654 for meta in reader.blocks() {
655 let impact = BlockImpact::from_meta(meta);
656 overall_max_tf = overall_max_tf.max(impact.max_tf);
657 overall_min_len = overall_min_len.min(impact.min_len);
658 }
659 let term = self.term_of(index);
660 let term_start = u32::try_from(retained.terms.len())
661 .map_err(|_| SealedSegmentError::Geometry("term bytes exceed u32"))?;
662 let term_len = u32::try_from(term.len())
663 .map_err(|_| SealedSegmentError::Geometry("term length exceeds u32"))?;
664 retained.terms.extend_from_slice(term);
665 retained.spans.push(ListSpan {
666 field: span.field,
667 term_start,
668 term_len,
669 meta_start,
670 docids_start: streams,
671 tfs_start: streams.saturating_add(docid_bytes),
672 block_count,
673 doc_freq: reader.document_frequency(),
674 overall_max_tf,
675 overall_min_len,
676 union_doc_freq: 0,
677 field_count: 0,
678 });
679 retained.blob.extend_from_slice(encoded.as_bytes());
680 lists.push(list);
681 }
682 stamp_retained_unions(&mut retained, &lists)?;
683 Ok(retained)
684 }
685
686 pub(crate) fn position_reader_bytes_controlled<E>(
687 &self,
688 term: &[u8],
689 field: FieldId,
690 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
691 ) -> Result<Option<usize>, E> {
692 work.check_now()?;
693 let (start, end) = self.term_span_range(term);
694 let Some(spans) = self.spans.get(start..end) else {
695 return Ok(None);
696 };
697 for span in spans {
698 work.step()?;
699 if span.field == field {
700 return Ok(
701 (span.block_count as usize).checked_mul(std::mem::size_of::<BlockMeta>())
702 );
703 }
704 }
705 Ok(Some(0))
706 }
707
708 pub(crate) fn position_reader_controlled<E: From<SealedSegmentError> + From<PostingsError>>(
709 &self,
710 term: &[u8],
711 field: FieldId,
712 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
713 ) -> Result<Option<PostingsReader<'_>>, E> {
714 work.check_now()?;
715 let (start, end) = self.term_span_range(term);
716 for index in start..end {
717 work.step()?;
718 if self
719 .spans
720 .get(index)
721 .is_some_and(|span| span.field == field)
722 {
723 return Ok(Some(PostingsReader::open_controlled(
724 self.list_bytes(index)?,
725 work,
726 )?));
727 }
728 }
729 Ok(None)
730 }
731
732 fn list_bytes(&self, index: usize) -> Result<&[u8], SealedSegmentError> {
733 let span = self
734 .spans
735 .get(index)
736 .ok_or(SealedSegmentError::Geometry("posting span is missing"))?;
737 let start = usize::try_from(span.meta_start)
738 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?
739 .checked_sub(HEADER_LEN)
740 .ok_or(SealedSegmentError::Geometry(
741 "posting header start underflows",
742 ))?;
743 let end = self
744 .spans
745 .get(index.saturating_add(1))
746 .map(|next| {
747 usize::try_from(next.meta_start)
748 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?
749 .checked_sub(HEADER_LEN)
750 .ok_or(SealedSegmentError::Geometry(
751 "posting header start underflows",
752 ))
753 })
754 .transpose()?
755 .unwrap_or(self.blob.len());
756 self.blob
757 .get(start..end)
758 .ok_or(SealedSegmentError::Geometry(
759 "posting list is out of bounds",
760 ))
761 }
762
763 fn term_of(&self, index: usize) -> &[u8] {
765 let Some(span) = self.spans.get(index) else {
766 return &[];
767 };
768 let start = usize::try_from(span.term_start).unwrap_or(usize::MAX);
769 let len = usize::try_from(span.term_len).unwrap_or(0);
770 self.terms
771 .get(start..start.saturating_add(len))
772 .unwrap_or(&[])
773 }
774
775 #[must_use]
777 pub const fn row_count(&self) -> u32 {
778 self.row_count
779 }
780
781 #[must_use]
783 pub const fn is_empty(&self) -> bool {
784 self.row_count == 0
785 }
786
787 #[must_use]
789 pub const fn postings_per_block(&self) -> u16 {
790 self.postings_per_block
791 }
792
793 #[must_use]
795 pub fn field_lengths(&self, field: FieldId) -> Option<&[u32]> {
796 self.lengths
797 .iter()
798 .find(|entry| entry.field == field)
799 .map(|entry| entry.lengths.as_slice())
800 }
801
802 pub(crate) fn field_lengths_controlled<E>(
803 &self,
804 field: FieldId,
805 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
806 ) -> Result<Option<&[u32]>, E> {
807 for entry in &self.lengths {
808 work.step()?;
809 if entry.field == field {
810 return Ok(Some(entry.lengths.as_slice()));
811 }
812 }
813 Ok(None)
814 }
815
816 #[must_use]
820 pub fn total_lengths(&self) -> &[u32] {
821 &self.total_lengths
822 }
823
824 pub fn fields(&self) -> impl Iterator<Item = FieldId> + '_ {
826 self.lengths.iter().map(|entry| entry.field)
827 }
828
829 #[must_use]
831 pub fn field_length(&self, row: u32, field: FieldId) -> u32 {
832 self.field_lengths(field)
833 .and_then(|lengths| {
834 usize::try_from(row)
835 .ok()
836 .and_then(|slot| lengths.get(slot).copied())
837 })
838 .unwrap_or(0)
839 }
840
841 #[must_use]
843 pub fn document_length(&self, row: u32) -> u32 {
844 let Ok(slot) = usize::try_from(row) else {
845 return 0;
846 };
847 self.lengths
848 .iter()
849 .filter_map(|entry| entry.lengths.get(slot).copied())
850 .fold(0_u32, u32::saturating_add)
851 }
852
853 #[must_use]
855 pub fn total_tokens(&self) -> u64 {
856 self.lengths
857 .iter()
858 .flat_map(|entry| entry.lengths.iter())
859 .map(|length| u64::from(*length))
860 .sum()
861 }
862
863 fn term_span_range(&self, term: &[u8]) -> (usize, usize) {
865 let start = self.spans.partition_point(|span| {
866 let bytes_start = usize::try_from(span.term_start).unwrap_or(usize::MAX);
867 let len = usize::try_from(span.term_len).unwrap_or(0);
868 self.terms
869 .get(bytes_start..bytes_start.saturating_add(len))
870 .unwrap_or(&[])
871 < term
872 });
873 let mut end = start;
874 while end < self.spans.len() && self.term_of(end) == term {
875 end = end.saturating_add(1);
876 }
877 (start, end)
878 }
879
880 #[must_use]
886 pub fn document_frequency(&self, term: &[u8], fields: &[FieldId]) -> u32 {
887 match self
888 .document_frequency_controlled(term, fields, || Ok::<(), std::convert::Infallible>(()))
889 {
890 Ok(frequency) => frequency,
891 Err(never) => match never {},
892 }
893 }
894
895 pub(crate) fn document_frequency_controlled<E>(
896 &self,
897 term: &[u8],
898 fields: &[FieldId],
899 check: impl FnMut() -> Result<(), E>,
900 ) -> Result<u32, E> {
901 let mut work = super::control::WorkCheck::new(check);
902 work.check_now()?;
903 let (start, end) = self.term_span_range(term);
904 let mut matched: Option<usize> = None;
905 let mut count = 0_usize;
906 for index in start..end {
907 work.step()?;
908 let Some(span) = self.spans.get(index) else {
909 continue;
910 };
911 if fields.contains(&span.field) {
912 matched = matched.or(Some(index));
913 count = count.saturating_add(1);
914 }
915 }
916 let frequency = match count {
917 0 => 0,
918 1 => matched
919 .and_then(|index| self.spans.get(index))
920 .map_or(0, |span| span.doc_freq),
921 _ => {
922 let total = u32::try_from(count).unwrap_or(u32::MAX);
923 let all_fields = self
924 .spans
925 .get(start)
926 .is_some_and(|span| span.field_count == total);
927 if all_fields {
928 self.spans.get(start).map_or(0, |span| span.union_doc_freq)
930 } else {
931 let mut cursors: Vec<ListCursor<'_>> = Vec::new();
933 for index in start..end {
934 work.step()?;
935 let Some(span) = self.spans.get(index) else {
936 continue;
937 };
938 if fields.contains(&span.field)
939 && let Some(cursor) = self.cursor_at(index, 1_000)
940 {
941 cursors.push(cursor);
942 }
943 }
944 union_count_controlled(&mut cursors, &mut work)?
945 }
946 }
947 };
948 work.check_now()?;
949 Ok(frequency)
950 }
951
952 pub(crate) fn single_posting_live_frequency(
957 &self,
958 term: &[u8],
959 fields: &[FieldId],
960 live_rows: &DocBitmap,
961 ) -> Option<u32> {
962 let (start, end) = self.term_span_range(term);
963 let mut matched = None;
964 for index in start..end {
965 let span = self.spans.get(index)?;
966 if !fields.contains(&span.field) || span.doc_freq == 0 {
967 continue;
968 }
969 if span.doc_freq > 1 || matched.is_some() {
970 return None;
971 }
972 matched = Some(index);
973 }
974 let Some(index) = matched else {
975 return Some(0);
976 };
977 let mut cursor = self.cursor_at(index, 1_000)?;
978 cursor.reset();
979 let row = cursor.current()?;
980 #[cfg(any(test, feature = "test-seams"))]
981 super::preparation_observer::live_df_walk(
982 1,
983 usize::try_from(cursor.blocks_decoded).unwrap_or(usize::MAX),
984 );
985 Some(u32::from(live_rows.contains(row)))
986 }
987
988 pub(crate) fn live_document_frequency(
989 &self,
990 term: &[u8],
991 fields: &[FieldId],
992 live_rows: &DocBitmap,
993 ) -> u32 {
994 match self.live_document_frequency_controlled(term, fields, live_rows, || {
995 Ok::<(), std::convert::Infallible>(())
996 }) {
997 Ok(frequency) => frequency,
998 Err(never) => match never {},
999 }
1000 }
1001
1002 pub(crate) fn live_document_frequency_controlled<E>(
1003 &self,
1004 term: &[u8],
1005 fields: &[FieldId],
1006 live_rows: &DocBitmap,
1007 check: impl FnMut() -> Result<(), E>,
1008 ) -> Result<u32, E> {
1009 let mut work = super::control::WorkCheck::new(check);
1010 work.check_now()?;
1011 let weights = FieldWeights::flat(fields);
1012 let Some(mut stream) = TermStream::open_controlled(self, term, &weights, &mut work)? else {
1013 return Ok(0);
1014 };
1015 let mut count = 0_u32;
1016 #[cfg(any(test, feature = "test-seams"))]
1017 let mut visited = 0_usize;
1018 let result = (|| {
1019 while let Some(row) = stream.current_row() {
1020 work.step()?;
1021 #[cfg(any(test, feature = "test-seams"))]
1022 {
1023 visited = visited.saturating_add(1);
1024 }
1025 if live_rows.contains(row) {
1026 count = count.saturating_add(1);
1027 }
1028 stream.advance_controlled(&mut work)?;
1029 }
1030 work.check_now()?;
1031 Ok(count)
1032 })();
1033 #[cfg(any(test, feature = "test-seams"))]
1034 super::preparation_observer::live_df_walk(
1035 visited,
1036 usize::try_from(stream.blocks_decoded()).unwrap_or(usize::MAX),
1037 );
1038 result
1039 }
1040
1041 fn cursor_at(&self, index: usize, weight: u64) -> Option<ListCursor<'_>> {
1043 let span = self.spans.get(index)?;
1044 let per_block = usize::from(self.postings_per_block);
1045 Some(ListCursor {
1046 meta: self.slice_from(span.meta_start),
1047 docids: self.slice_from(span.docids_start),
1048 tfs: self.slice_from(span.tfs_start),
1049 block_count: usize::try_from(span.block_count).unwrap_or(0),
1050 per_block,
1051 weight,
1052 block: NO_BLOCK,
1053 slot: 0,
1054 count: 0,
1055 impact: BlockImpact::default(),
1056 tfs_block: NO_BLOCK,
1057 tf_offset: 0,
1058 tf_bits: 0,
1059 overall: BlockImpact {
1060 max_tf: span.overall_max_tf,
1061 min_len: span.overall_min_len,
1062 },
1063 decoded_docids: vec![0; per_block],
1064 decoded_tfs: vec![0; per_block],
1065 blocks_decoded: 0,
1066 blocks_skipped: 0,
1067 tf_blocks_decoded: 0,
1068 })
1069 }
1070
1071 fn slice_from(&self, offset: u32) -> &[u8] {
1073 usize::try_from(offset)
1074 .ok()
1075 .and_then(|start| self.blob.get(start..))
1076 .unwrap_or(&[])
1077 }
1078}
1079
1080fn stamp_retained_unions(
1081 segment: &mut SealedSegment,
1082 lists: &[PostingList],
1083) -> Result<(), SealedSegmentError> {
1084 if segment.spans.len() != lists.len() {
1085 return Err(SealedSegmentError::Geometry(
1086 "retained posting list count differs from spans",
1087 ));
1088 }
1089 let mut start = 0_usize;
1090 while start < segment.spans.len() {
1091 let term = segment.term_of(start).to_vec();
1092 let mut end = start.saturating_add(1);
1093 while end < segment.spans.len() && segment.term_of(end) == term {
1094 end = end.saturating_add(1);
1095 }
1096 let mut documents = std::collections::BTreeSet::new();
1097 for list in lists.get(start..end).ok_or(SealedSegmentError::Geometry(
1098 "retained posting-list group is missing",
1099 ))? {
1100 documents.extend(list.postings().iter().map(|posting| posting.docid));
1101 }
1102 let union = u32::try_from(documents.len())
1103 .map_err(|_| SealedSegmentError::Geometry("union document count exceeds u32"))?;
1104 let field_count = u32::try_from(end.saturating_sub(start))
1105 .map_err(|_| SealedSegmentError::Geometry("field count exceeds u32"))?;
1106 let spans = segment
1107 .spans
1108 .get_mut(start..end)
1109 .ok_or(SealedSegmentError::Geometry(
1110 "retained span group is missing",
1111 ))?;
1112 for span in spans {
1113 span.union_doc_freq = union;
1114 span.field_count = field_count;
1115 }
1116 start = end;
1117 }
1118 Ok(())
1119}
1120
1121fn validate_decoded_spans(
1122 spans: &[ListSpan],
1123 terms: &[u8],
1124 blob: &[u8],
1125 postings_per_block: u16,
1126 row_count: u32,
1127) -> Result<(), SealedSegmentError> {
1128 let mut previous_key: Option<(&[u8], FieldId)> = None;
1129 let mut decoded_lists = Vec::with_capacity(spans.len());
1130 for (index, span) in spans.iter().enumerate() {
1131 let term_start = usize::try_from(span.term_start)
1132 .map_err(|_| SealedSegmentError::Geometry("term start exceeds usize"))?;
1133 let term_len = usize::try_from(span.term_len)
1134 .map_err(|_| SealedSegmentError::Geometry("term length exceeds usize"))?;
1135 let term_end = term_start
1136 .checked_add(term_len)
1137 .ok_or(SealedSegmentError::Geometry("term span overflows"))?;
1138 let term = terms
1139 .get(term_start..term_end)
1140 .ok_or(SealedSegmentError::Geometry("term span is out of bounds"))?;
1141 if previous_key.is_some_and(|(previous_term, previous_field)| {
1142 previous_term > term || (previous_term == term && previous_field >= span.field)
1143 }) {
1144 return Err(SealedSegmentError::Geometry(
1145 "term spans are not strictly ordered",
1146 ));
1147 }
1148 previous_key = Some((term, span.field));
1149 let meta_start = usize::try_from(span.meta_start)
1150 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?;
1151 let list_start = meta_start
1152 .checked_sub(HEADER_LEN)
1153 .ok_or(SealedSegmentError::Geometry(
1154 "posting header start underflows",
1155 ))?;
1156 let list_end = spans
1157 .get(index.saturating_add(1))
1158 .map(|next| {
1159 usize::try_from(next.meta_start)
1160 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?
1161 .checked_sub(HEADER_LEN)
1162 .ok_or(SealedSegmentError::Geometry(
1163 "posting header start underflows",
1164 ))
1165 })
1166 .transpose()?
1167 .unwrap_or(blob.len());
1168 let list = blob
1169 .get(list_start..list_end)
1170 .ok_or(SealedSegmentError::Geometry(
1171 "posting list is out of bounds",
1172 ))?;
1173 let reader = PostingsReader::open(list)?;
1174 let block_count = u32::try_from(reader.blocks().len())
1175 .map_err(|_| SealedSegmentError::Geometry("posting blocks exceed u32"))?;
1176 let expected_meta = u32::try_from(list_start)
1177 .map_err(|_| SealedSegmentError::Geometry("posting start exceeds u32"))?
1178 .checked_add(HEADER_LEN_U32)
1179 .ok_or(SealedSegmentError::Geometry(
1180 "posting metadata offset overflow",
1181 ))?;
1182 let expected_docids = expected_meta
1183 .checked_add(block_count.saturating_mul(META_LEN_U32))
1184 .ok_or(SealedSegmentError::Geometry(
1185 "posting stream offset overflow",
1186 ))?;
1187 let docid_bytes = reader
1188 .blocks()
1189 .iter()
1190 .map(|meta| stream_end(meta.docids_offset, meta.count, meta.docid_bits))
1191 .max()
1192 .unwrap_or(0);
1193 let expected_tfs = expected_docids
1194 .checked_add(docid_bytes)
1195 .ok_or(SealedSegmentError::Geometry("posting tf offset overflow"))?;
1196 let mut overall_max_tf = 0_u32;
1197 let mut overall_min_len = u16::MAX;
1198 for block in reader.blocks() {
1199 let impact = BlockImpact::from_meta(block);
1200 overall_max_tf = overall_max_tf.max(impact.max_tf);
1201 overall_min_len = overall_min_len.min(impact.min_len);
1202 }
1203 if reader.postings_per_block() != postings_per_block
1204 || reader.document_frequency() != span.doc_freq
1205 || block_count != span.block_count
1206 || span.meta_start != expected_meta
1207 || span.docids_start != expected_docids
1208 || span.tfs_start != expected_tfs
1209 || span.overall_max_tf != overall_max_tf
1210 || span.overall_min_len != overall_min_len
1211 || span.field_count == 0
1212 || span.union_doc_freq > row_count
1213 {
1214 return Err(SealedSegmentError::Geometry(
1215 "posting-list summary does not match span",
1216 ));
1217 }
1218 if reader
1219 .blocks()
1220 .iter()
1221 .any(|block| block.last_docid >= row_count)
1222 {
1223 return Err(SealedSegmentError::Geometry(
1224 "posting docid exceeds row count",
1225 ));
1226 }
1227 decoded_lists.push(reader.decode_all()?);
1228 }
1229 validate_union_summaries(spans, terms, &decoded_lists)?;
1230 Ok(())
1231}
1232
1233fn validate_union_summaries(
1234 spans: &[ListSpan],
1235 terms: &[u8],
1236 lists: &[PostingList],
1237) -> Result<(), SealedSegmentError> {
1238 if spans.len() != lists.len() {
1239 return Err(SealedSegmentError::Geometry(
1240 "decoded posting list count differs from spans",
1241 ));
1242 }
1243 let term_of = |span: &ListSpan| -> Result<&[u8], SealedSegmentError> {
1244 let start = usize::try_from(span.term_start)
1245 .map_err(|_| SealedSegmentError::Geometry("term start exceeds usize"))?;
1246 let length = usize::try_from(span.term_len)
1247 .map_err(|_| SealedSegmentError::Geometry("term length exceeds usize"))?;
1248 let end = start
1249 .checked_add(length)
1250 .ok_or(SealedSegmentError::Geometry("term span overflows"))?;
1251 terms
1252 .get(start..end)
1253 .ok_or(SealedSegmentError::Geometry("term span is out of bounds"))
1254 };
1255 let mut start = 0_usize;
1256 while start < spans.len() {
1257 let first = spans
1258 .get(start)
1259 .ok_or(SealedSegmentError::Geometry("term group start is missing"))?;
1260 let term = term_of(first)?;
1261 let mut end = start.saturating_add(1);
1262 while let Some(span) = spans.get(end) {
1263 if term_of(span)? != term {
1264 break;
1265 }
1266 end = end.saturating_add(1);
1267 }
1268 let field_count = u32::try_from(end.saturating_sub(start))
1269 .map_err(|_| SealedSegmentError::Geometry("term field count exceeds u32"))?;
1270 let mut documents = std::collections::BTreeSet::new();
1271 for list in lists.get(start..end).ok_or(SealedSegmentError::Geometry(
1272 "term posting-list group is missing",
1273 ))? {
1274 documents.extend(list.postings().iter().map(|posting| posting.docid));
1275 }
1276 let union = u32::try_from(documents.len())
1277 .map_err(|_| SealedSegmentError::Geometry("union document count exceeds u32"))?;
1278 if spans
1279 .get(start..end)
1280 .ok_or(SealedSegmentError::Geometry("term span group is missing"))?
1281 .iter()
1282 .any(|span| span.field_count != field_count || span.union_doc_freq != union)
1283 {
1284 return Err(SealedSegmentError::Geometry(
1285 "term union summary does not match posting lists",
1286 ));
1287 }
1288 start = end;
1289 }
1290 Ok(())
1291}
1292
1293struct RegionCursor<'a> {
1294 bytes: &'a [u8],
1295 offset: usize,
1296}
1297
1298impl<'a> RegionCursor<'a> {
1299 const fn new(bytes: &'a [u8]) -> Self {
1300 Self { bytes, offset: 0 }
1301 }
1302
1303 fn take(&mut self, length: usize) -> Result<&'a [u8], SealedSegmentError> {
1304 let end = self
1305 .offset
1306 .checked_add(length)
1307 .ok_or(SealedSegmentError::Geometry("cursor offset overflow"))?;
1308 let value = self
1309 .bytes
1310 .get(self.offset..end)
1311 .ok_or(SealedSegmentError::Geometry("region is truncated"))?;
1312 self.offset = end;
1313 Ok(value)
1314 }
1315
1316 fn take_array<const N: usize>(&mut self) -> Result<[u8; N], SealedSegmentError> {
1317 self.take(N)?
1318 .try_into()
1319 .map_err(|_| SealedSegmentError::Geometry("region is truncated"))
1320 }
1321
1322 fn read_u16(&mut self) -> Result<u16, SealedSegmentError> {
1323 self.take_array().map(u16::from_le_bytes)
1324 }
1325
1326 fn read_u32(&mut self) -> Result<u32, SealedSegmentError> {
1327 self.take_array().map(u32::from_le_bytes)
1328 }
1329
1330 fn read_u64(&mut self) -> Result<u64, SealedSegmentError> {
1331 self.take_array().map(u64::from_le_bytes)
1332 }
1333
1334 fn finish(self) -> Result<(), SealedSegmentError> {
1335 if self.offset == self.bytes.len() {
1336 Ok(())
1337 } else {
1338 Err(SealedSegmentError::Geometry("region has trailing bytes"))
1339 }
1340 }
1341}
1342
1343const HEADER_LEN_U32: u32 = HEADER_LEN as u32;
1345
1346const META_LEN_U32: u32 = BLOCK_META_LEN as u32;
1348
1349const NO_BLOCK: usize = usize::MAX;
1351
1352fn stream_end(offset: u32, count: u16, bits: u8) -> u32 {
1354 let packed = u32::from(count).saturating_mul(u32::from(bits)).div_ceil(8);
1355 offset.saturating_add(packed)
1356}
1357
1358fn finish_union_group<E>(
1365 spans: &mut [ListSpan],
1366 start: usize,
1367 lists: &[&PostingList],
1368 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1369) -> Result<(), E> {
1370 let end = spans.len();
1371 if end <= start {
1372 return Ok(());
1373 }
1374 let count = u32::try_from(end.saturating_sub(start)).unwrap_or(u32::MAX);
1375 let union = if count == 1 {
1376 spans.get(start).map_or(0, |span| span.doc_freq)
1377 } else {
1378 union_postings_controlled(lists, work)?
1379 };
1380 if let Some(range) = spans.get_mut(start..end) {
1381 for span in range {
1382 work.step()?;
1383 span.union_doc_freq = union;
1384 span.field_count = count;
1385 }
1386 }
1387 Ok(())
1388}
1389
1390fn union_postings_controlled<E>(
1392 lists: &[&PostingList],
1393 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1394) -> Result<u32, E> {
1395 work.check_now()?;
1396 let mut positions = vec![0_usize; lists.len()];
1397 let mut distinct = 0_u32;
1398 loop {
1399 let mut lowest: Option<u32> = None;
1400 for (slot, list) in lists.iter().enumerate() {
1401 work.step()?;
1402 #[cfg(test)]
1403 SEAL_UNION_PROBES.with(|value| value.set(value.get() + 1));
1404 let position = positions.get(slot).copied().unwrap_or(usize::MAX);
1405 if let Some(posting) = list.postings().get(position) {
1406 lowest = Some(lowest.map_or(posting.docid, |low| low.min(posting.docid)));
1407 }
1408 }
1409 let Some(low) = lowest else {
1410 break;
1411 };
1412 distinct = distinct.saturating_add(1);
1413 for (slot, list) in lists.iter().enumerate() {
1414 work.step()?;
1415 let Some(position) = positions.get_mut(slot) else {
1416 continue;
1417 };
1418 if list
1419 .postings()
1420 .get(*position)
1421 .is_some_and(|posting| posting.docid == low)
1422 {
1423 *position = position.saturating_add(1);
1424 }
1425 }
1426 }
1427 work.check_now()?;
1428 Ok(distinct)
1429}
1430
1431#[cfg(test)]
1432std::thread_local! {
1433 pub(crate) static SEAL_LENGTH_PROBES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1434 static SEAL_TOTAL_PROBES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1435 static SEAL_UNION_PROBES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1436}
1437
1438fn union_count_controlled<E>(
1440 cursors: &mut [ListCursor<'_>],
1441 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1442) -> Result<u32, E> {
1443 for cursor in cursors.iter_mut() {
1444 work.step()?;
1445 cursor.reset();
1446 }
1447 let mut distinct = 0_u32;
1448 loop {
1449 let mut first = None;
1450 for cursor in cursors.iter() {
1451 work.step()?;
1452 if let Some(row) = cursor.current() {
1453 first = Some(first.map_or(row, |held: u32| held.min(row)));
1454 }
1455 }
1456 let Some(row) = first else {
1457 break;
1458 };
1459 distinct = distinct.saturating_add(1);
1460 for cursor in cursors.iter_mut() {
1461 work.step()?;
1462 if cursor.current() == Some(row) {
1463 cursor.advance();
1464 }
1465 }
1466 }
1467 Ok(distinct)
1468}
1469
1470#[derive(Clone, Debug)]
1476pub struct ListCursor<'segment> {
1477 meta: &'segment [u8],
1478 docids: &'segment [u8],
1479 tfs: &'segment [u8],
1480 block_count: usize,
1481 per_block: usize,
1482 weight: u64,
1483 block: usize,
1484 slot: usize,
1485 count: usize,
1486 impact: BlockImpact,
1489 tfs_block: usize,
1496 tf_offset: u32,
1498 tf_bits: u8,
1500 overall: BlockImpact,
1502 decoded_docids: Vec<u32>,
1503 decoded_tfs: Vec<u32>,
1504 blocks_decoded: u64,
1505 blocks_skipped: u64,
1506 tf_blocks_decoded: u64,
1507}
1508
1509impl ListCursor<'_> {
1510 fn meta_at(&self, index: usize) -> Option<BlockMeta> {
1512 let start = index.checked_mul(BLOCK_META_LEN)?;
1513 let row = self.meta.get(start..start.checked_add(BLOCK_META_LEN)?)?;
1514 BlockMeta::read(row).ok()
1515 }
1516
1517 fn last_docid(&self, index: usize) -> Option<u32> {
1519 let start = index.checked_mul(BLOCK_META_LEN)?;
1520 let bytes = self.meta.get(start..start.checked_add(4)?)?;
1521 let mut buffer = [0_u8; 4];
1522 buffer.copy_from_slice(bytes);
1523 Some(u32::from_le_bytes(buffer))
1524 }
1525
1526 fn impact_at(&self, index: usize) -> BlockImpact {
1528 self.meta_at(index)
1529 .map_or_else(BlockImpact::default, |meta| BlockImpact::from_meta(&meta))
1530 }
1531
1532 fn load(&mut self, index: usize) -> bool {
1537 if self.block == index {
1538 return true;
1539 }
1540 let Some(meta) = self.meta_at(index) else {
1541 return false;
1542 };
1543 let count = usize::from(meta.count).min(self.per_block);
1544 let docid_start = usize::try_from(meta.docids_offset).unwrap_or(usize::MAX);
1545 let Some(packed) = self.docids.get(docid_start..) else {
1546 return false;
1547 };
1548 let base = if index == 0 {
1551 0
1552 } else {
1553 self.last_docid(index.saturating_sub(1)).unwrap_or(0)
1554 };
1555 let Some(target) = self.decoded_docids.get_mut(..count) else {
1556 return false;
1557 };
1558 if unpack(packed, meta.docid_bits, count, target).is_none() {
1559 return false;
1560 }
1561 prefix_sum(target, base);
1562
1563 self.tf_offset = meta.tfs_offset;
1566 self.tf_bits = meta.tf_bits;
1567 self.block = index;
1568 self.count = count;
1569 self.impact = BlockImpact::from_meta(&meta);
1570 self.blocks_decoded = self.blocks_decoded.saturating_add(1);
1571 true
1572 }
1573
1574 pub fn reset(&mut self) {
1580 self.slot = 0;
1581 if self.block == 0 {
1582 return;
1583 }
1584 self.block = NO_BLOCK;
1585 self.count = 0;
1586 if self.block_count > 0 {
1587 self.load(0);
1588 }
1589 }
1590
1591 #[must_use]
1593 pub fn current(&self) -> Option<u32> {
1594 if self.block == NO_BLOCK {
1595 return None;
1596 }
1597 self.decoded_docids.get(self.slot).copied()
1598 }
1599
1600 fn ensure_tfs(&mut self) -> bool {
1602 if self.tfs_block == self.block {
1603 return true;
1604 }
1605 let start = usize::try_from(self.tf_offset).unwrap_or(usize::MAX);
1606 let Some(packed) = self.tfs.get(start..) else {
1607 return false;
1608 };
1609 let Some(target) = self.decoded_tfs.get_mut(..self.count) else {
1610 return false;
1611 };
1612 if unpack(packed, self.tf_bits, self.count, target).is_none() {
1613 return false;
1614 }
1615 self.tfs_block = self.block;
1616 self.tf_blocks_decoded = self.tf_blocks_decoded.saturating_add(1);
1617 true
1618 }
1619
1620 #[must_use]
1624 pub fn current_tf(&mut self) -> Option<u32> {
1625 if self.block == NO_BLOCK || !self.ensure_tfs() {
1626 return None;
1627 }
1628 self.decoded_tfs.get(self.slot).copied()
1629 }
1630
1631 #[must_use]
1636 pub fn current_weighted_tf(&mut self) -> Option<u64> {
1637 self.current_tf()
1638 .map(|tf| u64::from(tf).saturating_mul(self.weight))
1639 }
1640
1641 pub fn advance(&mut self) {
1643 if self.block == NO_BLOCK {
1644 return;
1645 }
1646 self.slot = self.slot.saturating_add(1);
1647 if self.slot < self.count {
1648 return;
1649 }
1650 let next = self.block.saturating_add(1);
1651 if next >= self.block_count || !self.load(next) {
1652 self.block = NO_BLOCK;
1653 return;
1654 }
1655 self.slot = 0;
1656 }
1657
1658 pub fn seek(&mut self, row: u32) {
1665 if self.block == NO_BLOCK {
1666 return;
1667 }
1668 if self.current().is_some_and(|current| current >= row) {
1669 return;
1670 }
1671 if self.last_docid(self.block).is_some_and(|last| last < row) {
1672 let mut low = self.block.saturating_add(1);
1673 let mut high = self.block_count;
1674 while low < high {
1675 let middle = low.saturating_add(high.saturating_sub(low) / 2);
1676 if self.last_docid(middle).is_some_and(|last| last < row) {
1677 low = middle.saturating_add(1);
1678 } else {
1679 high = middle;
1680 }
1681 }
1682 let hopped = low.saturating_sub(self.block.saturating_add(1));
1683 self.blocks_skipped = self
1684 .blocks_skipped
1685 .saturating_add(u64::try_from(hopped).unwrap_or(u64::MAX));
1686 if low >= self.block_count || !self.load(low) {
1687 self.block = NO_BLOCK;
1688 return;
1689 }
1690 self.slot = 0;
1691 }
1692 loop {
1693 if self.block == NO_BLOCK {
1694 return;
1695 }
1696 let found = self
1699 .decoded_docids
1700 .get(self.slot..self.count)
1701 .map_or(0, |within| within.partition_point(|&docid| docid < row));
1702 let landed = self.slot.saturating_add(found);
1703 if landed < self.count {
1704 self.slot = landed;
1705 return;
1706 }
1707 self.slot = self.count.saturating_sub(1);
1710 self.advance();
1711 }
1712 }
1713
1714 #[must_use]
1716 pub const fn exhausted(&self) -> bool {
1717 self.block == NO_BLOCK
1718 }
1719
1720 #[must_use]
1724 pub const fn current_impact(&self) -> BlockImpact {
1725 if self.block == NO_BLOCK {
1726 return BlockImpact {
1727 max_tf: 0,
1728 min_len: 0,
1729 };
1730 }
1731 self.impact
1732 }
1733
1734 fn block_for(&self, row: u32) -> Option<usize> {
1740 if self.block == NO_BLOCK {
1741 return None;
1742 }
1743 if self.current().is_some_and(|current| current > row) {
1744 return None;
1747 }
1748 if self.last_docid(self.block).is_some_and(|last| last >= row) {
1749 return Some(self.block);
1750 }
1751 let mut low = self.block.saturating_add(1);
1752 let mut high = self.block_count;
1753 while low < high {
1754 let middle = low.saturating_add(high.saturating_sub(low) / 2);
1755 if self.last_docid(middle).is_some_and(|last| last < row) {
1756 low = middle.saturating_add(1);
1757 } else {
1758 high = middle;
1759 }
1760 }
1761 if low >= self.block_count {
1762 return None;
1763 }
1764 Some(low)
1765 }
1766
1767 #[must_use]
1775 pub fn impact_for(&self, row: u32) -> Option<BlockImpact> {
1776 let index = self.block_for(row)?;
1777 if index == self.block {
1778 return Some(self.current_impact());
1779 }
1780 Some(self.impact_at(index))
1781 }
1782
1783 fn impact_horizon_for(&self, row: u32) -> Option<u32> {
1794 if self.block == NO_BLOCK {
1795 return None;
1796 }
1797 if let Some(current) = self.current()
1798 && current > row
1799 {
1800 return Some(current.saturating_sub(1));
1801 }
1802 self.block_for(row).and_then(|index| self.last_docid(index))
1803 }
1804
1805 #[must_use]
1809 pub const fn overall_impact(&self) -> BlockImpact {
1810 self.overall
1811 }
1812}
1813
1814#[derive(Clone, Debug)]
1820pub struct TermStream<'segment> {
1821 runs: Vec<ListCursor<'segment>>,
1822 weights: Vec<u64>,
1823 unit_single: bool,
1827 heads_always_valid: bool,
1832 head: Option<u32>,
1833 head_tf: Option<u32>,
1835}
1836
1837#[cfg(test)]
1838thread_local! {
1839 static ZERO_TF_ROWS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
1840}
1841
1842impl<'segment> TermStream<'segment> {
1843 pub(crate) fn open_controlled<E>(
1844 segment: &'segment SealedSegment,
1845 term: &[u8],
1846 weights: &FieldWeights,
1847 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1848 ) -> Result<Option<Self>, E> {
1849 Self::open_with_capacity_controlled(segment, term, weights, 0, work)
1850 }
1851
1852 #[must_use]
1856 pub fn open(
1857 segment: &'segment SealedSegment,
1858 term: &[u8],
1859 weights: &FieldWeights,
1860 ) -> Option<Self> {
1861 Self::open_with_capacity(segment, term, weights, 0)
1862 }
1863
1864 pub(crate) fn allocation_bytes(
1867 segment: &SealedSegment,
1868 term: &[u8],
1869 weights: &FieldWeights,
1870 ) -> Option<usize> {
1871 Self::run_count(segment, term, weights).checked_mul(
1872 std::mem::size_of::<ListCursor<'_>>()
1873 .checked_add(std::mem::size_of::<u64>())?
1874 .checked_add(
1875 usize::from(segment.postings_per_block)
1876 .checked_mul(2 * std::mem::size_of::<u32>())?,
1877 )?,
1878 )
1879 }
1880
1881 fn run_count(segment: &SealedSegment, term: &[u8], weights: &FieldWeights) -> usize {
1882 let (start, end) = segment.term_span_range(term);
1883 segment
1884 .spans
1885 .get(start..end)
1886 .unwrap_or_default()
1887 .iter()
1888 .filter(|span| weights.weight(span.field) != 0)
1889 .count()
1890 }
1891
1892 #[cfg(test)]
1893 pub(crate) fn open_sized(
1894 segment: &'segment SealedSegment,
1895 term: &[u8],
1896 weights: &FieldWeights,
1897 ) -> Option<Self> {
1898 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
1899 match Self::open_sized_controlled(segment, term, weights, &mut work) {
1900 Ok(stream) => stream,
1901 Err(never) => match never {},
1902 }
1903 }
1904
1905 pub(crate) fn open_sized_controlled<E>(
1906 segment: &'segment SealedSegment,
1907 term: &[u8],
1908 weights: &FieldWeights,
1909 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1910 ) -> Result<Option<Self>, E> {
1911 let (start, end) = segment.term_span_range(term);
1912 let mut capacity = 0;
1913 for span in segment.spans.get(start..end).unwrap_or_default() {
1914 work.step()?;
1915 capacity += usize::from(weights.weight(span.field) != 0);
1916 }
1917 Self::open_with_capacity_controlled(segment, term, weights, capacity, work)
1918 }
1919
1920 fn open_with_capacity(
1921 segment: &'segment SealedSegment,
1922 term: &[u8],
1923 weights: &FieldWeights,
1924 capacity: usize,
1925 ) -> Option<Self> {
1926 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
1927 match Self::open_with_capacity_controlled(segment, term, weights, capacity, &mut work) {
1928 Ok(stream) => stream,
1929 Err(never) => match never {},
1930 }
1931 }
1932
1933 fn open_with_capacity_controlled<E>(
1934 segment: &'segment SealedSegment,
1935 term: &[u8],
1936 weights: &FieldWeights,
1937 capacity: usize,
1938 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1939 ) -> Result<Option<Self>, E> {
1940 work.step()?;
1941 let (start, end) = segment.term_span_range(term);
1942 let mut runs: Vec<ListCursor<'segment>> = Vec::with_capacity(capacity);
1943 let mut scales: Vec<u64> = Vec::with_capacity(capacity);
1944 for index in start..end {
1945 work.step()?;
1946 let Some(span) = segment.spans.get(index) else {
1947 return Ok(None);
1948 };
1949 let weight = u64::from(weights.weight(span.field));
1950 if weight == 0 {
1951 continue;
1952 }
1953 let Some(cursor) = segment.cursor_at(index, weight) else {
1954 return Ok(None);
1955 };
1956 runs.push(cursor);
1957 scales.push(weight);
1958 }
1959 if runs.is_empty() {
1960 return Ok(None);
1961 }
1962 let unit_single = runs.len() == 1 && scales.first().copied() == Some(1_000);
1963 let heads_always_valid = scales.iter().all(|weight| *weight >= 1_000);
1964 let mut stream = Self {
1965 runs,
1966 weights: scales,
1967 unit_single,
1968 heads_always_valid,
1969 head: None,
1970 head_tf: None,
1971 };
1972 stream.reset_controlled(work)?;
1973 Ok(Some(stream))
1974 }
1975
1976 pub fn reset(&mut self) {
1978 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
1979 match self.reset_controlled(&mut work) {
1980 Ok(()) => {}
1981 Err(never) => match never {},
1982 }
1983 }
1984
1985 pub(crate) fn reset_controlled<E>(
1986 &mut self,
1987 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1988 ) -> Result<(), E> {
1989 for run in &mut self.runs {
1990 work.step()?;
1991 run.reset();
1992 }
1993 self.refresh_controlled(work)
1994 }
1995
1996 fn refresh_controlled<E>(
2003 &mut self,
2004 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2005 ) -> Result<(), E> {
2006 self.head_tf = None;
2007 if self.heads_always_valid {
2008 self.head = if self.unit_single {
2009 self.runs.first().and_then(ListCursor::current)
2010 } else {
2011 self.first_row_controlled(work)?
2012 };
2013 return Ok(());
2014 }
2015 loop {
2016 let Some(row) = self.first_row_controlled(work)? else {
2017 self.head = None;
2018 return Ok(());
2019 };
2020 let mut total = 0_u64;
2021 for run in &mut self.runs {
2022 work.step()?;
2023 if run.current() == Some(row) {
2024 total = total.saturating_add(run.current_weighted_tf().unwrap_or(0));
2025 }
2026 }
2027 let tf = u32::try_from(total / 1_000).unwrap_or(u32::MAX);
2028 if tf > 0 {
2029 self.head = Some(row);
2030 self.head_tf = Some(tf);
2031 return Ok(());
2032 }
2033 #[cfg(test)]
2034 ZERO_TF_ROWS.with(|count| count.set(count.get() + 1));
2035 self.step_controlled(row, work)?;
2036 }
2037 }
2038
2039 fn first_row_controlled<E>(
2040 &self,
2041 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2042 ) -> Result<Option<u32>, E> {
2043 let mut first = None;
2044 for run in &self.runs {
2045 work.step()?;
2046 if let Some(row) = run.current() {
2047 first = Some(first.map_or(row, |held: u32| held.min(row)));
2048 }
2049 }
2050 Ok(first)
2051 }
2052
2053 fn step_controlled<E>(
2055 &mut self,
2056 row: u32,
2057 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2058 ) -> Result<(), E> {
2059 for run in &mut self.runs {
2060 work.step()?;
2061 if run.current() == Some(row) {
2062 run.advance();
2063 }
2064 }
2065 Ok(())
2066 }
2067
2068 #[must_use]
2070 pub const fn current_row(&self) -> Option<u32> {
2071 self.head
2072 }
2073
2074 #[must_use]
2077 pub fn current_tf(&mut self) -> Option<u32> {
2078 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2079 match self.current_tf_controlled(&mut work) {
2080 Ok(tf) => tf,
2081 Err(never) => match never {},
2082 }
2083 }
2084
2085 pub(crate) fn current_tf_controlled<E>(
2086 &mut self,
2087 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2088 ) -> Result<Option<u32>, E> {
2089 let Some(row) = self.head else {
2090 return Ok(None);
2091 };
2092 if let Some(tf) = self.head_tf {
2093 return Ok(Some(tf));
2094 }
2095 let mut total = 0_u64;
2096 for run in &mut self.runs {
2097 work.step()?;
2098 if run.current() == Some(row) {
2099 total = total.saturating_add(run.current_weighted_tf().unwrap_or(0));
2100 }
2101 }
2102 let tf = u32::try_from(total / 1_000).unwrap_or(u32::MAX);
2103 self.head_tf = Some(tf);
2104 Ok(Some(tf))
2105 }
2106
2107 pub fn advance(&mut self) {
2109 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2110 match self.advance_controlled(&mut work) {
2111 Ok(()) => {}
2112 Err(never) => match never {},
2113 }
2114 }
2115
2116 pub(crate) fn advance_controlled<E>(
2117 &mut self,
2118 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2119 ) -> Result<(), E> {
2120 let Some(row) = self.head else {
2121 return Ok(());
2122 };
2123 self.step_controlled(row, work)?;
2124 self.refresh_controlled(work)
2125 }
2126
2127 pub fn seek(&mut self, row: u32) {
2129 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2130 match self.seek_controlled(row, &mut work) {
2131 Ok(()) => {}
2132 Err(never) => match never {},
2133 }
2134 }
2135
2136 pub(crate) fn seek_controlled<E>(
2137 &mut self,
2138 row: u32,
2139 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2140 ) -> Result<(), E> {
2141 for run in &mut self.runs {
2142 work.step()?;
2143 run.seek(row);
2144 }
2145 self.refresh_controlled(work)
2146 }
2147
2148 #[must_use]
2150 pub const fn exhausted(&self) -> bool {
2151 self.head.is_none()
2152 }
2153
2154 #[must_use]
2156 pub fn blocks_decoded(&self) -> u64 {
2157 self.runs
2158 .iter()
2159 .map(|run| run.blocks_decoded)
2160 .fold(0_u64, u64::saturating_add)
2161 }
2162
2163 #[must_use]
2165 pub fn blocks_skipped(&self) -> u64 {
2166 self.runs
2167 .iter()
2168 .map(|run| run.blocks_skipped)
2169 .fold(0_u64, u64::saturating_add)
2170 }
2171
2172 #[must_use]
2177 pub fn tf_blocks_decoded(&self) -> u64 {
2178 self.runs
2179 .iter()
2180 .map(|run| run.tf_blocks_decoded)
2181 .fold(0_u64, u64::saturating_add)
2182 }
2183
2184 #[must_use]
2188 pub fn posting_count(&self) -> usize {
2189 self.runs
2190 .iter()
2191 .map(|run| {
2192 run.block_count
2193 .saturating_sub(1)
2194 .saturating_mul(run.per_block)
2195 .saturating_add(
2196 run.meta_at(run.block_count.saturating_sub(1))
2197 .map_or(0, |meta| usize::from(meta.count)),
2198 )
2199 })
2200 .fold(0_usize, usize::saturating_add)
2201 }
2202
2203 pub(crate) fn block_bound_controlled<E>(
2204 &self,
2205 scorer: &TermScorer,
2206 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2207 ) -> Result<f64, E> {
2208 self.bound_from_controlled(
2209 scorer,
2210 |run| (!run.exhausted()).then(|| run.current_impact()),
2211 work,
2212 )
2213 }
2214
2215 pub(crate) fn upper_bound_controlled<E>(
2216 &self,
2217 scorer: &TermScorer,
2218 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2219 ) -> Result<f64, E> {
2220 self.bound_from_controlled(
2221 scorer,
2222 |run| (!run.exhausted()).then(|| run.overall_impact()),
2223 work,
2224 )
2225 }
2226
2227 pub(crate) fn bound_for_controlled<E>(
2228 &self,
2229 row: u32,
2230 scorer: &TermScorer,
2231 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2232 ) -> Result<f64, E> {
2233 self.bound_from_controlled(scorer, |run| run.impact_for(row), work)
2234 }
2235
2236 pub(crate) fn bound_horizon_for_controlled<E>(
2237 &self,
2238 row: u32,
2239 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2240 ) -> Result<Option<u32>, E> {
2241 let mut edge = None;
2242 for run in &self.runs {
2243 work.step()?;
2244 if let Some(found) = run.impact_horizon_for(row) {
2245 edge = Some(edge.map_or(found, |held: u32| held.min(found)));
2246 }
2247 }
2248 Ok(edge)
2249 }
2250
2251 #[must_use]
2256 pub fn block_bound(&self, scorer: &TermScorer) -> f64 {
2257 self.bound_from(scorer, |run| {
2258 (!run.exhausted()).then(|| run.current_impact())
2259 })
2260 }
2261
2262 #[must_use]
2271 pub fn bound_for(&self, row: u32, scorer: &TermScorer) -> f64 {
2272 self.bound_from(scorer, |run| run.impact_for(row))
2273 }
2274
2275 #[must_use]
2294 pub fn bound_horizon_for(&self, row: u32) -> Option<u32> {
2295 self.runs
2296 .iter()
2297 .filter_map(|run| run.impact_horizon_for(row))
2298 .min()
2299 }
2300
2301 #[must_use]
2303 pub fn upper_bound(&self, scorer: &TermScorer) -> f64 {
2304 self.bound_from(scorer, |run| {
2305 (!run.exhausted()).then(|| run.overall_impact())
2306 })
2307 }
2308
2309 fn bound_from(
2314 &self,
2315 scorer: &TermScorer,
2316 pick: impl Fn(&ListCursor<'segment>) -> Option<BlockImpact>,
2317 ) -> f64 {
2318 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2319 match self.bound_from_controlled(scorer, pick, &mut work) {
2320 Ok(bound) => bound,
2321 Err(never) => match never {},
2322 }
2323 }
2324
2325 fn bound_from_controlled<E>(
2326 &self,
2327 scorer: &TermScorer,
2328 pick: impl Fn(&ListCursor<'segment>) -> Option<BlockImpact>,
2329 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2330 ) -> Result<f64, E> {
2331 let mut max_tf = 0_u64;
2332 let mut min_len = u64::MAX;
2333 for (run, weight) in self.runs.iter().zip(self.weights.iter()) {
2334 work.step()?;
2335 let Some(impact) = pick(run) else {
2336 continue;
2337 };
2338 max_tf = max_tf.saturating_add(
2341 u64::from(impact.max_tf)
2342 .saturating_mul(*weight)
2343 .div_ceil(1_000),
2344 );
2345 min_len = min_len.min(u64::from(impact.min_len).saturating_mul(*weight) / 1_000);
2346 }
2347 if min_len == u64::MAX {
2348 return Ok(0.0);
2349 }
2350 Ok(scorer.score(
2351 super::bm25::Tf(u32::try_from(max_tf).unwrap_or(u32::MAX)),
2352 super::bm25::DocLen(u32::try_from(min_len).unwrap_or(u32::MAX)),
2353 ))
2354 }
2355}
2356
2357#[cfg(test)]
2358#[allow(
2359 clippy::expect_used,
2360 clippy::indexing_slicing,
2361 clippy::panic,
2362 clippy::unwrap_used
2363)]
2364mod tests {
2365 use super::*;
2366 use crate::fts::index::{DEFAULT_FIELD, Document};
2367 use crate::fts::postings::{Posting, PostingList};
2368 use crate::fts::tokenizer::{Analyzer, Profile};
2369
2370 fn analyzer() -> Analyzer {
2371 Analyzer::new(Profile::Code.config()).expect("valid config")
2372 }
2373
2374 fn segment_of(texts: &[String]) -> SegmentIndex {
2375 let analyzer = analyzer();
2376 let mut segment = SegmentIndex::new();
2377 for text in texts {
2378 segment
2379 .push_document(&analyzer, &Document::with_text(text))
2380 .expect("indexable");
2381 }
2382 segment
2383 }
2384
2385 #[test]
2386 fn astra_18_sealing_cancels_inside_row_preparation() {
2387 let source = segment_of(&vec!["alpha beta".to_owned(); 4_096]);
2388 let reference = SealedSegment::seal(&source)
2389 .expect("reference")
2390 .encode_region()
2391 .expect("bytes");
2392 let mut observed = Vec::new();
2393 for (phase, probes) in [
2394 ("length copy", &SEAL_LENGTH_PROBES),
2395 ("total lengths", &SEAL_TOTAL_PROBES),
2396 ] {
2397 probes.with(|value| value.set(0));
2398 let mut work = super::super::control::WorkCheck::new(|| {
2399 if probes.with(std::cell::Cell::get) >= 64 {
2400 Err(PostingsError::ZeroBlockSize)
2401 } else {
2402 Ok(())
2403 }
2404 });
2405 let result = SealedSegment::seal_controlled(&source, &mut work);
2406 let count = probes.with(std::cell::Cell::get);
2407 println!("{phase}: canceled={}, rows={count}", result.is_err());
2408 observed.push((
2409 phase,
2410 matches!(result, Err(PostingsError::ZeroBlockSize)),
2411 count,
2412 ));
2413 }
2414 let mut work = super::super::control::WorkCheck::new(|| Ok::<(), PostingsError>(()));
2415 assert_eq!(
2416 SealedSegment::seal_controlled(&source, &mut work)
2417 .expect("clean")
2418 .encode_region()
2419 .expect("clean bytes"),
2420 reference
2421 );
2422 for (phase, canceled, count) in observed {
2423 assert!(canceled, "{phase}");
2424 assert!((64..=128).contains(&count), "{phase}: {count}");
2425 }
2426 }
2427
2428 #[test]
2429 fn astra_18_sealing_union_cancels_inside_merge() {
2430 let mut even = PostingList::new();
2431 let mut all = PostingList::new();
2432 for docid in 0..4_096 {
2433 let posting = Posting {
2434 docid,
2435 tf: 1,
2436 positions: vec![0],
2437 };
2438 all.push(posting.clone()).expect("all");
2439 if docid % 2 == 0 {
2440 even.push(posting).expect("even");
2441 }
2442 }
2443 SEAL_UNION_PROBES.with(|value| value.set(0));
2444 let mut work = super::super::control::WorkCheck::new(|| {
2445 if SEAL_UNION_PROBES.with(std::cell::Cell::get) >= 64 {
2446 Err("cancelled")
2447 } else {
2448 Ok(())
2449 }
2450 });
2451 let result = union_postings_controlled(&[&even, &all], &mut work);
2452 let probes = SEAL_UNION_PROBES.with(std::cell::Cell::get);
2453 println!("union probes={probes}");
2454 assert_eq!(result, Err("cancelled"));
2455 assert!((64..=128).contains(&probes));
2456 let mut work = super::super::control::WorkCheck::new(|| Ok::<(), ()>(()));
2457 assert_eq!(
2458 union_postings_controlled(&[&even, &all], &mut work),
2459 Ok(4_096)
2460 );
2461 }
2462
2463 #[test]
2464 fn astra_18_field_subset_df_cancellation_preserves_union() {
2465 let mut source = SegmentIndex::new();
2466 let analyzer = analyzer();
2467 for row in 0..4_096 {
2468 let mut document = Document::new();
2469 document.set(FieldId(0), if row % 2 == 0 { "alpha" } else { "beta" });
2470 document.set(FieldId(1), if row % 3 == 0 { "alpha" } else { "beta" });
2471 document.set(FieldId(2), "alpha");
2472 source.push_document(&analyzer, &document).unwrap();
2473 }
2474 let segment = SealedSegment::seal(&source).unwrap();
2475 let fields = [FieldId(0), FieldId(1)];
2476 let expected = (0..4_096)
2477 .filter(|row| row % 2 == 0 || row % 3 == 0)
2478 .count() as u32;
2479 assert_eq!(segment.document_frequency(b"alpha", &fields), expected);
2480 assert_eq!(
2481 segment.document_frequency_controlled(b"alpha", &fields, || Ok::<(), ()>(())),
2482 Ok(expected)
2483 );
2484 let mut checks = 0;
2485 assert_eq!(
2486 segment.document_frequency_controlled(b"alpha", &fields, || {
2487 checks += 1;
2488 if checks == 3 {
2489 Err("cancelled")
2490 } else {
2491 Ok(())
2492 }
2493 }),
2494 Err("cancelled")
2495 );
2496 assert_eq!(checks, 3);
2497 for (fields, expected) in [
2499 (vec![FieldId(0)], 2_048),
2500 (vec![FieldId(0), FieldId(1), FieldId(2)], 4_096),
2501 ] {
2502 assert_eq!(
2503 segment.document_frequency_controlled(b"alpha", &fields, || Ok::<(), ()>(())),
2504 Ok(expected)
2505 );
2506 }
2507 }
2508
2509 #[test]
2510 fn astra_18_fractional_stream_cancels_in_hidden_zero_tf_walk() {
2511 let analyzer = analyzer();
2512 let fields = FieldWeights::new(&[(DEFAULT_FIELD, 500)]);
2513 let mut observed = Vec::new();
2514 for action in ["open", "reset", "seek", "advance"] {
2515 let mut source = SegmentIndex::new();
2516 for row in 0..4_096 {
2517 let text = if row == 4_095 || (matches!(action, "advance" | "seek") && row == 0) {
2518 "alpha alpha"
2519 } else {
2520 "alpha"
2521 };
2522 source
2523 .push_document(&analyzer, &Document::with_text(text))
2524 .unwrap();
2525 }
2526 let segment = SealedSegment::seal(&source).unwrap();
2527 let mut stream = TermStream::open(&segment, b"alpha", &fields).unwrap();
2528 assert_eq!(
2529 stream.current_row(),
2530 Some(if matches!(action, "advance" | "seek") {
2531 0
2532 } else {
2533 4_095
2534 })
2535 );
2536 ZERO_TF_ROWS.with(|count| count.set(0));
2537 let mut checks = 0;
2538 let mut work = super::super::control::WorkCheck::new(|| {
2539 checks += 1;
2540 if checks == 3 {
2541 Err("cancelled")
2542 } else {
2543 Ok(())
2544 }
2545 });
2546 let result = match action {
2547 "open" => TermStream::open_controlled(&segment, b"alpha", &fields, &mut work)
2548 .map(|stream| stream.and_then(|stream| stream.current_row())),
2549 "reset" => stream
2550 .reset_controlled(&mut work)
2551 .map(|()| stream.current_row()),
2552 "seek" => stream
2553 .seek_controlled(1, &mut work)
2554 .map(|()| stream.current_row()),
2555 "advance" => stream
2556 .advance_controlled(&mut work)
2557 .map(|()| stream.current_row()),
2558 _ => unreachable!(),
2559 };
2560 let rows = ZERO_TF_ROWS.with(std::cell::Cell::get);
2561 println!("{action}: result={result:?}, checks={checks}, zero_tf_rows={rows}");
2562 observed.push((action, result, rows));
2563 }
2564 for (action, result, rows) in observed {
2567 assert_eq!(result, Err("cancelled"), "{action}");
2568 assert!(
2569 rows > 0 && rows <= 128,
2570 "{action} walked {rows} zero-TF rows"
2571 );
2572 }
2573 }
2574
2575 #[test]
2576 fn a_sealed_cursor_reproduces_every_posting_of_the_active_segment() {
2577 let texts: Vec<String> = (0..500)
2580 .map(|index| {
2581 if index % 3 == 0 {
2582 format!("alpha filler{index}")
2583 } else {
2584 format!("beta filler{index}")
2585 }
2586 })
2587 .collect();
2588 let active = segment_of(&texts);
2589 let sealed = SealedSegment::seal(&active).expect("seals");
2590
2591 for term in [b"alpha".as_slice(), b"beta".as_slice()] {
2592 let list = active.posting_list(term, DEFAULT_FIELD).expect("present");
2593 let mut cursor = sealed
2594 .cursor_at(sealed.term_span_range(term).0, 1_000)
2595 .expect("span");
2596 cursor.reset();
2597 for posting in list.postings() {
2598 assert_eq!(cursor.current(), Some(posting.docid));
2599 assert_eq!(
2600 cursor.current_weighted_tf(),
2601 Some(u64::from(posting.tf) * 1_000)
2602 );
2603 cursor.advance();
2604 }
2605 assert!(cursor.exhausted(), "the cursor must end with the list");
2606 }
2607 }
2608
2609 #[test]
2610 fn a_landed_block_pays_its_tf_decode_only_when_scored() {
2611 let texts: Vec<String> = (0..400).map(|index| format!("alpha d{index}")).collect();
2615 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
2616 let mut stream = TermStream::open(
2617 &sealed,
2618 b"alpha",
2619 &crate::fts::search::FieldWeights::flat(&[DEFAULT_FIELD]),
2620 )
2621 .expect("stream");
2622 assert_eq!(
2623 stream.tf_blocks_decoded(),
2624 0,
2625 "opening must not touch tf bytes"
2626 );
2627 stream.seek(200);
2628 stream.seek(383);
2629 assert!(stream.blocks_decoded() >= 3, "three landings were paid");
2630 assert_eq!(
2631 stream.tf_blocks_decoded(),
2632 0,
2633 "positioning must not touch tf bytes"
2634 );
2635 assert_eq!(stream.current_row(), Some(383));
2636 assert_eq!(stream.current_tf(), Some(1));
2637 assert_eq!(
2638 stream.tf_blocks_decoded(),
2639 1,
2640 "scoring pays exactly the one block it reads"
2641 );
2642 }
2643
2644 #[test]
2645 fn seeking_jumps_whole_blocks_without_decoding_them() {
2646 let texts: Vec<String> = (0..400).map(|index| format!("alpha d{index}")).collect();
2647 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
2648 let mut cursor = sealed
2649 .cursor_at(sealed.term_span_range(b"alpha").0, 1_000)
2650 .expect("span");
2651 cursor.reset();
2652 let decoded_before = cursor.blocks_decoded;
2653 cursor.seek(383);
2654 assert_eq!(cursor.current(), Some(383));
2655 assert!(
2656 cursor.blocks_skipped >= 4,
2657 "a 400-posting list at 64 per block must skip at least four"
2658 );
2659 assert_eq!(
2660 cursor.blocks_decoded,
2661 decoded_before + 1,
2662 "a seek must decode exactly the block it lands in"
2663 );
2664 }
2665
2666 #[test]
2667 fn seeking_past_the_end_exhausts_the_cursor() {
2668 let texts: Vec<String> = (0..100).map(|index| format!("alpha d{index}")).collect();
2669 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
2670 let mut cursor = sealed
2671 .cursor_at(sealed.term_span_range(b"alpha").0, 1_000)
2672 .expect("span");
2673 cursor.reset();
2674 cursor.seek(9_999);
2675 assert!(cursor.exhausted());
2676 assert_eq!(cursor.current(), None);
2677 assert_eq!(cursor.current_weighted_tf(), None);
2678 }
2679
2680 #[test]
2681 fn an_empty_segment_seals_and_answers_nothing() {
2682 let sealed = SealedSegment::seal(&SegmentIndex::new()).expect("seals");
2683 assert!(sealed.is_empty());
2684 assert_eq!(sealed.row_count(), 0);
2685 assert_eq!(sealed.total_tokens(), 0);
2686 assert_eq!(sealed.document_length(0), 0);
2687 assert_eq!(sealed.field_length(0, DEFAULT_FIELD), 0);
2688 assert!(sealed.field_lengths(DEFAULT_FIELD).is_none());
2689 assert_eq!(sealed.document_frequency(b"alpha", &[DEFAULT_FIELD]), 0);
2690 assert!(
2691 TermStream::open(&sealed, b"alpha", &FieldWeights::flat(&[DEFAULT_FIELD])).is_none()
2692 );
2693 }
2694
2695 #[test]
2696 fn a_zero_weight_field_contributes_no_run() {
2697 let texts: Vec<String> = (0..4).map(|index| format!("alpha d{index}")).collect();
2698 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
2699 let weights = FieldWeights::new(&[(DEFAULT_FIELD, 0)]);
2700 assert!(TermStream::open(&sealed, b"alpha", &weights).is_none());
2701 }
2702
2703 #[test]
2704 fn a_bound_over_one_field_is_the_stored_impact_pair() {
2705 let texts: Vec<String> = (0..300).map(|index| format!("alpha d{index}")).collect();
2706 let active = segment_of(&texts);
2707 let sealed = SealedSegment::seal(&active).expect("seals");
2708 let stats = crate::fts::bm25::CorpusStats::new(300, 600).expect("stats");
2709 let scorer = TermScorer::new(
2710 crate::fts::bm25::Df(300),
2711 &stats,
2712 crate::fts::bm25::Bm25Params::default(),
2713 );
2714 let stream = TermStream::open(&sealed, b"alpha", &FieldWeights::flat(&[DEFAULT_FIELD]))
2715 .expect("stream");
2716 let impact = stream.runs[0].current_impact();
2717 assert!((stream.block_bound(&scorer) - impact.bound(&scorer)).abs() < 1e-15);
2718 }
2719
2720 #[test]
2721 fn a_traversed_list_costs_a_couple_of_bytes_per_posting() {
2722 let texts: Vec<String> = (0..20_000)
2737 .map(|index| format!("alpha t{} filler{index}", index % 97))
2738 .collect();
2739 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
2740
2741 let mut long_postings = 0_usize;
2742 let mut long_bytes = 0_usize;
2743 let mut short_lists = 0_usize;
2744 for index in 0..sealed.spans.len() {
2745 let Some(span) = sealed.spans.get(index).copied() else {
2746 continue;
2747 };
2748 let cursor = sealed.cursor_at(index, 1_000).expect("span");
2749 if cursor.block_count <= 1 {
2750 short_lists = short_lists.saturating_add(1);
2751 continue;
2752 }
2753 long_postings =
2754 long_postings.saturating_add(usize::try_from(span.doc_freq).unwrap_or(0));
2755 let last = cursor.block_count.saturating_sub(1);
2760 let meta = cursor.meta_at(last).expect("last block");
2761 let tf_end = usize::try_from(meta.tfs_offset).unwrap_or(0)
2762 + (usize::from(meta.count) * usize::from(meta.tf_bits)).div_ceil(8);
2763 long_bytes = long_bytes
2764 .saturating_add(usize::try_from(span.tfs_start - span.docids_start).unwrap_or(0))
2765 .saturating_add(tf_end)
2766 .saturating_add(cursor.block_count.saturating_mul(BLOCK_META_LEN));
2767 }
2768
2769 let per_posting = long_bytes as f64 / long_postings as f64;
2770 println!(
2771 "SEALED DENSITY traversed_postings={long_postings} bytes={long_bytes} \
2772 per_posting={per_posting:.3} single_block_lists={short_lists}"
2773 );
2774 assert!(
2775 per_posting < 4.0,
2776 "a traversed list costs {per_posting:.3} bytes per posting; the \
2777 in-memory form cost a 32-byte stride plus an allocation, and a \
2778 regression here is a decode regression the wall clock would only \
2779 confirm later"
2780 );
2781 }
2782
2783 #[test]
2784 fn every_decoded_width_falls_in_the_narrow_kernel_path() {
2785 for (label, texts) in [
2794 (
2795 "zipf100k",
2796 (0..100_000)
2797 .map(|index| {
2798 let mut words: Vec<String> = Vec::new();
2799 for term in 0..12_usize {
2800 if index % (term + 1) == 0 {
2801 words.push(format!("t{term}"));
2802 }
2803 }
2804 if words.is_empty() {
2805 words.push(String::from("filler"));
2806 }
2807 words.join(" ")
2808 })
2809 .collect::<Vec<String>>(),
2810 ),
2811 (
2812 "textish",
2813 (0..50_000)
2814 .map(|index| format!("alpha t{} u{} filler{index}", index % 97, index % 1_009))
2815 .collect::<Vec<String>>(),
2816 ),
2817 ] {
2818 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
2819 let mut histogram = [0_u64; 33];
2820 let mut postings = 0_u64;
2821 let mut widest = 0_u8;
2822 for index in 0..sealed.spans.len() {
2823 let cursor = sealed.cursor_at(index, 1_000).expect("span");
2824 if cursor.block_count <= 1 {
2827 continue;
2828 }
2829 for block in 0..cursor.block_count {
2830 let meta = cursor.meta_at(block).expect("meta");
2831 postings = postings.saturating_add(u64::from(meta.count));
2832 for bits in [meta.docid_bits, meta.tf_bits] {
2833 widest = widest.max(bits);
2834 if let Some(slot) = histogram.get_mut(usize::from(bits)) {
2835 *slot = slot.saturating_add(u64::from(meta.count));
2836 }
2837 }
2838 }
2839 }
2840 let mut shape: Vec<String> = Vec::new();
2841 for (bits, count) in histogram.iter().enumerate() {
2842 if *count > 0 {
2843 shape.push(format!(
2844 "{bits}b={:.1}%",
2845 *count as f64 * 50.0 / postings as f64
2846 ));
2847 }
2848 }
2849 println!(
2850 "WIDTHS {label} postings={postings} widest={widest} {}",
2851 shape.join(" ")
2852 );
2853 assert!(postings > 0, "{label} produced no traversable list");
2854 assert!(
2855 widest <= crate::kernels::postings::NARROW_MAX_BITS,
2856 "{label} decodes at {widest} bits, outside the narrow kernel \
2857 path the histogram was used to justify. Re-measure the \
2858 histogram and re-decide which arm to tune; do not widen \
2859 this bound."
2860 );
2861 }
2862 }
2863
2864 #[test]
2865 fn a_list_whose_lengths_are_absent_still_seals() {
2866 let mut list = PostingList::new();
2870 list.push(Posting {
2871 docid: 0,
2872 tf: 2,
2873 positions: vec![0, 1],
2874 })
2875 .expect("ascending");
2876 let impacts = crate::fts::postings::block_impacts(&list, 64, &[]);
2877 assert_eq!(impacts[0].max_tf, 2);
2878 assert_eq!(impacts[0].min_len, 1);
2879 }
2880}
2881
2882#[cfg(test)]
2883#[allow(clippy::expect_used)]
2884mod astra_10_allocation_tests {
2885 use super::*;
2886
2887 #[test]
2888 fn astra_10_stream_reservation_matches_actual_owned_capacities() {
2889 use crate::fts::index::{Document, FieldId, SegmentIndex};
2890 use crate::fts::tokenizer::{Analyzer, Profile};
2891 let analyzer = Analyzer::new(Profile::Code.config()).expect("analyzer");
2892 let mut index = SegmentIndex::new();
2893 for _ in 0..130 {
2894 let mut doc = Document::new();
2895 doc.set(FieldId(0), "alpha beta");
2896 doc.set(FieldId(1), "alpha alpha");
2897 doc.set(FieldId(2), "beta");
2898 index.push_document(&analyzer, &doc).expect("document");
2899 }
2900 let segment = SealedSegment::seal(&index).expect("sealed");
2901 for weights in [
2902 FieldWeights::flat(&[FieldId(0)]),
2903 FieldWeights::new(&[(FieldId(0), 500), (FieldId(1), 1500), (FieldId(2), 0)]),
2904 ] {
2905 for term in [b"alpha".as_slice(), b"beta", b"absent"] {
2906 let planned =
2907 TermStream::allocation_bytes(&segment, term, &weights).expect("fits usize");
2908 let actual = TermStream::open_sized(&segment, term, &weights).map_or(0, |stream| {
2909 stream.runs.capacity() * std::mem::size_of::<ListCursor<'_>>()
2910 + stream.weights.capacity() * std::mem::size_of::<u64>()
2911 + stream
2912 .runs
2913 .iter()
2914 .map(|run| {
2915 (run.decoded_docids.capacity() + run.decoded_tfs.capacity())
2916 * std::mem::size_of::<u32>()
2917 })
2918 .sum::<usize>()
2919 });
2920 assert_eq!(
2921 planned, actual,
2922 "reserve the actual buffers, including absent and zero-weight fields"
2923 );
2924 }
2925 }
2926 }
2927}