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: InvertedIndex + ?Sized> {
82    index: &'a 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
93pub(crate) fn open<'a, T: InvertedIndex + ?Sized>(
94    index: &'a T,
95    field: &'a str,
96    term: &'a TokenTermKey,
97    control: &StorageReadControl,
98) -> StorageBackendResult<BudgetedPostingReadCursor<'a>> {
99    let mut first = None;
100    let mut previous = None;
101    let mut doc_freq = 0u64;
102    index.visit_score_clusters(field, term, None, usize::MAX, control, &mut |cluster| {
103        control.check()?;
104        if previous.is_some_and(|previous| previous >= cluster.cluster_id) {
105            return Err(corrupt("cluster identifiers are not strictly increasing"));
106        }
107        previous = Some(cluster.cluster_id);
108        let directory = directory(&cluster, control)?;
109        doc_freq = doc_freq
110            .checked_add(directory.count as u64)
111            .ok_or_else(|| corrupt("posting count overflow"))?;
112        if first.is_none() {
113            first = Some(retain(cluster, &directory, control)?);
114        }
115        Ok(())
116    })?;
117    control.check()?;
118    let mut cursor = ControlledCursor {
119        index,
120        field,
121        term,
122        control: control.clone(),
123        cluster: first,
124        entries: BudgetedVec::new(control.memory()),
125        doc_freq,
126        block: 0,
127        position: 0,
128    };
129    if let Some(cluster) = &cursor.cluster {
130        decode_block(cluster, 0, &mut cursor.entries, control)?;
131    }
132    BudgetedPostingReadCursor::new(cursor, control)
133}
134
135fn decode_block(
136    cluster: &EncodedCluster,
137    at: usize,
138    entries: &mut BudgetedVec<PostingScore>,
139    control: &StorageReadControl,
140) -> StorageBackendResult<()> {
141    // A fixed block staging area preserves the current cursor on cancellation or corrupt input.
142    let mut pending = [PostingScore {
143        doc_id: 0,
144        term_freq: 0,
145        doc_length: 0,
146    }; DEFAULT_BLOCK_SIZE];
147    let block = cluster.blocks[at];
148    let mut count = 0;
149    visit_score_block(
150        &cluster.bytes,
151        cluster.id,
152        block,
153        &mut || control.check(),
154        |entry, _| {
155            pending[count] = entry;
156            count += 1;
157            Ok(())
158        },
159    )?;
160    control.check()?;
161    entries.reserve(count.saturating_sub(entries.len()))?;
162    entries.clear();
163    for entry in &pending[..count] {
164        entries.push(*entry)?;
165    }
166    Ok(())
167}
168
169impl<T: InvertedIndex + ?Sized> ControlledCursor<'_, T> {
170    fn load_after(&mut self, after: u64) -> StorageBackendResult<bool> {
171        let mut next = None;
172        self.index.visit_score_clusters(
173            self.field,
174            self.term,
175            Some(after),
176            1,
177            &self.control,
178            &mut |cluster| {
179                if cluster.cluster_id <= after || next.is_some() {
180                    return Err(corrupt("score source violated the requested cluster range"));
181                }
182                let directory = directory(&cluster, &self.control)?;
183                next = Some(retain(cluster, &directory, &self.control)?);
184                Ok(())
185            },
186        )?;
187        self.control.check()?;
188        if let Some(cluster) = &next {
189            decode_block(cluster, 0, &mut self.entries, &self.control)?;
190        } else {
191            self.entries = BudgetedVec::new(self.control.memory());
192        }
193        self.cluster = next;
194        self.block = 0;
195        self.position = 0;
196        Ok(self.cluster.is_some())
197    }
198
199    fn load_block(&mut self, at: usize) -> StorageBackendResult<()> {
200        let cluster = self.cluster.as_ref().expect("live cluster");
201        decode_block(cluster, at, &mut self.entries, &self.control)?;
202        self.block = at;
203        self.position = 0;
204        Ok(())
205    }
206}
207
208impl<T: InvertedIndex + ?Sized> PostingReadCursor for ControlledCursor<'_, T> {
209    fn doc_freq(&self) -> u64 {
210        self.doc_freq
211    }
212    fn current(&self) -> Option<PostingScore> {
213        self.entries.get(self.position).copied()
214    }
215    fn advance(&mut self) -> StorageBackendResult<Option<PostingScore>> {
216        self.control.check()?;
217        let Some(cluster) = &self.cluster else {
218            return Ok(None);
219        };
220        if self.position + 1 < self.entries.len() {
221            self.position += 1;
222        } else if self.block + 1 < cluster.blocks.len() {
223            self.load_block(self.block + 1)?;
224        } else {
225            self.load_after(cluster.id)?;
226        }
227        Ok(self.current())
228    }
229    fn advance_to(&mut self, target: DocId) -> StorageBackendResult<Option<PostingScore>> {
230        self.control.check()?;
231        if self.current().is_none_or(|entry| entry.doc_id >= target) {
232            return Ok(self.current());
233        }
234        loop {
235            self.control.check()?;
236            let Some(cluster) = &self.cluster else {
237                return Ok(None);
238            };
239            let base = cluster_base(cluster.id)?;
240            let at = cluster
241                .blocks
242                .partition_point(|block| base + u64::from(block.last_offset) < target);
243            if at < cluster.blocks.len() {
244                if at != self.block {
245                    self.load_block(at)?;
246                }
247                self.position = self.entries.partition_point(|entry| entry.doc_id < target);
248                return Ok(self.current());
249            }
250            let after = cluster.id.max(cluster_id(target).saturating_sub(1));
251            if !self.load_after(after)? {
252                return Ok(None);
253            }
254        }
255    }
256}
uqa_storage/clustered_postings/controlled_cursor.rs

uqa_storage/clustered_postings/
controlled_cursor.rs
