Skip to main content

nodedb_lite/engine/array/
read.rs

1//! Read paths of the Lite array engine: point reads, slices, and the
2//! surrogate bitmap scan. Each read merges memtable and segment versions and
3//! resolves them through the bitemporal ceiling.
4
5use std::collections::HashMap;
6use std::sync::Arc;
7
8use nodedb_array::query::ceiling::{CeilingParams, CeilingResult, ceiling_resolve_cell};
9use nodedb_array::query::slice::{DimRange, Slice};
10use nodedb_array::schema::ArraySchema;
11use nodedb_array::segment::extract_cell_bytes;
12use nodedb_array::tile::cell_payload::CellPayload;
13use nodedb_array::tile::sparse_tile::SparseTile;
14use nodedb_array::types::TileId;
15use nodedb_array::types::coord::value::CoordValue;
16use nodedb_array::{TilePayload, tile_id_for_cell};
17
18use crate::engine::array::engine::ArrayEngineState;
19use crate::engine::array::memtable::ArrayMemtable;
20use crate::error::LiteError;
21use crate::storage::engine::StorageEngine;
22
23impl ArrayEngineState {
24    /// Read the most recent live payload for `coord` at or before `system_as_of`.
25    /// Returns `None` for NotFound, Tombstoned, or Erased results.
26    pub async fn read_coord<S: StorageEngine>(
27        &self,
28        storage: &Arc<S>,
29        name: &str,
30        coord: &[CoordValue],
31        system_as_of: i64,
32    ) -> Result<Option<CellPayload>, LiteError> {
33        let state = self.arrays.get(name).ok_or_else(|| LiteError::BadRequest {
34            detail: format!("array '{name}' not found"),
35        })?;
36
37        let schema = &state.schema;
38        let seg_ids: Vec<u64> = state.manifest.segments.iter().map(|s| s.id).collect();
39
40        let tile_id =
41            tile_id_for_cell(schema, coord, system_as_of).map_err(|e| LiteError::Storage {
42                detail: format!("tile_id_for_cell: {e}"),
43            })?;
44        let hilbert_prefix = tile_id.hilbert_prefix;
45
46        let mem_versions = state
47            .memtable
48            .iter_cell_versions(hilbert_prefix, system_as_of, coord);
49
50        let seg_versions = cell_versions_from_segments(
51            storage,
52            name,
53            &seg_ids,
54            hilbert_prefix,
55            system_as_of,
56            coord,
57        )
58        .await?;
59
60        let mut all_versions: Vec<(TileId, Vec<u8>)> =
61            mem_versions.into_iter().chain(seg_versions).collect();
62        // ceiling_resolve_cell requires newest-first ordering.
63        all_versions.sort_unstable_by_key(|v| std::cmp::Reverse(v.0.system_from_ms));
64
65        let params = CeilingParams {
66            system_as_of,
67            valid_at_ms: None,
68        };
69        let result = ceiling_resolve_cell(
70            all_versions.iter().map(|(id, b)| (*id, b.as_slice())),
71            coord,
72            &params,
73        )
74        .map_err(|e| LiteError::Storage {
75            detail: format!("ceiling_resolve_cell: {e}"),
76        })?;
77
78        match result {
79            CeilingResult::Live(payload) => Ok(Some(payload)),
80            CeilingResult::Tombstoned | CeilingResult::Erased | CeilingResult::NotFound => Ok(None),
81        }
82    }
83
84    /// Return all live cells whose coordinates fall within `ranges` at or
85    /// before `system_as_of`.
86    pub async fn slice<S: StorageEngine>(
87        &mut self,
88        storage: &Arc<S>,
89        name: &str,
90        ranges: Vec<Option<DimRange>>,
91        system_as_of: i64,
92    ) -> Result<Vec<CellPayload>, LiteError> {
93        let state = self
94            .arrays
95            .get_mut(name)
96            .ok_or_else(|| LiteError::BadRequest {
97                detail: format!("array '{name}' not found"),
98            })?;
99
100        let schema = state.schema.clone();
101        let seg_ids: Vec<u64> = state.manifest.segments.iter().map(|s| s.id).collect();
102        let slice_pred = Slice::new(ranges);
103
104        // coord_key → Vec<(TileId, cell_bytes)>.
105        let mut coord_versions: HashMap<Vec<u8>, Vec<(TileId, Vec<u8>)>> = HashMap::new();
106
107        // Segments (oldest-first iteration, index from end to get newest-first).
108        for &seg_id in &seg_ids {
109            let bytes = crate::engine::array::segments::load_segment(storage, name, seg_id).await?;
110            let reader = crate::engine::array::segments::open_reader(&bytes)?;
111            for idx in (0..reader.tile_count()).rev() {
112                let entry_tile_id = reader.tiles()[idx].tile_id;
113                if entry_tile_id.system_from_ms > system_as_of {
114                    continue;
115                }
116                let payload = reader.read_tile(idx).map_err(|e| LiteError::Storage {
117                    detail: format!("read_tile: {e}"),
118                })?;
119                if let TilePayload::Sparse(tile) = &payload {
120                    collect_tile_into_map(tile, entry_tile_id, &mut coord_versions)?;
121                }
122            }
123        }
124
125        // Memtable.
126        let mem_tiles = state
127            .memtable
128            .drain_all_tiles_read_only(system_as_of, &schema)?;
129        for (tile_id, tile) in &mem_tiles {
130            collect_tile_into_map(tile, *tile_id, &mut coord_versions)?;
131        }
132
133        // Ceiling + slice filter.
134        let mut results = Vec::new();
135        for (coord_key, mut versions) in coord_versions {
136            versions.sort_by_key(|v| std::cmp::Reverse(v.0.system_from_ms));
137
138            let coord: Vec<CoordValue> =
139                zerompk::from_msgpack(&coord_key).map_err(|e| LiteError::Serialization {
140                    detail: format!("decode coord: {e}"),
141                })?;
142
143            let params = CeilingParams {
144                system_as_of,
145                valid_at_ms: None,
146            };
147            let result = ceiling_resolve_cell(
148                versions.iter().map(|(id, b)| (*id, b.as_slice())),
149                &coord,
150                &params,
151            )
152            .map_err(|e| LiteError::Storage {
153                detail: format!("ceiling_resolve_cell: {e}"),
154            })?;
155
156            if let CeilingResult::Live(payload) = result
157                && cell_in_slice(&coord, &slice_pred)
158            {
159                results.push(payload);
160            }
161        }
162
163        Ok(results)
164    }
165
166    /// Return the set of surrogates for all live cells whose coordinates
167    /// fall within `ranges` at or before `system_as_of`.
168    ///
169    /// This is the cross-engine prefilter primitive: the returned bitmap
170    /// gates the HNSW candidate set in the vector search path so only
171    /// vector records whose array-cell counterpart matches the slice
172    /// predicate are considered.
173    pub async fn surrogate_bitmap_scan<S: StorageEngine>(
174        &mut self,
175        storage: &Arc<S>,
176        name: &str,
177        ranges: Vec<Option<DimRange>>,
178        system_as_of: i64,
179    ) -> Result<roaring::RoaringBitmap, LiteError> {
180        let cells = self.slice(storage, name, ranges, system_as_of).await?;
181        let mut bitmap = roaring::RoaringBitmap::new();
182        for payload in cells {
183            bitmap.insert(payload.surrogate.as_u32());
184        }
185        Ok(bitmap)
186    }
187}
188
189async fn cell_versions_from_segments<S: StorageEngine>(
190    storage: &Arc<S>,
191    name: &str,
192    seg_ids: &[u64],
193    hilbert_prefix: u64,
194    system_as_of: i64,
195    coord: &[CoordValue],
196) -> Result<Vec<(TileId, Vec<u8>)>, LiteError> {
197    let tile_payloads = crate::engine::array::segments::collect_tile_versions_across_segments(
198        storage,
199        name,
200        seg_ids,
201        hilbert_prefix,
202        system_as_of,
203    )
204    .await?;
205    let mut out = Vec::new();
206    for (tile_id, payload) in tile_payloads {
207        if let TilePayload::Sparse(tile) = &payload
208            && let Ok(Some(cell_bytes)) = extract_cell_bytes(tile, coord)
209        {
210            out.push((tile_id, cell_bytes));
211        }
212    }
213    Ok(out)
214}
215
216fn collect_tile_into_map(
217    tile: &SparseTile,
218    tile_id: TileId,
219    acc: &mut HashMap<Vec<u8>, Vec<(TileId, Vec<u8>)>>,
220) -> Result<(), LiteError> {
221    let n = tile.row_count();
222    for row in 0..n {
223        let coord = tile_row_coord(tile, row)?;
224        let key = zerompk::to_msgpack_vec(&coord).map_err(|e| LiteError::Serialization {
225            detail: format!("encode coord: {e}"),
226        })?;
227        if let Ok(Some(cell_bytes)) = extract_cell_bytes(tile, &coord) {
228            acc.entry(key).or_default().push((tile_id, cell_bytes));
229        }
230    }
231    Ok(())
232}
233
234fn tile_row_coord(tile: &SparseTile, row: usize) -> Result<Vec<CoordValue>, LiteError> {
235    let arity = tile.dim_dicts.len();
236    let mut coord = Vec::with_capacity(arity);
237    for dim_idx in 0..arity {
238        let dict = tile
239            .dim_dicts
240            .get(dim_idx)
241            .ok_or_else(|| LiteError::Storage {
242                detail: format!("tile_row_coord: dim {dim_idx} missing"),
243            })?;
244        let entry_idx = *dict.indices.get(row).ok_or_else(|| LiteError::Storage {
245            detail: format!("tile_row_coord: row {row} out of range"),
246        })? as usize;
247        let val = dict
248            .values
249            .get(entry_idx)
250            .ok_or_else(|| LiteError::Storage {
251                detail: format!("tile_row_coord: dict entry {entry_idx} missing"),
252            })?;
253        coord.push(val.clone());
254    }
255    Ok(coord)
256}
257
258fn cell_in_slice(coord: &[CoordValue], slice: &Slice) -> bool {
259    use CoordValue::*;
260    use nodedb_array::types::domain::DomainBound;
261    for (i, range) in slice.dim_ranges.iter().enumerate() {
262        let Some(r) = range else { continue };
263        let Some(c) = coord.get(i) else { return false };
264        let in_range = match (&r.lo, c, &r.hi) {
265            (DomainBound::Int64(l), Int64(v), DomainBound::Int64(h)) => l <= v && v <= h,
266            (DomainBound::Float64(l), Float64(v), DomainBound::Float64(h)) => l <= v && v <= h,
267            (DomainBound::String(l), String(v), DomainBound::String(h)) => l <= v && v <= h,
268            _ => false,
269        };
270        if !in_range {
271            return false;
272        }
273    }
274    true
275}
276
277impl ArrayMemtable {
278    /// Read-only snapshot of all tile-version entries at or before
279    /// `system_as_of`, materialised as `SparseTile`s.
280    pub fn drain_all_tiles_read_only(
281        &self,
282        system_as_of: i64,
283        schema: &ArraySchema,
284    ) -> Result<Vec<(TileId, SparseTile)>, LiteError> {
285        // Collect all prefixes in the memtable.
286        let mut prefixes: std::collections::HashSet<u64> = std::collections::HashSet::new();
287        for id in self.tiles_iter() {
288            if id.system_from_ms <= system_as_of {
289                prefixes.insert(id.hilbert_prefix);
290            }
291        }
292        let mut result = Vec::new();
293        for prefix in prefixes {
294            let mut tiles = self.iter_tiles_for_prefix(prefix, system_as_of, schema)?;
295            result.append(&mut tiles);
296        }
297        Ok(result)
298    }
299}