1use crate::{KevyError, KevyResult};
19use std::io;
20use std::sync::RwLock;
21
22use kevy_index::{
23 Catalog, Cursor, IndexKind, IndexSpec, IndexValue, Segment, SegmentStats, ValType,
24};
25
26use crate::store::{Store, lock_write};
27
28pub(crate) use crate::ops_index_sync::{each_written_key_pub, on_commit, sync_segs};
29
30pub type IndexPage = (Vec<(Vec<u8>, IndexValue)>, Option<Cursor>);
47
48#[cfg(feature = "text")]
50pub type FieldSpans = (Vec<u8>, Vec<(u32, u32)>);
51#[cfg(feature = "text")]
53pub type HighlightedHit = (Vec<u8>, f64, Vec<FieldSpans>);
54
55#[cfg(feature = "text")]
60#[path = "ops_index_highlight.rs"]
61pub(crate) mod highlight;
62
63#[path = "ops_index_claused.rs"]
67pub(crate) mod claused;
68
69#[cfg(feature = "text")]
71#[path = "ops_index_match_opts.rs"]
72pub(crate) mod match_opts;
73#[path = "ops_index_opts.rs"]
74pub(crate) mod opts;
75
76#[path = "ops_index_advise.rs"]
78pub(crate) mod advise;
79
80#[path = "ops_index_admin.rs"]
82mod admin;
83
84#[cfg(feature = "text")]
85#[path = "ops_index_text.rs"]
86mod text;
87
88#[cfg(feature = "text")]
90#[path = "ops_index_text_cold.rs"]
91pub(crate) mod text_cold;
92
93pub(crate) fn merge_page(mut all: Vec<(IndexValue, Vec<u8>)>, limit: usize) -> IndexPage {
99 all.sort();
100 all.truncate(limit);
101 let next = if all.len() == limit {
102 all.last().map(|(v, k)| Cursor::new(v.clone(), k.clone()))
103 } else {
104 None
105 };
106 (all.into_iter().map(|(v, k)| (k, v)).collect(), next)
107}
108
109pub(crate) fn built(b: kevy_index::IndexSpecBuilder) -> KevyResult<IndexSpec> {
111 Ok(b.build().map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?)
112}
113
114#[derive(Debug, Default)]
118pub(crate) struct IndexReg {
119 pub(crate) catalog: RwLock<(u64, Catalog)>,
120 pub(crate) usage:
121 RwLock<std::collections::HashMap<Vec<u8>, std::sync::Arc<kevy_index::UsageCell>>>,
122}
123
124#[cfg(not(target_arch = "wasm32"))]
127pub(crate) type WinRef<'a> = Option<&'a kevy_window::WindowRt>;
128#[cfg(target_arch = "wasm32")]
129pub(crate) type WinRef<'a> = Option<&'a core::convert::Infallible>;
130
131#[derive(Debug, Default)]
134pub(crate) struct ShardSegs {
135 pub(crate) version: u64,
136 pub(crate) segs: Vec<(IndexSpec, Segment)>,
137 #[cfg(feature = "text")]
140 pub(crate) text: Vec<(IndexSpec, kevy_text::TextSegment)>,
141 #[cfg(feature = "vector")]
143 pub(crate) ann: Vec<(IndexSpec, kevy_vector::Hnsw)>,
144 pub(crate) agg: Vec<(IndexSpec, kevy_index::AggSegment)>,
146 #[cfg(not(target_arch = "wasm32"))]
150 pub(crate) windows: Vec<(Vec<u8>, kevy_window::WindowRt)>,
151 #[cfg(all(feature = "text", not(target_arch = "wasm32")))]
155 pub(crate) cold_text: Vec<(Vec<u8>, kevy_window::TextColdDir)>,
156 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
163 pub(crate) stats_dirty: bool,
164 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
165 pub(crate) reserved_cache: u64,
166 pub(crate) drain: crate::ops_index_changes::Drain,
168}
169
170impl ShardSegs {
171 #[cfg(not(target_arch = "wasm32"))]
176 pub(crate) fn window_of(&self, name: &[u8]) -> Option<&kevy_window::WindowRt> {
177 self.windows.iter().find(|(n, _)| n == name).map(|(_, w)| w)
178 }
179
180 #[inline]
184 pub(crate) fn mark_stats_dirty(&mut self) {
185 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
186 {
187 self.stats_dirty = true;
188 }
189 }
190}
191
192impl Store {
193 pub fn idx_create(
196 &self,
197 name: &[u8],
198 prefix: &[u8],
199 field: &[u8],
200 ty: ValType,
201 kind: IndexKind,
202 ) -> KevyResult<()> {
203 if prefix.is_empty() {
204 return Err(KevyError::InvalidInput("empty prefix".into()));
205 }
206 #[cfg(not(feature = "text"))]
207 if kind == IndexKind::Text {
208 return Err(KevyError::Unsupported("text indexes need the `text` feature".into()));
209 }
210 #[cfg(not(feature = "vector"))]
211 if kind == IndexKind::Ann {
212 return Err(KevyError::Unsupported("vector indexes need the `vector` feature".into()));
213 }
214 let spec = built(IndexSpec::builder(name, prefix, kind, ty).with_field(field))?;
215 self.catalog_change(|| self.register_spec(spec))
216 }
217
218 pub(crate) fn register_spec(&self, spec: IndexSpec) -> KevyResult<()> {
221 #[cfg(all(feature = "tier", not(target_arch = "wasm32")))]
224 crate::ops_index_sync::tier_floor_check(&self.shards)?;
225 {
226 let mut g =
227 self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
228 let (ver, cat) = &mut *g;
229 cat.create(spec).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))?;
230 *ver += 1;
231 }
232 self.advise_clear();
233 self.usage_rekey();
234 for shard in self.shards.iter() {
236 let mut g = lock_write(shard);
237 let inner = &mut *g;
238 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
239 }
240 Ok(())
241 }
242
243 #[cfg(feature = "vector")]
246 pub fn idx_create_ann(
247 &self,
248 name: &[u8],
249 prefix: &[u8],
250 field: &[u8],
251 params: kevy_index::AnnSpec,
252 ) -> KevyResult<()> {
253 if params.dim == 0 || params.distance > 2 || params.m == 1 {
255 return Err(KevyError::InvalidInput("bad ann parameters".into()));
256 }
257 let ann = params
258 .with_m(if params.m == 0 { 16 } else { params.m })
259 .with_ef(if params.ef == 0 { 200 } else { params.ef });
260 let spec = IndexSpec::builder(name, prefix, IndexKind::Ann, ValType::Vector)
261 .with_field(field)
262 .with_ann(ann);
263 let spec = built(spec)?;
264 self.catalog_change(|| self.register_spec(spec))
265 }
266
267 pub fn idx_drop(&self, name: &[u8]) -> KevyResult<bool> {
270 self.catalog_change(|| Ok(self.drop_index(name)))
271 }
272
273 pub(crate) fn drop_index(&self, name: &[u8]) -> bool {
276 let hit = {
277 let mut g =
278 self.indexes.catalog.write().unwrap_or_else(std::sync::PoisonError::into_inner);
279 let (ver, cat) = &mut *g;
280 let hit = cat.drop_index(name);
281 if hit {
282 *ver += 1;
283 }
284 hit
285 };
286 if hit {
287 self.advise_clear();
288 self.usage_rekey();
289 }
290 hit
291 }
292
293 #[cfg(feature = "text")]
302 pub fn idx_match(
303 &self,
304 name: &[u8],
305 query: &[u8],
306 limit: usize,
307 ) -> KevyResult<Vec<(Vec<u8>, f64)>> {
308 Ok(self
309 .idx_match_with(name, query, limit, crate::MatchOpts::default())?
310 .into_iter()
311 .map(|(key, score, _)| (key, score))
312 .collect())
313 }
314
315 pub fn idx_create_agg(
318 &self,
319 name: &[u8],
320 prefix: &[u8],
321 field: &[u8],
322 ty: ValType,
323 group_by: &[u8],
324 ) -> KevyResult<()> {
325 if !matches!(ty, ValType::I64 | ValType::F64) || group_by.is_empty() {
326 return Err(KevyError::InvalidInput("agg requires numeric type + group field".into()));
327 }
328 let spec = IndexSpec::builder(name, prefix, IndexKind::Agg, ty)
329 .with_field(field)
330 .with_group_by(group_by);
331 let spec = built(spec)?;
332 self.catalog_change(|| self.register_spec(spec))
333 }
334
335 pub fn idx_group(&self, name: &[u8], group: &[u8]) -> KevyResult<kevy_index::GroupStats> {
337 let mut merged = kevy_index::GroupStats::default();
338 let mut found = false;
339 for shard in self.shards.iter() {
340 let mut g = lock_write(shard);
341 let inner = &mut *g;
342 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
343 if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name() == name) {
344 found = true;
345 merged.merge(&a.group(group));
346 }
347 }
348 if !found {
349 return Err(KevyError::NotFound("no such aggregate index".into()));
350 }
351 Ok(merged)
352 }
353
354 pub fn idx_groups(
356 &self,
357 name: &[u8],
358 by: kevy_index::AggBy,
359 limit: usize,
360 ) -> KevyResult<Vec<(Vec<u8>, kevy_index::GroupStats)>> {
361 let limit = limit.clamp(1, 1000);
362 let mut merged: std::collections::HashMap<Vec<u8>, kevy_index::GroupStats> =
365 std::collections::HashMap::new();
366 let mut found = false;
367 for shard in self.shards.iter() {
368 let mut g = lock_write(shard);
369 let inner = &mut *g;
370 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
371 if let Some((_, a)) = inner.idx_segs.agg.iter().find(|(s, _)| s.name() == name) {
372 found = true;
373 for (gk, st) in a.all_groups() {
374 match merged.get_mut(&gk) {
375 Some(m) => m.merge(&st),
376 None => {
377 merged.insert(gk, st);
378 }
379 }
380 }
381 }
382 }
383 if !found {
384 return Err(KevyError::NotFound("no such aggregate index".into()));
385 }
386 let mut ranked: Vec<(Vec<u8>, kevy_index::GroupStats)> = merged.into_iter().collect();
387 kevy_index::sort_groups(&mut ranked, by);
388 ranked.truncate(limit);
389 Ok(ranked)
390 }
391
392 #[cfg(feature = "vector")]
395 pub fn idx_knn(
396 &self,
397 name: &[u8],
398 query: &[f32],
399 k: usize,
400 ef: usize,
401 ) -> KevyResult<Vec<(Vec<u8>, f32)>> {
402 let k = k.clamp(1, 1000);
403 let mut all: Vec<(Vec<u8>, f32)> = Vec::new();
404 let mut found = false;
405 for shard in self.shards.iter() {
406 let mut g = lock_write(shard);
407 let inner = &mut *g;
408 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
409 if let Some((_, graph)) = inner.idx_segs.ann.iter().find(|(s, _)| s.name() == name) {
410 found = true;
411 all.extend(graph.knn(query, k, ef));
412 }
413 }
414 if !found {
415 return Err(KevyError::NotFound("no such vector index".into()));
416 }
417 all.sort_by(|a, b| a.1.total_cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
418 all.truncate(k);
419 Ok(all)
420 }
421
422 fn for_each_segment(&self, name: &[u8], mut f: impl FnMut(&Segment)) -> KevyResult<()> {
423 let mut found = false;
424 for shard in self.shards.iter() {
425 let mut g = lock_write(shard);
426 let inner = &mut *g;
427 sync_segs(&self.indexes, &mut inner.idx_segs, &mut inner.store);
428 if let Some((_, seg)) = inner.idx_segs.segs.iter().find(|(s, _)| s.name() == name) {
429 found = true;
430 f(seg);
431 }
432 }
433 if found { Ok(()) } else { Err(KevyError::NotFound("no such index".into())) }
434 }
435}