Skip to main content

rushdb_query/
index_access.rs

1//! Declared ordered indexes in the planner (SD-4.4): when every
2//! requested label has a `Ready` ordered index on a leaf's property, the leaf
3//! is answered from `INDEX_ENTRY` rows instead of a presence-narrowed walk.
4//! Each candidate entry is re-checked against the predicate, so the owner
5//! set is exact.
6
7use rushdb_lmpg::index;
8use rushdb_lmpg::reader::TryRead;
9use rushdb_types::encoding::{decode_order_preserving, encode_order_preserving, logical_tag};
10use rushdb_types::ids::TenantId;
11use rushdb_types::index::IndexDef;
12use rushdb_types::model::LogicalType;
13use rushdb_types::value::{normalize_name, Scalar};
14
15use crate::planner::{range_scalar_matches, scalar_variants, BoundedOwners, OwnerSet, RangeOp};
16use crate::{Expr, FieldPred, QueryError};
17
18type Json = serde_json::Value;
19
20/// The `Ready` ordered indexes covering property `name` for every label in
21/// `labels` (one per label), or `None` when some label has none or no label
22/// is given (an index belongs to a label).
23pub(crate) fn covering<R: TryRead>(
24    snap: &R,
25    tenant: TenantId,
26    name: &str,
27    labels: &[String],
28) -> Result<Option<Vec<IndexDef>>, QueryError> {
29    let norm = normalize_name(name);
30    if labels.is_empty() || norm.is_empty() || rushdb_types::value::is_reserved_name(&norm) {
31        return Ok(None);
32    }
33    let mut out = Vec::with_capacity(labels.len());
34    for label in labels {
35        match index::ready_ordered(snap, tenant, label, &norm)? {
36            Some(def) => out.push(def),
37            None => return Ok(None),
38        }
39    }
40    Ok(Some(out))
41}
42
43/// `[start, end)` covering every entry that begins with `prefix`.
44fn prefix_bounds(prefix: Vec<u8>) -> (Vec<u8>, Vec<u8>) {
45    let end = rushdb_storage::prefix_range(&prefix).1;
46    (prefix, end)
47}
48
49fn scan<R: TryRead>(
50    snap: &R,
51    tenant: TenantId,
52    defs: &[IndexDef],
53    (start, end): (&[u8], &[u8]),
54    keep: &dyn Fn(&[u8]) -> bool,
55    owners: &mut OwnerSet,
56    owner_limit: usize,
57) -> Result<(), QueryError> {
58    // Streamed straight into the owner set: an entry range of millions of
59    // owners never exists as a row list.
60    let mut over_budget = None;
61    let mut owners = BoundedOwners::new(owners, owner_limit);
62    for def in defs {
63        index::for_each_entry(snap, tenant, def.id, start, end, &mut |entry, owner| {
64            if !keep(entry) {
65                return true;
66            }
67            match owners.insert(owner) {
68                Ok(()) => true,
69                Err(error) => {
70                    over_budget = Some(error);
71                    false
72                }
73            }
74        })?;
75        if let Some(error) = over_budget {
76            return Err(error);
77        }
78    }
79    Ok(())
80}
81
82fn entry_of(ty: LogicalType, scalar: &Scalar) -> Option<Vec<u8>> {
83    let mut entry = vec![logical_tag(ty)];
84    entry.extend(encode_order_preserving(ty, scalar).ok()?);
85    Some(entry)
86}
87
88/// Owners whose value equals any of `values` (every type variant a value
89/// may be stored under, as the presence path probes).
90fn equal<R: TryRead>(
91    snap: &R,
92    tenant: TenantId,
93    defs: &[IndexDef],
94    values: &[Json],
95    owner_limit: usize,
96) -> Result<Option<OwnerSet>, QueryError> {
97    let mut owners = OwnerSet::new();
98    for value in values {
99        let Some(variants) = scalar_variants(value) else {
100            return Ok(None);
101        };
102        for (ty, scalar) in variants {
103            let Some(entry) = entry_of(ty, &scalar) else {
104                return Ok(None);
105            };
106            let (start, end) = prefix_bounds(entry.clone());
107            scan(
108                snap,
109                tenant,
110                defs,
111                (&start, &end),
112                &|found| found == entry.as_slice(),
113                &mut owners,
114                owner_limit,
115            )?;
116        }
117    }
118    Ok(Some(owners))
119}
120
121/// Owners holding a numeric value satisfying `op bound`.
122fn range<R: TryRead>(
123    snap: &R,
124    tenant: TenantId,
125    defs: &[IndexDef],
126    op: RangeOp,
127    bound: f64,
128    owner_limit: usize,
129) -> Result<Option<OwnerSet>, QueryError> {
130    let mut owners = OwnerSet::new();
131    let above = matches!(op, RangeOp::Gt | RangeOp::Gte);
132    for ty in [
133        LogicalType::Int64,
134        LogicalType::Float64,
135        LogicalType::DateTime,
136    ] {
137        let tag = logical_tag(ty);
138        // Bound the scan at the bound's encoding, rounded outward for the
139        // integer domains; the exact comparison happens on each entry.
140        let edge = match ty {
141            LogicalType::Float64 => Scalar::Float64(bound),
142            LogicalType::Int64 => Scalar::Int64(outward(bound, above)),
143            _ => Scalar::DateTime(outward(bound, above)),
144        };
145        let Some(edge) = entry_of(ty, &edge) else {
146            continue;
147        };
148        let (start, end) = if above {
149            (edge, vec![tag + 1])
150        } else {
151            (vec![tag], rushdb_storage::prefix_range(&edge).1)
152        };
153        scan(
154            snap,
155            tenant,
156            defs,
157            (&start, &end),
158            &|entry| {
159                entry.split_first().is_some_and(|(&found, payload)| {
160                    found == tag
161                        && decode_order_preserving(ty, payload)
162                            .is_ok_and(|scalar| range_scalar_matches(&scalar, op, bound))
163                })
164            },
165            &mut owners,
166            owner_limit,
167        )?;
168    }
169    Ok(Some(owners))
170}
171
172#[expect(
173    clippy::cast_possible_truncation,
174    reason = "floor/ceil made the bound integral; `as` saturates at the i64 edges"
175)]
176fn outward(bound: f64, above: bool) -> i64 {
177    if above {
178        bound.floor() as i64
179    } else {
180        bound.ceil() as i64
181    }
182}
183
184/// Owners holding a string starting with `prefix`.
185fn starts_with<R: TryRead>(
186    snap: &R,
187    tenant: TenantId,
188    defs: &[IndexDef],
189    prefix: &str,
190    owner_limit: usize,
191) -> Result<Option<OwnerSet>, QueryError> {
192    let mut entry = vec![logical_tag(LogicalType::Utf8)];
193    entry.extend_from_slice(prefix.as_bytes());
194    let (start, end) = prefix_bounds(entry);
195    let mut owners = OwnerSet::new();
196    scan(
197        snap,
198        tenant,
199        defs,
200        (&start, &end),
201        &|entry| {
202            decode_order_preserving(LogicalType::Utf8, &entry[1..]).is_ok_and(
203                |scalar| matches!(scalar, Scalar::Utf8(text) if text.starts_with(prefix)),
204            )
205        },
206        &mut owners,
207        owner_limit,
208    )?;
209    Ok(Some(owners))
210}
211
212/// The leaf's property name when an ordered index can answer it.
213fn indexed_property(expr: &Expr) -> Option<&str> {
214    match expr {
215        Expr::Field(
216            FieldPred::Eq(name, _)
217            | FieldPred::In(name, _)
218            | FieldPred::Gt(name, _)
219            | FieldPred::Gte(name, _)
220            | FieldPred::Lt(name, _)
221            | FieldPred::Lte(name, _)
222            | FieldPred::StartsWith(name, _),
223        ) => Some(name),
224        _ => None,
225    }
226}
227
228/// Exact owners of `expr` from declared indexes, or `None` when no `Ready`
229/// ordered index covers it for every label in `labels`.
230pub(crate) fn leaf_owners<R: TryRead>(
231    snap: &R,
232    tenant: TenantId,
233    expr: &Expr,
234    labels: &[String],
235    owner_limit: usize,
236) -> Result<Option<OwnerSet>, QueryError> {
237    let Some(name) = indexed_property(expr) else {
238        return Ok(None);
239    };
240    let Some(defs) = covering(snap, tenant, name, labels)? else {
241        return Ok(None);
242    };
243    // A leaf's set is a function of the snapshot: readers of one snapshot
244    // share it instead of each reading every index entry again.
245    let cache = snap.owner_sets();
246    let ids: Vec<u32> = defs.iter().map(|def| def.id).collect();
247    let key = format!("{tenant:?}|{ids:?}|{expr:?}").into_bytes();
248    if let Some(hit) = cache
249        .and_then(|(sets, ci)| sets.get(&key, ci))
250        .and_then(|set| set.downcast::<OwnerSet>().ok())
251    {
252        if hit.len() > u64::try_from(owner_limit).unwrap_or(u64::MAX) {
253            return Err(QueryError::ResourceLimit(
254                "aggregate candidate budget exceeded".into(),
255            ));
256        }
257        return Ok(Some(OwnerSet::clone(&hit)));
258    }
259    let owners = leaf_owners_uncached(snap, tenant, expr, &defs, owner_limit)?;
260    if let (Some(owners), Some((sets, ci))) = (owners.as_ref(), cache) {
261        sets.insert(
262            key,
263            ci,
264            std::sync::Arc::new(owners.clone()),
265            owners.serialized_size(),
266        );
267    }
268    Ok(owners)
269}
270
271fn leaf_owners_uncached<R: TryRead>(
272    snap: &R,
273    tenant: TenantId,
274    expr: &Expr,
275    defs: &[IndexDef],
276    owner_limit: usize,
277) -> Result<Option<OwnerSet>, QueryError> {
278    let defs = defs.to_vec();
279    match expr {
280        Expr::Field(FieldPred::Eq(_, value)) if !value.is_null() => equal(
281            snap,
282            tenant,
283            &defs,
284            std::slice::from_ref(value),
285            owner_limit,
286        ),
287        Expr::Field(FieldPred::In(_, values)) if values.iter().all(|v| !v.is_null()) => {
288            equal(snap, tenant, &defs, values, owner_limit)
289        }
290        Expr::Field(FieldPred::Gt(_, bound)) => {
291            range(snap, tenant, &defs, RangeOp::Gt, *bound, owner_limit)
292        }
293        Expr::Field(FieldPred::Gte(_, bound)) => {
294            range(snap, tenant, &defs, RangeOp::Gte, *bound, owner_limit)
295        }
296        Expr::Field(FieldPred::Lt(_, bound)) => {
297            range(snap, tenant, &defs, RangeOp::Lt, *bound, owner_limit)
298        }
299        Expr::Field(FieldPred::Lte(_, bound)) => {
300            range(snap, tenant, &defs, RangeOp::Lte, *bound, owner_limit)
301        }
302        Expr::Field(FieldPred::StartsWith(_, prefix)) => {
303            starts_with(snap, tenant, &defs, prefix, owner_limit)
304        }
305        _ => Ok(None),
306    }
307}
308
309/// Whether a declared index answers `expr` (a `startsWith` leaf is only
310/// eligible through one).
311pub(crate) fn answers<R: TryRead>(
312    snap: &R,
313    tenant: TenantId,
314    expr: &Expr,
315    labels: &[String],
316) -> Result<bool, QueryError> {
317    let Some(name) = indexed_property(expr) else {
318        return Ok(false);
319    };
320    Ok(covering(snap, tenant, name, labels)?.is_some())
321}
322
323/// Name of the first declared index the predicate's leaves can use, for
324/// `EXPLAIN`'s access path.
325///
326/// # Errors
327/// Returns the underlying storage error from an index-catalog read.
328pub fn explain_index<R: TryRead>(
329    snap: &R,
330    tenant: TenantId,
331    labels: &[String],
332    where_: &Expr,
333) -> Result<Option<String>, QueryError> {
334    let mut leaves = Vec::new();
335    crate::planner::collect_leaves(where_, &mut leaves);
336    for leaf in leaves {
337        if let Some(name) = indexed_property(leaf) {
338            if let Some(defs) = covering(snap, tenant, name, labels)? {
339                return Ok(Some(
340                    defs.iter()
341                        .map(|def| def.name.as_str())
342                        .collect::<Vec<_>>()
343                        .join(","),
344                ));
345            }
346        }
347    }
348    Ok(None)
349}