1use core::ops::Bound;
76
77use yo_common::{Code, Error, Result};
78use yo_kv::{Cursor, Elements, Full};
79use yo_shape::Metric;
80use yo_vector::{Match, Signature};
81
82use crate::head::{DEPTH_MAX, Kind};
83use crate::index::{self, IndexKind, Key, PathIndex};
84use crate::path::{Step, Steps};
85use crate::vector::{self, VectorIndex};
86use crate::{Builder, Keys, Value};
87
88#[derive(Debug)]
90pub struct Docs {
91 rows: Elements<()>,
93 keys: Keys,
95 build: Builder,
97 indexes: Vec<PathIndex>,
99 taken: Vec<Vec<u8>>,
106 vectors: Vec<VectorIndex>,
111 drawn: Vec<Vec<f32>>,
118}
119
120impl Default for Docs {
121 fn default() -> Docs {
124 Docs::new()
125 }
126}
127
128impl Docs {
129 #[must_use]
131 pub fn new() -> Docs {
132 Docs {
133 rows: Elements::tailed(0, 0),
134 keys: Keys::new(),
135 build: Builder::new(),
136 indexes: Vec::new(),
137 taken: Vec::new(),
138 vectors: Vec::new(),
139 drawn: Vec::new(),
140 }
141 }
142
143 #[must_use]
148 pub fn with_capacity(n: usize, each: usize) -> Docs {
149 Docs {
150 rows: Elements::tailed(n, n.saturating_mul(each)),
151 keys: Keys::new(),
152 build: Builder::with_capacity(each),
153 indexes: Vec::new(),
154 taken: Vec::new(),
155 vectors: Vec::new(),
156 drawn: Vec::new(),
157 }
158 }
159
160 pub fn put(&mut self, id: &[u8], value: Value<'_>) -> Result<bool> {
168 self.write(id, value, None)
169 }
170
171 pub fn put_bytes(&mut self, id: &[u8], doc: &[u8]) -> Result<bool> {
178 let value = Value::new(doc)
179 .ok_or_else(|| Error::new(Code::Corrupt, "the document is not a readable value"))?;
180 self.write(id, value, Some(doc))
181 }
182
183 fn write(&mut self, id: &[u8], value: Value<'_>, raw: Option<&[u8]>) -> Result<bool> {
194 let Docs {
195 rows,
196 keys,
197 build,
198 indexes,
199 taken,
200 vectors,
201 drawn,
202 } = self;
203
204 taken.resize(indexes.len(), Vec::new());
205 for (slot, index) in taken.iter_mut().zip(indexes.iter()) {
206 slot.clear();
207 let Some(at) = value.path_bytes(index.path())? else {
210 continue;
211 };
212 if index.keys_at(at, slot).is_err() {
213 return Err(Error::fmt(
214 Code::Full,
215 format_args!(
216 "a value at {} is longer than {} bytes and cannot be indexed",
217 String::from_utf8_lossy(index.path()),
218 index::KEY_MAX
219 ),
220 ));
221 }
222 }
223
224 drawn.resize(vectors.len(), Vec::new());
225 for (slot, index) in drawn.iter_mut().zip(vectors.iter()) {
226 slot.clear();
227 let Some(at) = value.path_bytes(index.path())? else {
228 continue;
229 };
230 vector::coordinates(at, index.dim(), index.path(), slot)?;
231 }
232 let tag = if vectors.is_empty() {
236 0
237 } else {
238 vector::tag_of(
239 indexes
240 .iter()
241 .map(PathIndex::path)
242 .zip(taken.iter().map(Vec::as_slice)),
243 )
244 };
245
246 unindex(rows, keys, indexes, id);
247
248 build.clear();
249 let fresh = if intern_into(keys, build, value, 0)? {
250 store(rows, id, build.finish()?)?
251 } else if let Some(raw) = raw {
252 store(rows, id, raw)?
255 } else {
256 build.clear();
257 build.embed(&value)?;
258 store(rows, id, build.finish()?)?
259 };
260
261 for (slot, index) in taken.iter().zip(indexes.iter_mut()) {
262 let mut filed = Ok(());
263 index::each_key(slot, |key| {
264 if filed.is_ok() {
265 filed = index.add(key, id);
266 }
267 });
268 filed?;
269 }
270 for (slot, index) in drawn.iter().zip(vectors.iter_mut()) {
273 if slot.is_empty() {
274 index.collection_mut().remove(id);
275 } else {
276 index.collection_mut().put_tagged(id, slot, tag)?;
277 }
278 }
279 Ok(fresh)
280 }
281
282 pub fn create_index(&mut self, path: &str) -> Result<()> {
295 self.create_index_bytes(path.as_bytes(), IndexKind::Equality)
296 }
297
298 pub fn create_ordered_index(&mut self, path: &str) -> Result<()> {
310 self.create_index_bytes(path.as_bytes(), IndexKind::Ordered)
311 }
312
313 pub fn create_array_index(&mut self, path: &str) -> Result<()> {
324 self.create_index_bytes(path.as_bytes(), IndexKind::Array)
325 }
326
327 pub fn create_text_index(&mut self, path: &str) -> Result<()> {
340 self.create_index_bytes(path.as_bytes(), IndexKind::Text)
341 }
342
343 pub fn create_index_bytes(&mut self, path: &[u8], kind: IndexKind) -> Result<()> {
346 for step in Steps::new(path) {
347 step?;
348 }
349 let old = match self.indexes.iter().position(|i| i.path() == path) {
355 Some(at) if self.indexes[at].kind() == kind => return Ok(()),
356 Some(at)
357 if kind == IndexKind::Equality && self.indexes[at].kind() == IndexKind::Ordered =>
358 {
359 return Ok(());
360 }
361 found => found,
362 };
363 let mut index = PathIndex::new(path, kind);
364 let mut list = Vec::new();
365 for (id, bytes) in self.rows.pairs() {
366 let Some(value) = Value::new(bytes) else {
367 continue;
368 };
369 let doc = Doc {
370 value,
371 keys: &self.keys,
372 };
373 let Some(at) = doc.path_bytes(path)? else {
374 continue;
375 };
376 list.clear();
377 if index.keys_at(at.value(), &mut list).is_err() {
378 return Err(Error::fmt(
379 Code::Full,
380 format_args!(
381 "a value at {} in {} is longer than {} bytes and cannot be indexed",
382 String::from_utf8_lossy(path),
383 String::from_utf8_lossy(id),
384 index::KEY_MAX
385 ),
386 ));
387 }
388 let mut filed = Ok(());
389 index::each_key(&list, |key| {
390 if filed.is_ok() {
391 filed = index.add(key, id);
392 }
393 });
394 filed?;
395 }
396 match old {
397 Some(at) => self.indexes[at] = index,
398 None => {
399 self.indexes.push(index);
400 self.taken.push(Vec::new());
401 }
402 }
403 self.retag();
404 Ok(())
405 }
406
407 pub fn drop_index(&mut self, path: &str) -> bool {
409 self.drop_index_bytes(path.as_bytes())
410 }
411
412 pub fn drop_index_bytes(&mut self, path: &[u8]) -> bool {
414 let Some(at) = self.indexes.iter().position(|i| i.path() == path) else {
415 return false;
416 };
417 self.indexes.remove(at);
418 self.taken.truncate(self.indexes.len());
419 self.retag();
420 true
421 }
422
423 #[must_use]
425 pub fn indexes(&self) -> &[PathIndex] {
426 &self.indexes
427 }
428
429 #[must_use]
431 pub fn index(&self, path: &str) -> Option<&PathIndex> {
432 self.indexes.iter().find(|i| i.path() == path.as_bytes())
433 }
434
435 pub fn create_vector_index(&mut self, path: &str, dim: usize) -> Result<()> {
449 self.create_vector_index_bytes(path.as_bytes(), dim, Metric::Cosine)
450 }
451
452 pub fn create_vector_index_with(
454 &mut self,
455 path: &str,
456 dim: usize,
457 metric: Metric,
458 ) -> Result<()> {
459 self.create_vector_index_bytes(path.as_bytes(), dim, metric)
460 }
461
462 pub fn create_vector_index_bytes(
471 &mut self,
472 path: &[u8],
473 dim: usize,
474 metric: Metric,
475 ) -> Result<()> {
476 for step in Steps::new(path) {
477 step?;
478 }
479 let old = match self.vectors.iter().position(|v| v.path() == path) {
483 Some(at) if self.vectors[at].dim() == dim && self.vectors[at].metric() == metric => {
484 return Ok(());
485 }
486 found => found,
487 };
488 let mut index = VectorIndex::new(path, dim, metric)?;
489 let mut list = Vec::new();
490 let mut v = Vec::new();
491 for (id, bytes) in self.rows.pairs() {
492 let Some(value) = Value::new(bytes) else {
493 continue;
494 };
495 let doc = Doc {
496 value,
497 keys: &self.keys,
498 };
499 let Some(at) = doc.path_bytes(path)? else {
500 continue;
501 };
502 vector::coordinates(at.value(), dim, path, &mut v)?;
503 let tag = tag_for(&doc, &self.indexes, &mut list);
504 index.collection_mut().put_tagged(id, &v, tag)?;
505 }
506 match old {
507 Some(at) => self.vectors[at] = index,
508 None => {
509 self.vectors.push(index);
510 self.drawn.push(Vec::new());
511 }
512 }
513 Ok(())
514 }
515
516 pub fn drop_vector_index(&mut self, path: &str) -> bool {
518 self.drop_vector_index_bytes(path.as_bytes())
519 }
520
521 pub fn drop_vector_index_bytes(&mut self, path: &[u8]) -> bool {
523 let Some(at) = self.vectors.iter().position(|v| v.path() == path) else {
524 return false;
525 };
526 self.vectors.remove(at);
527 self.drawn.truncate(self.vectors.len());
528 true
529 }
530
531 #[must_use]
534 pub fn vector_indexes(&self) -> &[VectorIndex] {
535 &self.vectors
536 }
537
538 #[must_use]
540 pub fn vector_index(&self, path: &str) -> Option<&VectorIndex> {
541 self.vectors.iter().find(|v| v.path() == path.as_bytes())
542 }
543
544 #[must_use]
551 pub fn embedding(&self, path: &str, id: &[u8]) -> Option<&[f32]> {
552 self.vector_index(path)?.collection().get(id)
553 }
554
555 pub fn nearest(
565 &self,
566 path: &str,
567 q: &[f32],
568 k: usize,
569 f: impl FnMut(&[u8], Doc<'_>, f32),
570 ) -> Result<usize> {
571 let hits = self.vector(path)?.collection().search(q, k, None)?;
572 Ok(self.answer(&hits, f))
573 }
574
575 pub fn nearest_where(
590 &self,
591 path: &str,
592 q: &[f32],
593 k: usize,
594 want: &[(&str, Key)],
595 f: impl FnMut(&[u8], Doc<'_>, f32),
596 ) -> Result<usize> {
597 let filter = self.wanted(want)?;
598 let hits = self
599 .vector(path)?
600 .collection()
601 .search_where(q, k, None, &filter)?;
602 Ok(self.answer(&hits, f))
603 }
604
605 pub fn nearest_to(
612 &self,
613 path: &str,
614 id: &[u8],
615 k: usize,
616 f: impl FnMut(&[u8], Doc<'_>, f32),
617 ) -> Result<usize> {
618 let index = self.vector(path)?;
619 let Some(q) = index.collection().get(id) else {
620 return Ok(0);
621 };
622 let hits = index.collection().search(q, k, Some(id))?;
623 Ok(self.answer(&hits, f))
624 }
625
626 fn vector(&self, path: &str) -> Result<&VectorIndex> {
628 self.vector_index(path).ok_or_else(|| {
629 Error::fmt(
630 Code::Invalid,
631 format_args!("there is no vector index on {path}, so this would be a scan"),
632 )
633 })
634 }
635
636 fn wanted(&self, want: &[(&str, Key)]) -> Result<Signature> {
638 let mut sig = Signature::default();
639 for (path, key) in want {
640 if self.index(path).is_none() {
641 return Err(Error::fmt(
642 Code::Invalid,
643 format_args!("there is no index on {path}, so a search cannot filter on it"),
644 ));
645 }
646 sig.insert(path, key.as_bytes());
647 }
648 Ok(sig)
649 }
650
651 fn answer(&self, hits: &[Match], mut f: impl FnMut(&[u8], Doc<'_>, f32)) -> usize {
653 let mut n = 0usize;
654 for hit in hits {
655 if let Some(doc) = self.get(&hit.key) {
656 f(&hit.key, doc, hit.distance);
657 n += 1;
658 }
659 }
660 n
661 }
662
663 fn retag(&mut self) {
671 let Docs {
672 rows,
673 keys,
674 indexes,
675 vectors,
676 ..
677 } = self;
678 if vectors.is_empty() {
679 return;
680 }
681 let mut list = Vec::new();
682 for (id, bytes) in rows.pairs() {
683 let Some(value) = Value::new(bytes) else {
684 continue;
685 };
686 let doc = Doc { value, keys };
687 let tag = tag_for(&doc, indexes, &mut list);
688 for index in vectors.iter_mut() {
689 index.collection_mut().retag(id, tag);
690 }
691 }
692 }
693
694 pub fn find(&self, path: &str, key: &Key, mut f: impl FnMut(&[u8], Doc<'_>)) -> Result<usize> {
703 let index = self.index(path).ok_or_else(|| {
704 Error::fmt(
705 Code::Invalid,
706 format_args!("there is no index on {path}, so this would be a scan"),
707 )
708 })?;
709 let Some(set) = index.get(key) else {
710 return Ok(0);
711 };
712 let mut n = 0usize;
713 index::each_id(set, |id| {
714 if let Some(doc) = self.get(id) {
715 f(id, doc);
716 n += 1;
717 }
718 });
719 Ok(n)
720 }
721
722 pub fn count(&self, path: &str, key: &Key) -> Result<usize> {
727 let index = self.index(path).ok_or_else(|| {
728 Error::fmt(
729 Code::Invalid,
730 format_args!("there is no index on {path}, so this would be a scan"),
731 )
732 })?;
733 Ok(index.count(key))
734 }
735
736 pub fn range(
748 &self,
749 path: &str,
750 lo: Bound<&Key>,
751 hi: Bound<&Key>,
752 mut f: impl FnMut(&[u8], Doc<'_>),
753 ) -> Result<usize> {
754 let index = self.ordered(path)?;
755 let mut n = 0usize;
756 for (_, set) in index.range(lo, hi) {
757 index::each_id(set, |id| {
758 if let Some(doc) = self.get(id) {
759 f(id, doc);
760 n += 1;
761 }
762 });
763 }
764 Ok(n)
765 }
766
767 pub fn range_rev(
769 &self,
770 path: &str,
771 lo: Bound<&Key>,
772 hi: Bound<&Key>,
773 mut f: impl FnMut(&[u8], Doc<'_>),
774 ) -> Result<usize> {
775 let index = self.ordered(path)?;
776 let mut n = 0usize;
777 for (_, set) in index.range_rev(lo, hi) {
778 index::each_id(set, |id| {
779 if let Some(doc) = self.get(id) {
780 f(id, doc);
781 n += 1;
782 }
783 });
784 }
785 Ok(n)
786 }
787
788 pub fn count_range(&self, path: &str, lo: Bound<&Key>, hi: Bound<&Key>) -> Result<usize> {
795 Ok(self.ordered(path)?.count_in(lo, hi))
796 }
797
798 fn ordered(&self, path: &str) -> Result<&PathIndex> {
800 match self.index(path) {
801 Some(index) if index.kind() == IndexKind::Ordered => Ok(index),
802 Some(_) => Err(Error::fmt(
803 Code::Invalid,
804 format_args!("the index on {path} answers equality and not ranges"),
805 )),
806 None => Err(Error::fmt(
807 Code::Invalid,
808 format_args!("there is no index on {path}, so this would be a scan"),
809 )),
810 }
811 }
812
813 #[must_use]
815 pub fn get(&self, id: &[u8]) -> Option<Doc<'_>> {
816 let value = Value::new(self.rows.tail(id)?)?;
817 Some(Doc {
818 value,
819 keys: &self.keys,
820 })
821 }
822
823 #[must_use]
828 pub fn bytes(&self, id: &[u8]) -> Option<&[u8]> {
829 self.rows.tail(id)
830 }
831
832 #[must_use]
834 pub fn contains(&self, id: &[u8]) -> bool {
835 self.rows.contains(id)
836 }
837
838 pub fn remove(&mut self, id: &[u8]) -> bool {
847 let Docs {
848 rows,
849 keys,
850 indexes,
851 vectors,
852 ..
853 } = self;
854 unindex(rows, keys, indexes, id);
855 for index in vectors.iter_mut() {
856 index.collection_mut().remove(id);
857 }
858 rows.remove(id).is_some()
859 }
860
861 #[must_use]
863 pub fn len(&self) -> usize {
864 self.rows.len()
865 }
866
867 #[must_use]
869 pub fn is_empty(&self) -> bool {
870 self.rows.is_empty()
871 }
872
873 #[must_use]
875 pub fn keys(&self) -> &Keys {
876 &self.keys
877 }
878
879 pub fn iter(&self) -> impl Iterator<Item = (&[u8], Doc<'_>)> {
881 let keys = &self.keys;
882 self.rows.pairs().filter_map(move |(id, bytes)| {
883 let value = Value::new(bytes)?;
884 Some((id, Doc { value, keys }))
885 })
886 }
887
888 pub fn scan<F>(&self, cursor: Cursor, count: usize, mut f: F) -> Cursor
891 where
892 F: FnMut(&[u8], Doc<'_>),
893 {
894 let keys = &self.keys;
895 self.rows.scan_pairs(cursor, count, |id, bytes| {
896 if let Some(value) = Value::new(bytes) {
897 f(id, Doc { value, keys });
898 }
899 })
900 }
901
902 pub fn clear(&mut self) {
913 self.rows.clear();
914 self.build.clear();
915 for index in &mut self.indexes {
916 index.clear();
917 }
918 for index in &mut self.vectors {
919 index.clear();
920 }
921 }
922
923 #[must_use]
925 pub fn memory_bytes(&self) -> usize {
926 self.rows.memory_bytes()
927 + self.keys.memory_bytes()
928 + self
929 .indexes
930 .iter()
931 .map(PathIndex::memory_bytes)
932 .sum::<usize>()
933 + self
934 .vectors
935 .iter()
936 .map(VectorIndex::memory_bytes)
937 .sum::<usize>()
938 }
939}
940
941fn tag_for(doc: &Doc<'_>, indexes: &[PathIndex], list: &mut Vec<u8>) -> u64 {
948 let mut sig = Signature::default();
949 for index in indexes {
950 let Ok(Some(at)) = doc.path_bytes(index.path()) else {
951 continue;
952 };
953 list.clear();
954 let _ = index.keys_at(at.value(), list);
955 vector::add_keys(&mut sig, index.path(), list);
956 }
957 sig.bits()
958}
959
960fn unindex(rows: &Elements<()>, keys: &Keys, indexes: &mut [PathIndex], id: &[u8]) {
970 if indexes.is_empty() {
971 return;
972 }
973 let Some(bytes) = rows.tail(id) else {
974 return;
975 };
976 let Some(value) = Value::new(bytes) else {
977 return;
978 };
979 let doc = Doc { value, keys };
980 let mut list = Vec::new();
981 for index in indexes {
982 let Ok(Some(at)) = doc.path_bytes(index.path()) else {
983 continue;
984 };
985 list.clear();
986 let _ = index.keys_at(at.value(), &mut list);
989 index::each_key(&list, |key| index.take(key, id));
990 }
991}
992
993fn store(rows: &mut Elements<()>, id: &[u8], bytes: &[u8]) -> Result<bool> {
996 match rows.set_tailed(id, bytes, ()) {
997 Ok((_, fresh)) => Ok(fresh),
998 Err(Full::Name) => Err(Error::fmt(
999 Code::Full,
1000 format_args!("a document id is at most {} bytes", yo_kv::NAME_MAX),
1001 )),
1002 Err(Full::Rows) => Err(Error::fmt(
1003 Code::Full,
1004 format_args!("a collection holds at most {} documents", yo_kv::MAX_ROWS),
1005 )),
1006 }
1007}
1008
1009fn intern_into(keys: &mut Keys, b: &mut Builder, value: Value<'_>, depth: usize) -> Result<bool> {
1019 let corrupt = || Error::new(Code::Corrupt, "the document is not readable at that point");
1020 match value.kind() {
1021 Kind::Null => b.null()?,
1022 Kind::Bool => b.bool(value.as_bool().ok_or_else(corrupt)?)?,
1023 Kind::Int => b.int(value.as_int().ok_or_else(corrupt)?)?,
1024 Kind::Float => b.float(value.as_float().ok_or_else(corrupt)?)?,
1025 Kind::Text => b.text_bytes(value.text_bytes().ok_or_else(corrupt)?)?,
1026 Kind::Array => {
1027 b.begin_array()?;
1028 for i in 0..value.len() {
1029 let child = value.at(i).ok_or_else(corrupt)?;
1030 if !intern_into(keys, b, child, depth + 1)? {
1031 return Ok(false);
1032 }
1033 }
1034 b.end_array()?;
1035 }
1036 Kind::Object => {
1037 if value.is_interned() {
1038 return Err(Error::new(
1039 Code::Invalid,
1040 "this document's keys are ids from another collection's key table",
1041 ));
1042 }
1043 b.begin_object_interned()?;
1044 for i in 0..value.len() {
1045 let name = value.key_at(i).ok_or_else(corrupt)?;
1046 let Some(id) = keys.intern(name) else {
1047 return Ok(false);
1048 };
1049 b.key_id(id)?;
1050 let child = value.at(i).ok_or_else(corrupt)?;
1051 if !intern_into(keys, b, child, depth + 1)? {
1052 return Ok(false);
1053 }
1054 }
1055 b.end_object()?;
1056 }
1057 }
1058 debug_assert!(depth <= DEPTH_MAX, "the builder caps the depth");
1059 Ok(true)
1060}
1061
1062#[derive(Clone, Copy)]
1070pub struct Doc<'a> {
1071 value: Value<'a>,
1072 keys: &'a Keys,
1073}
1074
1075impl<'a> Doc<'a> {
1076 #[must_use]
1078 pub fn new(value: Value<'a>, keys: &'a Keys) -> Doc<'a> {
1079 Doc { value, keys }
1080 }
1081
1082 #[must_use]
1084 pub fn value(&self) -> Value<'a> {
1085 self.value
1086 }
1087
1088 #[must_use]
1090 pub fn keys(&self) -> &'a Keys {
1091 self.keys
1092 }
1093
1094 #[must_use]
1096 pub fn kind(&self) -> Kind {
1097 self.value.kind()
1098 }
1099
1100 #[must_use]
1102 pub fn is_null(&self) -> bool {
1103 self.value.is_null()
1104 }
1105
1106 #[must_use]
1108 pub fn as_bool(&self) -> Option<bool> {
1109 self.value.as_bool()
1110 }
1111
1112 #[must_use]
1114 pub fn as_int(&self) -> Option<i64> {
1115 self.value.as_int()
1116 }
1117
1118 #[must_use]
1120 pub fn as_float(&self) -> Option<f64> {
1121 self.value.as_float()
1122 }
1123
1124 #[must_use]
1126 pub fn as_text(&self) -> Option<&'a str> {
1127 self.value.as_text()
1128 }
1129
1130 #[must_use]
1132 pub fn text_bytes(&self) -> Option<&'a [u8]> {
1133 self.value.text_bytes()
1134 }
1135
1136 #[must_use]
1138 pub fn len(&self) -> usize {
1139 self.value.len()
1140 }
1141
1142 #[must_use]
1144 pub fn is_empty(&self) -> bool {
1145 self.value.is_empty()
1146 }
1147
1148 #[must_use]
1155 pub fn get(&self, key: &[u8]) -> Option<Doc<'a>> {
1156 let value = if self.value.is_interned() {
1157 self.value.get_id(self.keys.id(key)?)?
1158 } else {
1159 self.value.get(key)?
1160 };
1161 Some(Doc {
1162 value,
1163 keys: self.keys,
1164 })
1165 }
1166
1167 #[must_use]
1169 pub fn at(&self, i: usize) -> Option<Doc<'a>> {
1170 Some(Doc {
1171 value: self.value.at(i)?,
1172 keys: self.keys,
1173 })
1174 }
1175
1176 #[must_use]
1178 pub fn key_at(&self, i: usize) -> Option<&'a [u8]> {
1179 if self.value.is_interned() {
1180 self.keys.name(self.value.key_id_at(i)?)
1181 } else {
1182 self.value.key_at(i)
1183 }
1184 }
1185
1186 #[must_use]
1193 pub fn members(&self) -> DocMembers<'a> {
1194 DocMembers { d: *self, i: 0 }
1195 }
1196
1197 #[must_use]
1199 pub fn iter(&self) -> DocElems<'a> {
1200 DocElems { d: *self, i: 0 }
1201 }
1202
1203 pub fn path(&self, path: &str) -> Result<Option<Doc<'a>>> {
1208 self.path_bytes(path.as_bytes())
1209 }
1210
1211 pub fn path_bytes(&self, path: &[u8]) -> Result<Option<Doc<'a>>> {
1213 let mut at = *self;
1214 for step in Steps::new(path) {
1215 let next = match step? {
1216 Step::Key(k) => at.get(k),
1217 Step::Index(_) if at.kind() != Kind::Array => None,
1218 Step::Index(i) => {
1219 let n = at.len();
1220 let i = if i < 0 {
1221 match n.checked_sub(i.unsigned_abs() as usize) {
1222 Some(i) => i,
1223 None => return Ok(None),
1224 }
1225 } else {
1226 i as usize
1227 };
1228 at.at(i)
1229 }
1230 };
1231 let Some(next) = next else {
1232 return Ok(None);
1233 };
1234 at = next;
1235 }
1236 Ok(Some(at))
1237 }
1238}
1239
1240impl core::fmt::Debug for Doc<'_> {
1241 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
1242 match self.kind() {
1243 Kind::Object => {
1244 let mut m = f.debug_map();
1245 for (k, v) in self.members() {
1246 m.entry(&String::from_utf8_lossy(k), &v);
1247 }
1248 m.finish()
1249 }
1250 Kind::Array => f.debug_list().entries(self.iter()).finish(),
1251 _ => self.value.fmt(f),
1252 }
1253 }
1254}
1255
1256#[derive(Clone)]
1258pub struct DocMembers<'a> {
1259 d: Doc<'a>,
1260 i: usize,
1261}
1262
1263impl<'a> Iterator for DocMembers<'a> {
1264 type Item = (&'a [u8], Doc<'a>);
1265
1266 fn next(&mut self) -> Option<(&'a [u8], Doc<'a>)> {
1267 let key = self.d.key_at(self.i)?;
1268 let val = self.d.at(self.i)?;
1269 self.i += 1;
1270 Some((key, val))
1271 }
1272
1273 fn size_hint(&self) -> (usize, Option<usize>) {
1274 let left = self.d.len().saturating_sub(self.i);
1275 (left, Some(left))
1276 }
1277}
1278
1279#[derive(Clone)]
1281pub struct DocElems<'a> {
1282 d: Doc<'a>,
1283 i: usize,
1284}
1285
1286impl<'a> Iterator for DocElems<'a> {
1287 type Item = Doc<'a>;
1288
1289 fn next(&mut self) -> Option<Doc<'a>> {
1290 let out = self.d.at(self.i)?;
1291 self.i += 1;
1292 Some(out)
1293 }
1294
1295 fn size_hint(&self) -> (usize, Option<usize>) {
1296 let left = self.d.len().saturating_sub(self.i);
1297 (left, Some(left))
1298 }
1299}
1300
1301#[cfg(test)]
1302mod tests {
1303 use super::*;
1304
1305 fn order(id: i64, status: &str, lines: usize) -> Vec<u8> {
1307 let mut b = Builder::new();
1308 b.begin_object().expect("open");
1309 b.key(b"id").expect("key");
1310 b.int(id).expect("value");
1311 b.key(b"customer").expect("key");
1312 b.int(id * 7).expect("value");
1313 b.key(b"status").expect("key");
1314 b.text(status).expect("value");
1315 b.key(b"lines").expect("key");
1316 b.begin_array().expect("open");
1317 for i in 0..lines {
1318 b.begin_object().expect("open");
1319 b.key(b"sku").expect("key");
1320 b.text(&format!("sku-{i}")).expect("value");
1321 b.key(b"qty").expect("key");
1322 b.int(i as i64 + 1).expect("value");
1323 b.end_object().expect("close");
1324 }
1325 b.end_array().expect("close");
1326 b.end_object().expect("close");
1327 b.finish().expect("finished").to_vec()
1328 }
1329
1330 #[test]
1331 fn a_document_reads_back_the_way_it_went_in() {
1332 let mut docs = Docs::new();
1333 assert!(
1334 docs.put_bytes(b"order:1", &order(1, "open", 3))
1335 .expect("put")
1336 );
1337 assert!(
1338 !docs
1339 .put_bytes(b"order:1", &order(1, "shut", 3))
1340 .expect("put")
1341 );
1342 assert_eq!(docs.len(), 1);
1343
1344 let d = docs.get(b"order:1").expect("stored");
1345 assert_eq!(d.get(b"id").and_then(|v| v.as_int()), Some(1));
1346 assert_eq!(d.get(b"status").and_then(|v| v.as_text()), Some("shut"));
1347 assert_eq!(d.get(b"lines").map(|v| v.len()), Some(3));
1348 assert_eq!(
1349 d.path("$.lines[1].sku")
1350 .expect("a path")
1351 .and_then(|v| v.as_text()),
1352 Some("sku-1")
1353 );
1354 assert_eq!(
1355 d.path("$.lines[-1].qty")
1356 .expect("a path")
1357 .and_then(|v| v.as_int()),
1358 Some(3)
1359 );
1360 assert!(d.get(b"missing").is_none());
1361 }
1362
1363 #[test]
1364 fn the_keys_are_interned_and_the_names_come_back() {
1365 let mut docs = Docs::new();
1366 docs.put_bytes(b"order:1", &order(1, "open", 2))
1367 .expect("put");
1368 let names: Vec<String> = docs
1369 .keys()
1370 .iter()
1371 .map(|(n, _)| String::from_utf8_lossy(n).into_owned())
1372 .collect();
1373 names.iter().for_each(|n| assert!(!n.is_empty()));
1374 assert_eq!(
1375 docs.keys().len(),
1376 6,
1377 "id customer status lines sku qty: {names:?}"
1378 );
1379
1380 let d = docs.get(b"order:1").expect("stored");
1381 assert!(d.value().is_interned());
1382 let mut got: Vec<&[u8]> = d.members().map(|(k, _)| k).collect();
1383 got.sort_unstable();
1384 assert_eq!(got, [&b"customer"[..], b"id", b"lines", b"status"]);
1385 let line = d.path("$.lines[0]").expect("a path").expect("there");
1386 assert!(line.value().is_interned());
1387 let mut inner: Vec<&[u8]> = line.members().map(|(k, _)| k).collect();
1388 inner.sort_unstable();
1389 assert_eq!(inner, [&b"qty"[..], b"sku"]);
1390 }
1391
1392 fn shrinkage(shape: impl Fn(i64) -> Vec<u8>) -> f64 {
1394 let mut docs = Docs::new();
1395 let mut plain = 0usize;
1396 for i in 0..256i64 {
1397 let bytes = shape(i);
1398 plain += bytes.len();
1399 docs.put_bytes(format!("d:{i}").as_bytes(), &bytes)
1400 .expect("put");
1401 }
1402 let stored: usize = (0..256i64)
1403 .map(|i| {
1404 docs.bytes(format!("d:{i}").as_bytes())
1405 .expect("stored")
1406 .len()
1407 })
1408 .sum();
1409 stored as f64 / plain as f64
1410 }
1411
1412 #[test]
1413 fn interning_makes_a_collection_of_the_same_shape_smaller() {
1414 let names = shrinkage(|i| {
1423 let mut b = Builder::new();
1424 b.begin_object().expect("open");
1425 for f in 0..20 {
1426 b.key(format!("some_field_name_{f:02}").as_bytes())
1427 .expect("key");
1428 b.int(i + f).expect("value");
1429 }
1430 b.end_object().expect("close");
1431 b.finish().expect("finished").to_vec()
1432 });
1433 assert!(names < 0.60, "a document of names kept {names}");
1434
1435 let orders = shrinkage(|i| order(i, "open", 2));
1440 assert!(orders < 0.80, "an order collection kept {orders}");
1441 }
1442
1443 #[test]
1444 fn a_document_whose_keys_are_already_ids_is_refused() {
1445 let mut b = Builder::new();
1446 b.begin_object_interned().expect("open");
1447 b.key_id(0).expect("key");
1448 b.int(1).expect("value");
1449 b.end_object().expect("close");
1450 let bytes = b.finish().expect("finished").to_vec();
1451
1452 let mut docs = Docs::new();
1453 let err = docs.put_bytes(b"x", &bytes).expect_err("refused");
1454 assert_eq!(err.code(), Code::Invalid);
1455 }
1456
1457 #[test]
1458 fn a_document_that_is_not_readable_is_refused() {
1459 let mut docs = Docs::new();
1460 let err = docs.put_bytes(b"x", &[2, 0, 0, 0]).expect_err("refused");
1461 assert_eq!(err.code(), Code::Corrupt);
1462 assert!(docs.is_empty());
1463 }
1464
1465 #[test]
1466 fn a_removal_leaves_every_other_document_where_it_was() {
1467 let mut docs = Docs::new();
1468 for i in 0..64i64 {
1469 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1470 .expect("put");
1471 }
1472 for i in (0..64i64).step_by(3) {
1473 assert!(docs.remove(format!("order:{i}").as_bytes()));
1474 }
1475 assert_eq!(docs.len(), 64 - 22);
1476 for i in 0..64i64 {
1477 let id = format!("order:{i}");
1478 match docs.get(id.as_bytes()) {
1479 Some(d) => {
1480 assert!(i % 3 != 0, "{id} was removed");
1481 assert_eq!(d.get(b"id").and_then(|v| v.as_int()), Some(i));
1482 }
1483 None => assert!(i % 3 == 0, "{id} was not removed"),
1484 }
1485 }
1486 assert_eq!(docs.keys().len(), 6, "a removal does not un-intern a name");
1487 }
1488
1489 #[test]
1490 fn a_walk_sees_every_document_once() {
1491 let mut docs = Docs::new();
1492 for i in 0..200i64 {
1493 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1494 .expect("put");
1495 }
1496
1497 let mut seen: Vec<i64> = docs
1498 .iter()
1499 .map(|(_, d)| d.get(b"id").and_then(|v| v.as_int()).expect("an id"))
1500 .collect();
1501 seen.sort_unstable();
1502 assert_eq!(seen, (0..200).collect::<Vec<i64>>());
1503
1504 let mut scanned = Vec::new();
1505 let mut cursor = Cursor::START;
1506 loop {
1507 cursor = docs.scan(cursor, 16, |id, _| scanned.push(id.to_vec()));
1508 if cursor.is_end() {
1509 break;
1510 }
1511 }
1512 scanned.sort_unstable();
1513 scanned.dedup();
1514 assert_eq!(scanned.len(), 200);
1515 }
1516
1517 #[test]
1518 fn an_empty_collection_answers_nothing_rather_than_failing() {
1519 let docs = Docs::new();
1520 assert!(docs.is_empty());
1521 assert!(docs.get(b"nothing").is_none());
1522 assert!(docs.bytes(b"nothing").is_none());
1523 assert!(!docs.contains(b"nothing"));
1524 assert_eq!(docs.iter().count(), 0);
1525 }
1526
1527 #[test]
1528 fn a_document_prints_with_its_names_back_on() {
1529 let mut docs = Docs::new();
1530 docs.put_bytes(b"order:1", &order(1, "open", 1))
1531 .expect("put");
1532 let text = format!("{:?}", docs.get(b"order:1").expect("stored"));
1533 assert!(text.contains("\"status\": \"open\""), "{text}");
1534 assert!(text.contains("\"sku\": \"sku-0\""), "{text}");
1535 }
1536
1537 fn found(docs: &Docs, path: &str, key: &Key) -> Vec<String> {
1539 let mut out = Vec::new();
1540 let n = docs
1541 .find(path, key, |id, d| {
1542 assert!(!d.is_empty(), "the document came back whole");
1543 out.push(String::from_utf8_lossy(id).into_owned());
1544 })
1545 .expect("indexed");
1546 assert_eq!(n, out.len(), "the count is what the callback saw");
1547 out.sort();
1548 out
1549 }
1550
1551 #[test]
1552 fn an_index_declared_after_the_documents_finds_them() {
1553 let mut docs = Docs::new();
1554 for i in 0..64i64 {
1555 let status = if i % 4 == 0 { "shut" } else { "open" };
1556 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, status, 1))
1557 .expect("put");
1558 }
1559 docs.create_index("$.status").expect("indexed");
1560 assert_eq!(docs.index("$.status").expect("there").len(), 2);
1561 assert_eq!(docs.count("$.status", &Key::text("shut")).expect("i"), 16);
1562 assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), 48);
1563 assert_eq!(found(&docs, "$.status", &Key::text("shut")).len(), 16);
1564 assert!(found(&docs, "$.status", &Key::text("gone")).is_empty());
1565
1566 docs.put_bytes(b"order:64", &order(64, "shut", 1))
1568 .expect("put");
1569 assert_eq!(docs.count("$.status", &Key::text("shut")).expect("i"), 17);
1570 }
1571
1572 #[test]
1573 fn an_overwrite_moves_a_document_from_one_key_to_the_other() {
1574 let mut docs = Docs::new();
1575 docs.create_index("$.status").expect("indexed");
1576 docs.put_bytes(b"order:1", &order(1, "open", 1))
1577 .expect("put");
1578 assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
1579
1580 docs.put_bytes(b"order:1", &order(1, "shut", 1))
1581 .expect("put");
1582 assert!(
1583 found(&docs, "$.status", &Key::text("open")).is_empty(),
1584 "the old key kept it"
1585 );
1586 assert_eq!(found(&docs, "$.status", &Key::text("shut")), ["order:1"]);
1587 assert_eq!(docs.index("$.status").expect("there").postings(), 1);
1588 }
1589
1590 #[test]
1591 fn a_removal_takes_a_document_out_of_every_index() {
1592 let mut docs = Docs::new();
1593 docs.create_index("$.status").expect("indexed");
1594 docs.create_index("$.customer").expect("indexed");
1595 for i in 0..8i64 {
1596 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1597 .expect("put");
1598 }
1599 assert!(docs.remove(b"order:3"));
1600 assert_eq!(found(&docs, "$.status", &Key::text("open")).len(), 7);
1601 assert_eq!(docs.count("$.customer", &Key::int(21)).expect("i"), 0);
1602 assert_eq!(docs.count("$.customer", &Key::int(28)).expect("i"), 1);
1603 for index in docs.indexes() {
1604 assert_eq!(index.postings(), 7);
1605 }
1606
1607 assert!(!docs.remove(b"order:3"), "it is already gone");
1608 assert_eq!(docs.index("$.status").expect("there").postings(), 7);
1609 }
1610
1611 #[test]
1612 fn a_path_that_names_a_container_or_nothing_is_simply_not_filed() {
1613 let mut docs = Docs::new();
1614 docs.create_index("$.lines").expect("indexed");
1615 docs.create_index("$.shipped").expect("indexed");
1616 docs.create_index("$.lines[0].qty").expect("indexed");
1617 for i in 0..4i64 {
1618 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 2))
1619 .expect("put");
1620 }
1621 assert_eq!(docs.len(), 4);
1622 assert!(
1623 docs.index("$.lines").expect("there").is_empty(),
1624 "an array has no equality key"
1625 );
1626 assert!(
1627 docs.index("$.shipped").expect("there").is_empty(),
1628 "no document has that path"
1629 );
1630 assert_eq!(
1631 docs.count("$.lines[0].qty", &Key::int(1)).expect("i"),
1632 4,
1633 "a path through an array reaches a scalar"
1634 );
1635 }
1636
1637 #[test]
1638 fn a_value_too_long_to_index_fails_the_write_and_stores_nothing() {
1639 let mut b = Builder::new();
1640 b.begin_object().expect("open");
1641 b.key(b"status").expect("key");
1642 b.text(&"x".repeat(crate::KEY_MAX)).expect("value");
1643 b.end_object().expect("close");
1644 let huge = b.finish().expect("finished").to_vec();
1645
1646 let mut docs = Docs::new();
1647 docs.create_index("$.status").expect("indexed");
1648 let err = docs.put_bytes(b"order:1", &huge).expect_err("refused");
1649 assert_eq!(err.code(), Code::Full);
1650 assert!(
1651 docs.is_empty(),
1652 "a write that cannot be indexed leaves nothing behind"
1653 );
1654
1655 assert!(docs.drop_index("$.status"));
1657 docs.put_bytes(b"order:1", &huge).expect("put");
1658 assert_eq!(docs.len(), 1);
1659 }
1660
1661 #[test]
1662 fn a_query_on_a_path_with_no_index_says_so_rather_than_scanning() {
1663 let mut docs = Docs::new();
1664 docs.put_bytes(b"order:1", &order(1, "open", 1))
1665 .expect("put");
1666 let err = docs
1667 .find("$.status", &Key::text("open"), |_, _| ())
1668 .expect_err("refused");
1669 assert_eq!(err.code(), Code::Invalid);
1670 assert_eq!(
1671 docs.count("$.status", &Key::text("open"))
1672 .expect_err("refused")
1673 .code(),
1674 Code::Invalid
1675 );
1676 assert!(docs.index("$.status").is_none());
1677 assert!(!docs.drop_index("$.status"));
1678 }
1679
1680 #[test]
1681 fn declaring_the_same_index_twice_leaves_the_first_one_alone() {
1682 let mut docs = Docs::new();
1683 docs.create_index("$.status").expect("indexed");
1684 docs.put_bytes(b"order:1", &order(1, "open", 1))
1685 .expect("put");
1686 docs.create_index("$.status").expect("indexed again");
1687 assert_eq!(docs.indexes().len(), 1);
1688 assert_eq!(
1689 docs.index("$.status").expect("there").postings(),
1690 1,
1691 "a redeclaration did not double file anything"
1692 );
1693 assert!(docs.create_index("$.[").is_err(), "the path has to parse");
1694 }
1695
1696 #[test]
1697 fn clearing_a_collection_empties_its_indexes_and_keeps_them() {
1698 let mut docs = Docs::new();
1699 docs.create_index("$.status").expect("indexed");
1700 for i in 0..8i64 {
1701 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1702 .expect("put");
1703 }
1704 docs.clear();
1705 assert!(docs.is_empty());
1706 assert!(docs.index("$.status").expect("still declared").is_empty());
1707 assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), 0);
1708
1709 docs.put_bytes(b"order:9", &order(9, "open", 1))
1710 .expect("put");
1711 assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:9"]);
1712 }
1713
1714 #[test]
1715 fn two_indexes_intersect_as_the_sets_they_are() {
1716 let mut docs = Docs::new();
1717 docs.create_index("$.status").expect("indexed");
1718 docs.create_index("$.customer").expect("indexed");
1719 for i in 0..32i64 {
1720 let status = if i % 2 == 0 { "open" } else { "shut" };
1721 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i % 4, status, 1))
1722 .expect("put");
1723 }
1724
1725 let open = Key::text("open");
1729 let customer = Key::int(14);
1730 let small = docs.count("$.customer", &customer).expect("indexed");
1731 let large = docs.count("$.status", &open).expect("indexed");
1732 assert_eq!((small, large), (8, 16));
1733
1734 let small = docs.index("$.customer").expect("there").get(&customer);
1735 let large = docs.index("$.status").expect("there").get(&open);
1736 let (Some(small), Some(large)) = (small, large) else {
1737 panic!("both keys are filed");
1738 };
1739 let mut both = Vec::new();
1740 index::each_id(small, |id| {
1741 if large.contains(id) {
1742 both.push(String::from_utf8_lossy(id).into_owned());
1743 }
1744 });
1745 both.sort();
1746 assert_eq!(
1747 both,
1748 [
1749 "order:10", "order:14", "order:18", "order:2", "order:22", "order:26", "order:30",
1750 "order:6"
1751 ]
1752 );
1753 }
1754
1755 fn ranged(docs: &Docs, lo: Bound<&Key>, hi: Bound<&Key>) -> Vec<i64> {
1757 let mut out = Vec::new();
1758 let n = docs
1759 .range("$.customer", lo, hi, |_, d| {
1760 out.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1761 })
1762 .expect("ordered");
1763 assert_eq!(n, out.len());
1764
1765 let mut back = Vec::new();
1766 docs.range_rev("$.customer", lo, hi, |_, d| {
1767 back.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1768 })
1769 .expect("ordered");
1770 back.reverse();
1771 assert_eq!(out, back, "backwards is forwards read the other way");
1772 assert_eq!(
1773 docs.count_range("$.customer", lo, hi).expect("ordered"),
1774 out.len()
1775 );
1776 out
1777 }
1778
1779 #[test]
1780 fn an_ordered_index_answers_a_range_in_order() {
1781 let mut docs = Docs::new();
1782 docs.create_ordered_index("$.customer").expect("ordered");
1783 for i in 0..64i64 {
1785 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1786 .expect("put");
1787 }
1788
1789 assert_eq!(
1790 ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1791 (0..64i64).map(|i| i * 7).collect::<Vec<i64>>()
1792 );
1793 let (lo, hi) = (Key::int(70), Key::int(105));
1794 assert_eq!(
1795 ranged(&docs, Bound::Included(&lo), Bound::Included(&hi)),
1796 [70, 77, 84, 91, 98, 105]
1797 );
1798 assert_eq!(
1799 ranged(&docs, Bound::Excluded(&lo), Bound::Excluded(&hi)),
1800 [77, 84, 91, 98]
1801 );
1802 assert_eq!(
1804 ranged(
1805 &docs,
1806 Bound::Included(&Key::int(71)),
1807 Bound::Excluded(&Key::int(90))
1808 ),
1809 [77, 84]
1810 );
1811 assert!(ranged(&docs, Bound::Included(&Key::int(442)), Bound::Unbounded).is_empty());
1812
1813 assert_eq!(docs.count("$.customer", &Key::int(70)).expect("i"), 1);
1815 assert_eq!(
1816 docs.index("$.customer").expect("there").kind(),
1817 IndexKind::Ordered
1818 );
1819 }
1820
1821 #[test]
1822 fn a_range_stays_right_through_writes_and_removals() {
1823 let mut docs = Docs::new();
1824 for i in 0..128i64 {
1825 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1826 .expect("put");
1827 }
1828 docs.create_ordered_index("$.customer").expect("ordered");
1831 assert_eq!(ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(), 128);
1832
1833 for i in (0..128i64).step_by(2) {
1836 assert!(docs.remove(format!("order:{i}").as_bytes()));
1837 }
1838 assert_eq!(
1839 ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1840 (0..128i64)
1841 .filter(|i| i % 2 == 1)
1842 .map(|i| i * 7)
1843 .collect::<Vec<i64>>()
1844 );
1845
1846 docs.put_bytes(b"order:1", &order(200, "open", 1))
1848 .expect("put");
1849 let after = ranged(&docs, Bound::Unbounded, Bound::Unbounded);
1850 assert_eq!(after.first(), Some(&21), "seven is gone");
1851 assert_eq!(after.last(), Some(&1400), "and it came back at the top");
1852 }
1853
1854 #[test]
1855 fn an_equality_index_refuses_a_range_rather_than_answering_nothing() {
1856 let mut docs = Docs::new();
1857 docs.create_index("$.customer").expect("indexed");
1858 docs.put_bytes(b"order:1", &order(1, "open", 1))
1859 .expect("put");
1860 let err = docs
1861 .range("$.customer", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1862 .expect_err("refused");
1863 assert_eq!(err.code(), Code::Invalid);
1864 assert!(err.to_string().contains("equality"), "{err}");
1865 assert_eq!(
1866 docs.range("$.status", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1867 .expect_err("refused")
1868 .code(),
1869 Code::Invalid
1870 );
1871 }
1872
1873 #[test]
1874 fn asking_for_an_order_on_an_equality_index_upgrades_it() {
1875 let mut docs = Docs::new();
1876 docs.create_index("$.customer").expect("indexed");
1877 for i in 0..8i64 {
1878 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1879 .expect("put");
1880 }
1881 assert_eq!(
1882 docs.index("$.customer").expect("there").kind(),
1883 IndexKind::Equality
1884 );
1885
1886 docs.create_ordered_index("$.customer").expect("upgraded");
1887 assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
1888 assert_eq!(ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(), 8);
1889
1890 docs.create_index("$.customer").expect("already there");
1893 assert_eq!(
1894 docs.index("$.customer").expect("there").kind(),
1895 IndexKind::Ordered
1896 );
1897 assert_eq!(docs.indexes().len(), 1);
1898 }
1899
1900 fn tagged(title: &str, tags: &[&str]) -> Vec<u8> {
1902 let mut b = Builder::new();
1903 b.begin_object().expect("open");
1904 b.key(b"title").expect("key");
1905 b.text(title).expect("value");
1906 b.key(b"tags").expect("key");
1907 b.begin_array().expect("open");
1908 for tag in tags {
1909 b.text(tag).expect("value");
1910 }
1911 b.end_array().expect("close");
1912 b.end_object().expect("close");
1913 b.finish().expect("finished").to_vec()
1914 }
1915
1916 #[test]
1917 fn an_array_index_files_a_document_under_every_element() {
1918 let mut docs = Docs::new();
1919 docs.create_array_index("$.tags").expect("indexed");
1920 docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1921 .expect("put");
1922 docs.put_bytes(b"b", &tagged("two", &["blue", "green"]))
1923 .expect("put");
1924 docs.put_bytes(b"c", &tagged("three", &[])).expect("put");
1925
1926 assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
1927 assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1928 assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["b"]);
1929 assert!(found(&docs, "$.tags", &Key::text("puce")).is_empty());
1930 assert_eq!(
1931 docs.index("$.tags").expect("there").len(),
1932 3,
1933 "three distinct tags over two documents"
1934 );
1935 }
1936
1937 #[test]
1938 fn an_array_index_takes_every_element_back_out_again() {
1939 let mut docs = Docs::new();
1940 docs.create_array_index("$.tags").expect("indexed");
1941 docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1942 .expect("put");
1943 docs.put_bytes(b"b", &tagged("two", &["blue"]))
1944 .expect("put");
1945
1946 docs.put_bytes(b"a", &tagged("one", &["blue", "green"]))
1948 .expect("put");
1949 assert!(found(&docs, "$.tags", &Key::text("red")).is_empty());
1950 assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1951 assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["a"]);
1952
1953 assert!(docs.remove(b"a"));
1954 assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["b"]);
1955 assert!(found(&docs, "$.tags", &Key::text("green")).is_empty());
1956 assert_eq!(
1957 docs.index("$.tags").expect("there").len(),
1958 1,
1959 "a tag nobody has left is not a key any more"
1960 );
1961 }
1962
1963 #[test]
1964 fn an_array_index_treats_one_value_as_a_list_of_one() {
1965 let mut docs = Docs::new();
1966 docs.create_array_index("$.status").expect("indexed");
1967 docs.put_bytes(b"order:1", &order(1, "open", 1))
1968 .expect("put");
1969 assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
1970 }
1971
1972 #[test]
1973 fn the_same_element_twice_is_one_posting() {
1974 let mut docs = Docs::new();
1975 docs.create_array_index("$.tags").expect("indexed");
1976 docs.put_bytes(b"a", &tagged("one", &["red", "red", "red"]))
1977 .expect("put");
1978 assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
1979 assert_eq!(docs.index("$.tags").expect("there").postings(), 1);
1980
1981 assert!(docs.remove(b"a"));
1983 assert_eq!(docs.index("$.tags").expect("there").postings(), 0);
1984 assert!(docs.index("$.tags").expect("there").is_empty());
1985 }
1986
1987 #[test]
1988 fn a_text_index_files_a_document_under_every_word() {
1989 let mut docs = Docs::new();
1990 docs.create_text_index("$.title").expect("indexed");
1991 docs.put_bytes(b"a", &tagged("A red bicycle", &[]))
1992 .expect("put");
1993 docs.put_bytes(b"b", &tagged("The red car, and a bicycle!", &[]))
1994 .expect("put");
1995
1996 assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a", "b"]);
1997 assert_eq!(found(&docs, "$.title", &word("car")), ["b"]);
1998 assert_eq!(
1999 found(&docs, "$.title", &word("RED")),
2000 ["a", "b"],
2001 "a search folds case the same way the write did"
2002 );
2003 assert!(found(&docs, "$.title", &word("lorry")).is_empty());
2004 }
2005
2006 #[test]
2007 fn a_text_index_follows_the_words_through_a_rewrite() {
2008 let mut docs = Docs::new();
2009 docs.create_text_index("$.title").expect("indexed");
2010 docs.put_bytes(b"a", &tagged("a red bicycle", &[]))
2011 .expect("put");
2012 docs.put_bytes(b"a", &tagged("a blue bicycle", &[]))
2013 .expect("put");
2014 assert!(found(&docs, "$.title", &word("red")).is_empty());
2015 assert_eq!(found(&docs, "$.title", &word("blue")), ["a"]);
2016 assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a"]);
2017
2018 assert!(docs.remove(b"a"));
2019 assert!(docs.index("$.title").expect("there").is_empty());
2020 }
2021
2022 #[test]
2023 fn a_text_index_declared_after_the_documents_finds_them() {
2024 let mut docs = Docs::new();
2025 for i in 0..16i64 {
2026 let title = if i % 2 == 0 {
2027 "a red one"
2028 } else {
2029 "a blue one"
2030 };
2031 docs.put_bytes(format!("t:{i}").as_bytes(), &tagged(title, &[]))
2032 .expect("put");
2033 }
2034 docs.create_text_index("$.title").expect("indexed");
2035 assert_eq!(docs.count("$.title", &word("red")).expect("i"), 8);
2036 assert_eq!(docs.count("$.title", &word("one")).expect("i"), 16);
2037 assert_eq!(
2038 docs.index("$.title").expect("there").len(),
2039 4,
2040 "a, red, blue and one"
2041 );
2042 }
2043
2044 #[test]
2045 fn changing_what_an_index_is_asked_rebuilds_it() {
2046 let mut docs = Docs::new();
2047 docs.create_index("$.tags").expect("indexed");
2048 docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
2049 .expect("put");
2050 assert!(
2051 found(&docs, "$.tags", &Key::text("red")).is_empty(),
2052 "an equality index over an array files nothing"
2053 );
2054
2055 docs.create_array_index("$.tags").expect("rebuilt");
2056 assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
2057 assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
2058
2059 docs.create_array_index("$.tags").expect("already there");
2060 assert_eq!(docs.indexes().len(), 1);
2061 }
2062
2063 fn word(w: &str) -> Key {
2065 Key::word(w).expect("one word")
2066 }
2067
2068 #[test]
2069 fn a_collection_whose_key_table_is_full_stores_the_rest_with_names() {
2070 let mut docs = Docs::new();
2073 for i in 0..crate::KEYS_MAX {
2074 let name = format!("filler{i}");
2075 assert!(docs.keys.intern(name.as_bytes()).is_some());
2076 }
2077 assert!(docs.keys().is_full());
2078
2079 docs.put_bytes(b"order:1", &order(1, "open", 1))
2080 .expect("put");
2081 let d = docs.get(b"order:1").expect("stored");
2082 assert!(!d.value().is_interned(), "there were no ids left to use");
2083 assert_eq!(d.get(b"status").and_then(|v| v.as_text()), Some("open"));
2084 assert_eq!(
2085 d.path("$.lines[0].sku")
2086 .expect("a path")
2087 .and_then(|v| v.as_text()),
2088 Some("sku-0")
2089 );
2090 }
2091
2092 fn item(lang: &str, v: &[f32]) -> Vec<u8> {
2097 let mut b = Builder::new();
2098 b.begin_object().expect("open");
2099 b.key(b"lang").expect("key");
2100 b.text(lang).expect("value");
2101 b.key(b"embedding").expect("key");
2102 b.begin_array().expect("open");
2103 for x in v {
2104 b.float(f64::from(*x)).expect("value");
2105 }
2106 b.end_array().expect("close");
2107 b.end_object().expect("close");
2108 b.finish().expect("finished").to_vec()
2109 }
2110
2111 fn bare(lang: &str) -> Vec<u8> {
2113 let mut b = Builder::new();
2114 b.begin_object().expect("open");
2115 b.key(b"lang").expect("key");
2116 b.text(lang).expect("value");
2117 b.end_object().expect("close");
2118 b.finish().expect("finished").to_vec()
2119 }
2120
2121 fn spread(n: u64) -> [f32; 8] {
2124 let mut s = n.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1;
2125 let mut v = [0.0f32; 8];
2126 for x in &mut v {
2127 s ^= s << 13;
2128 s ^= s >> 7;
2129 s ^= s << 17;
2130 *x = (s >> 40) as f32 / 4096.0 - 1.0;
2131 }
2132 v
2133 }
2134
2135 #[test]
2136 fn a_document_and_its_embedding_are_one_write() {
2137 let mut docs = Docs::new();
2138 docs.create_vector_index("$.embedding", 3).expect("index");
2139 for (id, v) in [
2140 ("a", [1.0, 0.0, 0.0]),
2141 ("b", [0.0, 1.0, 0.0]),
2142 ("c", [0.0, 0.0, 1.0]),
2143 ] {
2144 docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2145 }
2146 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 3);
2147
2148 let mut got = Vec::new();
2151 let n = docs
2152 .nearest("$.embedding", &[0.9, 0.1, 0.0], 3, |id, doc, d| {
2153 let lang = doc
2154 .get(b"lang")
2155 .and_then(|v| v.as_text())
2156 .map(str::to_owned);
2157 got.push((id.to_vec(), lang, d));
2158 })
2159 .expect("nearest");
2160 assert_eq!(n, 3);
2161 assert_eq!(got[0].0, b"a".to_vec());
2162 assert_eq!(got[0].1.as_deref(), Some("en"));
2163 assert!(got[0].2 <= got[1].2 && got[1].2 <= got[2].2);
2164
2165 assert!(
2167 docs.nearest("$.lang", &[1.0, 0.0, 0.0], 1, |_, _, _| {})
2168 .is_err()
2169 );
2170 }
2171
2172 #[test]
2173 fn an_embedding_of_the_wrong_shape_fails_the_write_and_stores_nothing() {
2174 let mut docs = Docs::new();
2175 docs.create_vector_index("$.embedding", 3).expect("index");
2176 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2177 .expect("put");
2178
2179 for wrong in [vec![1.0, 0.0], vec![1.0, 0.0, 0.0, 0.0]] {
2180 assert!(docs.put_bytes(b"b", &item("en", &wrong)).is_err());
2181 }
2182 assert!(docs.get(b"b").is_none(), "the write left nothing behind");
2183 assert_eq!(docs.len(), 1);
2184 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2185 }
2186
2187 #[test]
2188 fn a_document_with_nothing_at_the_path_is_not_in_the_index() {
2189 let mut docs = Docs::new();
2190 docs.create_vector_index("$.embedding", 3).expect("index");
2191 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2192 .expect("put");
2193 docs.put_bytes(b"b", &bare("en")).expect("put");
2194 assert_eq!(docs.len(), 2);
2195 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2196
2197 docs.put_bytes(b"a", &bare("en")).expect("put");
2200 assert!(
2201 docs.vector_index("$.embedding")
2202 .expect("declared")
2203 .is_empty()
2204 );
2205 assert!(docs.get(b"a").is_some());
2206
2207 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2209 .expect("put");
2210 assert!(docs.remove(b"a"));
2211 assert!(
2212 docs.vector_index("$.embedding")
2213 .expect("declared")
2214 .is_empty()
2215 );
2216 }
2217
2218 #[test]
2219 fn a_filter_finds_the_nearest_match_and_not_the_nearest_that_matches() {
2220 let mut docs = Docs::new();
2221 docs.create_index("$.lang").expect("index");
2222 docs.create_vector_index("$.embedding", 8).expect("index");
2223
2224 for n in 0..400u64 {
2227 let id = format!("en:{n}");
2228 docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2229 .expect("put");
2230 }
2231 docs.put_bytes(b"fr", &item("fr", &spread(9_999)))
2232 .expect("put");
2233
2234 let q = spread(3);
2235 let mut top = Vec::new();
2236 docs.nearest("$.embedding", &q, 20, |id, _, _| top.push(id.to_vec()))
2237 .expect("nearest");
2238 assert_eq!(top[0], b"en:3".to_vec());
2239 assert!(
2240 !top.iter().any(|id| id == b"fr"),
2241 "searching and then filtering would have answered nothing"
2242 );
2243
2244 let french = [("$.lang", Key::text("fr"))];
2246 let mut found = Vec::new();
2247 docs.nearest_where("$.embedding", &q, 5, &french, |id, _, _| {
2248 found.push(id.to_vec())
2249 })
2250 .expect("nearest");
2251 assert_eq!(found, [b"fr".to_vec()]);
2252
2253 let nothing = [("$.topic", Key::text("finance"))];
2255 assert!(
2256 docs.nearest_where("$.embedding", &q, 5, ¬hing, |_, _, _| {})
2257 .is_err()
2258 );
2259 }
2260
2261 #[test]
2262 fn declaring_either_index_last_gives_the_same_answers() {
2263 let q = spread(11);
2264 let french = [("$.lang", Key::text("fr"))];
2265
2266 let mut late = Docs::new();
2269 late.create_vector_index("$.embedding", 8).expect("index");
2270 for n in 0..200u64 {
2271 let lang = if n % 50 == 0 { "fr" } else { "en" };
2272 let id = format!("{n}");
2273 late.put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2274 .expect("put");
2275 }
2276 late.create_index("$.lang").expect("index");
2277
2278 let mut early = Docs::new();
2280 early.create_index("$.lang").expect("index");
2281 for n in 0..200u64 {
2282 let lang = if n % 50 == 0 { "fr" } else { "en" };
2283 let id = format!("{n}");
2284 early
2285 .put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2286 .expect("put");
2287 }
2288 early.create_vector_index("$.embedding", 8).expect("index");
2289
2290 let mut a = Vec::new();
2291 late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2292 a.push(id.to_vec())
2293 })
2294 .expect("nearest");
2295 let mut b = Vec::new();
2296 early
2297 .nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2298 b.push(id.to_vec())
2299 })
2300 .expect("nearest");
2301 assert_eq!(a.len(), 4, "there are four French documents to find");
2302 assert_eq!(a, b);
2303
2304 assert!(late.drop_index("$.lang"));
2306 late.create_index("$.lang").expect("index");
2307 let mut again = Vec::new();
2308 late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2309 again.push(id.to_vec());
2310 })
2311 .expect("nearest");
2312 assert_eq!(again, a);
2313 }
2314
2315 #[test]
2316 fn nearest_to_leaves_the_document_itself_out() {
2317 let mut docs = Docs::new();
2318 docs.create_vector_index("$.embedding", 3).expect("index");
2319 for (id, v) in [
2320 ("a", [1.0, 0.0, 0.0]),
2321 ("b", [0.9, 0.1, 0.0]),
2322 ("c", [0.0, 0.0, 1.0]),
2323 ] {
2324 docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2325 }
2326
2327 let mut like = Vec::new();
2328 docs.nearest_to("$.embedding", b"a", 2, |id, _, _| like.push(id.to_vec()))
2329 .expect("nearest");
2330 assert_eq!(like, [b"b".to_vec(), b"c".to_vec()]);
2331
2332 docs.put_bytes(b"d", &bare("en")).expect("put");
2334 let mut none = 0;
2335 assert_eq!(
2336 docs.nearest_to("$.embedding", b"d", 2, |_, _, _| none += 1)
2337 .expect("nearest"),
2338 0
2339 );
2340 }
2341
2342 #[test]
2343 fn declaring_the_same_vector_index_again_rebuilds_nothing() {
2344 let mut docs = Docs::new();
2345 for n in 0..8u64 {
2346 let id = format!("{n}");
2347 docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2348 .expect("put");
2349 }
2350 docs.create_vector_index("$.embedding", 8).expect("index");
2351 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 8);
2352
2353 docs.create_vector_index("$.embedding", 8).expect("again");
2355 assert_eq!(docs.vector_indexes().len(), 1);
2356
2357 assert!(docs.create_vector_index("$.embedding", 4).is_err());
2361 let still = docs.vector_index("$.embedding").expect("still declared");
2362 assert_eq!(still.dim(), 8);
2363 assert_eq!(still.len(), 8);
2364
2365 assert!(docs.drop_vector_index("$.embedding"));
2366 assert!(!docs.drop_vector_index("$.embedding"));
2367 assert!(docs.vector_indexes().is_empty());
2368 assert_eq!(docs.len(), 8, "the documents are untouched");
2369 }
2370
2371 #[test]
2372 fn clearing_a_collection_empties_the_vector_index_and_keeps_it_declared() {
2373 let mut docs = Docs::new();
2374 docs.create_vector_index("$.embedding", 3).expect("index");
2375 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2376 .expect("put");
2377 let full = docs.memory_bytes();
2378
2379 docs.clear();
2380 assert!(docs.is_empty());
2381 assert!(
2382 docs.vector_index("$.embedding")
2383 .expect("declared")
2384 .is_empty()
2385 );
2386 assert!(docs.memory_bytes() < full);
2387
2388 docs.put_bytes(b"b", &item("en", &[0.0, 1.0, 0.0]))
2389 .expect("put");
2390 let mut got = Vec::new();
2391 docs.nearest("$.embedding", &[0.0, 1.0, 0.0], 1, |id, _, _| {
2392 got.push(id.to_vec())
2393 })
2394 .expect("nearest");
2395 assert_eq!(got, [b"b".to_vec()]);
2396 }
2397}