uqa_storage/clustered_postings/
controlled_cursor.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Retain one encoded cluster and reuse one score block under the consumer's allocation owner.
8
9use super::{
10    cluster_base, cluster_id, corrupt,
11    scores::{visit_score_block, ScoreDirectory},
12    BudgetedPostingReadCursor, DocId, PostingReadCursor, PostingScore, ScoreBlock,
13    StorageBackendResult, DEFAULT_BLOCK_SIZE, OCCURRENCE_FORMAT_VERSION,
14};
15use crate::{read_control::StorageReadControl, InvertedIndex, TokenTermKey};
16use uqa_core::memory::BudgetedVec;
17
18#[derive(Clone, Copy)]
19pub struct EncodedScoreClusterRef<'a> {
20    pub cluster_id: u64,
21    pub stored_count: Option<u64>,
22    pub bytes: &'a [u8],
23}
24pub type ScoreClusterVisitor<'a> =
25    dyn FnMut(EncodedScoreClusterRef<'_>) -> StorageBackendResult<()> + 'a;
26
27struct EncodedCluster {
28    id: u64,
29    bytes: BudgetedVec<u8>,
30    blocks: BudgetedVec<ScoreBlock>,
31}
32
33fn directory<'a>(
34    cluster: &EncodedScoreClusterRef<'a>,
35    control: &StorageReadControl,
36) -> StorageBackendResult<ScoreDirectory<'a>> {
37    control.check()?;
38    cluster_base(cluster.cluster_id)?;
39    let directory = ScoreDirectory::new(cluster.bytes, &mut || control.check())?;
40    if cluster.bytes[4] != OCCURRENCE_FORMAT_VERSION {
41        return Err(corrupt("occurrence index contains a legacy score payload"));
42    }
43    if cluster
44        .stored_count
45        .is_some_and(|count| count != directory.count as u64)
46    {
47        return Err(corrupt(
48            "stored posting count disagrees with the score payload",
49        ));
50    }
51    Ok(directory)
52}
53
54fn retain(
55    cluster: EncodedScoreClusterRef<'_>,
56    directory: &ScoreDirectory<'_>,
57    control: &StorageReadControl,
58) -> StorageBackendResult<EncodedCluster> {
59    let mut bytes = BudgetedVec::new(control.memory());
60    bytes.reserve(cluster.bytes.len())?;
61    for (index, byte) in cluster.bytes.iter().copied().enumerate() {
62        if index % 1024 == 0 {
63            control.check()?;
64        }
65        bytes.push(byte)?;
66    }
67    let mut blocks = BudgetedVec::new(control.memory());
68    blocks.reserve(directory.blocks)?;
69    for index in 0..directory.blocks {
70        control.check()?;
71        blocks.push(directory.block(index)?)?;
72    }
73    control.check()?;
74    Ok(EncodedCluster {
75        id: cluster.cluster_id,
76        bytes,
77        blocks,
78    })
79}
80
81struct ControlledCursor<'a, T> {
82    index: T,
83    field: &'a str,
84    term: &'a TokenTermKey,
85    control: StorageReadControl,
86    cluster: Option<EncodedCluster>,
87    entries: BudgetedVec<PostingScore>,
88    doc_freq: u64,
89    block: usize,
90    position: usize,
91}
92
93/// Open a bounded score cursor over a borrowed index or an owned retained snapshot. The supplied handle must preserve one read boundary throughout traversal.
94pub fn open<'a, T>(
95    index: T,
96    field: &'a str,
97    term: &'a TokenTermKey,
98    control: &StorageReadControl,
99) -> StorageBackendResult<BudgetedPostingReadCursor<'a>>
100where
101    T: std::ops::Deref + Send + 'a,
102    T::Target: InvertedIndex,
103{
104    let mut first = None;
105    let mut previous = None;
106    let mut doc_freq = 0u64;
107    index.visit_score_clusters(field, term, None, usize::MAX, control, &mut |cluster| {
108        control.check()?;
109        if previous.is_some_and(|previous| previous >= cluster.cluster_id) {
110            return Err(corrupt("cluster identifiers are not strictly increasing"));
111        }
112        previous = Some(cluster.cluster_id);
113        let directory = directory(&cluster, control)?;
114        doc_freq = doc_freq
115            .checked_add(directory.count as u64)
116            .ok_or_else(|| corrupt("posting count overflow"))?;
117        if first.is_none() {
118            first = Some(retain(cluster, &directory, control)?);
119        }
120        Ok(())
121    })?;
122    control.check()?;
123    let mut cursor = ControlledCursor {
124        index,
125        field,
126        term,
127        control: control.clone(),
128        cluster: first,
129        entries: BudgetedVec::new(control.memory()),
130        doc_freq,
131        block: 0,
132        position: 0,
133    };
134    if let Some(cluster) = &cursor.cluster {
135        decode_block(cluster, 0, &mut cursor.entries, control)?;
136    }
137    BudgetedPostingReadCursor::new(cursor, control)
138}
139
140fn decode_block(
141    cluster: &EncodedCluster,
142    at: usize,
143    entries: &mut BudgetedVec<PostingScore>,
144    control: &StorageReadControl,
145) -> StorageBackendResult<()> {
146    // A fixed block staging area preserves the current cursor on cancellation or corrupt input.
147    let mut pending = [PostingScore {
148        doc_id: 0,
149        term_freq: 0,
150        doc_length: 0,
151    }; DEFAULT_BLOCK_SIZE];
152    let block = cluster.blocks[at];
153    let mut count = 0;
154    visit_score_block(
155        &cluster.bytes,
156        cluster.id,
157        block,
158        &mut || control.check(),
159        |entry, _| {
160            pending[count] = entry;
161            count += 1;
162            Ok(())
163        },
164    )?;
165    control.check()?;
166    entries.reserve(count.saturating_sub(entries.len()))?;
167    entries.clear();
168    for entry in &pending[..count] {
169        entries.push(*entry)?;
170    }
171    Ok(())
172}
173
174impl<T> ControlledCursor<'_, T>
175where
176    T: std::ops::Deref,
177    T::Target: InvertedIndex,
178{
179    fn load_after(&mut self, after: u64) -> StorageBackendResult<bool> {
180        let mut next = None;
181        self.index.visit_score_clusters(
182            self.field,
183            self.term,
184            Some(after),
185            1,
186            &self.control,
187            &mut |cluster| {
188                if cluster.cluster_id <= after || next.is_some() {
189                    return Err(corrupt("score source violated the requested cluster range"));
190                }
191                let directory = directory(&cluster, &self.control)?;
192                next = Some(retain(cluster, &directory, &self.control)?);
193                Ok(())
194            },
195        )?;
196        self.control.check()?;
197        if let Some(cluster) = &next {
198            decode_block(cluster, 0, &mut self.entries, &self.control)?;
199        } else {
200            self.entries = BudgetedVec::new(self.control.memory());
201        }
202        self.cluster = next;
203        self.block = 0;
204        self.position = 0;
205        Ok(self.cluster.is_some())
206    }
207
208    fn load_block(&mut self, at: usize) -> StorageBackendResult<()> {
209        let cluster = self.cluster.as_ref().expect("live cluster");
210        decode_block(cluster, at, &mut self.entries, &self.control)?;
211        self.block = at;
212        self.position = 0;
213        Ok(())
214    }
215}
216
217impl<T> PostingReadCursor for ControlledCursor<'_, T>
218where
219    T: std::ops::Deref + Send,
220    T::Target: InvertedIndex,
221{
222    fn doc_freq(&self) -> u64 {
223        self.doc_freq
224    }
225    fn current(&self) -> Option<PostingScore> {
226        self.entries.get(self.position).copied()
227    }
228    fn advance(&mut self) -> StorageBackendResult<Option<PostingScore>> {
229        self.control.check()?;
230        let Some(cluster) = &self.cluster else {
231            return Ok(None);
232        };
233        if self.position + 1 < self.entries.len() {
234            self.position += 1;
235        } else if self.block + 1 < cluster.blocks.len() {
236            self.load_block(self.block + 1)?;
237        } else {
238            self.load_after(cluster.id)?;
239        }
240        Ok(self.current())
241    }
242    fn advance_to(&mut self, target: DocId) -> StorageBackendResult<Option<PostingScore>> {
243        self.control.check()?;
244        if self.current().is_none_or(|entry| entry.doc_id >= target) {
245            return Ok(self.current());
246        }
247        loop {
248            self.control.check()?;
249            let Some(cluster) = &self.cluster else {
250                return Ok(None);
251            };
252            let base = cluster_base(cluster.id)?;
253            let at = cluster
254                .blocks
255                .partition_point(|block| base + u64::from(block.last_offset) < target);
256            if at < cluster.blocks.len() {
257                if at != self.block {
258                    self.load_block(at)?;
259                }
260                self.position = self.entries.partition_point(|entry| entry.doc_id < target);
261                return Ok(self.current());
262            }
263            let after = cluster.id.max(cluster_id(target).saturating_sub(1));
264            if !self.load_after(after)? {
265                return Ok(None);
266            }
267        }
268    }
269}
uqa_storage/clustered_postings/controlled_cursor.rs

uqa_storage/clustered_postings/
controlled_cursor.rs
