1use 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
20pub(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
43fn 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 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
88fn 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
121fn 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 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
184fn 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
212fn 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
228pub(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 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
309pub(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
323pub 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}