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
156pub(crate) struct GraphSealed<'m, C> {
159 sealed: SealedSegment,
160 terms_charge: C,
161 spans_charge: C,
162 blob_charge: C,
163 outer_lengths_charge: C,
164 inner_lengths_charge: C,
165 totals_charge: C,
166 marker: std::marker::PhantomData<&'m ()>,
167}
168
169impl<C: super::control::CapacityCharge> GraphSealed<'_, C> {
170 pub(crate) const fn sealed(&self) -> &SealedSegment {
171 &self.sealed
172 }
173
174 pub(crate) fn owned_bytes(&self) -> usize {
175 self.terms_charge
176 .bytes()
177 .saturating_add(self.spans_charge.bytes())
178 .saturating_add(self.blob_charge.bytes())
179 .saturating_add(self.outer_lengths_charge.bytes())
180 .saturating_add(self.inner_lengths_charge.bytes())
181 .saturating_add(self.totals_charge.bytes())
182 }
183
184 #[cfg(test)]
185 pub(crate) fn actual_capacity_bytes(&self) -> usize {
186 self.sealed
187 .terms
188 .capacity()
189 .saturating_add(self.sealed.blob.capacity())
190 .saturating_add(
191 self.sealed
192 .spans
193 .capacity()
194 .saturating_mul(std::mem::size_of::<ListSpan>()),
195 )
196 .saturating_add(
197 self.sealed
198 .lengths
199 .capacity()
200 .saturating_mul(std::mem::size_of::<FieldLengths>()),
201 )
202 .saturating_add(
203 self.sealed
204 .lengths
205 .iter()
206 .map(|entry| {
207 entry
208 .lengths
209 .capacity()
210 .saturating_mul(std::mem::size_of::<u32>())
211 })
212 .fold(0_usize, usize::saturating_add),
213 )
214 .saturating_add(
215 self.sealed
216 .total_lengths
217 .capacity()
218 .saturating_mul(std::mem::size_of::<u32>()),
219 )
220 }
221}
222
223impl<'m, C> GraphSealed<'m, C> {
224 pub(crate) fn seal_policy<'a, P, I>(
225 row_lengths: &[u32],
226 postings: I,
227 policy: &mut P,
228 ) -> Result<Self, P::Error>
229 where
230 P: super::control::BuildPolicy<'m, Charge = C>,
231 P::Error: From<PostingsError> + From<SealedSegmentError>,
232 I: IntoIterator<Item = (&'a str, &'a [super::postings::OwnedPosting<'m, C>])>,
233 'm: 'a,
234 C: 'a,
235 {
236 use super::control::GuardedVec;
237 policy.checkpoint()?;
238 let row_count = u32::try_from(row_lengths.len())
239 .map_err(|_| P::Error::from(SealedSegmentError::Geometry("row count exceeds u32")))?;
240 let mut terms = GuardedVec::with_capacity(policy, 0)?;
241 let mut spans = GuardedVec::with_capacity(policy, 0)?;
242 let mut blob = GuardedVec::with_capacity(policy, 0)?;
243 let mut inner_lengths = GuardedVec::with_capacity(policy, row_lengths.len())?;
244 let mut totals = GuardedVec::with_capacity(policy, row_lengths.len())?;
245 for length in row_lengths {
246 policy.step(1)?;
247 if *length == 0 {
248 return Err(P::Error::from(SealedSegmentError::Geometry(
249 "graph lexical row length is zero",
250 )));
251 }
252 inner_lengths.push(*length, policy)?;
253 totals.push(*length, policy)?;
254 }
255 let mut outer_lengths =
256 GuardedVec::with_capacity(policy, usize::from(!row_lengths.is_empty()))?;
257
258 let mut previous_term: Option<&[u8]> = None;
259 for (term, list) in postings {
260 policy.step(1)?;
261 if term.is_empty() || list.is_empty() {
262 return Err(P::Error::from(SealedSegmentError::Geometry(
263 "graph posting term/list is empty",
264 )));
265 }
266 if let Some(previous) = previous_term
267 && compare_bytes_policy(previous, term.as_bytes(), policy)?
268 != std::cmp::Ordering::Less
269 {
270 return Err(P::Error::from(SealedSegmentError::Geometry(
271 "graph terms are not strictly ordered",
272 )));
273 }
274 previous_term = Some(term.as_bytes());
275 let encoded = super::postings::encode_graph_policy(list, row_lengths, policy)?;
276 let summary = super::postings::validate_graph_policy(
277 encoded.bytes.as_slice(),
278 row_count,
279 policy,
280 )?;
281 if encoded.summary != summary {
282 return Err(P::Error::from(SealedSegmentError::Geometry(
283 "encoded posting-list summary does not match validation",
284 )));
285 }
286 let base = u32::try_from(blob.len()).map_err(|_| {
287 P::Error::from(SealedSegmentError::Geometry("posting blob exceeds u32"))
288 })?;
289 let meta_start = base.checked_add(HEADER_LEN_U32).ok_or_else(|| {
290 P::Error::from(SealedSegmentError::Geometry(
291 "posting metadata offset overflow",
292 ))
293 })?;
294 let streams =
295 meta_start
296 .checked_add(summary.block_count.checked_mul(META_LEN_U32).ok_or_else(
297 || {
298 P::Error::from(SealedSegmentError::Geometry(
299 "posting stream offset overflow",
300 ))
301 },
302 )?)
303 .ok_or_else(|| {
304 P::Error::from(SealedSegmentError::Geometry(
305 "posting stream offset overflow",
306 ))
307 })?;
308 let term_start = u32::try_from(terms.len()).map_err(|_| {
309 P::Error::from(SealedSegmentError::Geometry("term bytes exceed u32"))
310 })?;
311 terms.extend_from_slice(term.as_bytes(), policy)?;
312 spans.push(
313 ListSpan {
314 field: super::index::DEFAULT_FIELD,
315 term_start,
316 term_len: u32::try_from(term.len()).map_err(|_| {
317 P::Error::from(SealedSegmentError::Geometry("term length exceeds u32"))
318 })?,
319 meta_start,
320 docids_start: streams,
321 tfs_start: streams.checked_add(summary.docid_bytes).ok_or_else(|| {
322 P::Error::from(SealedSegmentError::Geometry("posting tf offset overflow"))
323 })?,
324 block_count: summary.block_count,
325 doc_freq: summary.doc_freq,
326 overall_max_tf: summary.overall_max_tf,
327 overall_min_len: summary.overall_min_len,
328 union_doc_freq: summary.doc_freq,
329 field_count: 1,
330 },
331 policy,
332 )?;
333 blob.extend_from_slice(encoded.bytes.as_slice(), policy)?;
334 }
335 let (inner_values, inner_lengths_charge) = inner_lengths.into_parts();
336 if !row_lengths.is_empty() {
337 outer_lengths.push(
338 FieldLengths {
339 field: super::index::DEFAULT_FIELD,
340 lengths: inner_values,
341 },
342 policy,
343 )?;
344 }
345 let (terms, terms_charge) = terms.into_parts();
346 let (spans, spans_charge) = spans.into_parts();
347 let (blob, blob_charge) = blob.into_parts();
348 let (lengths, outer_lengths_charge) = outer_lengths.into_parts();
349 let (total_lengths, totals_charge) = totals.into_parts();
350 let sealed = SealedSegment {
351 terms,
352 spans,
353 blob,
354 lengths,
355 total_lengths,
356 row_count,
357 postings_per_block: DEFAULT_POSTINGS_PER_BLOCK,
358 };
359 policy.checkpoint()?;
360 Ok(Self {
361 sealed,
362 terms_charge,
363 spans_charge,
364 blob_charge,
365 outer_lengths_charge,
366 inner_lengths_charge,
367 totals_charge,
368 marker: std::marker::PhantomData,
369 })
370 }
371
372 pub(crate) fn encode_region_policy<P>(
373 &self,
374 policy: &mut P,
375 ) -> Result<super::control::GuardedVec<'m, u8, C>, P::Error>
376 where
377 P: super::control::BuildPolicy<'m, Charge = C>,
378 P::Error: From<SealedSegmentError>,
379 {
380 use super::control::GuardedVec;
381 let span_count = u32::try_from(self.sealed.spans.len())
382 .map_err(|_| P::Error::from(SealedSegmentError::Geometry("span count exceeds u32")))?;
383 let field_count = u32::try_from(self.sealed.lengths.len())
384 .map_err(|_| P::Error::from(SealedSegmentError::Geometry("field count exceeds u32")))?;
385 let spans_len = self
386 .sealed
387 .spans
388 .len()
389 .checked_mul(REGION_SPAN_LEN)
390 .ok_or_else(|| P::Error::from(SealedSegmentError::Geometry("span bytes overflow")))?;
391 let fields_len = self
392 .sealed
393 .lengths
394 .len()
395 .checked_mul(REGION_FIELD_HEADER_LEN)
396 .and_then(|value| value.checked_add(self.sealed.total_lengths.len().checked_mul(4)?))
397 .ok_or_else(|| P::Error::from(SealedSegmentError::Geometry("field bytes overflow")))?;
398 let capacity = REGION_HEADER_LEN
399 .checked_add(spans_len)
400 .and_then(|v| v.checked_add(fields_len))
401 .and_then(|v| v.checked_add(self.sealed.terms.len()))
402 .and_then(|v| v.checked_add(self.sealed.blob.len()))
403 .ok_or_else(|| {
404 P::Error::from(SealedSegmentError::Geometry(
405 "lexical region length overflow",
406 ))
407 })?;
408 let mut bytes = GuardedVec::with_capacity(policy, capacity)?;
409 bytes.extend_from_slice(®ION_MAGIC, policy)?;
410 #[cfg(test)]
411 super::control::test_stage_probe(super::control::TestStage::RegionEncode);
412 bytes.extend_from_slice(®ION_VERSION.to_le_bytes(), policy)?;
413 bytes.extend_from_slice(&self.sealed.postings_per_block.to_le_bytes(), policy)?;
414 bytes.extend_from_slice(&self.sealed.row_count.to_le_bytes(), policy)?;
415 bytes.extend_from_slice(&span_count.to_le_bytes(), policy)?;
416 bytes.extend_from_slice(&field_count.to_le_bytes(), policy)?;
417 bytes.extend_from_slice(&0_u32.to_le_bytes(), policy)?;
418 bytes.extend_from_slice(
419 &u64::try_from(self.sealed.terms.len())
420 .unwrap_or(u64::MAX)
421 .to_le_bytes(),
422 policy,
423 )?;
424 bytes.extend_from_slice(
425 &u64::try_from(self.sealed.blob.len())
426 .unwrap_or(u64::MAX)
427 .to_le_bytes(),
428 policy,
429 )?;
430 for span in &self.sealed.spans {
431 for value in [span.field.0, 0] {
432 bytes.extend_from_slice(&value.to_le_bytes(), policy)?;
433 }
434 for value in [
435 span.term_start,
436 span.term_len,
437 span.meta_start,
438 span.docids_start,
439 span.tfs_start,
440 span.block_count,
441 span.doc_freq,
442 span.overall_max_tf,
443 ] {
444 bytes.extend_from_slice(&value.to_le_bytes(), policy)?;
445 }
446 bytes.extend_from_slice(&span.overall_min_len.to_le_bytes(), policy)?;
447 bytes.extend_from_slice(&0_u16.to_le_bytes(), policy)?;
448 bytes.extend_from_slice(&span.union_doc_freq.to_le_bytes(), policy)?;
449 bytes.extend_from_slice(&span.field_count.to_le_bytes(), policy)?;
450 }
451 for entry in &self.sealed.lengths {
452 bytes.extend_from_slice(&entry.field.0.to_le_bytes(), policy)?;
453 bytes.extend_from_slice(&0_u16.to_le_bytes(), policy)?;
454 bytes.extend_from_slice(&self.sealed.row_count.to_le_bytes(), policy)?;
455 for length in &entry.lengths {
456 bytes.extend_from_slice(&length.to_le_bytes(), policy)?;
457 }
458 }
459 bytes.extend_from_slice(&self.sealed.terms, policy)?;
460 bytes.extend_from_slice(&self.sealed.blob, policy)?;
461 Ok(bytes)
462 }
463
464 pub(crate) fn decode_policy<P>(bytes: &[u8], policy: &mut P) -> Result<(Self, u64), P::Error>
465 where
466 P: super::control::BuildPolicy<'m, Charge = C>,
467 P::Error: From<PostingsError> + From<SealedSegmentError>,
468 {
469 use super::control::GuardedVec;
470 policy.checkpoint()?;
471 let mut cursor = RegionCursor::new(bytes);
472 if cursor.take_array::<4>().map_err(P::Error::from)? != REGION_MAGIC {
473 return Err(P::Error::from(SealedSegmentError::Geometry(
474 "magic did not match",
475 )));
476 }
477 if cursor.read_u16().map_err(P::Error::from)? != REGION_VERSION {
478 return Err(P::Error::from(SealedSegmentError::Geometry(
479 "version is not readable",
480 )));
481 }
482 let postings_per_block = cursor.read_u16().map_err(P::Error::from)?;
483 if postings_per_block == 0 {
484 return Err(P::Error::from(SealedSegmentError::Geometry(
485 "postings per block is zero",
486 )));
487 }
488 let row_count = cursor.read_u32().map_err(P::Error::from)?;
489 let span_count =
490 usize::try_from(cursor.read_u32().map_err(P::Error::from)?).map_err(|_| {
491 P::Error::from(SealedSegmentError::Geometry("span count exceeds usize"))
492 })?;
493 let field_count =
494 usize::try_from(cursor.read_u32().map_err(P::Error::from)?).map_err(|_| {
495 P::Error::from(SealedSegmentError::Geometry("field count exceeds usize"))
496 })?;
497 if cursor.read_u32().map_err(P::Error::from)? != 0 {
498 return Err(P::Error::from(SealedSegmentError::Geometry(
499 "header reserved field is nonzero",
500 )));
501 }
502 let terms_len = usize::try_from(cursor.read_u64().map_err(P::Error::from)?)
503 .map_err(|_| P::Error::from(SealedSegmentError::Geometry("term bytes exceed usize")))?;
504 let blob_len =
505 usize::try_from(cursor.read_u64().map_err(P::Error::from)?).map_err(|_| {
506 P::Error::from(SealedSegmentError::Geometry("postings bytes exceed usize"))
507 })?;
508 let rows = usize::try_from(row_count)
509 .map_err(|_| P::Error::from(SealedSegmentError::Geometry("row count exceeds usize")))?;
510 let span_bytes = span_count
511 .checked_mul(REGION_SPAN_LEN)
512 .ok_or_else(|| P::Error::from(SealedSegmentError::Geometry("span bytes overflow")))?;
513 let one_field = REGION_FIELD_HEADER_LEN
514 .checked_add(rows.checked_mul(4).ok_or_else(|| {
515 P::Error::from(SealedSegmentError::Geometry("field bytes overflow"))
516 })?)
517 .ok_or_else(|| P::Error::from(SealedSegmentError::Geometry("field bytes overflow")))?;
518 let required = REGION_HEADER_LEN
519 .checked_add(span_bytes)
520 .and_then(|value| value.checked_add(field_count.checked_mul(one_field)?))
521 .and_then(|value| value.checked_add(terms_len))
522 .and_then(|value| value.checked_add(blob_len))
523 .ok_or_else(|| {
524 P::Error::from(SealedSegmentError::Geometry(
525 "lexical region length overflow",
526 ))
527 })?;
528 if required != bytes.len() {
529 return Err(P::Error::from(SealedSegmentError::Geometry(
530 "region is truncated",
531 )));
532 }
533 if (row_count == 0
534 && (field_count != 0 || span_count != 0 || terms_len != 0 || blob_len != 0))
535 || (row_count != 0 && field_count != 1)
536 {
537 return Err(P::Error::from(SealedSegmentError::Geometry(
538 "graph lexical region is not the default single field",
539 )));
540 }
541
542 let mut spans = GuardedVec::with_capacity(policy, span_count)?;
543 for _ in 0..span_count {
544 policy.step(1)?;
545 let field = FieldId(cursor.read_u16().map_err(P::Error::from)?);
546 if cursor.read_u16().map_err(P::Error::from)? != 0 {
547 return Err(P::Error::from(SealedSegmentError::Geometry(
548 "span reserved field is nonzero",
549 )));
550 }
551 let span = ListSpan {
552 field,
553 term_start: cursor.read_u32().map_err(P::Error::from)?,
554 term_len: cursor.read_u32().map_err(P::Error::from)?,
555 meta_start: cursor.read_u32().map_err(P::Error::from)?,
556 docids_start: cursor.read_u32().map_err(P::Error::from)?,
557 tfs_start: cursor.read_u32().map_err(P::Error::from)?,
558 block_count: cursor.read_u32().map_err(P::Error::from)?,
559 doc_freq: cursor.read_u32().map_err(P::Error::from)?,
560 overall_max_tf: cursor.read_u32().map_err(P::Error::from)?,
561 overall_min_len: cursor.read_u16().map_err(P::Error::from)?,
562 union_doc_freq: {
563 if cursor.read_u16().map_err(P::Error::from)? != 0 {
564 return Err(P::Error::from(SealedSegmentError::Geometry(
565 "span trailing reserved field is nonzero",
566 )));
567 }
568 cursor.read_u32().map_err(P::Error::from)?
569 },
570 field_count: cursor.read_u32().map_err(P::Error::from)?,
571 };
572 spans.push(span, policy)?;
573 #[cfg(test)]
574 super::control::test_stage_probe(super::control::TestStage::RegionDecode);
575 }
576 let mut outer_lengths = GuardedVec::with_capacity(policy, field_count)?;
577 let mut inner_lengths = GuardedVec::with_capacity(policy, rows)?;
578 let mut totals = GuardedVec::with_capacity(policy, rows)?;
579 let mut total_tokens = 0_u64;
580 if field_count == 1 {
581 let field = FieldId(cursor.read_u16().map_err(P::Error::from)?);
582 if field != super::index::DEFAULT_FIELD
583 || cursor.read_u16().map_err(P::Error::from)? != 0
584 || cursor.read_u32().map_err(P::Error::from)? != row_count
585 {
586 return Err(P::Error::from(SealedSegmentError::Geometry(
587 "field header is invalid",
588 )));
589 }
590 for _ in 0..rows {
591 policy.step(1)?;
592 let length = cursor.read_u32().map_err(P::Error::from)?;
593 if length == 0 {
594 return Err(P::Error::from(SealedSegmentError::Geometry(
595 "graph lexical row lengths are invalid",
596 )));
597 }
598 total_tokens = total_tokens.checked_add(u64::from(length)).ok_or_else(|| {
599 P::Error::from(SealedSegmentError::Geometry("total token count overflow"))
600 })?;
601 inner_lengths.push(length, policy)?;
602 totals.push(length, policy)?;
603 }
604 }
605 let mut terms = GuardedVec::with_capacity(policy, terms_len)?;
606 terms.extend_from_slice(cursor.take(terms_len).map_err(P::Error::from)?, policy)?;
607 let mut blob = GuardedVec::with_capacity(policy, blob_len)?;
608 blob.extend_from_slice(cursor.take(blob_len).map_err(P::Error::from)?, policy)?;
609 cursor.finish().map_err(P::Error::from)?;
610
611 let mut previous: Option<(&[u8], FieldId)> = None;
612 for (index, span) in spans.as_slice().iter().enumerate() {
613 policy.step(1)?;
614 let start = usize::try_from(span.term_start).map_err(|_| {
615 P::Error::from(SealedSegmentError::Geometry("term start exceeds usize"))
616 })?;
617 let end = start
618 .checked_add(usize::try_from(span.term_len).unwrap_or(usize::MAX))
619 .ok_or_else(|| {
620 P::Error::from(SealedSegmentError::Geometry("term span overflows"))
621 })?;
622 let term = terms.as_slice().get(start..end).ok_or_else(|| {
623 P::Error::from(SealedSegmentError::Geometry("term span is out of bounds"))
624 })?;
625 let out_of_order = if let Some((prior, field)) = previous {
626 match compare_bytes_policy(prior, term, policy)? {
627 std::cmp::Ordering::Greater => true,
628 std::cmp::Ordering::Equal => field >= span.field,
629 std::cmp::Ordering::Less => false,
630 }
631 } else {
632 false
633 };
634 if term.is_empty() || out_of_order {
635 return Err(P::Error::from(SealedSegmentError::Geometry(
636 "term spans are not strictly ordered",
637 )));
638 }
639 previous = Some((term, span.field));
640 if span.field != super::index::DEFAULT_FIELD
641 || span.field_count != 1
642 || span.union_doc_freq != span.doc_freq
643 {
644 return Err(P::Error::from(SealedSegmentError::Geometry(
645 "graph lexical posting shape is invalid",
646 )));
647 }
648 let list_start = usize::try_from(span.meta_start)
649 .unwrap_or(usize::MAX)
650 .checked_sub(HEADER_LEN)
651 .ok_or_else(|| {
652 P::Error::from(SealedSegmentError::Geometry(
653 "posting header start underflows",
654 ))
655 })?;
656 let list_end = spans
657 .as_slice()
658 .get(index.saturating_add(1))
659 .and_then(|next| {
660 usize::try_from(next.meta_start)
661 .unwrap_or(usize::MAX)
662 .checked_sub(HEADER_LEN)
663 })
664 .unwrap_or(blob.len());
665 let list = blob.as_slice().get(list_start..list_end).ok_or_else(|| {
666 P::Error::from(SealedSegmentError::Geometry(
667 "posting list is out of bounds",
668 ))
669 })?;
670 let summary = super::postings::validate_graph_policy(list, row_count, policy)?;
671 let expected_meta = u32::try_from(list_start)
672 .unwrap_or(u32::MAX)
673 .checked_add(HEADER_LEN_U32)
674 .ok_or_else(|| {
675 P::Error::from(SealedSegmentError::Geometry(
676 "posting metadata offset overflow",
677 ))
678 })?;
679 let expected_docids = expected_meta
680 .checked_add(summary.block_count.saturating_mul(META_LEN_U32))
681 .ok_or_else(|| {
682 P::Error::from(SealedSegmentError::Geometry(
683 "posting stream offset overflow",
684 ))
685 })?;
686 if summary.postings_per_block != postings_per_block
687 || span.meta_start != expected_meta
688 || span.docids_start != expected_docids
689 || span.tfs_start != expected_docids.saturating_add(summary.docid_bytes)
690 || span.block_count != summary.block_count
691 || span.doc_freq != summary.doc_freq
692 || span.overall_max_tf != summary.overall_max_tf
693 || span.overall_min_len != summary.overall_min_len
694 {
695 return Err(P::Error::from(SealedSegmentError::Geometry(
696 "posting-list summary does not match span",
697 )));
698 }
699 }
700 let (inner_values, inner_lengths_charge) = inner_lengths.into_parts();
701 if field_count == 1 {
702 outer_lengths.push(
703 FieldLengths {
704 field: super::index::DEFAULT_FIELD,
705 lengths: inner_values,
706 },
707 policy,
708 )?;
709 }
710 let (terms, terms_charge) = terms.into_parts();
711 let (spans, spans_charge) = spans.into_parts();
712 let (blob, blob_charge) = blob.into_parts();
713 let (lengths, outer_lengths_charge) = outer_lengths.into_parts();
714 let (total_lengths, totals_charge) = totals.into_parts();
715 let sealed = SealedSegment {
716 terms,
717 spans,
718 blob,
719 lengths,
720 total_lengths,
721 row_count,
722 postings_per_block,
723 };
724 policy.checkpoint()?;
725 Ok((
726 Self {
727 sealed,
728 terms_charge,
729 spans_charge,
730 blob_charge,
731 outer_lengths_charge,
732 inner_lengths_charge,
733 totals_charge,
734 marker: std::marker::PhantomData,
735 },
736 total_tokens,
737 ))
738 }
739}
740
741fn compare_bytes_policy<'m, P: super::control::BuildPolicy<'m>>(
742 left: &[u8],
743 right: &[u8],
744 policy: &mut P,
745) -> Result<std::cmp::Ordering, P::Error> {
746 for (left, right) in left.iter().zip(right) {
747 policy.step(1)?;
748 let order = left.cmp(right);
749 if order != std::cmp::Ordering::Equal {
750 return Ok(order);
751 }
752 }
753 Ok(left.len().cmp(&right.len()))
754}
755
756impl SealedSegment {
757 pub(crate) fn resident_bytes(&self) -> Result<usize, SealedSegmentError> {
758 self.resident_bytes_controlled(&mut super::control::WorkCheck::new(|| {
759 Ok::<(), SealedSegmentError>(())
760 }))
761 }
762
763 pub(crate) fn resident_bytes_controlled<E: From<SealedSegmentError>>(
764 &self,
765 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
766 ) -> Result<usize, E> {
767 work.check_now()?;
768 let direct = self
769 .terms
770 .capacity()
771 .checked_add(self.blob.capacity())
772 .and_then(|bytes| {
773 self.spans
774 .capacity()
775 .checked_mul(std::mem::size_of::<ListSpan>())
776 .and_then(|spans| bytes.checked_add(spans))
777 })
778 .and_then(|bytes| {
779 self.lengths
780 .capacity()
781 .checked_mul(std::mem::size_of::<FieldLengths>())
782 .and_then(|lengths| bytes.checked_add(lengths))
783 })
784 .and_then(|bytes| {
785 self.total_lengths
786 .capacity()
787 .checked_mul(std::mem::size_of::<u32>())
788 .and_then(|lengths| bytes.checked_add(lengths))
789 })
790 .ok_or(SealedSegmentError::Geometry(
791 "sealed lexical resident bytes overflow",
792 ))?;
793 self.lengths.iter().try_fold(direct, |total, entry| {
794 work.step()?;
795 entry
796 .lengths
797 .capacity()
798 .checked_mul(std::mem::size_of::<u32>())
799 .and_then(|bytes| total.checked_add(bytes))
800 .ok_or(SealedSegmentError::Geometry(
801 "sealed lexical resident bytes overflow",
802 ))
803 .map_err(E::from)
804 })
805 }
806
807 pub fn terms(&self) -> impl Iterator<Item = &[u8]> {
813 (0..self.spans.len()).filter_map(move |index| {
814 let term = self.term_of(index);
815 (index == 0 || self.term_of(index.saturating_sub(1)) != term).then_some(term)
816 })
817 }
818
819 pub(crate) fn vocabulary_terms(&self) -> impl Iterator<Item = (&[u8], FieldId)> + Clone {
820 self.spans
821 .iter()
822 .enumerate()
823 .map(move |(index, span)| (self.term_of(index), span.field))
824 }
825
826 pub fn seal(segment: &SegmentIndex) -> Result<Self, PostingsError> {
835 Self::seal_controlled(
836 segment,
837 &mut super::control::WorkCheck::new(|| Ok::<(), PostingsError>(())),
838 )
839 }
840
841 pub(crate) fn seal_controlled<E: From<PostingsError>>(
842 segment: &SegmentIndex,
843 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
844 ) -> Result<Self, E> {
845 work.check_now()?;
846 let per_block = DEFAULT_POSTINGS_PER_BLOCK;
847 let mut sealed = Self {
848 terms: Vec::new(),
849 spans: Vec::new(),
850 blob: Vec::new(),
851 lengths: Vec::new(),
852 total_lengths: Vec::new(),
853 row_count: segment.row_count(),
854 postings_per_block: per_block,
855 };
856 for field in segment.fields() {
857 work.step()?;
858 let source = segment.field_lengths(field).unwrap_or(&[]);
859 let mut lengths = Vec::with_capacity(source.len());
860 for length in source {
861 work.step()?;
862 #[cfg(test)]
863 SEAL_LENGTH_PROBES.with(|value| value.set(value.get() + 1));
864 lengths.push(*length);
865 }
866 sealed.lengths.push(FieldLengths { field, lengths });
867 }
868
869 let rows = usize::try_from(sealed.row_count).unwrap_or(0);
870 sealed.total_lengths = vec![0_u32; rows];
871 for entry in &sealed.lengths {
872 work.step()?;
873 for (slot, length) in entry.lengths.iter().enumerate() {
874 work.step()?;
875 #[cfg(test)]
876 SEAL_TOTAL_PROBES.with(|value| value.set(value.get() + 1));
877 if let Some(total) = sealed.total_lengths.get_mut(slot) {
878 *total = total.saturating_add(*length);
879 }
880 }
881 }
882
883 let mut group_term: Option<&[u8]> = None;
889 let mut group_start = 0_usize;
890 let mut group_lists: Vec<&PostingList> = Vec::new();
891 for (key, list) in segment.postings() {
892 work.step()?;
893 if list.is_empty() {
894 continue;
895 }
896 if group_term != Some(key.term.as_slice()) {
897 finish_union_group(&mut sealed.spans, group_start, &group_lists, work)?;
898 group_term = Some(key.term.as_slice());
899 group_start = sealed.spans.len();
900 group_lists.clear();
901 }
902 group_lists.push(list);
903 let mut field_lengths = &[][..];
904 for entry in &sealed.lengths {
905 work.step()?;
906 if entry.field == key.field {
907 field_lengths = entry.lengths.as_slice();
908 break;
909 }
910 }
911 let impacts =
912 super::postings::block_impacts_controlled(list, per_block, field_lengths, work)?;
913 let encoded =
914 super::postings::encode_v2_controlled(list, per_block, &[], &impacts, work)?;
915 let bytes = encoded.as_bytes();
916 let reader = PostingsReader::open_controlled(bytes, work)?;
919 let base = u32::try_from(sealed.blob.len()).unwrap_or(u32::MAX);
920 let block_count = u32::try_from(reader.blocks().len()).unwrap_or(u32::MAX);
921 let meta_start = base.saturating_add(HEADER_LEN_U32);
922 let streams = meta_start.saturating_add(block_count.saturating_mul(META_LEN_U32));
923 let mut docid_bytes = 0;
924 for meta in reader.blocks() {
925 work.step()?;
926 docid_bytes =
927 docid_bytes.max(stream_end(meta.docids_offset, meta.count, meta.docid_bits));
928 }
929 let mut overall_max_tf = 0_u32;
930 let mut overall_min_len = u16::MAX;
931 for meta in reader.blocks() {
932 work.step()?;
933 let impact = BlockImpact::from_meta(meta);
934 overall_max_tf = overall_max_tf.max(impact.max_tf);
935 overall_min_len = overall_min_len.min(impact.min_len);
936 }
937
938 let term_start = u32::try_from(sealed.terms.len()).unwrap_or(u32::MAX);
939 super::control::extend_bytes(&mut sealed.terms, &key.term, work)?;
940 sealed.spans.push(ListSpan {
941 field: key.field,
942 term_start,
943 term_len: u32::try_from(key.term.len()).unwrap_or(u32::MAX),
944 meta_start,
945 docids_start: streams,
946 tfs_start: streams.saturating_add(docid_bytes),
947 block_count,
948 doc_freq: reader.document_frequency(),
949 overall_max_tf,
950 overall_min_len,
951 union_doc_freq: 0,
953 field_count: 0,
954 });
955 super::control::extend_bytes(&mut sealed.blob, bytes, work)?;
956 }
957 finish_union_group(&mut sealed.spans, group_start, &group_lists, work)?;
958 work.check_now()?;
959 Ok(sealed)
960 }
961
962 pub fn encode_region(&self) -> Result<Vec<u8>, SealedSegmentError> {
964 let span_count = u32::try_from(self.spans.len())
965 .map_err(|_| SealedSegmentError::Geometry("span count exceeds u32"))?;
966 let field_count = u32::try_from(self.lengths.len())
967 .map_err(|_| SealedSegmentError::Geometry("field count exceeds u32"))?;
968 let terms_len = u64::try_from(self.terms.len())
969 .map_err(|_| SealedSegmentError::Geometry("term bytes exceed u64"))?;
970 let blob_len = u64::try_from(self.blob.len())
971 .map_err(|_| SealedSegmentError::Geometry("postings bytes exceed u64"))?;
972 let spans_len = self
973 .spans
974 .len()
975 .checked_mul(REGION_SPAN_LEN)
976 .ok_or(SealedSegmentError::Geometry("span bytes overflow"))?;
977 let fields_len = self.lengths.iter().try_fold(0_usize, |total, entry| {
978 if entry.lengths.len() != self.row_count as usize {
979 return Err(SealedSegmentError::Geometry(
980 "field lengths do not cover every row",
981 ));
982 }
983 let body = entry
984 .lengths
985 .len()
986 .checked_mul(4)
987 .and_then(|bytes| bytes.checked_add(REGION_FIELD_HEADER_LEN))
988 .ok_or(SealedSegmentError::Geometry("field bytes overflow"))?;
989 total
990 .checked_add(body)
991 .ok_or(SealedSegmentError::Geometry("field bytes overflow"))
992 })?;
993 let capacity = REGION_HEADER_LEN
994 .checked_add(spans_len)
995 .and_then(|value| value.checked_add(fields_len))
996 .and_then(|value| value.checked_add(self.terms.len()))
997 .and_then(|value| value.checked_add(self.blob.len()))
998 .ok_or(SealedSegmentError::Geometry(
999 "lexical region length overflow",
1000 ))?;
1001 let mut bytes = Vec::with_capacity(capacity);
1002 bytes.extend_from_slice(®ION_MAGIC);
1003 bytes.extend_from_slice(®ION_VERSION.to_le_bytes());
1004 bytes.extend_from_slice(&self.postings_per_block.to_le_bytes());
1005 bytes.extend_from_slice(&self.row_count.to_le_bytes());
1006 bytes.extend_from_slice(&span_count.to_le_bytes());
1007 bytes.extend_from_slice(&field_count.to_le_bytes());
1008 bytes.extend_from_slice(&0_u32.to_le_bytes());
1009 bytes.extend_from_slice(&terms_len.to_le_bytes());
1010 bytes.extend_from_slice(&blob_len.to_le_bytes());
1011 for span in &self.spans {
1012 bytes.extend_from_slice(&span.field.0.to_le_bytes());
1013 bytes.extend_from_slice(&0_u16.to_le_bytes());
1014 bytes.extend_from_slice(&span.term_start.to_le_bytes());
1015 bytes.extend_from_slice(&span.term_len.to_le_bytes());
1016 bytes.extend_from_slice(&span.meta_start.to_le_bytes());
1017 bytes.extend_from_slice(&span.docids_start.to_le_bytes());
1018 bytes.extend_from_slice(&span.tfs_start.to_le_bytes());
1019 bytes.extend_from_slice(&span.block_count.to_le_bytes());
1020 bytes.extend_from_slice(&span.doc_freq.to_le_bytes());
1021 bytes.extend_from_slice(&span.overall_max_tf.to_le_bytes());
1022 bytes.extend_from_slice(&span.overall_min_len.to_le_bytes());
1023 bytes.extend_from_slice(&0_u16.to_le_bytes());
1024 bytes.extend_from_slice(&span.union_doc_freq.to_le_bytes());
1025 bytes.extend_from_slice(&span.field_count.to_le_bytes());
1026 }
1027 for entry in &self.lengths {
1028 bytes.extend_from_slice(&entry.field.0.to_le_bytes());
1029 bytes.extend_from_slice(&0_u16.to_le_bytes());
1030 bytes.extend_from_slice(&self.row_count.to_le_bytes());
1031 for length in &entry.lengths {
1032 bytes.extend_from_slice(&length.to_le_bytes());
1033 }
1034 }
1035 bytes.extend_from_slice(&self.terms);
1036 bytes.extend_from_slice(&self.blob);
1037 Ok(bytes)
1038 }
1039
1040 pub fn decode_region(bytes: &[u8]) -> Result<Self, SealedSegmentError> {
1042 let mut cursor = RegionCursor::new(bytes);
1043 if cursor.take_array::<4>()? != REGION_MAGIC {
1044 return Err(SealedSegmentError::Geometry("magic did not match"));
1045 }
1046 if cursor.read_u16()? != REGION_VERSION {
1047 return Err(SealedSegmentError::Geometry("version is not readable"));
1048 }
1049 let postings_per_block = cursor.read_u16()?;
1050 if postings_per_block == 0 {
1051 return Err(SealedSegmentError::Geometry("postings per block is zero"));
1052 }
1053 let row_count = cursor.read_u32()?;
1054 let span_count = usize::try_from(cursor.read_u32()?)
1055 .map_err(|_| SealedSegmentError::Geometry("span count exceeds usize"))?;
1056 let field_count = usize::try_from(cursor.read_u32()?)
1057 .map_err(|_| SealedSegmentError::Geometry("field count exceeds usize"))?;
1058 if cursor.read_u32()? != 0 {
1059 return Err(SealedSegmentError::Geometry(
1060 "header reserved field is nonzero",
1061 ));
1062 }
1063 let terms_len = usize::try_from(cursor.read_u64()?)
1064 .map_err(|_| SealedSegmentError::Geometry("term bytes exceed usize"))?;
1065 let blob_len = usize::try_from(cursor.read_u64()?)
1066 .map_err(|_| SealedSegmentError::Geometry("postings bytes exceed usize"))?;
1067 let mut spans = Vec::with_capacity(span_count);
1068 for _ in 0..span_count {
1069 let field = FieldId(cursor.read_u16()?);
1070 if cursor.read_u16()? != 0 {
1071 return Err(SealedSegmentError::Geometry(
1072 "span reserved field is nonzero",
1073 ));
1074 }
1075 let span = ListSpan {
1076 field,
1077 term_start: cursor.read_u32()?,
1078 term_len: cursor.read_u32()?,
1079 meta_start: cursor.read_u32()?,
1080 docids_start: cursor.read_u32()?,
1081 tfs_start: cursor.read_u32()?,
1082 block_count: cursor.read_u32()?,
1083 doc_freq: cursor.read_u32()?,
1084 overall_max_tf: cursor.read_u32()?,
1085 overall_min_len: cursor.read_u16()?,
1086 union_doc_freq: {
1087 if cursor.read_u16()? != 0 {
1088 return Err(SealedSegmentError::Geometry(
1089 "span trailing reserved field is nonzero",
1090 ));
1091 }
1092 cursor.read_u32()?
1093 },
1094 field_count: cursor.read_u32()?,
1095 };
1096 spans.push(span);
1097 }
1098 let mut lengths = Vec::with_capacity(field_count);
1099 let rows = usize::try_from(row_count)
1100 .map_err(|_| SealedSegmentError::Geometry("row count exceeds usize"))?;
1101 let mut previous_field: Option<FieldId> = None;
1102 for _ in 0..field_count {
1103 let field = FieldId(cursor.read_u16()?);
1104 if previous_field.is_some_and(|previous| previous >= field) {
1105 return Err(SealedSegmentError::Geometry(
1106 "fields are not strictly ascending",
1107 ));
1108 }
1109 previous_field = Some(field);
1110 if cursor.read_u16()? != 0 || cursor.read_u32()? != row_count {
1111 return Err(SealedSegmentError::Geometry("field header is invalid"));
1112 }
1113 let mut field_lengths = Vec::with_capacity(rows);
1114 for _ in 0..rows {
1115 field_lengths.push(cursor.read_u32()?);
1116 }
1117 lengths.push(FieldLengths {
1118 field,
1119 lengths: field_lengths,
1120 });
1121 }
1122 let terms = cursor.take(terms_len)?.to_vec();
1123 let blob = cursor.take(blob_len)?.to_vec();
1124 cursor.finish()?;
1125 validate_decoded_spans(&spans, &terms, &blob, postings_per_block, row_count)?;
1126 let mut total_lengths = vec![0_u32; rows];
1127 for entry in &lengths {
1128 for (slot, length) in entry.lengths.iter().enumerate() {
1129 if let Some(total) = total_lengths.get_mut(slot) {
1130 *total = total.saturating_add(*length);
1131 }
1132 }
1133 }
1134 Ok(Self {
1135 terms,
1136 spans,
1137 blob,
1138 lengths,
1139 total_lengths,
1140 row_count,
1141 postings_per_block,
1142 })
1143 }
1144
1145 pub fn retain_rows(&self, survivors: &[usize]) -> Result<Self, SealedSegmentError> {
1150 let old_rows = usize::try_from(self.row_count)
1151 .map_err(|_| SealedSegmentError::Geometry("row count exceeds usize"))?;
1152 let mut mapping = vec![None; old_rows];
1153 let mut previous = None;
1154 for (new_row, old_row) in survivors.iter().copied().enumerate() {
1155 if old_row >= old_rows || previous.is_some_and(|value| value >= old_row) {
1156 return Err(SealedSegmentError::Geometry(
1157 "survivor rows are not ascending and in bounds",
1158 ));
1159 }
1160 previous = Some(old_row);
1161 let new_row = u32::try_from(new_row)
1162 .map_err(|_| SealedSegmentError::Geometry("survivor row exceeds u32"))?;
1163 let target = mapping
1164 .get_mut(old_row)
1165 .ok_or(SealedSegmentError::Geometry(
1166 "survivor row is out of bounds",
1167 ))?;
1168 *target = Some(new_row);
1169 }
1170 let row_count = u32::try_from(survivors.len())
1171 .map_err(|_| SealedSegmentError::Geometry("survivor count exceeds u32"))?;
1172 let mut retained = Self {
1173 terms: Vec::new(),
1174 spans: Vec::new(),
1175 blob: Vec::new(),
1176 lengths: Vec::with_capacity(self.lengths.len()),
1177 total_lengths: vec![0; survivors.len()],
1178 row_count,
1179 postings_per_block: self.postings_per_block,
1180 };
1181 for entry in &self.lengths {
1182 let mut lengths = Vec::with_capacity(survivors.len());
1183 for row in survivors {
1184 lengths.push(entry.lengths.get(*row).copied().ok_or(
1185 SealedSegmentError::Geometry("survivor field length is missing"),
1186 )?);
1187 }
1188 retained.lengths.push(FieldLengths {
1189 field: entry.field,
1190 lengths,
1191 });
1192 }
1193 for entry in &retained.lengths {
1194 for (row, length) in entry.lengths.iter().copied().enumerate() {
1195 let total =
1196 retained
1197 .total_lengths
1198 .get_mut(row)
1199 .ok_or(SealedSegmentError::Geometry(
1200 "retained total-length row is missing",
1201 ))?;
1202 *total = total.saturating_add(length);
1203 }
1204 }
1205
1206 let mut lists = Vec::new();
1207 for (index, span) in self.spans.iter().enumerate() {
1208 let encoded = self.list_bytes(index)?;
1209 let decoded = PostingsReader::open(encoded)?.decode_all()?;
1210 let mut list = PostingList::new();
1211 for posting in decoded.postings() {
1212 let old_row = usize::try_from(posting.docid)
1213 .map_err(|_| SealedSegmentError::Geometry("posting row exceeds usize"))?;
1214 let Some(new_row) = mapping.get(old_row).copied().flatten() else {
1215 continue;
1216 };
1217 let mut posting = posting.clone();
1218 posting.docid = new_row;
1219 list.push(posting)?;
1220 }
1221 if list.is_empty() {
1222 continue;
1223 }
1224 let field_lengths = retained
1225 .lengths
1226 .iter()
1227 .find(|entry| entry.field == span.field)
1228 .map_or(&[][..], |entry| entry.lengths.as_slice());
1229 let impacts = block_impacts(&list, retained.postings_per_block, field_lengths);
1230 let encoded = encode_v2(&list, retained.postings_per_block, &[], &impacts)?;
1231 let reader = PostingsReader::open(encoded.as_bytes())?;
1232 let base = u32::try_from(retained.blob.len())
1233 .map_err(|_| SealedSegmentError::Geometry("postings blob exceeds u32"))?;
1234 let block_count = u32::try_from(reader.blocks().len())
1235 .map_err(|_| SealedSegmentError::Geometry("posting blocks exceed u32"))?;
1236 let meta_start =
1237 base.checked_add(HEADER_LEN_U32)
1238 .ok_or(SealedSegmentError::Geometry(
1239 "posting metadata offset overflow",
1240 ))?;
1241 let streams = meta_start
1242 .checked_add(block_count.saturating_mul(META_LEN_U32))
1243 .ok_or(SealedSegmentError::Geometry(
1244 "posting stream offset overflow",
1245 ))?;
1246 let docid_bytes = reader
1247 .blocks()
1248 .iter()
1249 .map(|meta| stream_end(meta.docids_offset, meta.count, meta.docid_bits))
1250 .max()
1251 .unwrap_or(0);
1252 let mut overall_max_tf = 0_u32;
1253 let mut overall_min_len = u16::MAX;
1254 for meta in reader.blocks() {
1255 let impact = BlockImpact::from_meta(meta);
1256 overall_max_tf = overall_max_tf.max(impact.max_tf);
1257 overall_min_len = overall_min_len.min(impact.min_len);
1258 }
1259 let term = self.term_of(index);
1260 let term_start = u32::try_from(retained.terms.len())
1261 .map_err(|_| SealedSegmentError::Geometry("term bytes exceed u32"))?;
1262 let term_len = u32::try_from(term.len())
1263 .map_err(|_| SealedSegmentError::Geometry("term length exceeds u32"))?;
1264 retained.terms.extend_from_slice(term);
1265 retained.spans.push(ListSpan {
1266 field: span.field,
1267 term_start,
1268 term_len,
1269 meta_start,
1270 docids_start: streams,
1271 tfs_start: streams.saturating_add(docid_bytes),
1272 block_count,
1273 doc_freq: reader.document_frequency(),
1274 overall_max_tf,
1275 overall_min_len,
1276 union_doc_freq: 0,
1277 field_count: 0,
1278 });
1279 retained.blob.extend_from_slice(encoded.as_bytes());
1280 lists.push(list);
1281 }
1282 stamp_retained_unions(&mut retained, &lists)?;
1283 Ok(retained)
1284 }
1285
1286 pub(crate) fn position_reader_bytes_controlled<E>(
1287 &self,
1288 term: &[u8],
1289 field: FieldId,
1290 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1291 ) -> Result<Option<usize>, E> {
1292 work.check_now()?;
1293 let (start, end) = self.term_span_range(term);
1294 let Some(spans) = self.spans.get(start..end) else {
1295 return Ok(None);
1296 };
1297 for span in spans {
1298 work.step()?;
1299 if span.field == field {
1300 return Ok(
1301 (span.block_count as usize).checked_mul(std::mem::size_of::<BlockMeta>())
1302 );
1303 }
1304 }
1305 Ok(Some(0))
1306 }
1307
1308 pub(crate) fn position_reader_controlled<E: From<SealedSegmentError> + From<PostingsError>>(
1309 &self,
1310 term: &[u8],
1311 field: FieldId,
1312 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1313 ) -> Result<Option<PostingsReader<'_>>, E> {
1314 work.check_now()?;
1315 let (start, end) = self.term_span_range(term);
1316 for index in start..end {
1317 work.step()?;
1318 if self
1319 .spans
1320 .get(index)
1321 .is_some_and(|span| span.field == field)
1322 {
1323 return Ok(Some(PostingsReader::open_controlled(
1324 self.list_bytes(index)?,
1325 work,
1326 )?));
1327 }
1328 }
1329 Ok(None)
1330 }
1331
1332 fn list_bytes(&self, index: usize) -> Result<&[u8], SealedSegmentError> {
1333 let span = self
1334 .spans
1335 .get(index)
1336 .ok_or(SealedSegmentError::Geometry("posting span is missing"))?;
1337 let start = usize::try_from(span.meta_start)
1338 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?
1339 .checked_sub(HEADER_LEN)
1340 .ok_or(SealedSegmentError::Geometry(
1341 "posting header start underflows",
1342 ))?;
1343 let end = self
1344 .spans
1345 .get(index.saturating_add(1))
1346 .map(|next| {
1347 usize::try_from(next.meta_start)
1348 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?
1349 .checked_sub(HEADER_LEN)
1350 .ok_or(SealedSegmentError::Geometry(
1351 "posting header start underflows",
1352 ))
1353 })
1354 .transpose()?
1355 .unwrap_or(self.blob.len());
1356 self.blob
1357 .get(start..end)
1358 .ok_or(SealedSegmentError::Geometry(
1359 "posting list is out of bounds",
1360 ))
1361 }
1362
1363 fn term_of(&self, index: usize) -> &[u8] {
1365 let Some(span) = self.spans.get(index) else {
1366 return &[];
1367 };
1368 let start = usize::try_from(span.term_start).unwrap_or(usize::MAX);
1369 let len = usize::try_from(span.term_len).unwrap_or(0);
1370 self.terms
1371 .get(start..start.saturating_add(len))
1372 .unwrap_or(&[])
1373 }
1374
1375 #[must_use]
1377 pub const fn row_count(&self) -> u32 {
1378 self.row_count
1379 }
1380
1381 #[must_use]
1383 pub const fn is_empty(&self) -> bool {
1384 self.row_count == 0
1385 }
1386
1387 #[must_use]
1389 pub const fn postings_per_block(&self) -> u16 {
1390 self.postings_per_block
1391 }
1392
1393 #[must_use]
1395 pub fn field_lengths(&self, field: FieldId) -> Option<&[u32]> {
1396 self.lengths
1397 .iter()
1398 .find(|entry| entry.field == field)
1399 .map(|entry| entry.lengths.as_slice())
1400 }
1401
1402 pub(crate) fn field_lengths_controlled<E>(
1403 &self,
1404 field: FieldId,
1405 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1406 ) -> Result<Option<&[u32]>, E> {
1407 for entry in &self.lengths {
1408 work.step()?;
1409 if entry.field == field {
1410 return Ok(Some(entry.lengths.as_slice()));
1411 }
1412 }
1413 Ok(None)
1414 }
1415
1416 #[must_use]
1420 pub fn total_lengths(&self) -> &[u32] {
1421 &self.total_lengths
1422 }
1423
1424 pub fn fields(&self) -> impl Iterator<Item = FieldId> + '_ {
1426 self.lengths.iter().map(|entry| entry.field)
1427 }
1428
1429 #[must_use]
1431 pub fn field_length(&self, row: u32, field: FieldId) -> u32 {
1432 self.field_lengths(field)
1433 .and_then(|lengths| {
1434 usize::try_from(row)
1435 .ok()
1436 .and_then(|slot| lengths.get(slot).copied())
1437 })
1438 .unwrap_or(0)
1439 }
1440
1441 #[must_use]
1443 pub fn document_length(&self, row: u32) -> u32 {
1444 let Ok(slot) = usize::try_from(row) else {
1445 return 0;
1446 };
1447 self.lengths
1448 .iter()
1449 .filter_map(|entry| entry.lengths.get(slot).copied())
1450 .fold(0_u32, u32::saturating_add)
1451 }
1452
1453 #[must_use]
1455 pub fn total_tokens(&self) -> u64 {
1456 self.lengths
1457 .iter()
1458 .flat_map(|entry| entry.lengths.iter())
1459 .map(|length| u64::from(*length))
1460 .sum()
1461 }
1462
1463 fn term_span_range(&self, term: &[u8]) -> (usize, usize) {
1465 let start = self.spans.partition_point(|span| {
1466 let bytes_start = usize::try_from(span.term_start).unwrap_or(usize::MAX);
1467 let len = usize::try_from(span.term_len).unwrap_or(0);
1468 self.terms
1469 .get(bytes_start..bytes_start.saturating_add(len))
1470 .unwrap_or(&[])
1471 < term
1472 });
1473 let mut end = start;
1474 while end < self.spans.len() && self.term_of(end) == term {
1475 end = end.saturating_add(1);
1476 }
1477 (start, end)
1478 }
1479
1480 #[must_use]
1486 pub fn document_frequency(&self, term: &[u8], fields: &[FieldId]) -> u32 {
1487 match self
1488 .document_frequency_controlled(term, fields, || Ok::<(), std::convert::Infallible>(()))
1489 {
1490 Ok(frequency) => frequency,
1491 Err(never) => match never {},
1492 }
1493 }
1494
1495 pub(crate) fn document_frequency_controlled<E>(
1496 &self,
1497 term: &[u8],
1498 fields: &[FieldId],
1499 check: impl FnMut() -> Result<(), E>,
1500 ) -> Result<u32, E> {
1501 let mut work = super::control::WorkCheck::new(check);
1502 work.check_now()?;
1503 let (start, end) = self.term_span_range(term);
1504 let mut matched: Option<usize> = None;
1505 let mut count = 0_usize;
1506 for index in start..end {
1507 work.step()?;
1508 let Some(span) = self.spans.get(index) else {
1509 continue;
1510 };
1511 if fields.contains(&span.field) {
1512 matched = matched.or(Some(index));
1513 count = count.saturating_add(1);
1514 }
1515 }
1516 let frequency = match count {
1517 0 => 0,
1518 1 => matched
1519 .and_then(|index| self.spans.get(index))
1520 .map_or(0, |span| span.doc_freq),
1521 _ => {
1522 let total = u32::try_from(count).unwrap_or(u32::MAX);
1523 let all_fields = self
1524 .spans
1525 .get(start)
1526 .is_some_and(|span| span.field_count == total);
1527 if all_fields {
1528 self.spans.get(start).map_or(0, |span| span.union_doc_freq)
1530 } else {
1531 let mut cursors: Vec<ListCursor<'_>> = Vec::new();
1533 for index in start..end {
1534 work.step()?;
1535 let Some(span) = self.spans.get(index) else {
1536 continue;
1537 };
1538 if fields.contains(&span.field)
1539 && let Some(cursor) = self.cursor_at(index, 1_000)
1540 {
1541 cursors.push(cursor);
1542 }
1543 }
1544 union_count_controlled(&mut cursors, &mut work)?
1545 }
1546 }
1547 };
1548 work.check_now()?;
1549 Ok(frequency)
1550 }
1551
1552 pub(crate) fn single_posting_live_frequency(
1557 &self,
1558 term: &[u8],
1559 fields: &[FieldId],
1560 live_rows: &DocBitmap,
1561 ) -> Option<u32> {
1562 let (start, end) = self.term_span_range(term);
1563 let mut matched = None;
1564 for index in start..end {
1565 let span = self.spans.get(index)?;
1566 if !fields.contains(&span.field) || span.doc_freq == 0 {
1567 continue;
1568 }
1569 if span.doc_freq > 1 || matched.is_some() {
1570 return None;
1571 }
1572 matched = Some(index);
1573 }
1574 let Some(index) = matched else {
1575 return Some(0);
1576 };
1577 let mut cursor = self.cursor_at(index, 1_000)?;
1578 cursor.reset();
1579 let row = cursor.current()?;
1580 #[cfg(any(test, feature = "test-seams"))]
1581 super::preparation_observer::live_df_walk(
1582 1,
1583 usize::try_from(cursor.blocks_decoded).unwrap_or(usize::MAX),
1584 );
1585 Some(u32::from(live_rows.contains(row)))
1586 }
1587
1588 pub(crate) fn live_document_frequency(
1589 &self,
1590 term: &[u8],
1591 fields: &[FieldId],
1592 live_rows: &DocBitmap,
1593 ) -> u32 {
1594 match self.live_document_frequency_controlled(term, fields, live_rows, || {
1595 Ok::<(), std::convert::Infallible>(())
1596 }) {
1597 Ok(frequency) => frequency,
1598 Err(never) => match never {},
1599 }
1600 }
1601
1602 pub(crate) fn live_document_frequency_controlled<E>(
1603 &self,
1604 term: &[u8],
1605 fields: &[FieldId],
1606 live_rows: &DocBitmap,
1607 check: impl FnMut() -> Result<(), E>,
1608 ) -> Result<u32, E> {
1609 let mut work = super::control::WorkCheck::new(check);
1610 work.check_now()?;
1611 let weights = FieldWeights::flat(fields);
1612 let Some(mut stream) = TermStream::open_controlled(self, term, &weights, &mut work)? else {
1613 return Ok(0);
1614 };
1615 let mut count = 0_u32;
1616 #[cfg(any(test, feature = "test-seams"))]
1617 let mut visited = 0_usize;
1618 let result = (|| {
1619 while let Some(row) = stream.current_row() {
1620 work.step()?;
1621 #[cfg(any(test, feature = "test-seams"))]
1622 {
1623 visited = visited.saturating_add(1);
1624 }
1625 if live_rows.contains(row) {
1626 count = count.saturating_add(1);
1627 }
1628 stream.advance_controlled(&mut work)?;
1629 }
1630 work.check_now()?;
1631 Ok(count)
1632 })();
1633 #[cfg(any(test, feature = "test-seams"))]
1634 super::preparation_observer::live_df_walk(
1635 visited,
1636 usize::try_from(stream.blocks_decoded()).unwrap_or(usize::MAX),
1637 );
1638 result
1639 }
1640
1641 fn cursor_at(&self, index: usize, weight: u64) -> Option<ListCursor<'_>> {
1643 let span = self.spans.get(index)?;
1644 let per_block = usize::from(self.postings_per_block);
1645 Some(ListCursor {
1646 meta: self.slice_from(span.meta_start),
1647 docids: self.slice_from(span.docids_start),
1648 tfs: self.slice_from(span.tfs_start),
1649 block_count: usize::try_from(span.block_count).unwrap_or(0),
1650 per_block,
1651 weight,
1652 block: NO_BLOCK,
1653 slot: 0,
1654 count: 0,
1655 impact: BlockImpact::default(),
1656 tfs_block: NO_BLOCK,
1657 tf_offset: 0,
1658 tf_bits: 0,
1659 overall: BlockImpact {
1660 max_tf: span.overall_max_tf,
1661 min_len: span.overall_min_len,
1662 },
1663 decoded_docids: vec![0; per_block],
1664 decoded_tfs: vec![0; per_block],
1665 blocks_decoded: 0,
1666 blocks_skipped: 0,
1667 tf_blocks_decoded: 0,
1668 })
1669 }
1670
1671 fn slice_from(&self, offset: u32) -> &[u8] {
1673 usize::try_from(offset)
1674 .ok()
1675 .and_then(|start| self.blob.get(start..))
1676 .unwrap_or(&[])
1677 }
1678}
1679
1680fn stamp_retained_unions(
1681 segment: &mut SealedSegment,
1682 lists: &[PostingList],
1683) -> Result<(), SealedSegmentError> {
1684 if segment.spans.len() != lists.len() {
1685 return Err(SealedSegmentError::Geometry(
1686 "retained posting list count differs from spans",
1687 ));
1688 }
1689 let mut start = 0_usize;
1690 while start < segment.spans.len() {
1691 let term = segment.term_of(start).to_vec();
1692 let mut end = start.saturating_add(1);
1693 while end < segment.spans.len() && segment.term_of(end) == term {
1694 end = end.saturating_add(1);
1695 }
1696 let mut documents = std::collections::BTreeSet::new();
1697 for list in lists.get(start..end).ok_or(SealedSegmentError::Geometry(
1698 "retained posting-list group is missing",
1699 ))? {
1700 documents.extend(list.postings().iter().map(|posting| posting.docid));
1701 }
1702 let union = u32::try_from(documents.len())
1703 .map_err(|_| SealedSegmentError::Geometry("union document count exceeds u32"))?;
1704 let field_count = u32::try_from(end.saturating_sub(start))
1705 .map_err(|_| SealedSegmentError::Geometry("field count exceeds u32"))?;
1706 let spans = segment
1707 .spans
1708 .get_mut(start..end)
1709 .ok_or(SealedSegmentError::Geometry(
1710 "retained span group is missing",
1711 ))?;
1712 for span in spans {
1713 span.union_doc_freq = union;
1714 span.field_count = field_count;
1715 }
1716 start = end;
1717 }
1718 Ok(())
1719}
1720
1721fn validate_decoded_spans(
1722 spans: &[ListSpan],
1723 terms: &[u8],
1724 blob: &[u8],
1725 postings_per_block: u16,
1726 row_count: u32,
1727) -> Result<(), SealedSegmentError> {
1728 let mut previous_key: Option<(&[u8], FieldId)> = None;
1729 let mut decoded_lists = Vec::with_capacity(spans.len());
1730 for (index, span) in spans.iter().enumerate() {
1731 let term_start = usize::try_from(span.term_start)
1732 .map_err(|_| SealedSegmentError::Geometry("term start exceeds usize"))?;
1733 let term_len = usize::try_from(span.term_len)
1734 .map_err(|_| SealedSegmentError::Geometry("term length exceeds usize"))?;
1735 let term_end = term_start
1736 .checked_add(term_len)
1737 .ok_or(SealedSegmentError::Geometry("term span overflows"))?;
1738 let term = terms
1739 .get(term_start..term_end)
1740 .ok_or(SealedSegmentError::Geometry("term span is out of bounds"))?;
1741 if previous_key.is_some_and(|(previous_term, previous_field)| {
1742 previous_term > term || (previous_term == term && previous_field >= span.field)
1743 }) {
1744 return Err(SealedSegmentError::Geometry(
1745 "term spans are not strictly ordered",
1746 ));
1747 }
1748 previous_key = Some((term, span.field));
1749 let meta_start = usize::try_from(span.meta_start)
1750 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?;
1751 let list_start = meta_start
1752 .checked_sub(HEADER_LEN)
1753 .ok_or(SealedSegmentError::Geometry(
1754 "posting header start underflows",
1755 ))?;
1756 let list_end = spans
1757 .get(index.saturating_add(1))
1758 .map(|next| {
1759 usize::try_from(next.meta_start)
1760 .map_err(|_| SealedSegmentError::Geometry("posting offset exceeds usize"))?
1761 .checked_sub(HEADER_LEN)
1762 .ok_or(SealedSegmentError::Geometry(
1763 "posting header start underflows",
1764 ))
1765 })
1766 .transpose()?
1767 .unwrap_or(blob.len());
1768 let list = blob
1769 .get(list_start..list_end)
1770 .ok_or(SealedSegmentError::Geometry(
1771 "posting list is out of bounds",
1772 ))?;
1773 let reader = PostingsReader::open(list)?;
1774 let block_count = u32::try_from(reader.blocks().len())
1775 .map_err(|_| SealedSegmentError::Geometry("posting blocks exceed u32"))?;
1776 let expected_meta = u32::try_from(list_start)
1777 .map_err(|_| SealedSegmentError::Geometry("posting start exceeds u32"))?
1778 .checked_add(HEADER_LEN_U32)
1779 .ok_or(SealedSegmentError::Geometry(
1780 "posting metadata offset overflow",
1781 ))?;
1782 let expected_docids = expected_meta
1783 .checked_add(block_count.saturating_mul(META_LEN_U32))
1784 .ok_or(SealedSegmentError::Geometry(
1785 "posting stream offset overflow",
1786 ))?;
1787 let docid_bytes = reader
1788 .blocks()
1789 .iter()
1790 .map(|meta| stream_end(meta.docids_offset, meta.count, meta.docid_bits))
1791 .max()
1792 .unwrap_or(0);
1793 let expected_tfs = expected_docids
1794 .checked_add(docid_bytes)
1795 .ok_or(SealedSegmentError::Geometry("posting tf offset overflow"))?;
1796 let mut overall_max_tf = 0_u32;
1797 let mut overall_min_len = u16::MAX;
1798 for block in reader.blocks() {
1799 let impact = BlockImpact::from_meta(block);
1800 overall_max_tf = overall_max_tf.max(impact.max_tf);
1801 overall_min_len = overall_min_len.min(impact.min_len);
1802 }
1803 if reader.postings_per_block() != postings_per_block
1804 || reader.document_frequency() != span.doc_freq
1805 || block_count != span.block_count
1806 || span.meta_start != expected_meta
1807 || span.docids_start != expected_docids
1808 || span.tfs_start != expected_tfs
1809 || span.overall_max_tf != overall_max_tf
1810 || span.overall_min_len != overall_min_len
1811 || span.field_count == 0
1812 || span.union_doc_freq > row_count
1813 {
1814 return Err(SealedSegmentError::Geometry(
1815 "posting-list summary does not match span",
1816 ));
1817 }
1818 if reader
1819 .blocks()
1820 .iter()
1821 .any(|block| block.last_docid >= row_count)
1822 {
1823 return Err(SealedSegmentError::Geometry(
1824 "posting docid exceeds row count",
1825 ));
1826 }
1827 decoded_lists.push(reader.decode_all()?);
1828 }
1829 validate_union_summaries(spans, terms, &decoded_lists)?;
1830 Ok(())
1831}
1832
1833fn validate_union_summaries(
1834 spans: &[ListSpan],
1835 terms: &[u8],
1836 lists: &[PostingList],
1837) -> Result<(), SealedSegmentError> {
1838 if spans.len() != lists.len() {
1839 return Err(SealedSegmentError::Geometry(
1840 "decoded posting list count differs from spans",
1841 ));
1842 }
1843 let term_of = |span: &ListSpan| -> Result<&[u8], SealedSegmentError> {
1844 let start = usize::try_from(span.term_start)
1845 .map_err(|_| SealedSegmentError::Geometry("term start exceeds usize"))?;
1846 let length = usize::try_from(span.term_len)
1847 .map_err(|_| SealedSegmentError::Geometry("term length exceeds usize"))?;
1848 let end = start
1849 .checked_add(length)
1850 .ok_or(SealedSegmentError::Geometry("term span overflows"))?;
1851 terms
1852 .get(start..end)
1853 .ok_or(SealedSegmentError::Geometry("term span is out of bounds"))
1854 };
1855 let mut start = 0_usize;
1856 while start < spans.len() {
1857 let first = spans
1858 .get(start)
1859 .ok_or(SealedSegmentError::Geometry("term group start is missing"))?;
1860 let term = term_of(first)?;
1861 let mut end = start.saturating_add(1);
1862 while let Some(span) = spans.get(end) {
1863 if term_of(span)? != term {
1864 break;
1865 }
1866 end = end.saturating_add(1);
1867 }
1868 let field_count = u32::try_from(end.saturating_sub(start))
1869 .map_err(|_| SealedSegmentError::Geometry("term field count exceeds u32"))?;
1870 let mut documents = std::collections::BTreeSet::new();
1871 for list in lists.get(start..end).ok_or(SealedSegmentError::Geometry(
1872 "term posting-list group is missing",
1873 ))? {
1874 documents.extend(list.postings().iter().map(|posting| posting.docid));
1875 }
1876 let union = u32::try_from(documents.len())
1877 .map_err(|_| SealedSegmentError::Geometry("union document count exceeds u32"))?;
1878 if spans
1879 .get(start..end)
1880 .ok_or(SealedSegmentError::Geometry("term span group is missing"))?
1881 .iter()
1882 .any(|span| span.field_count != field_count || span.union_doc_freq != union)
1883 {
1884 return Err(SealedSegmentError::Geometry(
1885 "term union summary does not match posting lists",
1886 ));
1887 }
1888 start = end;
1889 }
1890 Ok(())
1891}
1892
1893struct RegionCursor<'a> {
1894 bytes: &'a [u8],
1895 offset: usize,
1896}
1897
1898impl<'a> RegionCursor<'a> {
1899 const fn new(bytes: &'a [u8]) -> Self {
1900 Self { bytes, offset: 0 }
1901 }
1902
1903 fn take(&mut self, length: usize) -> Result<&'a [u8], SealedSegmentError> {
1904 let end = self
1905 .offset
1906 .checked_add(length)
1907 .ok_or(SealedSegmentError::Geometry("cursor offset overflow"))?;
1908 let value = self
1909 .bytes
1910 .get(self.offset..end)
1911 .ok_or(SealedSegmentError::Geometry("region is truncated"))?;
1912 self.offset = end;
1913 Ok(value)
1914 }
1915
1916 fn take_array<const N: usize>(&mut self) -> Result<[u8; N], SealedSegmentError> {
1917 self.take(N)?
1918 .try_into()
1919 .map_err(|_| SealedSegmentError::Geometry("region is truncated"))
1920 }
1921
1922 fn read_u16(&mut self) -> Result<u16, SealedSegmentError> {
1923 self.take_array().map(u16::from_le_bytes)
1924 }
1925
1926 fn read_u32(&mut self) -> Result<u32, SealedSegmentError> {
1927 self.take_array().map(u32::from_le_bytes)
1928 }
1929
1930 fn read_u64(&mut self) -> Result<u64, SealedSegmentError> {
1931 self.take_array().map(u64::from_le_bytes)
1932 }
1933
1934 fn finish(self) -> Result<(), SealedSegmentError> {
1935 if self.offset == self.bytes.len() {
1936 Ok(())
1937 } else {
1938 Err(SealedSegmentError::Geometry("region has trailing bytes"))
1939 }
1940 }
1941}
1942
1943const HEADER_LEN_U32: u32 = HEADER_LEN as u32;
1945
1946const META_LEN_U32: u32 = BLOCK_META_LEN as u32;
1948
1949const NO_BLOCK: usize = usize::MAX;
1951
1952fn stream_end(offset: u32, count: u16, bits: u8) -> u32 {
1954 let packed = u32::from(count).saturating_mul(u32::from(bits)).div_ceil(8);
1955 offset.saturating_add(packed)
1956}
1957
1958fn finish_union_group<E>(
1965 spans: &mut [ListSpan],
1966 start: usize,
1967 lists: &[&PostingList],
1968 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1969) -> Result<(), E> {
1970 let end = spans.len();
1971 if end <= start {
1972 return Ok(());
1973 }
1974 let count = u32::try_from(end.saturating_sub(start)).unwrap_or(u32::MAX);
1975 let union = if count == 1 {
1976 spans.get(start).map_or(0, |span| span.doc_freq)
1977 } else {
1978 union_postings_controlled(lists, work)?
1979 };
1980 if let Some(range) = spans.get_mut(start..end) {
1981 for span in range {
1982 work.step()?;
1983 span.union_doc_freq = union;
1984 span.field_count = count;
1985 }
1986 }
1987 Ok(())
1988}
1989
1990fn union_postings_controlled<E>(
1992 lists: &[&PostingList],
1993 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
1994) -> Result<u32, E> {
1995 work.check_now()?;
1996 let mut positions = vec![0_usize; lists.len()];
1997 let mut distinct = 0_u32;
1998 loop {
1999 let mut lowest: Option<u32> = None;
2000 for (slot, list) in lists.iter().enumerate() {
2001 work.step()?;
2002 #[cfg(test)]
2003 SEAL_UNION_PROBES.with(|value| value.set(value.get() + 1));
2004 let position = positions.get(slot).copied().unwrap_or(usize::MAX);
2005 if let Some(posting) = list.postings().get(position) {
2006 lowest = Some(lowest.map_or(posting.docid, |low| low.min(posting.docid)));
2007 }
2008 }
2009 let Some(low) = lowest else {
2010 break;
2011 };
2012 distinct = distinct.saturating_add(1);
2013 for (slot, list) in lists.iter().enumerate() {
2014 work.step()?;
2015 let Some(position) = positions.get_mut(slot) else {
2016 continue;
2017 };
2018 if list
2019 .postings()
2020 .get(*position)
2021 .is_some_and(|posting| posting.docid == low)
2022 {
2023 *position = position.saturating_add(1);
2024 }
2025 }
2026 }
2027 work.check_now()?;
2028 Ok(distinct)
2029}
2030
2031#[cfg(test)]
2032std::thread_local! {
2033 pub(crate) static SEAL_LENGTH_PROBES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
2034 static SEAL_TOTAL_PROBES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
2035 static SEAL_UNION_PROBES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
2036}
2037
2038fn union_count_controlled<E>(
2040 cursors: &mut [ListCursor<'_>],
2041 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2042) -> Result<u32, E> {
2043 for cursor in cursors.iter_mut() {
2044 work.step()?;
2045 cursor.reset();
2046 }
2047 let mut distinct = 0_u32;
2048 loop {
2049 let mut first = None;
2050 for cursor in cursors.iter() {
2051 work.step()?;
2052 if let Some(row) = cursor.current() {
2053 first = Some(first.map_or(row, |held: u32| held.min(row)));
2054 }
2055 }
2056 let Some(row) = first else {
2057 break;
2058 };
2059 distinct = distinct.saturating_add(1);
2060 for cursor in cursors.iter_mut() {
2061 work.step()?;
2062 if cursor.current() == Some(row) {
2063 cursor.advance();
2064 }
2065 }
2066 }
2067 Ok(distinct)
2068}
2069
2070#[derive(Clone, Debug)]
2076pub struct ListCursor<'segment> {
2077 meta: &'segment [u8],
2078 docids: &'segment [u8],
2079 tfs: &'segment [u8],
2080 block_count: usize,
2081 per_block: usize,
2082 weight: u64,
2083 block: usize,
2084 slot: usize,
2085 count: usize,
2086 impact: BlockImpact,
2089 tfs_block: usize,
2096 tf_offset: u32,
2098 tf_bits: u8,
2100 overall: BlockImpact,
2102 decoded_docids: Vec<u32>,
2103 decoded_tfs: Vec<u32>,
2104 blocks_decoded: u64,
2105 blocks_skipped: u64,
2106 tf_blocks_decoded: u64,
2107}
2108
2109impl ListCursor<'_> {
2110 fn meta_at(&self, index: usize) -> Option<BlockMeta> {
2112 let start = index.checked_mul(BLOCK_META_LEN)?;
2113 let row = self.meta.get(start..start.checked_add(BLOCK_META_LEN)?)?;
2114 BlockMeta::read(row).ok()
2115 }
2116
2117 fn last_docid(&self, index: usize) -> Option<u32> {
2119 let start = index.checked_mul(BLOCK_META_LEN)?;
2120 let bytes = self.meta.get(start..start.checked_add(4)?)?;
2121 let mut buffer = [0_u8; 4];
2122 buffer.copy_from_slice(bytes);
2123 Some(u32::from_le_bytes(buffer))
2124 }
2125
2126 fn impact_at(&self, index: usize) -> BlockImpact {
2128 self.meta_at(index)
2129 .map_or_else(BlockImpact::default, |meta| BlockImpact::from_meta(&meta))
2130 }
2131
2132 fn load(&mut self, index: usize) -> bool {
2137 if self.block == index {
2138 return true;
2139 }
2140 let Some(meta) = self.meta_at(index) else {
2141 return false;
2142 };
2143 let count = usize::from(meta.count).min(self.per_block);
2144 let docid_start = usize::try_from(meta.docids_offset).unwrap_or(usize::MAX);
2145 let Some(packed) = self.docids.get(docid_start..) else {
2146 return false;
2147 };
2148 let base = if index == 0 {
2151 0
2152 } else {
2153 self.last_docid(index.saturating_sub(1)).unwrap_or(0)
2154 };
2155 let Some(target) = self.decoded_docids.get_mut(..count) else {
2156 return false;
2157 };
2158 if unpack(packed, meta.docid_bits, count, target).is_none() {
2159 return false;
2160 }
2161 prefix_sum(target, base);
2162
2163 self.tf_offset = meta.tfs_offset;
2166 self.tf_bits = meta.tf_bits;
2167 self.block = index;
2168 self.count = count;
2169 self.impact = BlockImpact::from_meta(&meta);
2170 self.blocks_decoded = self.blocks_decoded.saturating_add(1);
2171 true
2172 }
2173
2174 pub fn reset(&mut self) {
2180 self.slot = 0;
2181 if self.block == 0 {
2182 return;
2183 }
2184 self.block = NO_BLOCK;
2185 self.count = 0;
2186 if self.block_count > 0 {
2187 self.load(0);
2188 }
2189 }
2190
2191 #[must_use]
2193 pub fn current(&self) -> Option<u32> {
2194 if self.block == NO_BLOCK {
2195 return None;
2196 }
2197 self.decoded_docids.get(self.slot).copied()
2198 }
2199
2200 fn ensure_tfs(&mut self) -> bool {
2202 if self.tfs_block == self.block {
2203 return true;
2204 }
2205 let start = usize::try_from(self.tf_offset).unwrap_or(usize::MAX);
2206 let Some(packed) = self.tfs.get(start..) else {
2207 return false;
2208 };
2209 let Some(target) = self.decoded_tfs.get_mut(..self.count) else {
2210 return false;
2211 };
2212 if unpack(packed, self.tf_bits, self.count, target).is_none() {
2213 return false;
2214 }
2215 self.tfs_block = self.block;
2216 self.tf_blocks_decoded = self.tf_blocks_decoded.saturating_add(1);
2217 true
2218 }
2219
2220 #[must_use]
2224 pub fn current_tf(&mut self) -> Option<u32> {
2225 if self.block == NO_BLOCK || !self.ensure_tfs() {
2226 return None;
2227 }
2228 self.decoded_tfs.get(self.slot).copied()
2229 }
2230
2231 #[must_use]
2236 pub fn current_weighted_tf(&mut self) -> Option<u64> {
2237 self.current_tf()
2238 .map(|tf| u64::from(tf).saturating_mul(self.weight))
2239 }
2240
2241 pub fn advance(&mut self) {
2243 if self.block == NO_BLOCK {
2244 return;
2245 }
2246 self.slot = self.slot.saturating_add(1);
2247 if self.slot < self.count {
2248 return;
2249 }
2250 let next = self.block.saturating_add(1);
2251 if next >= self.block_count || !self.load(next) {
2252 self.block = NO_BLOCK;
2253 return;
2254 }
2255 self.slot = 0;
2256 }
2257
2258 pub fn seek(&mut self, row: u32) {
2265 if self.block == NO_BLOCK {
2266 return;
2267 }
2268 if self.current().is_some_and(|current| current >= row) {
2269 return;
2270 }
2271 if self.last_docid(self.block).is_some_and(|last| last < row) {
2272 let mut low = self.block.saturating_add(1);
2273 let mut high = self.block_count;
2274 while low < high {
2275 let middle = low.saturating_add(high.saturating_sub(low) / 2);
2276 if self.last_docid(middle).is_some_and(|last| last < row) {
2277 low = middle.saturating_add(1);
2278 } else {
2279 high = middle;
2280 }
2281 }
2282 let hopped = low.saturating_sub(self.block.saturating_add(1));
2283 self.blocks_skipped = self
2284 .blocks_skipped
2285 .saturating_add(u64::try_from(hopped).unwrap_or(u64::MAX));
2286 if low >= self.block_count || !self.load(low) {
2287 self.block = NO_BLOCK;
2288 return;
2289 }
2290 self.slot = 0;
2291 }
2292 loop {
2293 if self.block == NO_BLOCK {
2294 return;
2295 }
2296 let found = self
2299 .decoded_docids
2300 .get(self.slot..self.count)
2301 .map_or(0, |within| within.partition_point(|&docid| docid < row));
2302 let landed = self.slot.saturating_add(found);
2303 if landed < self.count {
2304 self.slot = landed;
2305 return;
2306 }
2307 self.slot = self.count.saturating_sub(1);
2310 self.advance();
2311 }
2312 }
2313
2314 #[must_use]
2316 pub const fn exhausted(&self) -> bool {
2317 self.block == NO_BLOCK
2318 }
2319
2320 #[must_use]
2324 pub const fn current_impact(&self) -> BlockImpact {
2325 if self.block == NO_BLOCK {
2326 return BlockImpact {
2327 max_tf: 0,
2328 min_len: 0,
2329 };
2330 }
2331 self.impact
2332 }
2333
2334 fn block_for(&self, row: u32) -> Option<usize> {
2340 if self.block == NO_BLOCK {
2341 return None;
2342 }
2343 if self.current().is_some_and(|current| current > row) {
2344 return None;
2347 }
2348 if self.last_docid(self.block).is_some_and(|last| last >= row) {
2349 return Some(self.block);
2350 }
2351 let mut low = self.block.saturating_add(1);
2352 let mut high = self.block_count;
2353 while low < high {
2354 let middle = low.saturating_add(high.saturating_sub(low) / 2);
2355 if self.last_docid(middle).is_some_and(|last| last < row) {
2356 low = middle.saturating_add(1);
2357 } else {
2358 high = middle;
2359 }
2360 }
2361 if low >= self.block_count {
2362 return None;
2363 }
2364 Some(low)
2365 }
2366
2367 #[must_use]
2375 pub fn impact_for(&self, row: u32) -> Option<BlockImpact> {
2376 let index = self.block_for(row)?;
2377 if index == self.block {
2378 return Some(self.current_impact());
2379 }
2380 Some(self.impact_at(index))
2381 }
2382
2383 fn impact_horizon_for(&self, row: u32) -> Option<u32> {
2394 if self.block == NO_BLOCK {
2395 return None;
2396 }
2397 if let Some(current) = self.current()
2398 && current > row
2399 {
2400 return Some(current.saturating_sub(1));
2401 }
2402 self.block_for(row).and_then(|index| self.last_docid(index))
2403 }
2404
2405 #[must_use]
2409 pub const fn overall_impact(&self) -> BlockImpact {
2410 self.overall
2411 }
2412}
2413
2414#[derive(Clone, Debug)]
2420pub struct TermStream<'segment> {
2421 runs: Vec<ListCursor<'segment>>,
2422 weights: Vec<u64>,
2423 unit_single: bool,
2427 heads_always_valid: bool,
2432 head: Option<u32>,
2433 head_tf: Option<u32>,
2435}
2436
2437#[cfg(test)]
2438thread_local! {
2439 static ZERO_TF_ROWS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
2440}
2441
2442impl<'segment> TermStream<'segment> {
2443 pub(crate) fn open_controlled<E>(
2444 segment: &'segment SealedSegment,
2445 term: &[u8],
2446 weights: &FieldWeights,
2447 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2448 ) -> Result<Option<Self>, E> {
2449 Self::open_with_capacity_controlled(segment, term, weights, 0, work)
2450 }
2451
2452 #[must_use]
2456 pub fn open(
2457 segment: &'segment SealedSegment,
2458 term: &[u8],
2459 weights: &FieldWeights,
2460 ) -> Option<Self> {
2461 Self::open_with_capacity(segment, term, weights, 0)
2462 }
2463
2464 pub(crate) fn allocation_bytes(
2467 segment: &SealedSegment,
2468 term: &[u8],
2469 weights: &FieldWeights,
2470 ) -> Option<usize> {
2471 Self::run_count(segment, term, weights).checked_mul(
2472 std::mem::size_of::<ListCursor<'_>>()
2473 .checked_add(std::mem::size_of::<u64>())?
2474 .checked_add(
2475 usize::from(segment.postings_per_block)
2476 .checked_mul(2 * std::mem::size_of::<u32>())?,
2477 )?,
2478 )
2479 }
2480
2481 fn run_count(segment: &SealedSegment, term: &[u8], weights: &FieldWeights) -> usize {
2482 let (start, end) = segment.term_span_range(term);
2483 segment
2484 .spans
2485 .get(start..end)
2486 .unwrap_or_default()
2487 .iter()
2488 .filter(|span| weights.weight(span.field) != 0)
2489 .count()
2490 }
2491
2492 #[cfg(test)]
2493 pub(crate) fn open_sized(
2494 segment: &'segment SealedSegment,
2495 term: &[u8],
2496 weights: &FieldWeights,
2497 ) -> Option<Self> {
2498 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2499 match Self::open_sized_controlled(segment, term, weights, &mut work) {
2500 Ok(stream) => stream,
2501 Err(never) => match never {},
2502 }
2503 }
2504
2505 pub(crate) fn open_sized_controlled<E>(
2506 segment: &'segment SealedSegment,
2507 term: &[u8],
2508 weights: &FieldWeights,
2509 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2510 ) -> Result<Option<Self>, E> {
2511 let (start, end) = segment.term_span_range(term);
2512 let mut capacity = 0;
2513 for span in segment.spans.get(start..end).unwrap_or_default() {
2514 work.step()?;
2515 capacity += usize::from(weights.weight(span.field) != 0);
2516 }
2517 Self::open_with_capacity_controlled(segment, term, weights, capacity, work)
2518 }
2519
2520 fn open_with_capacity(
2521 segment: &'segment SealedSegment,
2522 term: &[u8],
2523 weights: &FieldWeights,
2524 capacity: usize,
2525 ) -> Option<Self> {
2526 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2527 match Self::open_with_capacity_controlled(segment, term, weights, capacity, &mut work) {
2528 Ok(stream) => stream,
2529 Err(never) => match never {},
2530 }
2531 }
2532
2533 fn open_with_capacity_controlled<E>(
2534 segment: &'segment SealedSegment,
2535 term: &[u8],
2536 weights: &FieldWeights,
2537 capacity: usize,
2538 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2539 ) -> Result<Option<Self>, E> {
2540 work.step()?;
2541 let (start, end) = segment.term_span_range(term);
2542 let mut runs: Vec<ListCursor<'segment>> = Vec::with_capacity(capacity);
2543 let mut scales: Vec<u64> = Vec::with_capacity(capacity);
2544 for index in start..end {
2545 work.step()?;
2546 let Some(span) = segment.spans.get(index) else {
2547 return Ok(None);
2548 };
2549 let weight = u64::from(weights.weight(span.field));
2550 if weight == 0 {
2551 continue;
2552 }
2553 let Some(cursor) = segment.cursor_at(index, weight) else {
2554 return Ok(None);
2555 };
2556 runs.push(cursor);
2557 scales.push(weight);
2558 }
2559 if runs.is_empty() {
2560 return Ok(None);
2561 }
2562 let unit_single = runs.len() == 1 && scales.first().copied() == Some(1_000);
2563 let heads_always_valid = scales.iter().all(|weight| *weight >= 1_000);
2564 let mut stream = Self {
2565 runs,
2566 weights: scales,
2567 unit_single,
2568 heads_always_valid,
2569 head: None,
2570 head_tf: None,
2571 };
2572 stream.reset_controlled(work)?;
2573 Ok(Some(stream))
2574 }
2575
2576 pub fn reset(&mut self) {
2578 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2579 match self.reset_controlled(&mut work) {
2580 Ok(()) => {}
2581 Err(never) => match never {},
2582 }
2583 }
2584
2585 pub(crate) fn reset_controlled<E>(
2586 &mut self,
2587 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2588 ) -> Result<(), E> {
2589 for run in &mut self.runs {
2590 work.step()?;
2591 run.reset();
2592 }
2593 self.refresh_controlled(work)
2594 }
2595
2596 fn refresh_controlled<E>(
2603 &mut self,
2604 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2605 ) -> Result<(), E> {
2606 self.head_tf = None;
2607 if self.heads_always_valid {
2608 self.head = if self.unit_single {
2609 self.runs.first().and_then(ListCursor::current)
2610 } else {
2611 self.first_row_controlled(work)?
2612 };
2613 return Ok(());
2614 }
2615 loop {
2616 let Some(row) = self.first_row_controlled(work)? else {
2617 self.head = None;
2618 return Ok(());
2619 };
2620 let mut total = 0_u64;
2621 for run in &mut self.runs {
2622 work.step()?;
2623 if run.current() == Some(row) {
2624 total = total.saturating_add(run.current_weighted_tf().unwrap_or(0));
2625 }
2626 }
2627 let tf = u32::try_from(total / 1_000).unwrap_or(u32::MAX);
2628 if tf > 0 {
2629 self.head = Some(row);
2630 self.head_tf = Some(tf);
2631 return Ok(());
2632 }
2633 #[cfg(test)]
2634 ZERO_TF_ROWS.with(|count| count.set(count.get() + 1));
2635 self.step_controlled(row, work)?;
2636 }
2637 }
2638
2639 fn first_row_controlled<E>(
2640 &self,
2641 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2642 ) -> Result<Option<u32>, E> {
2643 let mut first = None;
2644 for run in &self.runs {
2645 work.step()?;
2646 if let Some(row) = run.current() {
2647 first = Some(first.map_or(row, |held: u32| held.min(row)));
2648 }
2649 }
2650 Ok(first)
2651 }
2652
2653 fn step_controlled<E>(
2655 &mut self,
2656 row: u32,
2657 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2658 ) -> Result<(), E> {
2659 for run in &mut self.runs {
2660 work.step()?;
2661 if run.current() == Some(row) {
2662 run.advance();
2663 }
2664 }
2665 Ok(())
2666 }
2667
2668 #[must_use]
2670 pub const fn current_row(&self) -> Option<u32> {
2671 self.head
2672 }
2673
2674 #[must_use]
2677 pub fn current_tf(&mut self) -> Option<u32> {
2678 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2679 match self.current_tf_controlled(&mut work) {
2680 Ok(tf) => tf,
2681 Err(never) => match never {},
2682 }
2683 }
2684
2685 pub(crate) fn current_tf_controlled<E>(
2686 &mut self,
2687 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2688 ) -> Result<Option<u32>, E> {
2689 let Some(row) = self.head else {
2690 return Ok(None);
2691 };
2692 if let Some(tf) = self.head_tf {
2693 return Ok(Some(tf));
2694 }
2695 let mut total = 0_u64;
2696 for run in &mut self.runs {
2697 work.step()?;
2698 if run.current() == Some(row) {
2699 total = total.saturating_add(run.current_weighted_tf().unwrap_or(0));
2700 }
2701 }
2702 let tf = u32::try_from(total / 1_000).unwrap_or(u32::MAX);
2703 self.head_tf = Some(tf);
2704 Ok(Some(tf))
2705 }
2706
2707 pub fn advance(&mut self) {
2709 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2710 match self.advance_controlled(&mut work) {
2711 Ok(()) => {}
2712 Err(never) => match never {},
2713 }
2714 }
2715
2716 pub(crate) fn advance_controlled<E>(
2717 &mut self,
2718 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2719 ) -> Result<(), E> {
2720 let Some(row) = self.head else {
2721 return Ok(());
2722 };
2723 self.step_controlled(row, work)?;
2724 self.refresh_controlled(work)
2725 }
2726
2727 pub fn seek(&mut self, row: u32) {
2729 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2730 match self.seek_controlled(row, &mut work) {
2731 Ok(()) => {}
2732 Err(never) => match never {},
2733 }
2734 }
2735
2736 pub(crate) fn seek_controlled<E>(
2737 &mut self,
2738 row: u32,
2739 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2740 ) -> Result<(), E> {
2741 for run in &mut self.runs {
2742 work.step()?;
2743 run.seek(row);
2744 }
2745 self.refresh_controlled(work)
2746 }
2747
2748 #[must_use]
2750 pub const fn exhausted(&self) -> bool {
2751 self.head.is_none()
2752 }
2753
2754 #[must_use]
2756 pub fn blocks_decoded(&self) -> u64 {
2757 self.runs
2758 .iter()
2759 .map(|run| run.blocks_decoded)
2760 .fold(0_u64, u64::saturating_add)
2761 }
2762
2763 #[must_use]
2765 pub fn blocks_skipped(&self) -> u64 {
2766 self.runs
2767 .iter()
2768 .map(|run| run.blocks_skipped)
2769 .fold(0_u64, u64::saturating_add)
2770 }
2771
2772 #[must_use]
2777 pub fn tf_blocks_decoded(&self) -> u64 {
2778 self.runs
2779 .iter()
2780 .map(|run| run.tf_blocks_decoded)
2781 .fold(0_u64, u64::saturating_add)
2782 }
2783
2784 #[must_use]
2788 pub fn posting_count(&self) -> usize {
2789 self.runs
2790 .iter()
2791 .map(|run| {
2792 run.block_count
2793 .saturating_sub(1)
2794 .saturating_mul(run.per_block)
2795 .saturating_add(
2796 run.meta_at(run.block_count.saturating_sub(1))
2797 .map_or(0, |meta| usize::from(meta.count)),
2798 )
2799 })
2800 .fold(0_usize, usize::saturating_add)
2801 }
2802
2803 pub(crate) fn block_bound_controlled<E>(
2804 &self,
2805 scorer: &TermScorer,
2806 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2807 ) -> Result<f64, E> {
2808 self.bound_from_controlled(
2809 scorer,
2810 |run| (!run.exhausted()).then(|| run.current_impact()),
2811 work,
2812 )
2813 }
2814
2815 pub(crate) fn upper_bound_controlled<E>(
2816 &self,
2817 scorer: &TermScorer,
2818 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2819 ) -> Result<f64, E> {
2820 self.bound_from_controlled(
2821 scorer,
2822 |run| (!run.exhausted()).then(|| run.overall_impact()),
2823 work,
2824 )
2825 }
2826
2827 pub(crate) fn bound_for_controlled<E>(
2828 &self,
2829 row: u32,
2830 scorer: &TermScorer,
2831 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2832 ) -> Result<f64, E> {
2833 self.bound_from_controlled(scorer, |run| run.impact_for(row), work)
2834 }
2835
2836 pub(crate) fn bound_horizon_for_controlled<E>(
2837 &self,
2838 row: u32,
2839 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2840 ) -> Result<Option<u32>, E> {
2841 let mut edge = None;
2842 for run in &self.runs {
2843 work.step()?;
2844 if let Some(found) = run.impact_horizon_for(row) {
2845 edge = Some(edge.map_or(found, |held: u32| held.min(found)));
2846 }
2847 }
2848 Ok(edge)
2849 }
2850
2851 #[must_use]
2856 pub fn block_bound(&self, scorer: &TermScorer) -> f64 {
2857 self.bound_from(scorer, |run| {
2858 (!run.exhausted()).then(|| run.current_impact())
2859 })
2860 }
2861
2862 #[must_use]
2871 pub fn bound_for(&self, row: u32, scorer: &TermScorer) -> f64 {
2872 self.bound_from(scorer, |run| run.impact_for(row))
2873 }
2874
2875 #[must_use]
2894 pub fn bound_horizon_for(&self, row: u32) -> Option<u32> {
2895 self.runs
2896 .iter()
2897 .filter_map(|run| run.impact_horizon_for(row))
2898 .min()
2899 }
2900
2901 #[must_use]
2903 pub fn upper_bound(&self, scorer: &TermScorer) -> f64 {
2904 self.bound_from(scorer, |run| {
2905 (!run.exhausted()).then(|| run.overall_impact())
2906 })
2907 }
2908
2909 fn bound_from(
2914 &self,
2915 scorer: &TermScorer,
2916 pick: impl Fn(&ListCursor<'segment>) -> Option<BlockImpact>,
2917 ) -> f64 {
2918 let mut work = super::control::WorkCheck::new(|| Ok::<(), std::convert::Infallible>(()));
2919 match self.bound_from_controlled(scorer, pick, &mut work) {
2920 Ok(bound) => bound,
2921 Err(never) => match never {},
2922 }
2923 }
2924
2925 fn bound_from_controlled<E>(
2926 &self,
2927 scorer: &TermScorer,
2928 pick: impl Fn(&ListCursor<'segment>) -> Option<BlockImpact>,
2929 work: &mut super::control::WorkCheck<impl FnMut() -> Result<(), E>>,
2930 ) -> Result<f64, E> {
2931 let mut max_tf = 0_u64;
2932 let mut min_len = u64::MAX;
2933 for (run, weight) in self.runs.iter().zip(self.weights.iter()) {
2934 work.step()?;
2935 let Some(impact) = pick(run) else {
2936 continue;
2937 };
2938 max_tf = max_tf.saturating_add(
2941 u64::from(impact.max_tf)
2942 .saturating_mul(*weight)
2943 .div_ceil(1_000),
2944 );
2945 min_len = min_len.min(u64::from(impact.min_len).saturating_mul(*weight) / 1_000);
2946 }
2947 if min_len == u64::MAX {
2948 return Ok(0.0);
2949 }
2950 Ok(scorer.score(
2951 super::bm25::Tf(u32::try_from(max_tf).unwrap_or(u32::MAX)),
2952 super::bm25::DocLen(u32::try_from(min_len).unwrap_or(u32::MAX)),
2953 ))
2954 }
2955}
2956
2957#[cfg(test)]
2958#[allow(
2959 clippy::expect_used,
2960 clippy::indexing_slicing,
2961 clippy::panic,
2962 clippy::unwrap_used
2963)]
2964mod tests {
2965 use super::*;
2966 use crate::fts::index::{DEFAULT_FIELD, Document};
2967 use crate::fts::postings::{Posting, PostingList};
2968 use crate::fts::tokenizer::{Analyzer, Profile};
2969
2970 fn analyzer() -> Analyzer {
2971 Analyzer::new(Profile::Code.config()).expect("valid config")
2972 }
2973
2974 fn segment_of(texts: &[String]) -> SegmentIndex {
2975 let analyzer = analyzer();
2976 let mut segment = SegmentIndex::new();
2977 for text in texts {
2978 segment
2979 .push_document(&analyzer, &Document::with_text(text))
2980 .expect("indexable");
2981 }
2982 segment
2983 }
2984
2985 #[test]
2986 fn astra_18_sealing_cancels_inside_row_preparation() {
2987 let source = segment_of(&vec!["alpha beta".to_owned(); 4_096]);
2988 let reference = SealedSegment::seal(&source)
2989 .expect("reference")
2990 .encode_region()
2991 .expect("bytes");
2992 let mut observed = Vec::new();
2993 for (phase, probes) in [
2994 ("length copy", &SEAL_LENGTH_PROBES),
2995 ("total lengths", &SEAL_TOTAL_PROBES),
2996 ] {
2997 probes.with(|value| value.set(0));
2998 let mut work = super::super::control::WorkCheck::new(|| {
2999 if probes.with(std::cell::Cell::get) >= 64 {
3000 Err(PostingsError::ZeroBlockSize)
3001 } else {
3002 Ok(())
3003 }
3004 });
3005 let result = SealedSegment::seal_controlled(&source, &mut work);
3006 let count = probes.with(std::cell::Cell::get);
3007 println!("{phase}: canceled={}, rows={count}", result.is_err());
3008 observed.push((
3009 phase,
3010 matches!(result, Err(PostingsError::ZeroBlockSize)),
3011 count,
3012 ));
3013 }
3014 let mut work = super::super::control::WorkCheck::new(|| Ok::<(), PostingsError>(()));
3015 assert_eq!(
3016 SealedSegment::seal_controlled(&source, &mut work)
3017 .expect("clean")
3018 .encode_region()
3019 .expect("clean bytes"),
3020 reference
3021 );
3022 for (phase, canceled, count) in observed {
3023 assert!(canceled, "{phase}");
3024 assert!((64..=128).contains(&count), "{phase}: {count}");
3025 }
3026 }
3027
3028 #[test]
3029 fn astra_18_sealing_union_cancels_inside_merge() {
3030 let mut even = PostingList::new();
3031 let mut all = PostingList::new();
3032 for docid in 0..4_096 {
3033 let posting = Posting {
3034 docid,
3035 tf: 1,
3036 positions: vec![0],
3037 };
3038 all.push(posting.clone()).expect("all");
3039 if docid % 2 == 0 {
3040 even.push(posting).expect("even");
3041 }
3042 }
3043 SEAL_UNION_PROBES.with(|value| value.set(0));
3044 let mut work = super::super::control::WorkCheck::new(|| {
3045 if SEAL_UNION_PROBES.with(std::cell::Cell::get) >= 64 {
3046 Err("cancelled")
3047 } else {
3048 Ok(())
3049 }
3050 });
3051 let result = union_postings_controlled(&[&even, &all], &mut work);
3052 let probes = SEAL_UNION_PROBES.with(std::cell::Cell::get);
3053 println!("union probes={probes}");
3054 assert_eq!(result, Err("cancelled"));
3055 assert!((64..=128).contains(&probes));
3056 let mut work = super::super::control::WorkCheck::new(|| Ok::<(), ()>(()));
3057 assert_eq!(
3058 union_postings_controlled(&[&even, &all], &mut work),
3059 Ok(4_096)
3060 );
3061 }
3062
3063 #[test]
3064 fn astra_18_field_subset_df_cancellation_preserves_union() {
3065 let mut source = SegmentIndex::new();
3066 let analyzer = analyzer();
3067 for row in 0..4_096 {
3068 let mut document = Document::new();
3069 document.set(FieldId(0), if row % 2 == 0 { "alpha" } else { "beta" });
3070 document.set(FieldId(1), if row % 3 == 0 { "alpha" } else { "beta" });
3071 document.set(FieldId(2), "alpha");
3072 source.push_document(&analyzer, &document).unwrap();
3073 }
3074 let segment = SealedSegment::seal(&source).unwrap();
3075 let fields = [FieldId(0), FieldId(1)];
3076 let expected = (0..4_096)
3077 .filter(|row| row % 2 == 0 || row % 3 == 0)
3078 .count() as u32;
3079 assert_eq!(segment.document_frequency(b"alpha", &fields), expected);
3080 assert_eq!(
3081 segment.document_frequency_controlled(b"alpha", &fields, || Ok::<(), ()>(())),
3082 Ok(expected)
3083 );
3084 let mut checks = 0;
3085 assert_eq!(
3086 segment.document_frequency_controlled(b"alpha", &fields, || {
3087 checks += 1;
3088 if checks == 3 {
3089 Err("cancelled")
3090 } else {
3091 Ok(())
3092 }
3093 }),
3094 Err("cancelled")
3095 );
3096 assert_eq!(checks, 3);
3097 for (fields, expected) in [
3099 (vec![FieldId(0)], 2_048),
3100 (vec![FieldId(0), FieldId(1), FieldId(2)], 4_096),
3101 ] {
3102 assert_eq!(
3103 segment.document_frequency_controlled(b"alpha", &fields, || Ok::<(), ()>(())),
3104 Ok(expected)
3105 );
3106 }
3107 }
3108
3109 #[test]
3110 fn astra_18_fractional_stream_cancels_in_hidden_zero_tf_walk() {
3111 let analyzer = analyzer();
3112 let fields = FieldWeights::new(&[(DEFAULT_FIELD, 500)]);
3113 let mut observed = Vec::new();
3114 for action in ["open", "reset", "seek", "advance"] {
3115 let mut source = SegmentIndex::new();
3116 for row in 0..4_096 {
3117 let text = if row == 4_095 || (matches!(action, "advance" | "seek") && row == 0) {
3118 "alpha alpha"
3119 } else {
3120 "alpha"
3121 };
3122 source
3123 .push_document(&analyzer, &Document::with_text(text))
3124 .unwrap();
3125 }
3126 let segment = SealedSegment::seal(&source).unwrap();
3127 let mut stream = TermStream::open(&segment, b"alpha", &fields).unwrap();
3128 assert_eq!(
3129 stream.current_row(),
3130 Some(if matches!(action, "advance" | "seek") {
3131 0
3132 } else {
3133 4_095
3134 })
3135 );
3136 ZERO_TF_ROWS.with(|count| count.set(0));
3137 let mut checks = 0;
3138 let mut work = super::super::control::WorkCheck::new(|| {
3139 checks += 1;
3140 if checks == 3 {
3141 Err("cancelled")
3142 } else {
3143 Ok(())
3144 }
3145 });
3146 let result = match action {
3147 "open" => TermStream::open_controlled(&segment, b"alpha", &fields, &mut work)
3148 .map(|stream| stream.and_then(|stream| stream.current_row())),
3149 "reset" => stream
3150 .reset_controlled(&mut work)
3151 .map(|()| stream.current_row()),
3152 "seek" => stream
3153 .seek_controlled(1, &mut work)
3154 .map(|()| stream.current_row()),
3155 "advance" => stream
3156 .advance_controlled(&mut work)
3157 .map(|()| stream.current_row()),
3158 _ => unreachable!(),
3159 };
3160 let rows = ZERO_TF_ROWS.with(std::cell::Cell::get);
3161 println!("{action}: result={result:?}, checks={checks}, zero_tf_rows={rows}");
3162 observed.push((action, result, rows));
3163 }
3164 for (action, result, rows) in observed {
3167 assert_eq!(result, Err("cancelled"), "{action}");
3168 assert!(
3169 rows > 0 && rows <= 128,
3170 "{action} walked {rows} zero-TF rows"
3171 );
3172 }
3173 }
3174
3175 #[test]
3176 fn a_sealed_cursor_reproduces_every_posting_of_the_active_segment() {
3177 let texts: Vec<String> = (0..500)
3180 .map(|index| {
3181 if index % 3 == 0 {
3182 format!("alpha filler{index}")
3183 } else {
3184 format!("beta filler{index}")
3185 }
3186 })
3187 .collect();
3188 let active = segment_of(&texts);
3189 let sealed = SealedSegment::seal(&active).expect("seals");
3190
3191 for term in [b"alpha".as_slice(), b"beta".as_slice()] {
3192 let list = active.posting_list(term, DEFAULT_FIELD).expect("present");
3193 let mut cursor = sealed
3194 .cursor_at(sealed.term_span_range(term).0, 1_000)
3195 .expect("span");
3196 cursor.reset();
3197 for posting in list.postings() {
3198 assert_eq!(cursor.current(), Some(posting.docid));
3199 assert_eq!(
3200 cursor.current_weighted_tf(),
3201 Some(u64::from(posting.tf) * 1_000)
3202 );
3203 cursor.advance();
3204 }
3205 assert!(cursor.exhausted(), "the cursor must end with the list");
3206 }
3207 }
3208
3209 #[test]
3210 fn a_landed_block_pays_its_tf_decode_only_when_scored() {
3211 let texts: Vec<String> = (0..400).map(|index| format!("alpha d{index}")).collect();
3215 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
3216 let mut stream = TermStream::open(
3217 &sealed,
3218 b"alpha",
3219 &crate::fts::search::FieldWeights::flat(&[DEFAULT_FIELD]),
3220 )
3221 .expect("stream");
3222 assert_eq!(
3223 stream.tf_blocks_decoded(),
3224 0,
3225 "opening must not touch tf bytes"
3226 );
3227 stream.seek(200);
3228 stream.seek(383);
3229 assert!(stream.blocks_decoded() >= 3, "three landings were paid");
3230 assert_eq!(
3231 stream.tf_blocks_decoded(),
3232 0,
3233 "positioning must not touch tf bytes"
3234 );
3235 assert_eq!(stream.current_row(), Some(383));
3236 assert_eq!(stream.current_tf(), Some(1));
3237 assert_eq!(
3238 stream.tf_blocks_decoded(),
3239 1,
3240 "scoring pays exactly the one block it reads"
3241 );
3242 }
3243
3244 #[test]
3245 fn seeking_jumps_whole_blocks_without_decoding_them() {
3246 let texts: Vec<String> = (0..400).map(|index| format!("alpha d{index}")).collect();
3247 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
3248 let mut cursor = sealed
3249 .cursor_at(sealed.term_span_range(b"alpha").0, 1_000)
3250 .expect("span");
3251 cursor.reset();
3252 let decoded_before = cursor.blocks_decoded;
3253 cursor.seek(383);
3254 assert_eq!(cursor.current(), Some(383));
3255 assert!(
3256 cursor.blocks_skipped >= 4,
3257 "a 400-posting list at 64 per block must skip at least four"
3258 );
3259 assert_eq!(
3260 cursor.blocks_decoded,
3261 decoded_before + 1,
3262 "a seek must decode exactly the block it lands in"
3263 );
3264 }
3265
3266 #[test]
3267 fn seeking_past_the_end_exhausts_the_cursor() {
3268 let texts: Vec<String> = (0..100).map(|index| format!("alpha d{index}")).collect();
3269 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
3270 let mut cursor = sealed
3271 .cursor_at(sealed.term_span_range(b"alpha").0, 1_000)
3272 .expect("span");
3273 cursor.reset();
3274 cursor.seek(9_999);
3275 assert!(cursor.exhausted());
3276 assert_eq!(cursor.current(), None);
3277 assert_eq!(cursor.current_weighted_tf(), None);
3278 }
3279
3280 #[test]
3281 fn an_empty_segment_seals_and_answers_nothing() {
3282 let sealed = SealedSegment::seal(&SegmentIndex::new()).expect("seals");
3283 assert!(sealed.is_empty());
3284 assert_eq!(sealed.row_count(), 0);
3285 assert_eq!(sealed.total_tokens(), 0);
3286 assert_eq!(sealed.document_length(0), 0);
3287 assert_eq!(sealed.field_length(0, DEFAULT_FIELD), 0);
3288 assert!(sealed.field_lengths(DEFAULT_FIELD).is_none());
3289 assert_eq!(sealed.document_frequency(b"alpha", &[DEFAULT_FIELD]), 0);
3290 assert!(
3291 TermStream::open(&sealed, b"alpha", &FieldWeights::flat(&[DEFAULT_FIELD])).is_none()
3292 );
3293 }
3294
3295 #[test]
3296 fn a_zero_weight_field_contributes_no_run() {
3297 let texts: Vec<String> = (0..4).map(|index| format!("alpha d{index}")).collect();
3298 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
3299 let weights = FieldWeights::new(&[(DEFAULT_FIELD, 0)]);
3300 assert!(TermStream::open(&sealed, b"alpha", &weights).is_none());
3301 }
3302
3303 #[test]
3304 fn a_bound_over_one_field_is_the_stored_impact_pair() {
3305 let texts: Vec<String> = (0..300).map(|index| format!("alpha d{index}")).collect();
3306 let active = segment_of(&texts);
3307 let sealed = SealedSegment::seal(&active).expect("seals");
3308 let stats = crate::fts::bm25::CorpusStats::new(300, 600).expect("stats");
3309 let scorer = TermScorer::new(
3310 crate::fts::bm25::Df(300),
3311 &stats,
3312 crate::fts::bm25::Bm25Params::default(),
3313 );
3314 let stream = TermStream::open(&sealed, b"alpha", &FieldWeights::flat(&[DEFAULT_FIELD]))
3315 .expect("stream");
3316 let impact = stream.runs[0].current_impact();
3317 assert!((stream.block_bound(&scorer) - impact.bound(&scorer)).abs() < 1e-15);
3318 }
3319
3320 #[test]
3321 fn a_traversed_list_costs_a_couple_of_bytes_per_posting() {
3322 let texts: Vec<String> = (0..20_000)
3337 .map(|index| format!("alpha t{} filler{index}", index % 97))
3338 .collect();
3339 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
3340
3341 let mut long_postings = 0_usize;
3342 let mut long_bytes = 0_usize;
3343 let mut short_lists = 0_usize;
3344 for index in 0..sealed.spans.len() {
3345 let Some(span) = sealed.spans.get(index).copied() else {
3346 continue;
3347 };
3348 let cursor = sealed.cursor_at(index, 1_000).expect("span");
3349 if cursor.block_count <= 1 {
3350 short_lists = short_lists.saturating_add(1);
3351 continue;
3352 }
3353 long_postings =
3354 long_postings.saturating_add(usize::try_from(span.doc_freq).unwrap_or(0));
3355 let last = cursor.block_count.saturating_sub(1);
3360 let meta = cursor.meta_at(last).expect("last block");
3361 let tf_end = usize::try_from(meta.tfs_offset).unwrap_or(0)
3362 + (usize::from(meta.count) * usize::from(meta.tf_bits)).div_ceil(8);
3363 long_bytes = long_bytes
3364 .saturating_add(usize::try_from(span.tfs_start - span.docids_start).unwrap_or(0))
3365 .saturating_add(tf_end)
3366 .saturating_add(cursor.block_count.saturating_mul(BLOCK_META_LEN));
3367 }
3368
3369 let per_posting = long_bytes as f64 / long_postings as f64;
3370 println!(
3371 "SEALED DENSITY traversed_postings={long_postings} bytes={long_bytes} \
3372 per_posting={per_posting:.3} single_block_lists={short_lists}"
3373 );
3374 assert!(
3375 per_posting < 4.0,
3376 "a traversed list costs {per_posting:.3} bytes per posting; the \
3377 in-memory form cost a 32-byte stride plus an allocation, and a \
3378 regression here is a decode regression the wall clock would only \
3379 confirm later"
3380 );
3381 }
3382
3383 #[test]
3384 fn every_decoded_width_falls_in_the_narrow_kernel_path() {
3385 for (label, texts) in [
3394 (
3395 "zipf100k",
3396 (0..100_000)
3397 .map(|index| {
3398 let mut words: Vec<String> = Vec::new();
3399 for term in 0..12_usize {
3400 if index % (term + 1) == 0 {
3401 words.push(format!("t{term}"));
3402 }
3403 }
3404 if words.is_empty() {
3405 words.push(String::from("filler"));
3406 }
3407 words.join(" ")
3408 })
3409 .collect::<Vec<String>>(),
3410 ),
3411 (
3412 "textish",
3413 (0..50_000)
3414 .map(|index| format!("alpha t{} u{} filler{index}", index % 97, index % 1_009))
3415 .collect::<Vec<String>>(),
3416 ),
3417 ] {
3418 let sealed = SealedSegment::seal(&segment_of(&texts)).expect("seals");
3419 let mut histogram = [0_u64; 33];
3420 let mut postings = 0_u64;
3421 let mut widest = 0_u8;
3422 for index in 0..sealed.spans.len() {
3423 let cursor = sealed.cursor_at(index, 1_000).expect("span");
3424 if cursor.block_count <= 1 {
3427 continue;
3428 }
3429 for block in 0..cursor.block_count {
3430 let meta = cursor.meta_at(block).expect("meta");
3431 postings = postings.saturating_add(u64::from(meta.count));
3432 for bits in [meta.docid_bits, meta.tf_bits] {
3433 widest = widest.max(bits);
3434 if let Some(slot) = histogram.get_mut(usize::from(bits)) {
3435 *slot = slot.saturating_add(u64::from(meta.count));
3436 }
3437 }
3438 }
3439 }
3440 let mut shape: Vec<String> = Vec::new();
3441 for (bits, count) in histogram.iter().enumerate() {
3442 if *count > 0 {
3443 shape.push(format!(
3444 "{bits}b={:.1}%",
3445 *count as f64 * 50.0 / postings as f64
3446 ));
3447 }
3448 }
3449 println!(
3450 "WIDTHS {label} postings={postings} widest={widest} {}",
3451 shape.join(" ")
3452 );
3453 assert!(postings > 0, "{label} produced no traversable list");
3454 assert!(
3455 widest <= crate::kernels::postings::NARROW_MAX_BITS,
3456 "{label} decodes at {widest} bits, outside the narrow kernel \
3457 path the histogram was used to justify. Re-measure the \
3458 histogram and re-decide which arm to tune; do not widen \
3459 this bound."
3460 );
3461 }
3462 }
3463
3464 #[test]
3465 fn a_list_whose_lengths_are_absent_still_seals() {
3466 let mut list = PostingList::new();
3470 list.push(Posting {
3471 docid: 0,
3472 tf: 2,
3473 positions: vec![0, 1],
3474 })
3475 .expect("ascending");
3476 let impacts = crate::fts::postings::block_impacts(&list, 64, &[]);
3477 assert_eq!(impacts[0].max_tf, 2);
3478 assert_eq!(impacts[0].min_len, 1);
3479 }
3480}
3481
3482#[cfg(test)]
3483#[allow(clippy::expect_used)]
3484mod astra_10_allocation_tests {
3485 use super::*;
3486
3487 #[test]
3488 fn astra_10_stream_reservation_matches_actual_owned_capacities() {
3489 use crate::fts::index::{Document, FieldId, SegmentIndex};
3490 use crate::fts::tokenizer::{Analyzer, Profile};
3491 let analyzer = Analyzer::new(Profile::Code.config()).expect("analyzer");
3492 let mut index = SegmentIndex::new();
3493 for _ in 0..130 {
3494 let mut doc = Document::new();
3495 doc.set(FieldId(0), "alpha beta");
3496 doc.set(FieldId(1), "alpha alpha");
3497 doc.set(FieldId(2), "beta");
3498 index.push_document(&analyzer, &doc).expect("document");
3499 }
3500 let segment = SealedSegment::seal(&index).expect("sealed");
3501 for weights in [
3502 FieldWeights::flat(&[FieldId(0)]),
3503 FieldWeights::new(&[(FieldId(0), 500), (FieldId(1), 1500), (FieldId(2), 0)]),
3504 ] {
3505 for term in [b"alpha".as_slice(), b"beta", b"absent"] {
3506 let planned =
3507 TermStream::allocation_bytes(&segment, term, &weights).expect("fits usize");
3508 let actual = TermStream::open_sized(&segment, term, &weights).map_or(0, |stream| {
3509 stream.runs.capacity() * std::mem::size_of::<ListCursor<'_>>()
3510 + stream.weights.capacity() * std::mem::size_of::<u64>()
3511 + stream
3512 .runs
3513 .iter()
3514 .map(|run| {
3515 (run.decoded_docids.capacity() + run.decoded_tfs.capacity())
3516 * std::mem::size_of::<u32>()
3517 })
3518 .sum::<usize>()
3519 });
3520 assert_eq!(
3521 planned, actual,
3522 "reserve the actual buffers, including absent and zero-weight fields"
3523 );
3524 }
3525 }
3526 }
3527}