1use kevy_index::{
8 Cursor, FacetBucket, IndexSpec, IndexValue, ScalarClauses, ScalarHit, ValType, ValueTest,
9 fold_facets, merge_claused, sort_facets,
10};
11
12use super::sync_segs;
13use crate::store::{Store, lock_write};
14use crate::{KevyError, KevyResult};
15
16#[derive(Clone, Copy)]
23pub enum ValueFilter<'a> {
24 Range {
26 field: &'a [u8],
28 min: &'a [u8],
30 max: &'a [u8],
32 },
33 Eq {
35 field: &'a [u8],
37 value: &'a [u8],
39 },
40}
41
42impl ValueFilter<'_> {
43 pub(crate) fn field(&self) -> &[u8] {
44 match self {
45 ValueFilter::Range { field, .. } | ValueFilter::Eq { field, .. } => field,
46 }
47 }
48}
49
50pub(crate) fn value_test(spec: &IndexSpec, f: &ValueFilter<'_>) -> KevyResult<(usize, ValueTest)> {
53 let stored: Vec<&[u8]> = spec.values.iter().map(|v| v.name.as_slice()).collect();
54 let pos = spec
55 .values
56 .iter()
57 .position(|v| v.name == f.field())
58 .ok_or_else(|| unknown_field("FILTER", f.field(), "store", &stored))?;
59 let ty = spec.values[pos].ty;
60 let now = (kevy_store::now_unix_ms() / 1000) as i64;
64 let (test, raw) = match f {
65 ValueFilter::Range { min, max, .. } => (ValueTest::range_at(ty, min, max, now), *min),
66 ValueFilter::Eq { value, .. } => (ValueTest::eq_at(ty, value, now), *value),
67 };
68 let test = test.ok_or_else(|| {
69 KevyError::InvalidInput(format!(
70 "FILTER bound '{}' is not a valid {}, which is how this index declares '{}'",
71 String::from_utf8_lossy(raw),
72 ty.tag(),
73 String::from_utf8_lossy(f.field()),
74 ))
75 })?;
76 Ok((pos, test))
77}
78
79fn unstored_field(spec: &IndexSpec, opts: &ScalarQueryOpts<'_>) -> Option<Vec<u8>> {
83 let stored = |f: &[u8]| spec.values.iter().any(|v| v.name == f);
84 for f in opts.filters {
85 if !stored(f.field()) {
86 return Some(f.field().to_vec());
87 }
88 }
89 if let Some((f, _)) = opts.sort
90 && !stored(f)
91 {
92 return Some(f.to_vec());
93 }
94 if let Some(f) = opts.distinct
95 && !stored(f)
96 {
97 return Some(f.to_vec());
98 }
99 opts.facets.iter().find(|f| !stored(f)).cloned()
100}
101
102pub(crate) fn unknown_field(clause: &str, bad: &[u8], verb: &str, offered: &[&[u8]]) -> KevyError {
104 let names: Vec<String> =
105 offered.iter().map(|n| String::from_utf8_lossy(n).into_owned()).collect();
106 KevyError::InvalidInput(format!(
107 "{clause} names field '{}', which this index does not {verb} — it {}: {}",
108 String::from_utf8_lossy(bad),
109 if verb.ends_with(['s', 'x', 'z']) { format!("{verb}es") } else { format!("{verb}s") },
112 names.join(", ")
113 ))
114}
115
116#[derive(Clone, Copy, Default)]
120pub struct ScalarQueryOpts<'a> {
121 pub filters: &'a [ValueFilter<'a>],
124 pub sort: Option<(&'a [u8], bool)>,
127 pub distinct: Option<&'a [u8]>,
130 pub facets: &'a [Vec<u8>],
133 pub offset: usize,
135}
136
137impl ScalarQueryOpts<'_> {
138 pub(crate) fn selects(&self) -> bool {
141 self.sort.is_some() || self.distinct.is_some() || !self.facets.is_empty() || self.offset > 0
142 }
143}
144
145#[derive(Debug)]
149pub struct ScalarPage {
150 pub rows: Vec<(Vec<u8>, IndexValue)>,
152 pub facets: Vec<Vec<(Vec<u8>, u64)>>,
154 pub cursor: Option<Cursor>,
157}
158
159type GatheredPages = (Vec<(ScalarHit, ())>, Vec<Vec<FacetBucket>>);
163
164struct Resolved {
166 filters: Vec<(usize, ValueTest)>,
167 sort: Option<(usize, bool, ValType)>,
168 distinct: Option<(usize, ValType)>,
169 facets: Vec<(usize, ValType)>,
170}
171
172fn value_field(spec: &IndexSpec, clause: &str, field: &[u8]) -> KevyResult<(usize, ValType)> {
174 let stored: Vec<&[u8]> = spec.values.iter().map(|v| v.name.as_slice()).collect();
175 let pos = spec
176 .values
177 .iter()
178 .position(|v| v.name == field)
179 .ok_or_else(|| unknown_field(clause, field, "store", &stored))?;
180 Ok((pos, spec.values[pos].ty))
181}
182
183fn resolve(spec: &IndexSpec, opts: &ScalarQueryOpts<'_>) -> KevyResult<Resolved> {
184 let filters =
185 opts.filters.iter().map(|f| value_test(spec, f)).collect::<KevyResult<Vec<_>>>()?;
186 let sort = match opts.sort {
187 Some((field, desc)) => {
188 let (pos, ty) = value_field(spec, "SORT", field)?;
189 Some((pos, desc, ty))
190 }
191 None => None,
192 };
193 let distinct = match opts.distinct {
194 Some(field) => Some(value_field(spec, "DISTINCT", field)?),
195 None => None,
196 };
197 let facets = opts
198 .facets
199 .iter()
200 .map(|f| value_field(spec, "FACET", f))
201 .collect::<KevyResult<Vec<_>>>()?;
202 Ok(Resolved { filters, sort, distinct, facets })
203}
204
205impl Store {
206 pub fn idx_create_with_values(
210 &self,
211 name: &[u8],
212 prefix: &[u8],
213 field: &[u8],
214 ty: ValType,
215 kind: kevy_index::IndexKind,
216 values: &[(&[u8], ValType)],
217 ) -> KevyResult<()> {
218 if prefix.is_empty() {
219 return Err(KevyError::InvalidInput("empty prefix".into()));
220 }
221 let spec = IndexSpec {
222 name: name.to_vec(),
223 prefix: prefix.to_vec(),
224 fields: vec![kevy_index::FieldSpec::new(field.to_vec())],
225 ty,
226 kind,
227 max_bytes: 0,
228 ann: None,
229 group_by: None,
230 with_positions: false,
231 values: values
232 .iter()
233 .map(|(n, t)| kevy_index::ValueSpec { name: n.to_vec(), ty: *t })
234 .collect(),
235 composite: None,
236 };
237 self.register_spec(spec)
238 }
239
240 pub fn idx_query_claused(
247 &self,
248 name: &[u8],
249 min: &IndexValue,
250 max: &IndexValue,
251 cursor: Option<&Cursor>,
252 limit: usize,
253 opts: ScalarQueryOpts<'_>,
254 ) -> KevyResult<ScalarPage> {
255 let limit = limit.clamp(1, 100_000);
256 let offset = opts.offset.min(10_000);
257 let spec = self.claused_spec(name)?;
258 let r = self.claused_resolve(name, &spec, &opts)?;
259 let clauses = ScalarClauses {
260 filters: &r.filters,
261 sort: r.sort,
262 distinct: r.distinct,
263 facets: &r.facets,
264 fetch: limit + offset,
265 };
266 let (all, mut facets) = self.gather_claused(name, min, max, cursor, &clauses)?;
267 let sort_desc = r.sort.map(|(_, desc, _)| desc);
268 let all = merge_claused(all, sort_desc, r.distinct.is_some(), offset, limit);
269 sort_facets(&mut facets);
270 let next = (!opts.selects() && all.len() == limit)
271 .then(|| {
272 all.last().map(|(h, ())| Cursor { value: h.value.clone(), key: h.key.clone() })
273 })
274 .flatten();
275 self.observe_hit(name);
276 Ok(ScalarPage {
277 rows: all.into_iter().map(|(h, ())| (h.key, h.value)).collect(),
278 facets: facets
279 .into_iter()
280 .map(|f| f.into_iter().map(|(_, label, n)| (label, n)).collect())
281 .collect(),
282 cursor: next,
283 })
284 }
285
286 pub fn idx_count_claused(
291 &self,
292 name: &[u8],
293 min: &IndexValue,
294 max: &IndexValue,
295 filters: &[ValueFilter<'_>],
296 ) -> KevyResult<u64> {
297 let spec = self.claused_spec(name)?;
298 let opts = ScalarQueryOpts { filters, ..ScalarQueryOpts::default() };
299 let r = self.claused_resolve(name, &spec, &opts)?;
300 let mut total = 0u64;
301 #[cfg(not(target_arch = "wasm32"))]
302 let probe = self.usage_cell(name);
303 self.for_each_segment_windowed(name, |spec, seg, win| {
304 total += seg.count_claused(min, max, &r.filters);
305 #[cfg(not(target_arch = "wasm32"))]
306 if let Some(w) = win {
307 crate::ops_index::advise::probe_window(&probe, w, min);
308 }
309 #[cfg(not(target_arch = "wasm32"))]
312 if let Some(w) = win.filter(|w| w.has_cold()) {
313 total += w
314 .cold_claused_count(spec.ty, min, max, &r.filters)
315 .map_err(|e| KevyError::Io(std::io::Error::other(e)))?;
316 }
317 #[cfg(target_arch = "wasm32")]
318 let _ = (spec, win);
319 Ok(())
320 })?;
321 self.observe_hit(name);
322 Ok(total)
323 }
324
325 fn claused_spec(&self, name: &[u8]) -> KevyResult<IndexSpec> {
328 let spec = {
329 let g = self.indexes.catalog.read().unwrap_or_else(std::sync::PoisonError::into_inner);
330 g.1.get(name).map(|(s, _)| s.clone())
331 };
332 spec.ok_or_else(|| {
333 self.observe_refused(name, kevy_index::AdviseShape::Range);
334 KevyError::NotFound("no such index".into())
335 })
336 }
337
338 fn claused_resolve(
341 &self,
342 name: &[u8],
343 spec: &IndexSpec,
344 opts: &ScalarQueryOpts<'_>,
345 ) -> KevyResult<Resolved> {
346 resolve(spec, opts).inspect_err(|_| {
347 if let Some(f) = unstored_field(spec, opts) {
348 self.observe_refused(name, kevy_index::AdviseShape::Filter(f));
349 }
350 })
351 }
352
353 fn gather_claused(
356 &self,
357 name: &[u8],
358 min: &IndexValue,
359 max: &IndexValue,
360 cursor: Option<&Cursor>,
361 clauses: &ScalarClauses<'_>,
362 ) -> KevyResult<GatheredPages> {
363 let mut all: Vec<(ScalarHit, ())> = Vec::new();
364 let mut facets: Vec<Vec<FacetBucket>> = vec![Vec::new(); clauses.facets.len()];
365 let mut found = false;
366 #[cfg(not(target_arch = "wasm32"))]
367 let probe = self.usage_cell(name);
368 for shard in self.shards.iter() {
369 let mut g = lock_write(shard);
370 let inner = &mut *g;
371 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
372 if let Some((_spec, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name == name) {
373 found = true;
374 #[cfg(not(target_arch = "wasm32"))]
375 if let Some(w) = inner.idx_segs.window_of(name) {
376 crate::ops_index::advise::probe_window(&probe, w, min);
377 }
378 let page = seg.query_claused(min, max, cursor, clauses);
379 all.extend(page.hits.into_iter().map(|h| (h, ())));
380 fold_facets(&mut facets, page.facets);
381 #[cfg(not(target_arch = "wasm32"))]
386 if let Some(w) = inner.idx_segs.window_of(name).filter(|w| w.has_cold()) {
387 let (chits, cfacets) = w
388 .cold_claused(_spec.ty, min, max, cursor, clauses)
389 .map_err(|e| KevyError::Io(std::io::Error::other(e)))?;
390 all.extend(chits.into_iter().map(|h| (h, ())));
391 fold_facets(&mut facets, cfacets);
392 }
393 }
394 }
395 if !found {
396 return Err(KevyError::NotFound("no such index".into()));
397 }
398 Ok((all, facets))
399 }
400}