uqa_storage/
clustered_postings.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Backend-neutral clustered posting codec and lazy score cursor.
8//!
9//! A posting key identifies `(table, field, term, cluster_id)`, where one
10//! cluster covers 2^16 consecutive document identifiers. Score data and
11//! positional payloads are encoded separately: ranking reads only the
12//! columnar document-offset, term-frequency, and document-length streams,
13//! while phrase/highlight consumers opt into the positions blob.
14
15use std::sync::Arc;
16
17use uqa_core::{DocId, TokenOccurrence};
18
19use crate::{StorageBackendError, StorageBackendResult, DEFAULT_BLOCK_SIZE};
20
21pub const POSTING_CLUSTER_DOCS: u64 = 1 << 16;
22
23const SCORE_MAGIC: &[u8; 4] = b"UQCS";
24const POSITIONS_MAGIC: &[u8; 4] = b"UQCP";
25const TERMS_MAGIC: &[u8; 4] = b"UQCT";
26const FORMAT_VERSION: u8 = 1;
27pub const OCCURRENCE_FORMAT_VERSION: u8 = 2;
28const HEADER_LEN: usize = 16;
29const SCORE_DIRECTORY_ENTRY_LEN: usize = 28;
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct PostingScore {
33    pub doc_id: DocId,
34    pub term_freq: u64,
35    pub doc_length: u64,
36}
37
38#[derive(Debug, Clone, PartialEq, Eq)]
39pub struct ClusterPosting {
40    pub doc_id: DocId,
41    pub term_freq: u64,
42    pub doc_length: u64,
43    pub positions: Vec<u32>,
44}
45
46pub trait PostingCursor: Send {
47    fn doc_freq(&self) -> u64;
48    fn ordinal(&self) -> u64;
49    fn current(&self) -> Option<PostingScore>;
50    fn advance(&mut self) -> StorageBackendResult<Option<PostingScore>>;
51    fn advance_to(&mut self, target: DocId) -> StorageBackendResult<Option<PostingScore>>;
52    fn boxed_clone(&self) -> Box<dyn PostingCursor>;
53}
54
55mod controlled_cursor;
56mod positions;
57mod read_cursor;
58pub(crate) use controlled_cursor::open as open_controlled_cursor;
59pub use controlled_cursor::{EncodedScoreClusterRef, ScoreClusterVisitor};
60pub use read_cursor::{BudgetedPostingReadCursor, OwnedPostingReadCursor, PostingReadCursor};
61
62impl Clone for Box<dyn PostingCursor> {
63    fn clone(&self) -> Self {
64        self.boxed_clone()
65    }
66}
67
68#[derive(Clone)]
69pub struct MaterializedPostingCursor {
70    entries: Arc<[PostingScore]>,
71    position: usize,
72}
73
74impl MaterializedPostingCursor {
75    pub fn new(entries: Vec<PostingScore>) -> StorageBackendResult<Self> {
76        validate_scores(&entries)?;
77        Ok(Self {
78            entries: entries.into(),
79            position: 0,
80        })
81    }
82}
83
84impl PostingCursor for MaterializedPostingCursor {
85    fn doc_freq(&self) -> u64 {
86        self.entries.len() as u64
87    }
88
89    fn ordinal(&self) -> u64 {
90        self.position as u64
91    }
92
93    fn current(&self) -> Option<PostingScore> {
94        self.entries.get(self.position).copied()
95    }
96
97    fn advance(&mut self) -> StorageBackendResult<Option<PostingScore>> {
98        self.position = self.position.saturating_add(1).min(self.entries.len());
99        Ok(self.current())
100    }
101
102    fn advance_to(&mut self, target: DocId) -> StorageBackendResult<Option<PostingScore>> {
103        if self.current().is_some_and(|entry| entry.doc_id >= target) {
104            return Ok(self.current());
105        }
106        let relative = self.entries[self.position..].partition_point(|entry| entry.doc_id < target);
107        self.position = self
108            .position
109            .saturating_add(relative)
110            .min(self.entries.len());
111        Ok(self.current())
112    }
113
114    fn boxed_clone(&self) -> Box<dyn PostingCursor> {
115        Box::new(self.clone())
116    }
117}
118
119#[derive(Debug, Clone)]
120pub struct EncodedScoreCluster {
121    pub cluster_id: u64,
122    pub bytes: Vec<u8>,
123}
124
125#[derive(Debug, Clone, Copy)]
126struct ScoreBlock {
127    version: u8,
128    count: usize,
129    last_offset: u16,
130    docs_start: usize,
131    docs_end: usize,
132    term_freqs_start: usize,
133    term_freqs_end: usize,
134    doc_lengths_start: usize,
135    doc_lengths_end: usize,
136}
137
138#[derive(Debug, Clone, Copy)]
139struct CursorBlock {
140    cluster_index: usize,
141    score_block: ScoreBlock,
142    first_ordinal: u64,
143    last_doc_id: DocId,
144}
145
146#[derive(Clone)]
147pub struct ClusteredPostingCursor {
148    clusters: Arc<[EncodedScoreCluster]>,
149    blocks: Arc<[CursorBlock]>,
150    doc_freq: u64,
151    block_index: usize,
152    entries: Vec<PostingScore>,
153    position_in_block: usize,
154}
155
156impl ClusteredPostingCursor {
157    pub fn new(clusters: Vec<EncodedScoreCluster>) -> StorageBackendResult<Self> {
158        let mut previous_cluster = None;
159        let mut blocks = Vec::new();
160        let mut doc_freq = 0_u64;
161        for (cluster_index, cluster) in clusters.iter().enumerate() {
162            if previous_cluster.is_some_and(|previous| previous >= cluster.cluster_id) {
163                return Err(corrupt("cluster identifiers are not strictly increasing"));
164            }
165            previous_cluster = Some(cluster.cluster_id);
166            let (_, score_blocks) = parse_score_blob(&cluster.bytes)?;
167            for score_block in score_blocks {
168                let first_ordinal = doc_freq;
169                doc_freq = doc_freq
170                    .checked_add(score_block.count as u64)
171                    .ok_or_else(|| corrupt("posting count overflow"))?;
172                blocks.push(CursorBlock {
173                    cluster_index,
174                    score_block,
175                    first_ordinal,
176                    last_doc_id: cluster_base(cluster.cluster_id)?
177                        .checked_add(u64::from(score_block.last_offset))
178                        .ok_or_else(|| corrupt("block document id overflow"))?,
179                });
180            }
181        }
182
183        let mut cursor = Self {
184            clusters: clusters.into(),
185            blocks: blocks.into(),
186            doc_freq,
187            block_index: 0,
188            entries: Vec::new(),
189            position_in_block: 0,
190        };
191        if !cursor.blocks.is_empty() {
192            cursor.load_block(0)?;
193        }
194        Ok(cursor)
195    }
196
197    fn load_block(&mut self, block_index: usize) -> StorageBackendResult<()> {
198        if block_index >= self.blocks.len() {
199            self.block_index = self.blocks.len();
200            self.entries.clear();
201            self.position_in_block = 0;
202            return Ok(());
203        }
204        let block = self.blocks[block_index];
205        let cluster = &self.clusters[block.cluster_index];
206        self.entries.clear();
207        self.entries.reserve(block.score_block.count);
208        decode_score_block_into(
209            &cluster.bytes,
210            cluster.cluster_id,
211            block.score_block,
212            &mut self.entries,
213        )?;
214        self.block_index = block_index;
215        self.position_in_block = 0;
216        Ok(())
217    }
218
219    fn exhausted(&self) -> bool {
220        self.block_index >= self.blocks.len()
221    }
222}
223
224impl PostingCursor for ClusteredPostingCursor {
225    fn doc_freq(&self) -> u64 {
226        self.doc_freq
227    }
228
229    fn ordinal(&self) -> u64 {
230        if self.exhausted() {
231            return self.doc_freq;
232        }
233        self.blocks[self.block_index]
234            .first_ordinal
235            .saturating_add(self.position_in_block as u64)
236    }
237
238    fn current(&self) -> Option<PostingScore> {
239        self.entries.get(self.position_in_block).copied()
240    }
241
242    fn advance(&mut self) -> StorageBackendResult<Option<PostingScore>> {
243        if self.exhausted() {
244            return Ok(None);
245        }
246        self.position_in_block += 1;
247        if self.position_in_block < self.entries.len() {
248            return Ok(self.current());
249        }
250        self.load_block(self.block_index + 1)?;
251        Ok(self.current())
252    }
253
254    fn advance_to(&mut self, target: DocId) -> StorageBackendResult<Option<PostingScore>> {
255        if self.current().is_some_and(|entry| entry.doc_id >= target) {
256            return Ok(self.current());
257        }
258        if self.exhausted() {
259            return Ok(None);
260        }
261        let relative =
262            self.blocks[self.block_index..].partition_point(|block| block.last_doc_id < target);
263        let block_index = self.block_index + relative;
264        if block_index >= self.blocks.len() {
265            self.load_block(self.blocks.len())?;
266            return Ok(None);
267        }
268        if block_index != self.block_index {
269            self.load_block(block_index)?;
270        }
271        self.position_in_block = self.entries.partition_point(|entry| entry.doc_id < target);
272        if self.position_in_block < self.entries.len() {
273            return Ok(self.current());
274        }
275        self.load_block(block_index + 1)?;
276        Ok(self.current())
277    }
278
279    fn boxed_clone(&self) -> Box<dyn PostingCursor> {
280        Box::new(self.clone())
281    }
282}
283
284pub fn cluster_id(doc_id: DocId) -> u64 {
285    doc_id / POSTING_CLUSTER_DOCS
286}
287
288fn cluster_base(cluster_id: u64) -> StorageBackendResult<DocId> {
289    cluster_id
290        .checked_mul(POSTING_CLUSTER_DOCS)
291        .ok_or_else(|| corrupt("cluster base document id overflow"))
292}
293
294fn cluster_offset(doc_id: DocId) -> StorageBackendResult<u16> {
295    u16::try_from(doc_id % POSTING_CLUSTER_DOCS)
296        .map_err(|_| corrupt("document offset exceeds clustered format"))
297}
298
299pub fn encode_cluster(entries: &[ClusterPosting]) -> StorageBackendResult<(Vec<u8>, Vec<u8>)> {
300    if entries.is_empty() {
301        return Err(corrupt("cannot encode an empty posting cluster"));
302    }
303    validate_cluster_entries(entries)?;
304    let scores: Vec<_> = entries
305        .iter()
306        .map(|entry| PostingScore {
307            doc_id: entry.doc_id,
308            term_freq: entry.term_freq,
309            doc_length: entry.doc_length,
310        })
311        .collect();
312    Ok((
313        encode_scores(&scores, FORMAT_VERSION)?,
314        encode_positions(entries)?,
315    ))
316}
317
318pub fn decode_cluster(
319    cluster_id: u64,
320    score_blob: &[u8],
321    positions_blob: &[u8],
322) -> StorageBackendResult<Vec<ClusterPosting>> {
323    if score_blob.get(4) != Some(&FORMAT_VERSION) {
324        return Err(corrupt("legacy cluster reader requires format version 1"));
325    }
326    let scores = decode_all_scores(cluster_id, score_blob)?;
327    let positions = decode_positions(positions_blob, &scores)?;
328    Ok(scores
329        .into_iter()
330        .zip(positions)
331        .map(|(score, positions)| ClusterPosting {
332            doc_id: score.doc_id,
333            term_freq: score.term_freq,
334            doc_length: score.doc_length,
335            positions,
336        })
337        .collect())
338}
339
340pub fn score_count(score_blob: &[u8]) -> StorageBackendResult<u64> {
341    score_count_with_control(score_blob, || Ok(()))
342}
343
344/// Validate the borrowed score directory and count its entries without allocating.
345pub fn score_count_with_control(
346    score_blob: &[u8],
347    mut poll: impl FnMut() -> StorageBackendResult<()>,
348) -> StorageBackendResult<u64> {
349    Ok(scores::ScoreDirectory::new(score_blob, &mut poll)?.count as u64)
350}
351
352pub fn decode_all_scores(
353    cluster_id: u64,
354    score_blob: &[u8],
355) -> StorageBackendResult<Vec<PostingScore>> {
356    let (count, blocks) = parse_score_blob(score_blob)?;
357    let mut entries = Vec::with_capacity(count);
358    for block in blocks {
359        decode_score_block_into(score_blob, cluster_id, block, &mut entries)?;
360    }
361    validate_scores(&entries)?;
362    Ok(entries)
363}
364
365fn position_entries(blob: &[u8], expected_count: usize) -> StorageBackendResult<Vec<&[u8]>> {
366    let directory = positions::PositionDirectory::new(blob, expected_count, &mut || Ok(()))?;
367    (0..expected_count)
368        .map(|index| directory.entry(index))
369        .collect()
370}
371
372fn validate_cluster_entries(entries: &[ClusterPosting]) -> StorageBackendResult<()> {
373    if entries.is_empty() {
374        return Ok(());
375    }
376    let expected_cluster = cluster_id(entries[0].doc_id);
377    let mut previous = None;
378    for entry in entries {
379        if cluster_id(entry.doc_id) != expected_cluster {
380            return Err(corrupt("one encoded value spans multiple clusters"));
381        }
382        if previous.is_some_and(|doc_id| doc_id >= entry.doc_id) {
383            return Err(corrupt("posting document ids are not strictly increasing"));
384        }
385        if entry.term_freq == 0
386            || entry.doc_length < entry.term_freq
387            || entry.term_freq != entry.positions.len() as u64
388        {
389            return Err(corrupt(
390                "posting frequency, document length, and positions disagree",
391            ));
392        }
393        if entry.positions.windows(2).any(|pair| pair[0] >= pair[1]) {
394            return Err(corrupt("term positions are not strictly increasing"));
395        }
396        previous = Some(entry.doc_id);
397    }
398    Ok(())
399}
400
401fn validate_scores(entries: &[PostingScore]) -> StorageBackendResult<()> {
402    let mut previous = None;
403    for entry in entries {
404        if previous.is_some_and(|doc_id| doc_id >= entry.doc_id) {
405            return Err(corrupt("posting scores are not strictly ordered"));
406        }
407        if entry.term_freq == 0 || entry.doc_length == 0 {
408            return Err(corrupt("invalid posting score frequencies"));
409        }
410        previous = Some(entry.doc_id);
411    }
412    Ok(())
413}
414
415fn validate_header(blob: &[u8], magic: [u8; 4]) -> StorageBackendResult<()> {
416    if blob.len() < HEADER_LEN || blob.get(..4) != Some(magic.as_slice()) {
417        return Err(corrupt("missing clustered posting header"));
418    }
419    if ![FORMAT_VERSION, OCCURRENCE_FORMAT_VERSION].contains(&blob[4]) {
420        return Err(corrupt("unsupported clustered posting version"));
421    }
422    if blob[5..8] != [0; 3] {
423        return Err(corrupt("clustered posting reserved header bits are set"));
424    }
425    Ok(())
426}
427
428fn put_varint(output: &mut Vec<u8>, mut value: u64) {
429    while value >= 0x80 {
430        output.push((value as u8 & 0x7f) | 0x80);
431        value >>= 7;
432    }
433    output.push(value as u8);
434}
435
436fn read_varint(input: &mut &[u8]) -> StorageBackendResult<u64> {
437    let mut value = 0_u64;
438    for shift in (0..=63).step_by(7) {
439        let Some((&byte, rest)) = input.split_first() else {
440            return Err(corrupt("truncated varint"));
441        };
442        *input = rest;
443        if shift == 63 && byte > 1 {
444            return Err(corrupt("varint overflow"));
445        }
446        value |= u64::from(byte & 0x7f) << shift;
447        if byte & 0x80 == 0 {
448            if shift > 0 && byte == 0 {
449                return Err(corrupt("noncanonical varint"));
450            }
451            return Ok(value);
452        }
453    }
454    Err(corrupt("unterminated varint"))
455}
456
457fn put_u32(output: &mut Vec<u8>, value: usize, field: &str) -> StorageBackendResult<()> {
458    let value = u32::try_from(value)
459        .map_err(|_| corrupt(format!("{field} exceeds the u32 on-disk format")))?;
460    output.extend_from_slice(&value.to_le_bytes());
461    Ok(())
462}
463
464fn read_u16(input: &[u8], offset: usize) -> StorageBackendResult<u16> {
465    let bytes: [u8; 2] = input
466        .get(offset..offset.saturating_add(2))
467        .ok_or_else(|| corrupt("truncated u16 field"))?
468        .try_into()
469        .map_err(|_| corrupt("invalid u16 field"))?;
470    Ok(u16::from_le_bytes(bytes))
471}
472
473fn read_u32(input: &[u8], offset: usize) -> StorageBackendResult<u32> {
474    let bytes: [u8; 4] = input
475        .get(offset..offset.saturating_add(4))
476        .ok_or_else(|| corrupt("truncated u32 field"))?
477        .try_into()
478        .map_err(|_| corrupt("invalid u32 field"))?;
479    Ok(u32::from_le_bytes(bytes))
480}
481
482fn corrupt(message: impl Into<String>) -> StorageBackendError {
483    StorageBackendError::Other(format!("corrupt clustered posting: {}", message.into()))
484}
485
486mod legacy;
487mod occurrences;
488mod scores;
489mod term_keys;
490
491use legacy::{decode_positions, encode_positions};
492pub use legacy::{decode_terms, encode_terms};
493pub use occurrences::{
494    decode_occurrence_cluster, decode_occurrence_cluster_budgeted,
495    decode_occurrence_document_budgeted, encode_occurrence_cluster, OccurrencePosting,
496};
497use scores::{decode_score_block_into, encode_scores, parse_score_blob};
498pub use term_keys::{decode_term_keys, encode_term_keys};
499
500#[cfg(test)]
501mod tests;
uqa_storage/clustered_postings.rs

uqa_storage/
clustered_postings.rs
