Skip to main content

nodedb_array/query/
slice.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Coordinate-range slice over a sparse tile.
4//!
5//! A [`Slice`] carries one optional [`DimRange`] per schema dimension
6//! (None = unconstrained). Bounds are typed [`DomainBound`]s so the
7//! filter dispatches by dim type without re-parsing.
8//!
9//! Two entry points:
10//! * [`tile_overlaps_slice`] — cheap MBR-only check; the query
11//!   executor calls this against each segment tile entry to skip
12//!   non-overlapping tiles before any payload decode.
13//! * [`slice_sparse`] — per-cell filter, returns a new [`SparseTile`]
14//!   containing only the surviving cells (preserving column order).
15
16use crate::error::ArrayResult;
17use crate::schema::ArraySchema;
18use crate::tile::sparse_tile::{SparseRow, SparseTile, SparseTileBuilder};
19use crate::types::cell_value::value::CellValue;
20use crate::types::coord::value::CoordValue;
21use crate::types::domain::DomainBound;
22
23/// Inclusive-by-default closed range on one dim. `lo`/`hi` variants
24/// must match the dim's type — mismatches are silently treated as
25/// "does not intersect" (the cell drops out).
26#[derive(
27    Debug,
28    Clone,
29    PartialEq,
30    serde::Serialize,
31    serde::Deserialize,
32    zerompk::ToMessagePack,
33    zerompk::FromMessagePack,
34)]
35pub struct DimRange {
36    pub lo: DomainBound,
37    pub hi: DomainBound,
38}
39
40impl DimRange {
41    pub fn new(lo: DomainBound, hi: DomainBound) -> Self {
42        Self { lo, hi }
43    }
44}
45
46/// Slice predicate. `dim_ranges.len()` must equal `schema.arity()`.
47#[derive(
48    Debug,
49    Clone,
50    Default,
51    serde::Serialize,
52    serde::Deserialize,
53    zerompk::ToMessagePack,
54    zerompk::FromMessagePack,
55)]
56pub struct Slice {
57    pub dim_ranges: Vec<Option<DimRange>>,
58}
59
60impl Slice {
61    pub fn new(dim_ranges: Vec<Option<DimRange>>) -> Self {
62        Self { dim_ranges }
63    }
64}
65
66/// True if a tile's per-dim MBR overlaps every constrained dim range.
67/// An empty `dim_ranges` (no constraints) trivially overlaps.
68pub fn tile_overlaps_slice(
69    dim_mins: &[DomainBound],
70    dim_maxs: &[DomainBound],
71    slice: &Slice,
72) -> bool {
73    for (i, range) in slice.dim_ranges.iter().enumerate() {
74        let Some(r) = range else { continue };
75        let Some(tile_min) = dim_mins.get(i) else {
76            return false;
77        };
78        let Some(tile_max) = dim_maxs.get(i) else {
79            return false;
80        };
81        // Disjoint when tile_max < range.lo OR tile_min > range.hi.
82        if bound_lt(tile_max, &r.lo) || bound_lt(&r.hi, tile_min) {
83            return false;
84        }
85    }
86    true
87}
88
89/// Filter cells in `tile` to those whose coords pass every dim range.
90/// Result is a freshly-built [`SparseTile`] — dictionaries shrink to
91/// surviving values, MBR/attr_stats are recomputed.
92pub fn slice_sparse(
93    schema: &ArraySchema,
94    tile: &SparseTile,
95    slice: &Slice,
96) -> ArrayResult<SparseTile> {
97    use crate::tile::sparse_tile::RowKind;
98    let mut b = SparseTileBuilder::new(schema);
99    let n = tile.row_count();
100    let mut live_idx = 0usize;
101    for row in 0..n {
102        // Sentinel rows carry no payload and must not be emitted into slice results.
103        if tile.row_kind(row)? != RowKind::Live {
104            continue;
105        }
106        let coord: Vec<CoordValue> = tile
107            .dim_dicts
108            .iter()
109            .map(|d| d.values[d.indices[row] as usize].clone())
110            .collect();
111        if !cell_in_slice(&coord, slice) {
112            live_idx += 1;
113            continue;
114        }
115        let attr_row = live_idx;
116        live_idx += 1;
117        let attrs: Vec<CellValue> = tile
118            .attr_cols
119            .iter()
120            .map(|col| col[attr_row].clone())
121            .collect();
122        // A stored cell keeps its identity; a derived row keeps none.
123        let surrogate = tile.row_surrogate(row)?;
124        let valid_from_ms = tile.valid_from_ms.get(row).copied().unwrap_or(0);
125        let valid_until_ms = tile
126            .valid_until_ms
127            .get(row)
128            .copied()
129            .unwrap_or(nodedb_types::OPEN_UPPER);
130        b.push_row(SparseRow {
131            coord: &coord,
132            attrs: &attrs,
133            surrogate,
134            valid_from_ms,
135            valid_until_ms,
136            kind: crate::tile::sparse_tile::RowKind::Live,
137        })?;
138    }
139    Ok(b.build())
140}
141
142fn cell_in_slice(coord: &[CoordValue], slice: &Slice) -> bool {
143    for (i, range) in slice.dim_ranges.iter().enumerate() {
144        let Some(r) = range else { continue };
145        let Some(c) = coord.get(i) else {
146            return false;
147        };
148        let cb = coord_to_bound(c);
149        if bound_lt(&cb, &r.lo) || bound_lt(&r.hi, &cb) {
150            return false;
151        }
152    }
153    true
154}
155
156fn bound_lt(a: &DomainBound, b: &DomainBound) -> bool {
157    match (a, b) {
158        (DomainBound::Int64(x), DomainBound::Int64(y)) => x < y,
159        (DomainBound::Float64(x), DomainBound::Float64(y)) => x < y,
160        (DomainBound::TimestampMs(x), DomainBound::TimestampMs(y)) => x < y,
161        (DomainBound::String(x), DomainBound::String(y)) => x < y,
162        // Type-mismatched bounds never satisfy a less-than relation.
163        _ => false,
164    }
165}
166
167fn coord_to_bound(c: &CoordValue) -> DomainBound {
168    match c {
169        CoordValue::Int64(v) => DomainBound::Int64(*v),
170        CoordValue::Float64(v) => DomainBound::Float64(*v),
171        CoordValue::TimestampMs(v) => DomainBound::TimestampMs(*v),
172        CoordValue::String(v) => DomainBound::String(v.clone()),
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use crate::schema::ArraySchemaBuilder;
180    use crate::schema::attr_spec::{AttrSpec, AttrType};
181    use crate::schema::dim_spec::{DimSpec, DimType};
182    use crate::types::domain::Domain;
183
184    fn schema() -> ArraySchema {
185        ArraySchemaBuilder::new("g")
186            .dim(DimSpec::new(
187                "x",
188                DimType::Int64,
189                Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)),
190            ))
191            .dim(DimSpec::new(
192                "y",
193                DimType::Int64,
194                Domain::new(DomainBound::Int64(0), DomainBound::Int64(99)),
195            ))
196            .attr(AttrSpec::new("v", AttrType::Int64, true))
197            .tile_extents(vec![10, 10])
198            .build()
199            .unwrap()
200    }
201
202    fn tile(rows: &[(i64, i64, i64)]) -> SparseTile {
203        let s = schema();
204        let mut b = SparseTileBuilder::new(&s);
205        for (x, y, v) in rows {
206            b.push(
207                &[CoordValue::Int64(*x), CoordValue::Int64(*y)],
208                &[CellValue::Int64(*v)],
209            )
210            .unwrap();
211        }
212        b.build()
213    }
214
215    #[test]
216    fn slice_keeps_only_in_range_cells() {
217        let s = schema();
218        let t = tile(&[(0, 0, 1), (5, 5, 2), (9, 9, 3)]);
219        let sl = Slice::new(vec![
220            Some(DimRange::new(DomainBound::Int64(4), DomainBound::Int64(10))),
221            None,
222        ]);
223        let out = slice_sparse(&s, &t, &sl).unwrap();
224        assert_eq!(out.nnz(), 2);
225    }
226
227    #[test]
228    fn empty_slice_keeps_everything() {
229        let s = schema();
230        let t = tile(&[(0, 0, 1), (5, 5, 2)]);
231        let sl = Slice::new(vec![None, None]);
232        let out = slice_sparse(&s, &t, &sl).unwrap();
233        assert_eq!(out.nnz(), 2);
234    }
235
236    #[test]
237    fn slice_over_never_written_region_is_empty() {
238        // A sparse array only materializes written cells. Slicing a coordinate
239        // region that overlaps none of them returns an empty result cleanly —
240        // no crash, no error.
241        let s = schema();
242        let t = tile(&[(0, 0, 1), (5, 5, 2), (9, 9, 3)]);
243        let sl = Slice::new(vec![
244            Some(DimRange::new(
245                DomainBound::Int64(50),
246                DomainBound::Int64(60),
247            )),
248            Some(DimRange::new(
249                DomainBound::Int64(50),
250                DomainBound::Int64(60),
251            )),
252        ]);
253        let out = slice_sparse(&s, &t, &sl).unwrap();
254        assert_eq!(out.nnz(), 0);
255    }
256
257    #[test]
258    fn mbr_overlap_skip_disjoint() {
259        let mins = vec![DomainBound::Int64(0), DomainBound::Int64(0)];
260        let maxs = vec![DomainBound::Int64(9), DomainBound::Int64(9)];
261        let sl = Slice::new(vec![
262            Some(DimRange::new(
263                DomainBound::Int64(20),
264                DomainBound::Int64(30),
265            )),
266            None,
267        ]);
268        assert!(!tile_overlaps_slice(&mins, &maxs, &sl));
269    }
270
271    #[test]
272    fn mbr_overlap_accepts_intersection() {
273        let mins = vec![DomainBound::Int64(0), DomainBound::Int64(0)];
274        let maxs = vec![DomainBound::Int64(9), DomainBound::Int64(9)];
275        let sl = Slice::new(vec![
276            Some(DimRange::new(DomainBound::Int64(5), DomainBound::Int64(15))),
277            None,
278        ]);
279        assert!(tile_overlaps_slice(&mins, &maxs, &sl));
280    }
281}