1use crate::keys;
13use crate::store::Store;
14use crate::Result;
15
16const CAT_NEXT_ID: u64 = 1;
18const CAT_VEC: u64 = 2;
21pub(crate) const CAT_NAV_PENDING: u64 = 3;
25
26#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
47pub struct VecMeta {
48 pub dim: u64,
49 pub seed: u64,
50 pub bits: u64,
51 pub max: u64,
52 pub medoid: u64,
53 pub watermark: u64,
54}
55
56impl VecMeta {
57 pub const LEN: usize = 48;
61
62 pub fn encode(&self) -> Vec<u8> {
63 let mut v = Vec::with_capacity(Self::LEN);
64 for w in [self.dim, self.seed, self.bits, self.max, self.medoid, self.watermark] {
65 v.extend_from_slice(&w.to_be_bytes());
66 }
67 v
68 }
69
70 pub fn decode(b: &[u8]) -> Option<VecMeta> {
71 if b.len() != Self::LEN { return None; }
72 let w = |i: usize| u64::from_be_bytes(b[i * 8..i * 8 + 8].try_into().unwrap());
73 Some(VecMeta { dim: w(0), seed: w(1), bits: w(2), max: w(3), medoid: w(4), watermark: w(5) })
74 }
75}
76
77pub struct Graph {
78 store: Store,
79 next_id: u64,
80 allocator_verified: bool,
84 vec_meta: std::collections::HashMap<u64, VecMeta>,
87 vec_dirty: std::collections::HashSet<u64>,
90 pub(crate) nav_fold_every: u64,
92 encoders: std::cell::RefCell<std::collections::HashMap<u64, std::sync::Arc<crate::vecquant::Encoder>>>,
96 pub redge: bool,
100}
101
102impl Graph {
103 pub fn new(store: Store) -> Result<Graph> {
104 let counter = store.get(&keys::catalog(CAT_NEXT_ID))?;
105 let (next_id, allocator_verified) = match counter {
106 Some(v) if v.len() == 8 => {
107 let next = u64::from_be_bytes(v.try_into().unwrap());
108 (next, next >= 1)
109 }
110 Some(_) => (0, false),
111 None => {
112 let mut rows = store.scan(&[keys::TAG_NODE])?;
116 let empty = match rows.next() {
117 None => true,
118 Some(Ok((key, _))) => key.first() != Some(&keys::TAG_NODE),
119 Some(Err(e)) => return Err(e),
120 };
121 if empty { (1, true) } else { (0, false) }
122 }
123 };
124 let mut vec_meta = std::collections::HashMap::new();
127 let vpre = keys::catalog_field_prefix(CAT_VEC);
128 store.scan(&vpre)?.for_each_ref(|k, v| {
129 if !k.starts_with(&vpre) { return false; }
130 if k.len() == 17 {
131 if let Some(m) = VecMeta::decode(v) { vec_meta.insert(keys::u64_at(k, 9), m); }
132 }
133 true
134 })?;
135 Ok(Graph { store, next_id, allocator_verified, vec_meta, vec_dirty: Default::default(),
136 nav_fold_every: crate::nav::NAV_FOLD_CHECKPOINT_EVERY,
137 encoders: Default::default(), redge: true })
138 }
139
140 pub fn store(&mut self) -> &mut Store { &mut self.store }
141 pub fn store_ref(&self) -> &Store { &self.store }
142
143 pub fn commit(&mut self) -> Result<()> {
146 if !self.allocator_verified {
147 return Err(crate::Error::Corrupt {
148 page_no: 0,
149 why: "id allocator counter is missing or malformed",
150 });
151 }
152 self.store.put(&keys::catalog(CAT_NEXT_ID), &self.next_id.to_be_bytes())?;
153 for field in std::mem::take(&mut self.vec_dirty) {
154 if let Some(m) = self.vec_meta.get(&field).copied() {
155 self.store.put(&keys::catalog_field(CAT_VEC, field), &m.encode())?;
156 }
157 }
158 self.store.commit()
159 }
160
161 pub fn checkpoint(&mut self) -> Result<()> { self.store.checkpoint() }
162
163 pub fn set_nav_fold_interval(&mut self, every: u64) { self.nav_fold_every = every.max(1); }
165
166 pub fn add_node(&mut self, ext: Option<&[u8]>, label: u64, props: &[u8]) -> Result<u64> {
171 if !self.allocator_verified {
172 return Err(crate::Error::Corrupt {
173 page_no: 0,
174 why: "id allocator counter is missing or malformed",
175 });
176 }
177 let id = self.next_id;
178 self.next_id = self.next_id.checked_add(1).ok_or(crate::Error::Corrupt {
179 page_no: 0,
180 why: "id allocator counter is exhausted",
181 })?;
182 let mut v = Vec::with_capacity(8 + props.len());
183 v.extend_from_slice(&label.to_be_bytes());
184 v.extend_from_slice(props);
185 self.store.put(&keys::node(id), &v)?;
186 self.store.put(&keys::label(label, id), &[])?;
187 if let Some(e) = ext {
188 let mut ev = Vec::with_capacity(8 + e.len());
189 ev.extend_from_slice(&id.to_be_bytes());
190 ev.extend_from_slice(e);
191 self.store.put(&keys::extkey(keys::ext_hash(e)), &ev)?;
192 }
193 Ok(id)
194 }
195
196 pub fn resolve(&self, ext: &[u8]) -> Result<Option<u64>> {
199 Ok(match self.store.get(&keys::extkey(keys::ext_hash(ext)))? {
200 Some(v) if v.len() >= 8 && &v[8..] == ext => {
201 Some(u64::from_be_bytes(v[..8].try_into().unwrap()))
202 }
203 _ => None,
204 })
205 }
206
207 pub fn get_node(&self, id: u64) -> Result<Option<(u64, Vec<u8>)>> {
208 let Some(v) = self.store.get(&keys::node(id))? else { return Ok(None); };
209 let label = v.get(..8).ok_or(crate::Error::Corrupt {
210 page_no: 0,
211 why: "node row is shorter than its eight-byte label",
212 })?;
213 Ok(Some((u64::from_be_bytes(label.try_into().unwrap()), v[8..].to_vec())))
214 }
215
216 pub fn add_edge(&mut self, ctx: u64, src: u64, ty: u64, dst: u64, props: &[u8]) -> Result<()> {
228 self.store.put(&keys::edge(ctx, src, ty, dst), props)?;
229 if self.redge {
230 self.store.put(&keys::redge(ctx, dst, ty, src), &[])?;
231 }
232 Ok(())
233 }
234
235 pub fn out_edges(&self, ctx: u64, src: u64, ty: Option<u64>) -> Result<EdgeIter<'_>> {
238 let prefix = match ty {
239 Some(t) => keys::edge_type_prefix(ctx, src, t),
240 None => keys::edge_prefix(ctx, src),
241 };
242 Ok(EdgeIter { inner: self.store.scan(&prefix)?, prefix, rev: false, base: ctx == 0 })
243 }
244
245 pub fn in_edges(&self, ctx: u64, dst: u64, ty: Option<u64>) -> Result<EdgeIter<'_>> {
246 assert!(self.redge, "in_edges requires the redge keyspace");
247 let prefix = match ty {
248 Some(t) => keys::redge_type_prefix(ctx, dst, t),
249 None => keys::redge_prefix(ctx, dst),
250 };
251 Ok(EdgeIter { inner: self.store.scan(&prefix)?, prefix, rev: true, base: ctx == 0 })
252 }
253
254 pub fn nodes_with_label(&self, l: u64) -> Result<LabelIter<'_>> {
256 let prefix = keys::label_prefix(l);
257 Ok(LabelIter { inner: self.store.scan(&prefix)?, prefix })
258 }
259
260 pub fn set_prop(&mut self, prop: u64, value: u64, id: u64) -> Result<()> {
266 self.store.put(&keys::prop(prop, value, id), &[])
267 }
268
269 pub fn update_prop(&mut self, prop: u64, old: u64, new: u64, id: u64) -> Result<()> {
275 self.store.delete(&keys::prop(prop, old, id))?;
276 self.store.put(&keys::prop(prop, new, id), &[])
277 }
278
279 pub fn prop_range(&self, prop: u64, lo: u64, hi: u64) -> Result<PropIter<'_>> {
281 let from = keys::prop_value_prefix(prop, lo);
282 Ok(PropIter { inner: self.store.scan(&from)?, prop, hi })
283 }
284
285 pub fn prop_eq(&self, prop: u64, value: u64) -> Result<PropIter<'_>> {
287 self.prop_range(prop, value, value)
288 }
289
290 pub fn count_label(&self, l: u64) -> Result<usize> {
292 let prefix = keys::label_prefix(l);
293 let mut n = 0;
294 self.store.scan(&prefix)?.for_each_ref(|k, _| {
295 if k.len() == 17 && k.starts_with(&prefix) { n += 1; true } else { false }
296 })?;
297 Ok(n)
298 }
299
300 pub fn count_prop_range(&self, prop: u64, lo: u64, hi: u64) -> Result<usize> {
302 let from = keys::prop_value_prefix(prop, lo);
303 let mut n = 0;
304 self.store.scan(&from)?.for_each_ref(|k, _| {
305 if k.len() == 25 && k[0] == keys::TAG_PROP && keys::u64_at(&k, 1) == prop
306 && keys::u64_at(&k, 9) <= hi { n += 1; true } else { false }
307 })?;
308 Ok(n)
309 }
310
311 pub fn vec_seed(dim: u64) -> u64 {
318 0xC0FF_EE00_2600_u64 ^ dim.wrapping_mul(0x9E37_79B9_7F4A_7C15)
319 }
320
321 pub fn vec_meta(&self, field: u64) -> Option<VecMeta> { self.vec_meta.get(&field).copied() }
323
324 pub fn set_vec_meta(&mut self, field: u64, m: VecMeta) -> Result<()> {
328 self.vec_meta.insert(field, m);
329 self.vec_dirty.remove(&field);
330 self.store.put(&keys::catalog_field(CAT_VEC, field), &m.encode())
331 }
332
333 pub fn clear_nav(&mut self, field: u64) -> Result<()> {
337 self.store.delete_prefix(&keys::nav_prefix(field))?;
338 self.store.delete_prefix(&keys::catalog_field(CAT_NAV_PENDING, field))?;
339 if let Some(mut meta) = self.vec_meta.get(&field).copied() {
340 meta.medoid = 0;
341 meta.watermark = 0;
342 self.set_vec_meta(field, meta)?;
343 }
344 Ok(())
345 }
346
347 pub fn set_vec(&mut self, field: u64, id: u64, v: &[f32]) -> Result<()> {
354 let dim = v.len() as u64;
355 let meta = match self.vec_meta.get(&field).copied() {
356 Some(m) if m.dim != dim => return Err(crate::Error::TooLarge), Some(m) => m,
358 None => {
359 let m = VecMeta { dim, seed: Self::vec_seed(dim),
360 bits: crate::vecquant::DEFAULT_BITS as u64,
361 ..Default::default() };
362 self.set_vec_meta(field, m)?;
363 m
364 }
365 };
366 let mut bytes = Vec::with_capacity(v.len() * 4);
367 for x in v {
368 bytes.extend_from_slice(&x.to_le_bytes());
369 }
370 self.store.put(&keys::vec_key(field, id), &bytes)?;
371 let enc = self.encoder(field).expect("recipe was just established");
375 let (norm, code) = enc.encode(v);
376 let mut row = Vec::with_capacity(4 + code.len());
377 row.extend_from_slice(&norm.to_le_bytes());
378 row.extend_from_slice(&code);
379 if id > meta.max {
380 self.vec_meta.entry(field).and_modify(|m| m.max = id);
381 self.vec_dirty.insert(field);
382 }
383 self.store.put(&keys::vcode_key(field, id), &row)?;
384 if id <= meta.watermark {
388 self.store.put(&keys::catalog_field_item(CAT_NAV_PENDING, field, id), &[])?;
389 }
390 Ok(())
391 }
392
393 pub fn delete_vec(&mut self, field: u64, id: u64) -> Result<bool> {
396 let had = self.store.delete(&keys::vec_key(field, id))?;
397 self.store.delete(&keys::vcode_key(field, id))?;
398 self.store.delete(&keys::catalog_field_item(CAT_NAV_PENDING, field, id))?;
399 Ok(had)
400 }
401
402 pub fn vec_dock_rows(enc: &crate::vecquant::Encoder, field: u64, id: u64, v: &[f32])
408 -> [(Vec<u8>, Vec<u8>); 2]
409 {
410 let mut bytes = Vec::with_capacity(v.len() * 4);
411 for x in v { bytes.extend_from_slice(&x.to_le_bytes()); }
412 let (norm, code) = enc.encode(v);
413 let mut row = Vec::with_capacity(4 + code.len());
414 row.extend_from_slice(&norm.to_le_bytes());
415 row.extend_from_slice(&code);
416 [(keys::vec_key(field, id), bytes), (keys::vcode_key(field, id), row)]
417 }
418
419 pub fn vec_meta_row(field: u64, m: VecMeta) -> (Vec<u8>, Vec<u8>) {
422 (keys::catalog_field(CAT_VEC, field), m.encode())
423 }
424
425 pub fn get_vec(&self, field: u64, id: u64) -> Result<Option<Vec<f32>>> {
426 self.store.get(&keys::vec_key(field, id))?.map(decode_f32s).transpose()
427 }
428
429 pub fn vec_distance(&self, field: u64, id: u64, query: &[f32], metric: Metric)
433 -> Result<Option<f32>>
434 {
435 let key = keys::vec_key(field, id);
436 let mut found = None;
437 let mut malformed = false;
438 self.store.scan(&key)?.for_each_ref(|candidate, bytes| {
439 if candidate != key.as_slice() { return false; }
440 if bytes.len() != query.len() * 4 {
441 malformed = true;
442 } else {
443 found = Some(metric.distance_bytes(bytes, query));
444 }
445 false
446 })?;
447 if malformed { return Err(crate::Error::TooLarge); }
448 Ok(found)
449 }
450
451 pub fn vec_dim(&self, field: u64) -> u64 { self.vec_meta.get(&field).map_or(0, |m| m.dim) }
453 pub fn vec_bits_pub(&self, field: u64) -> usize {
454 self.vec_meta.get(&field).map_or(crate::vecquant::DEFAULT_BITS, |m| m.bits as usize)
455 }
456 pub(crate) fn encoder(&self, field: u64) -> Option<std::sync::Arc<crate::vecquant::Encoder>> {
460 if let Some(e) = self.encoders.borrow().get(&field) { return Some(e.clone()); }
461 let m = self.vec_meta.get(&field).copied()?;
462 let e = std::sync::Arc::new(
463 crate::vecquant::Encoder::new(m.dim as usize, m.seed, m.bits as usize));
464 self.encoders.borrow_mut().insert(field, e.clone());
465 Some(e)
466 }
467
468 pub fn rescore(
476 &self,
477 field: u64,
478 candidates: &[u64],
479 query: &[f32],
480 metric: Metric,
481 k: usize,
482 ) -> Result<Vec<(u64, f32)>> {
483 if k == 0 { return Ok(Vec::new()); }
484 let mut heap: std::collections::BinaryHeap<Scored> = std::collections::BinaryHeap::new();
485 for &id in candidates {
486 let Some(d) = self.vec_distance(field, id, query, metric)? else { continue };
487 offer_score(&mut heap, Scored { d, id }, k);
488 }
489 let mut out: Vec<(u64, f32)> = heap.into_iter().map(|s| (s.id, s.d)).collect();
490 out.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
491 Ok(out)
492 }
493
494 pub fn rescore_all(&self, field: u64, query: &[f32], metric: Metric, k: usize)
498 -> Result<Vec<(u64, f32)>>
499 {
500 if k == 0 { return Ok(Vec::new()); }
501 let prefix = keys::vec_prefix(field);
502 let query_norm = metric.query_norm(query);
503 let mut heap: std::collections::BinaryHeap<Scored> = std::collections::BinaryHeap::new();
504 let mut malformed = false;
505 self.store.scan(&prefix)?.for_each_ref(|key, bytes| {
506 if !key.starts_with(&prefix) { return false; }
507 if key.len() != 17 || bytes.len() != query.len() * 4 {
508 malformed = true;
509 return false;
510 }
511 let id = keys::u64_at(key, 9);
512 let d = metric.distance_bytes_prepared(bytes, query, query_norm);
513 offer_score(&mut heap, Scored { d, id }, k);
514 true
515 })?;
516 if malformed { return Err(crate::Error::TooLarge); }
517 let mut out: Vec<(u64, f32)> = heap.into_iter().map(|s| (s.id, s.d)).collect();
518 out.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
519 Ok(out)
520 }
521
522 pub fn rescore_sorted(&self, field: u64, candidates: &[u64], query: &[f32],
526 metric: Metric, k: usize) -> Result<Vec<(u64, f32)>> {
527 if k == 0 || candidates.is_empty() { return Ok(Vec::new()); }
528 debug_assert!(candidates.windows(2).all(|pair| pair[0] <= pair[1]));
529 let prefix = keys::vec_prefix(field);
530 let query_norm = metric.query_norm(query);
531 let mut at = 0usize;
532 let mut heap: std::collections::BinaryHeap<Scored> = std::collections::BinaryHeap::new();
533 let mut malformed = false;
534 self.store.scan(&prefix)?.for_each_ref(|key, bytes| {
535 if !key.starts_with(&prefix) || at == candidates.len() { return false; }
536 if key.len() != 17 { malformed = true; return false; }
537 let id = keys::u64_at(key, 9);
538 while at < candidates.len() && candidates[at] < id { at += 1; }
539 if at == candidates.len() { return false; }
540 if candidates[at] != id { return true; }
541 while at + 1 < candidates.len() && candidates[at + 1] == id { at += 1; }
542 at += 1;
543 if bytes.len() != query.len() * 4 { malformed = true; return false; }
544 offer_score(&mut heap, Scored {
545 d: metric.distance_bytes_prepared(bytes, query, query_norm), id
546 }, k);
547 true
548 })?;
549 if malformed { return Err(crate::Error::TooLarge); }
550 let mut out: Vec<(u64, f32)> = heap.into_iter().map(|score| (score.id, score.d)).collect();
551 out.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
552 Ok(out)
553 }
554
555 pub fn distances_sorted(&self, field: u64, candidates: &[u64], query: &[f32], metric: Metric)
558 -> Result<Vec<(u64, f32)>>
559 {
560 if candidates.is_empty() { return Ok(Vec::new()); }
561 debug_assert!(candidates.windows(2).all(|pair| pair[0] <= pair[1]));
562 let prefix = keys::vec_prefix(field);
563 let query_norm = metric.query_norm(query);
564 let mut at = 0usize;
565 let mut out = Vec::with_capacity(candidates.len());
566 let mut malformed = false;
567 self.store.scan(&prefix)?.for_each_ref(|key, bytes| {
568 if !key.starts_with(&prefix) || at == candidates.len() { return false; }
569 if key.len() != 17 { malformed = true; return false; }
570 let id = keys::u64_at(key, 9);
571 while at < candidates.len() && candidates[at] < id { at += 1; }
572 if at == candidates.len() { return false; }
573 if candidates[at] != id { return true; }
574 if bytes.len() != query.len() * 4 { malformed = true; return false; }
575 out.push((id, metric.distance_bytes_prepared(bytes, query, query_norm)));
576 while at < candidates.len() && candidates[at] == id { at += 1; }
577 true
578 })?;
579 if malformed { return Err(crate::Error::TooLarge); }
580 Ok(out)
581 }
582
583 pub fn nearest(&self, field: u64, q: &[f32], k: usize, metric: Metric, oversample: usize)
592 -> Result<Vec<(u64, f32)>>
593 {
594 let Some(meta) = self.vec_meta.get(&field).copied() else {
595 return Err(crate::Error::TooLarge);
596 };
597 if meta.dim == 0 || q.len() as u64 != meta.dim {
598 return Err(crate::Error::TooLarge);
599 }
600 let enc = self.encoder(field).expect("a field with a recipe has an encoder");
601 let bits = meta.bits as usize;
602 let expected_code_len = 4usize
603 .checked_add(crate::vecquant::code_len(crate::vecquant::pad_dim(meta.dim as usize), bits))
604 .ok_or(crate::Error::Corrupt {
605 page_no: 0,
606 why: "vector code row length overflows",
607 })?;
608 let aq = enc.affine_query(q);
609 let proxy_boost = if matches!(metric, Metric::L1) { 4 } else { 1 };
614 let cap = k.saturating_mul(oversample).saturating_mul(proxy_boost).max(k);
615
616 let mut heap: std::collections::BinaryHeap<Scored> = std::collections::BinaryHeap::new();
617 let prefix = crate::keys::vcode_prefix(field);
621 let it = self.store.scan(&prefix)?;
622 let mut malformed = false;
623 it.for_each_ref(|key, val| {
624 if key.len() != 17 {
627 malformed = key.starts_with(&prefix);
628 return false;
629 }
630 if key[0] != crate::keys::TAG_VCODE || crate::keys::u64_at(key, 1) != field {
631 return false; }
633 if val.len() != expected_code_len {
634 malformed = true;
635 return false;
636 }
637 let id = crate::keys::u64_at(key, 9);
638 let norm = f32::from_le_bytes(val[0..4].try_into().unwrap());
639 let dot = if bits == 2 {
640 crate::vecquant::dot_est_affine2(norm, &val[4..], &aq)
641 } else {
642 crate::vecquant::dot_est_affine(norm, &val[4..], &aq)
643 };
644 let d = match metric {
651 Metric::L2 | Metric::L1 => norm * norm - 2.0 * dot,
652 Metric::Dot => -dot,
653 Metric::Cosine => if norm > 0.0 { -dot / norm } else { 0.0 },
654 };
655 heap.push(Scored { d, id });
656 if heap.len() > cap { heap.pop(); }
657 true
658 })?;
659 if malformed {
660 return Err(crate::Error::Corrupt {
661 page_no: 0,
662 why: "vector code row has an invalid length",
663 });
664 }
665
666 let cands: Vec<u64> = heap.into_iter().map(|s| s.id).collect();
667 self.rescore(field, &cands, q, metric, k)
668 }
669
670 pub fn nearest_par(&self, field: u64, q: &[f32], k: usize, metric: Metric, oversample: usize,
678 threads: usize) -> Result<Vec<(u64, f32)>>
679 {
680 if threads <= 1 { return self.nearest(field, q, k, metric, oversample); }
681 let mut ps = self.parallel_searcher(field, threads)?;
682 ps.nearest(q, k, metric, oversample)
683 }
684
685 pub fn parallel_searcher(&self, field: u64, threads: usize) -> Result<ParallelSearcher> {
690 let hi_id = self.vec_meta.get(&field).map_or(0, |m| m.max).max(self.next_id).max(1);
691 let chunk = hi_id.div_ceil(threads as u64).max(1);
692 let mut readers = Vec::new();
693 for t in 0..threads as u64 {
694 let (lo, hi) = (1 + t * chunk, (1 + (t + 1) * chunk).min(hi_id + 1));
695 if lo >= hi { break; }
696 let s = crate::store::Store::open_snapshot(self.store.dir(), crate::store::Config::default())?;
697 readers.push((Graph::new(s)?, lo, hi));
698 }
699 Ok(ParallelSearcher { field, readers })
700 }
701
702 fn geo_rows(field: u64, id: u64, g: &crate::spatial::Geom)
703 -> Result<(Vec<Vec<u8>>, Vec<u8>, Vec<u8>)>
704 {
705 use crate::spatial as sp;
706 let Some((xmin, xmax, ymin, ymax)) = g.bbox() else {
707 return Err(crate::Error::TooLarge); };
709 if !(xmin >= -180.0 && xmax <= 180.0 && ymin >= -90.0 && ymax <= 90.0)
714 || !xmin.is_finite() || !xmax.is_finite()
715 || !ymin.is_finite() || !ymax.is_finite() {
716 return Err(crate::Error::TooLarge);
717 }
718 let bx = sp::BoxF::from_f64(xmin, xmax, ymin, ymax);
719 let (level, cells) = match sp::cover_cells(xmin, xmax, ymin, ymax,
720 sp::LEVEL_FINE, sp::MAX_CELLS) {
721 Some(c) => (sp::LEVEL_FINE, c),
722 None => match sp::cover_cells(xmin, xmax, ymin, ymax,
723 sp::LEVEL_COARSE, sp::MAX_CELLS) {
724 Some(c) => (sp::LEVEL_COARSE, c),
725 None => (sp::LEVEL_WORLD, vec![(0, 0)]),
727 }
728 };
729 let new_keys: Vec<Vec<u8>> = cells.into_iter().map(|(cx, cy)| {
730 let h = sp::cell_hilbert(cx, cy, level);
731 keys::spat_key(field, level, h, id)
732 }).collect();
733 Ok((new_keys, bx.encode().to_vec(), g.encode()))
734 }
735
736 pub fn insert_geo_new(&mut self, field: u64, id: u64, g: &crate::spatial::Geom) -> Result<()> {
740 let (new_keys, encoded_box, encoded_geom) = Self::geo_rows(field, id, g)?;
741 for key in &new_keys { self.store.put(key, &encoded_box)?; }
742 self.store.put(&keys::geom_key(field, id), &encoded_geom)?;
743 Ok(())
744 }
745
746 pub fn insert_geo_postings_new(&mut self, field: u64, id: u64,
750 g: &crate::spatial::Geom) -> Result<()> {
751 let (new_keys, encoded_box, _) = Self::geo_rows(field, id, g)?;
752 for key in &new_keys { self.store.put(key, &encoded_box)?; }
753 Ok(())
754 }
755
756 pub fn geo_posting_rows(field: u64, id: u64, g: &crate::spatial::Geom)
761 -> Result<Vec<(Vec<u8>, Vec<u8>)>>
762 {
763 let (keys, bbox, _) = Self::geo_rows(field, id, g)?;
764 Ok(keys.into_iter().map(|key| (key, bbox.clone())).collect())
765 }
766
767 fn possible_geo_keys(field: u64, id: u64, g: &crate::spatial::Geom) -> Vec<Vec<u8>> {
768 use crate::spatial as sp;
769 let mut old_keys = Vec::new();
770 if let Some((xmin, xmax, ymin, ymax)) = g.bbox() {
771 for level in [sp::LEVEL_FINE, sp::LEVEL_COARSE] {
772 if let Some(cells) = sp::cover_cells(xmin, xmax, ymin, ymax, level, sp::MAX_CELLS) {
773 old_keys.extend(cells.into_iter().map(|(cx, cy)| {
774 keys::spat_key(field, level, sp::cell_hilbert(cx, cy, level), id)
775 }));
776 }
777 }
778 old_keys.push(keys::spat_key(field, sp::LEVEL_WORLD, 0, id));
779 }
780 old_keys
781 }
782
783 pub fn replace_geo_postings(&mut self, field: u64, id: u64,
786 old: Option<&crate::spatial::Geom>,
787 new: Option<&crate::spatial::Geom>) -> Result<()> {
788 let (new_keys, encoded_box) = match new {
789 Some(g) => {
790 let (keys, bbox, _) = Self::geo_rows(field, id, g)?;
791 for key in &keys { self.store.put(key, &bbox)?; }
792 (keys, Some(bbox))
793 }
794 None => (Vec::new(), None),
795 };
796 if old.is_some() {
797 if let Some(encoded_box) = encoded_box.as_deref() {
798 for key in &new_keys {
799 if self.store_ref().get(key)?.as_deref() != Some(encoded_box) {
800 return Err(crate::Error::Corrupt {
801 page_no: 0, why: "new spatial posting did not verify",
802 });
803 }
804 }
805 }
806 }
807 if let Some(old) = old {
808 for key in Self::possible_geo_keys(field, id, old) {
809 if !new_keys.contains(&key) { self.store.delete(&key)?; }
810 }
811 }
812 Ok(())
813 }
814
815 pub fn set_geo(&mut self, field: u64, id: u64, g: &crate::spatial::Geom) -> Result<()> {
819 use crate::spatial as sp;
820 let old = self.store_ref().get(&keys::geom_key(field, id))?
821 .and_then(|bytes| sp::Geom::decode(&bytes));
822 let (new_keys, encoded_box, encoded_geom) = Self::geo_rows(field, id, g)?;
823 for key in &new_keys { self.store.put(key, &encoded_box)?; }
826 self.store.put(&keys::geom_key(field, id), &encoded_geom)?;
827 if let Some(old) = old {
830 for key in &new_keys {
831 if self.store_ref().get(key)?.as_deref() != Some(encoded_box.as_slice()) {
832 return Err(crate::Error::Corrupt { page_no: 0, why: "new spatial posting did not verify" });
833 }
834 }
835 if self.store_ref().get(&keys::geom_key(field, id))?.as_deref()
836 != Some(encoded_geom.as_slice())
837 {
838 return Err(crate::Error::Corrupt { page_no: 0, why: "new geometry row did not verify" });
839 }
840 for key in Self::possible_geo_keys(field, id, &old) {
841 if !new_keys.contains(&key) { self.store.delete(&key)?; }
842 }
843 }
844 Ok(())
845 }
846
847 pub fn clear_unpublished_geo(&mut self, field: u64) -> Result<()> {
851 use crate::spatial as sp;
852 for level in [sp::LEVEL_FINE, sp::LEVEL_COARSE, sp::LEVEL_WORLD] {
853 self.store.delete_prefix(&keys::spat_prefix(field, level))?;
854 }
855 self.store.delete_prefix(&keys::geom_prefix(field))?;
856 Ok(())
857 }
858
859 pub fn delete_geo(&mut self, field: u64, id: u64) -> Result<bool> {
862 use crate::spatial as sp;
863 let Some(raw) = self.store_ref().get(&keys::geom_key(field, id))? else {
864 return Ok(false);
865 };
866 if let Some(g) = sp::Geom::decode(&raw) {
867 if let Some((xmin, xmax, ymin, ymax)) = g.bbox() {
868 for level in [sp::LEVEL_FINE, sp::LEVEL_COARSE] {
869 if let Some(cells) = sp::cover_cells(xmin, xmax, ymin, ymax, level, sp::MAX_CELLS) {
870 for (cx, cy) in cells {
871 let h = sp::cell_hilbert(cx, cy, level);
872 self.store.delete(&keys::spat_key(field, level, h, id))?;
873 }
874 }
875 }
876 self.store.delete(&keys::spat_key(field, sp::LEVEL_WORLD, 0, id))?;
877 }
878 }
879 self.store.delete(&keys::geom_key(field, id))?;
880 Ok(true)
881 }
882
883 pub fn get_geo(&self, field: u64, id: u64) -> Result<Option<crate::spatial::Geom>> {
884 Ok(self.store_ref().get(&keys::geom_key(field, id))?
885 .and_then(|b| crate::spatial::Geom::decode(&b)))
886 }
887
888 fn geo_candidates(&self, field: u64, xmin: f64, xmax: f64, ymin: f64, ymax: f64)
893 -> Result<Vec<(u64, crate::spatial::BoxF)>>
894 {
895 let diag = std::env::var_os("GEO_DIAG").is_some();
896 let t0 = std::time::Instant::now();
897 use crate::spatial as sp;
898 let qbox = sp::BoxF::from_f64(xmin, xmax, ymin, ymax);
899 let cover_started = std::time::Instant::now();
900 let covers: Vec<(u8, usize, Vec<(u64, u64)>)> =
901 [sp::LEVEL_FINE, sp::LEVEL_COARSE].into_iter().map(|level| {
902 let (x0, y0) = sp::cell_of(xmin, ymin, level);
903 let (x1, y1) = sp::cell_of(xmax, ymax, level);
904 let cells = (x1 - x0 + 1) as usize * (y1 - y0 + 1) as usize;
905 (level, cells, sp::cover_ranges(xmin, xmax, ymin, ymax, level, 256))
906 }).collect();
907 let cover_elapsed = cover_started.elapsed();
908 let pages_before = self.store_ref().pool_stats();
909 let gather_started = std::time::Instant::now();
910 let mut out: Vec<(u64, sp::BoxF)> = Vec::new();
911 let mut seen: std::collections::HashSet<u64> = std::collections::HashSet::new();
912 let mut world_rows = 0u64;
913 let mut level_rows = Vec::with_capacity(covers.len());
914 {
916 let prefix = keys::spat_prefix(field, sp::LEVEL_WORLD);
917 let it = self.store_ref().scan(&prefix)?;
918 it.for_each_ref(|key, val| {
919 if !key.starts_with(&prefix) || key.len() != prefix.len() + 16 {
920 return false;
921 }
922 world_rows += 1;
923 let id = u64::from_be_bytes(key[prefix.len() + 8..].try_into().unwrap());
924 if let Some(bx) = sp::BoxF::decode(val) {
925 if bx.intersects(&qbox) && seen.insert(id) {
926 out.push((id, bx));
927 }
928 }
929 true
930 })?;
931 }
932 for (level, cells, ranges) in covers {
933 let prefix = keys::spat_prefix(field, level);
934 let range_count = ranges.len();
935 let mut posting_rows = 0u64;
936 for (lo, hi) in ranges {
937 let mut from = prefix.clone();
938 from.extend_from_slice(&lo.to_be_bytes());
939 let it = self.store_ref().scan(&from)?;
940 it.for_each_ref(|key, val| {
941 if !key.starts_with(&prefix) || key.len() != prefix.len() + 16 {
942 return false;
943 }
944 let h = u64::from_be_bytes(key[prefix.len()..prefix.len() + 8].try_into().unwrap());
945 if h > hi { return false; }
946 posting_rows += 1;
947 let id = u64::from_be_bytes(key[prefix.len() + 8..].try_into().unwrap());
948 if let Some(bx) = sp::BoxF::decode(val) {
949 if bx.intersects(&qbox) && seen.insert(id) {
950 out.push((id, bx));
951 }
952 }
953 true
954 })?;
955 }
956 level_rows.push((level, cells, range_count, posting_rows));
957 }
958 if diag {
959 let pages_after = self.store_ref().pool_stats();
960 let logical = pages_after.hits.saturating_add(pages_after.misses)
961 .saturating_sub(pages_before.hits.saturating_add(pages_before.misses));
962 let misses = pages_after.misses.saturating_sub(pages_before.misses);
963 eprintln!(
964 "GEO_DIAG cover_ms={:.3} gather_ms={:.3} total_ms={:.3} world_rows={} candidates={} pool_logical={} pool_misses={} levels={:?}",
965 cover_elapsed.as_secs_f64() * 1000.0,
966 gather_started.elapsed().as_secs_f64() * 1000.0,
967 t0.elapsed().as_secs_f64() * 1000.0,
968 world_rows, out.len(), logical, misses, level_rows,
969 );
970 }
971 Ok(out)
972 }
973
974 pub fn st_distance(&self, field: u64, id: u64, lat: f64, lon: f64) -> Result<Option<f64>> {
980 use crate::{geomath as gm, spatial::Geom};
981 let Some(g) = self.get_geo(field, id)? else { return Ok(None) };
982 let d = match &g {
983 Geom::Point(x, y) => gm::geodesic_distance_m(lat, lon, *y, *x),
984 _ => {
985 let inside = g.rings_latlon().iter()
986 .any(|r| gm::point_in_polygon(lat, lon, r));
987 if inside { 0.0 } else {
988 let mut min = f64::MAX;
989 let mut visit = |x: f64, y: f64| {
990 let d = gm::geodesic_distance_m(lat, lon, y, x);
991 if d < min { min = d; }
992 };
993 match &g {
994 Geom::LineString(c) | Geom::MultiPoint(c) =>
995 c.iter().for_each(|p| visit(p[0], p[1])),
996 Geom::Polygon(rs) | Geom::MultiLineString(rs) =>
997 rs.iter().flatten().for_each(|p| visit(p[0], p[1])),
998 Geom::MultiPolygon(ps) =>
999 ps.iter().flatten().flatten().for_each(|p| visit(p[0], p[1])),
1000 Geom::Point(..) => unreachable!(),
1001 }
1002 min
1003 }
1004 }
1005 };
1006 Ok(Some(d))
1007 }
1008
1009 pub fn within_radius(&self, field: u64, lat: f64, lon: f64, meters: f64, k: usize)
1014 -> Result<Vec<(u64, f64)>>
1015 {
1016 use crate::geomath as gm;
1017 let dlat = meters / 110_574.0;
1019 let dlon = meters / (111_320.0 * lat.to_radians().cos().abs().max(0.01));
1020 let cands = self.geo_candidates(field, lon - dlon, lon + dlon,
1021 lat - dlat, lat + dlat)?;
1022 let mut hits: Vec<(u64, f64)> = Vec::new();
1023 let band = meters * 6e-3;
1028 for (id, bx) in cands {
1029 let d = match bx.as_point() {
1030 Some((x, y)) => {
1031 let h = gm::haversine_km(lat, lon, y, x) * 1000.0;
1032 if h > meters + band { continue; }
1033 if h < meters - band { h }
1034 else { gm::geodesic_distance_m(lat, lon, y, x) }
1035 }
1036 None => match self.st_distance(field, id, lat, lon)? {
1037 Some(d) => d,
1038 None => continue,
1039 },
1040 };
1041 if d <= meters { hits.push((id, d)); }
1042 }
1043 hits.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
1044 hits.truncate(k);
1045 Ok(hits)
1046 }
1047
1048 pub fn in_bbox(&self, field: u64, xmin: f64, xmax: f64, ymin: f64, ymax: f64)
1052 -> Result<Vec<u64>>
1053 {
1054 let mut ids: Vec<u64> = self.geo_candidates(field, xmin, xmax, ymin, ymax)?
1055 .into_iter().map(|(id, _)| id).collect();
1056 ids.sort_unstable();
1057 Ok(ids)
1058 }
1059
1060 pub fn in_bbox_with_boxes(&self, field: u64, xmin: f64, xmax: f64,
1064 ymin: f64, ymax: f64)
1065 -> Result<Vec<(u64, crate::spatial::BoxF)>>
1066 {
1067 self.geo_candidates(field, xmin, xmax, ymin, ymax)
1068 }
1069
1070 pub fn for_each_bbox_candidate(
1074 &self,
1075 field: u64,
1076 xmin: f64,
1077 xmax: f64,
1078 ymin: f64,
1079 ymax: f64,
1080 mut visit: impl FnMut(u64, crate::spatial::BoxF) -> Result<bool>,
1081 ) -> Result<()> {
1082 use crate::spatial as sp;
1083 let diag = std::env::var_os("GEO_DIAG").is_some();
1084 let started = std::time::Instant::now();
1085 let cover_started = std::time::Instant::now();
1086 let covers: Vec<(u8, usize, Vec<(u64, u64)>)> =
1087 [sp::LEVEL_FINE, sp::LEVEL_COARSE].into_iter().map(|level| {
1088 let (x0, y0) = sp::cell_of(xmin, ymin, level);
1089 let (x1, y1) = sp::cell_of(xmax, ymax, level);
1090 let cells = (x1 - x0 + 1) as usize * (y1 - y0 + 1) as usize;
1091 (level, cells, sp::cover_ranges(xmin, xmax, ymin, ymax, level, 256))
1092 }).collect();
1093 let cover_elapsed = cover_started.elapsed();
1094 let qbox = sp::BoxF::from_f64(xmin, xmax, ymin, ymax);
1095 let pages_before = self.store_ref().pool_stats();
1096 let mut keep_going = true;
1097 let mut failure = None;
1098 let mut callback_elapsed = std::time::Duration::ZERO;
1099 let mut candidates = 0u64;
1100 let mut world_rows = 0u64;
1101 let mut level_rows = Vec::with_capacity(covers.len());
1102 let world = keys::spat_prefix(field, sp::LEVEL_WORLD);
1103 self.store_ref().scan(&world)?.for_each_ref(|key, value| {
1104 if !keep_going || failure.is_some() { return false }
1105 if !key.starts_with(&world) || key.len() != world.len() + 16 { return false }
1106 world_rows += 1;
1107 let id = u64::from_be_bytes(key[world.len() + 8..].try_into().unwrap());
1108 if let Some(bx) = sp::BoxF::decode(value).filter(|bx| bx.intersects(&qbox)) {
1109 candidates += 1;
1110 let callback_started = std::time::Instant::now();
1111 match visit(id, bx) {
1112 Ok(keep) => keep_going = keep,
1113 Err(error) => failure = Some(error),
1114 }
1115 callback_elapsed += callback_started.elapsed();
1116 }
1117 keep_going && failure.is_none()
1118 })?;
1119 if let Some(error) = failure.take() { return Err(error) }
1120 if !keep_going { return Ok(()) }
1121
1122 for (level, cells, ranges) in covers {
1123 let prefix = keys::spat_prefix(field, level);
1124 let range_count = ranges.len();
1125 let mut posting_rows = 0u64;
1126 for &(lo, hi) in &ranges {
1127 let mut from = prefix.clone();
1128 from.extend_from_slice(&lo.to_be_bytes());
1129 self.store_ref().scan(&from)?.for_each_ref(|key, value| {
1130 if !keep_going || failure.is_some() { return false }
1131 if !key.starts_with(&prefix) || key.len() != prefix.len() + 16 {
1132 return false;
1133 }
1134 let h = u64::from_be_bytes(
1135 key[prefix.len()..prefix.len() + 8].try_into().unwrap());
1136 if h > hi { return false }
1137 posting_rows += 1;
1138 let id = u64::from_be_bytes(key[prefix.len() + 8..].try_into().unwrap());
1139 let Some(bx) = sp::BoxF::decode(value).filter(|bx| bx.intersects(&qbox))
1140 else { return true };
1141
1142 let cells = sp::cover_cells(
1147 bx.xmin as f64, bx.xmax as f64,
1148 bx.ymin as f64, bx.ymax as f64,
1149 level, sp::MAX_CELLS * 4);
1150 if let Some(cells) = cells {
1151 for (cx, cy) in cells {
1152 let other = sp::cell_hilbert(cx, cy, level);
1153 if other >= h || !ranges.iter().any(|&(a, b)| a <= other && other <= b) {
1154 continue;
1155 }
1156 match self.store_ref().get(&keys::spat_key(field, level, other, id)) {
1157 Ok(Some(_)) => return true,
1158 Ok(None) => {}
1159 Err(error) => { failure = Some(error); return false }
1160 }
1161 }
1162 }
1163 candidates += 1;
1164 let callback_started = std::time::Instant::now();
1165 match visit(id, bx) {
1166 Ok(keep) => keep_going = keep,
1167 Err(error) => failure = Some(error),
1168 }
1169 callback_elapsed += callback_started.elapsed();
1170 keep_going && failure.is_none()
1171 })?;
1172 if let Some(error) = failure.take() { return Err(error) }
1173 if !keep_going { return Ok(()) }
1174 }
1175 level_rows.push((level, cells, range_count, posting_rows));
1176 }
1177 if diag {
1178 let after = self.store_ref().pool_stats();
1179 let total = started.elapsed();
1180 let gather = total.saturating_sub(cover_elapsed).saturating_sub(callback_elapsed);
1181 eprintln!(
1182 "GEO_DIAG streaming=1 cover_ms={:.3} gather_ms={:.3} exact_callback_ms={:.3} total_ms={:.3} world_rows={} candidates={} pool_logical={} pool_misses={} levels={:?}",
1183 cover_elapsed.as_secs_f64() * 1000.0,
1184 gather.as_secs_f64() * 1000.0,
1185 callback_elapsed.as_secs_f64() * 1000.0,
1186 total.as_secs_f64() * 1000.0,
1187 world_rows, candidates,
1188 after.hits.saturating_add(after.misses)
1189 .saturating_sub(pages_before.hits.saturating_add(pages_before.misses)),
1190 after.misses.saturating_sub(pages_before.misses), level_rows,
1191 );
1192 }
1193 Ok(())
1194 }
1195
1196 pub fn knn_geo(&self, field: u64, lat: f64, lon: f64, k: usize)
1201 -> Result<Vec<(u64, f64)>>
1202 {
1203 let mut radius = 700.0; loop {
1205 let hits = self.within_radius(field, lat, lon, radius, k)?;
1206 if hits.len() >= k || radius > 21_000_000.0 {
1207 return Ok(hits);
1208 }
1209 radius *= 2.0;
1210 }
1211 }
1212
1213 pub fn contains_point(&self, field: u64, lat: f64, lon: f64) -> Result<Vec<u64>> {
1215 use crate::geomath as gm;
1216 let eps = 1e-9;
1217 let cands = self.geo_candidates(field, lon - eps, lon + eps, lat - eps, lat + eps)?;
1218 let mut out = Vec::new();
1219 for (id, _) in cands {
1220 let Some(g) = self.get_geo(field, id)? else { continue };
1221 if g.rings_latlon().iter().any(|r| gm::point_in_polygon(lat, lon, r)) {
1222 out.push(id);
1223 }
1224 }
1225 out.sort_unstable();
1226 Ok(out)
1227 }
1228
1229 pub fn perspective(&self, ctx: u64) -> Result<EdgeIter<'_>> {
1239 let prefix = keys::ctx_prefix(ctx);
1240 Ok(EdgeIter { inner: self.store.scan(&prefix)?, prefix, rev: false, base: ctx == 0 })
1241 }
1242
1243 pub fn bfs(&self, ctx: u64, from: u64, ty: Option<u64>, depth: usize) -> Result<Vec<u64>> {
1244 let mut seen = std::collections::HashSet::new();
1245 let mut order = Vec::new();
1246 let mut frontier = vec![from];
1247 seen.insert(from);
1248 for _ in 0..depth {
1249 frontier.sort_unstable();
1250 let mut next = Vec::new();
1251 for src in frontier.drain(..) {
1252 for e in self.out_edges(ctx, src, ty)? {
1253 let (_, _, dst, _) = e?;
1254 if seen.insert(dst) {
1255 order.push(dst);
1256 next.push(dst);
1257 }
1258 }
1259 }
1260 if next.is_empty() { break; }
1261 frontier = next;
1262 }
1263 Ok(order)
1264 }
1265}
1266
1267pub struct ParallelSearcher {
1271 field: u64,
1275 readers: Vec<(Graph, u64, u64)>,
1276}
1277
1278impl ParallelSearcher {
1279 pub fn nearest(&mut self, q: &[f32], k: usize, metric: Metric, oversample: usize)
1283 -> Result<Vec<(u64, f32)>>
1284 {
1285 let field = self.field;
1286 let Some((g0, _, _)) = self.readers.first() else { return Ok(Vec::new()) };
1287 let enc_bits = g0.vec_bits_pub(field);
1288 let Some(enc) = g0.encoder(field) else { return Ok(Vec::new()) };
1289 let expected_code_len = 4usize
1290 .checked_add(crate::vecquant::code_len(
1291 crate::vecquant::pad_dim(g0.vec_dim(field) as usize), enc_bits))
1292 .ok_or(crate::Error::Corrupt {
1293 page_no: 0,
1294 why: "vector code row length overflows",
1295 })?;
1296 let prefix = crate::keys::vcode_prefix(field);
1297 let aq = enc.affine_query(q);
1298 let proxy_boost = if matches!(metric, Metric::L1) { 4 } else { 1 };
1303 let cap = k.saturating_mul(oversample).saturating_mul(proxy_boost).max(k);
1304 let cands: Vec<u64> = std::thread::scope(|sc| -> Result<Vec<u64>> {
1305 let mut handles = Vec::new();
1306 for (g, lo, hi) in &mut self.readers {
1307 let (aq, prefix, lo, hi) = (&aq, &prefix, *lo, *hi);
1308 let g: &mut Graph = g;
1309 handles.push(sc.spawn(move || -> Result<Vec<(f32, u64)>> {
1310 let mut heap: std::collections::BinaryHeap<Scored> = std::collections::BinaryHeap::new();
1311 let mut malformed = false;
1312 let from = crate::keys::vcode_key(field, lo);
1313 let it = g.store.scan(&from)?;
1314 it.for_each_ref(|key, val| {
1315 if key.len() != 17 {
1316 malformed = key.starts_with(prefix);
1317 return false;
1318 }
1319 if key[0] != crate::keys::TAG_VCODE || crate::keys::u64_at(key, 1) != field {
1320 return false;
1321 }
1322 if val.len() != expected_code_len {
1323 malformed = true;
1324 return false;
1325 }
1326 let id = crate::keys::u64_at(key, 9);
1327 if id >= hi { return false; }
1328 let norm = f32::from_le_bytes(val[0..4].try_into().unwrap());
1329 let dot = if enc_bits == 2 {
1330 crate::vecquant::dot_est_affine2(norm, &val[4..], aq)
1331 } else {
1332 crate::vecquant::dot_est_affine(norm, &val[4..], aq)
1333 };
1334 let d = match metric {
1335 Metric::L2 | Metric::L1 => norm * norm - 2.0 * dot,
1336 Metric::Dot => -dot,
1337 Metric::Cosine => if norm > 0.0 { -dot / norm } else { 0.0 },
1338 };
1339 heap.push(Scored { d, id });
1340 if heap.len() > cap { heap.pop(); }
1341 true
1342 })?;
1343 if malformed {
1344 return Err(crate::Error::Corrupt {
1345 page_no: 0,
1346 why: "vector code row has an invalid length",
1347 });
1348 }
1349 Ok(heap.into_iter().map(|s| (s.d, s.id)).collect())
1350 }));
1351 }
1352 let mut all: Vec<(f32, u64)> = Vec::new();
1353 for h in handles { all.extend(h.join().unwrap()?); }
1354 all.sort_by(|a, b| a.0.total_cmp(&b.0));
1355 all.truncate(cap);
1356 Ok(all.into_iter().map(|(_, id)| id).collect())
1357 })?;
1358 self.readers[0].0.rescore(field, &cands, q, metric, k)
1359 }
1360}
1361
1362pub struct EdgeIter<'p> {
1365 inner: crate::btree::RangeIter<'p>,
1366 prefix: Vec<u8>,
1367 rev: bool,
1368 base: bool,
1370}
1371
1372impl Iterator for EdgeIter<'_> {
1373 type Item = Result<(u64, u64, u64, Vec<u8>)>;
1374 fn next(&mut self) -> Option<Self::Item> {
1375 let (k, v) = match self.inner.next()? {
1376 Ok(kv) => kv,
1377 Err(e) => return Some(Err(e)),
1378 };
1379 let want = if self.base { 25 } else { 33 };
1380 if !k.starts_with(&self.prefix) || k.len() != want {
1381 return None;
1382 }
1383 let o = if self.base { 1 } else { 9 };
1384 let (a, ty, b) = (keys::u64_at(&k, o), keys::u64_at(&k, o + 8), keys::u64_at(&k, o + 16));
1385 Some(Ok(if self.rev { (b, ty, a, v) } else { (a, ty, b, v) }))
1387 }
1388}
1389
1390pub struct LabelIter<'p> {
1391 inner: crate::btree::RangeIter<'p>,
1392 prefix: Vec<u8>,
1393}
1394
1395impl Iterator for LabelIter<'_> {
1396 type Item = Result<u64>;
1397 fn next(&mut self) -> Option<Self::Item> {
1398 let (k, _) = match self.inner.next()? {
1399 Ok(kv) => kv,
1400 Err(e) => return Some(Err(e)),
1401 };
1402 if !k.starts_with(&self.prefix) || k.len() != 17 {
1403 return None;
1404 }
1405 Some(Ok(keys::u64_at(&k, 9)))
1406 }
1407}
1408
1409pub struct PropIter<'p> {
1412 inner: crate::btree::RangeIter<'p>,
1413 prop: u64,
1414 hi: u64,
1415}
1416
1417impl Iterator for PropIter<'_> {
1418 type Item = Result<(u64, u64)>;
1419 fn next(&mut self) -> Option<Self::Item> {
1420 let (k, _) = match self.inner.next()? {
1421 Ok(kv) => kv,
1422 Err(e) => return Some(Err(e)),
1423 };
1424 if k.len() != 25 || k[0] != keys::TAG_PROP || keys::u64_at(&k, 1) != self.prop {
1425 return None;
1426 }
1427 let value = keys::u64_at(&k, 9);
1428 if value > self.hi {
1429 return None;
1430 }
1431 Some(Ok((value, keys::u64_at(&k, 17))))
1432 }
1433}
1434
1435fn decode_f32s(bytes: Vec<u8>) -> Result<Vec<f32>> {
1436 if bytes.len() % 4 != 0 {
1437 return Err(crate::Error::Corrupt {
1438 page_no: 0,
1439 why: "stored f32 list has trailing bytes",
1440 });
1441 }
1442 Ok(bytes.chunks_exact(4)
1443 .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
1444 .collect())
1445}
1446
1447#[derive(Clone, Copy, PartialEq, Eq, Debug)]
1450pub enum Metric { L2, Cosine, Dot, L1 }
1451
1452impl Metric {
1453 pub fn distance_bytes(&self, stored: &[u8], q: &[f32]) -> f32 {
1456 self.distance_bytes_prepared(stored, q, self.query_norm(q))
1457 }
1458
1459 #[inline]
1460 fn query_norm(&self, q: &[f32]) -> f32 {
1461 if *self != Metric::Cosine { return 0.0; }
1462 let mut norm2 = 0.0;
1463 for &v in q { norm2 += v * v; }
1464 norm2.sqrt()
1465 }
1466
1467 #[inline]
1468 fn distance_bytes_prepared(&self, stored: &[u8], q: &[f32], query_norm: f32) -> f32 {
1469 if stored.len() != q.len() * 4 { return f32::INFINITY; }
1470 #[inline(always)]
1471 fn lane(bytes: &[u8], i: usize) -> f32 {
1472 let bits = unsafe {
1476 std::ptr::read_unaligned(bytes.as_ptr().add(i * 4).cast::<u32>())
1477 };
1478 f32::from_bits(u32::from_le(bits))
1479 }
1480 match self {
1481 Metric::L2 => {
1482 let mut sum = 0.0;
1483 for (i, &b) in q.iter().enumerate() {
1484 let d = lane(stored, i) - b;
1485 sum += d * d;
1486 }
1487 sum
1488 }
1489 Metric::L1 => {
1490 let mut sum = 0.0;
1491 for (i, &b) in q.iter().enumerate() { sum += (lane(stored, i) - b).abs(); }
1492 sum
1493 }
1494 Metric::Dot => {
1495 let mut dot = 0.0;
1496 for (i, &b) in q.iter().enumerate() { dot += lane(stored, i) * b; }
1497 -dot
1498 }
1499 Metric::Cosine => {
1500 let mut dot = 0.0f32;
1501 let mut na = 0.0f32;
1502 for (i, &b) in q.iter().enumerate() {
1503 let a = lane(stored, i);
1504 dot += a * b; na += a * a;
1505 }
1506 let denom = (na.sqrt() * query_norm).max(f32::MIN_POSITIVE);
1507 -(dot / denom)
1508 }
1509 }
1510 }
1511}
1512
1513struct Scored { d: f32, id: u64 }
1515impl PartialEq for Scored { fn eq(&self, o: &Self) -> bool { self.d == o.d && self.id == o.id } }
1516impl Eq for Scored {}
1517impl Ord for Scored {
1518 fn cmp(&self, o: &Self) -> std::cmp::Ordering {
1519 self.d.total_cmp(&o.d).then(self.id.cmp(&o.id))
1520 }
1521}
1522impl PartialOrd for Scored { fn partial_cmp(&self, o: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(o)) } }
1523
1524#[inline]
1527fn offer_score(heap: &mut std::collections::BinaryHeap<Scored>, score: Scored, k: usize) {
1528 if heap.len() < k {
1529 heap.push(score);
1530 } else if score.cmp(heap.peek().expect("a full top-k heap is non-empty")).is_lt() {
1531 *heap.peek_mut().expect("a full top-k heap is non-empty") = score;
1532 }
1533}