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 let n = if cfg!(miri) { 16i64 } else { 256 };
1401 for i in 0..n {
1402 let bytes = shape(i);
1403 plain += bytes.len();
1404 docs.put_bytes(format!("d:{i}").as_bytes(), &bytes)
1405 .expect("put");
1406 }
1407 let stored: usize = (0..n)
1408 .map(|i| {
1409 docs.bytes(format!("d:{i}").as_bytes())
1410 .expect("stored")
1411 .len()
1412 })
1413 .sum();
1414 stored as f64 / plain as f64
1415 }
1416
1417 #[test]
1418 fn interning_makes_a_collection_of_the_same_shape_smaller() {
1419 let names = shrinkage(|i| {
1428 let mut b = Builder::new();
1429 b.begin_object().expect("open");
1430 for f in 0..20 {
1431 b.key(format!("some_field_name_{f:02}").as_bytes())
1432 .expect("key");
1433 b.int(i + f).expect("value");
1434 }
1435 b.end_object().expect("close");
1436 b.finish().expect("finished").to_vec()
1437 });
1438 assert!(names < 0.60, "a document of names kept {names}");
1439
1440 let orders = shrinkage(|i| order(i, "open", 2));
1445 assert!(orders < 0.80, "an order collection kept {orders}");
1446 }
1447
1448 #[test]
1449 fn a_document_whose_keys_are_already_ids_is_refused() {
1450 let mut b = Builder::new();
1451 b.begin_object_interned().expect("open");
1452 b.key_id(0).expect("key");
1453 b.int(1).expect("value");
1454 b.end_object().expect("close");
1455 let bytes = b.finish().expect("finished").to_vec();
1456
1457 let mut docs = Docs::new();
1458 let err = docs.put_bytes(b"x", &bytes).expect_err("refused");
1459 assert_eq!(err.code(), Code::Invalid);
1460 }
1461
1462 #[test]
1463 fn a_document_that_is_not_readable_is_refused() {
1464 let mut docs = Docs::new();
1465 let err = docs.put_bytes(b"x", &[2, 0, 0, 0]).expect_err("refused");
1466 assert_eq!(err.code(), Code::Corrupt);
1467 assert!(docs.is_empty());
1468 }
1469
1470 #[test]
1471 fn a_removal_leaves_every_other_document_where_it_was() {
1472 let mut docs = Docs::new();
1473 let n = if cfg!(miri) { 24i64 } else { 64 };
1477 for i in 0..n {
1478 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1479 .expect("put");
1480 }
1481 let gone = (0..n).step_by(3).count();
1482 for i in (0..n).step_by(3) {
1483 assert!(docs.remove(format!("order:{i}").as_bytes()));
1484 }
1485 assert_eq!(docs.len(), n as usize - gone);
1486 for i in 0..n {
1487 let id = format!("order:{i}");
1488 match docs.get(id.as_bytes()) {
1489 Some(d) => {
1490 assert!(i % 3 != 0, "{id} was removed");
1491 assert_eq!(d.get(b"id").and_then(|v| v.as_int()), Some(i));
1492 }
1493 None => assert!(i % 3 == 0, "{id} was not removed"),
1494 }
1495 }
1496 assert_eq!(docs.keys().len(), 6, "a removal does not un-intern a name");
1497 }
1498
1499 #[test]
1500 fn a_walk_sees_every_document_once() {
1501 let n = if cfg!(miri) { 48i64 } else { 200 };
1505 let mut docs = Docs::new();
1506 for i in 0..n {
1507 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1508 .expect("put");
1509 }
1510
1511 let mut seen: Vec<i64> = docs
1512 .iter()
1513 .map(|(_, d)| d.get(b"id").and_then(|v| v.as_int()).expect("an id"))
1514 .collect();
1515 seen.sort_unstable();
1516 assert_eq!(seen, (0..n).collect::<Vec<i64>>());
1517
1518 let mut scanned = Vec::new();
1519 let mut cursor = Cursor::START;
1520 loop {
1521 cursor = docs.scan(cursor, 16, |id, _| scanned.push(id.to_vec()));
1522 if cursor.is_end() {
1523 break;
1524 }
1525 }
1526 scanned.sort_unstable();
1527 scanned.dedup();
1528 assert_eq!(scanned.len(), n as usize);
1529 }
1530
1531 #[test]
1532 fn an_empty_collection_answers_nothing_rather_than_failing() {
1533 let docs = Docs::new();
1534 assert!(docs.is_empty());
1535 assert!(docs.get(b"nothing").is_none());
1536 assert!(docs.bytes(b"nothing").is_none());
1537 assert!(!docs.contains(b"nothing"));
1538 assert_eq!(docs.iter().count(), 0);
1539 }
1540
1541 #[test]
1542 fn a_document_prints_with_its_names_back_on() {
1543 let mut docs = Docs::new();
1544 docs.put_bytes(b"order:1", &order(1, "open", 1))
1545 .expect("put");
1546 let text = format!("{:?}", docs.get(b"order:1").expect("stored"));
1547 assert!(text.contains("\"status\": \"open\""), "{text}");
1548 assert!(text.contains("\"sku\": \"sku-0\""), "{text}");
1549 }
1550
1551 fn found(docs: &Docs, path: &str, key: &Key) -> Vec<String> {
1553 let mut out = Vec::new();
1554 let n = docs
1555 .find(path, key, |id, d| {
1556 assert!(!d.is_empty(), "the document came back whole");
1557 out.push(String::from_utf8_lossy(id).into_owned());
1558 })
1559 .expect("indexed");
1560 assert_eq!(n, out.len(), "the count is what the callback saw");
1561 out.sort();
1562 out
1563 }
1564
1565 #[test]
1566 fn an_index_declared_after_the_documents_finds_them() {
1567 let n = if cfg!(miri) { 24i64 } else { 64 };
1571 let (shut, open) = ((n / 4) as usize, (n - n / 4) as usize);
1572 let mut docs = Docs::new();
1573 for i in 0..n {
1574 let status = if i % 4 == 0 { "shut" } else { "open" };
1575 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, status, 1))
1576 .expect("put");
1577 }
1578 docs.create_index("$.status").expect("indexed");
1579 assert_eq!(docs.index("$.status").expect("there").len(), 2);
1580 assert_eq!(docs.count("$.status", &Key::text("shut")).expect("i"), shut);
1581 assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), open);
1582 assert_eq!(found(&docs, "$.status", &Key::text("shut")).len(), shut);
1583 assert!(found(&docs, "$.status", &Key::text("gone")).is_empty());
1584
1585 docs.put_bytes(format!("order:{n}").as_bytes(), &order(n, "shut", 1))
1587 .expect("put");
1588 assert_eq!(
1589 docs.count("$.status", &Key::text("shut")).expect("i"),
1590 shut + 1
1591 );
1592 }
1593
1594 #[test]
1595 fn an_overwrite_moves_a_document_from_one_key_to_the_other() {
1596 let mut docs = Docs::new();
1597 docs.create_index("$.status").expect("indexed");
1598 docs.put_bytes(b"order:1", &order(1, "open", 1))
1599 .expect("put");
1600 assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
1601
1602 docs.put_bytes(b"order:1", &order(1, "shut", 1))
1603 .expect("put");
1604 assert!(
1605 found(&docs, "$.status", &Key::text("open")).is_empty(),
1606 "the old key kept it"
1607 );
1608 assert_eq!(found(&docs, "$.status", &Key::text("shut")), ["order:1"]);
1609 assert_eq!(docs.index("$.status").expect("there").postings(), 1);
1610 }
1611
1612 #[test]
1613 fn a_removal_takes_a_document_out_of_every_index() {
1614 let mut docs = Docs::new();
1615 docs.create_index("$.status").expect("indexed");
1616 docs.create_index("$.customer").expect("indexed");
1617 for i in 0..8i64 {
1618 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1619 .expect("put");
1620 }
1621 assert!(docs.remove(b"order:3"));
1622 assert_eq!(found(&docs, "$.status", &Key::text("open")).len(), 7);
1623 assert_eq!(docs.count("$.customer", &Key::int(21)).expect("i"), 0);
1624 assert_eq!(docs.count("$.customer", &Key::int(28)).expect("i"), 1);
1625 for index in docs.indexes() {
1626 assert_eq!(index.postings(), 7);
1627 }
1628
1629 assert!(!docs.remove(b"order:3"), "it is already gone");
1630 assert_eq!(docs.index("$.status").expect("there").postings(), 7);
1631 }
1632
1633 #[test]
1634 fn a_path_that_names_a_container_or_nothing_is_simply_not_filed() {
1635 let mut docs = Docs::new();
1636 docs.create_index("$.lines").expect("indexed");
1637 docs.create_index("$.shipped").expect("indexed");
1638 docs.create_index("$.lines[0].qty").expect("indexed");
1639 for i in 0..4i64 {
1640 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 2))
1641 .expect("put");
1642 }
1643 assert_eq!(docs.len(), 4);
1644 assert!(
1645 docs.index("$.lines").expect("there").is_empty(),
1646 "an array has no equality key"
1647 );
1648 assert!(
1649 docs.index("$.shipped").expect("there").is_empty(),
1650 "no document has that path"
1651 );
1652 assert_eq!(
1653 docs.count("$.lines[0].qty", &Key::int(1)).expect("i"),
1654 4,
1655 "a path through an array reaches a scalar"
1656 );
1657 }
1658
1659 #[test]
1660 fn a_value_too_long_to_index_fails_the_write_and_stores_nothing() {
1661 let mut b = Builder::new();
1662 b.begin_object().expect("open");
1663 b.key(b"status").expect("key");
1664 b.text(&"x".repeat(crate::KEY_MAX)).expect("value");
1665 b.end_object().expect("close");
1666 let huge = b.finish().expect("finished").to_vec();
1667
1668 let mut docs = Docs::new();
1669 docs.create_index("$.status").expect("indexed");
1670 let err = docs.put_bytes(b"order:1", &huge).expect_err("refused");
1671 assert_eq!(err.code(), Code::Full);
1672 assert!(
1673 docs.is_empty(),
1674 "a write that cannot be indexed leaves nothing behind"
1675 );
1676
1677 assert!(docs.drop_index("$.status"));
1679 docs.put_bytes(b"order:1", &huge).expect("put");
1680 assert_eq!(docs.len(), 1);
1681 }
1682
1683 #[test]
1684 fn a_query_on_a_path_with_no_index_says_so_rather_than_scanning() {
1685 let mut docs = Docs::new();
1686 docs.put_bytes(b"order:1", &order(1, "open", 1))
1687 .expect("put");
1688 let err = docs
1689 .find("$.status", &Key::text("open"), |_, _| ())
1690 .expect_err("refused");
1691 assert_eq!(err.code(), Code::Invalid);
1692 assert_eq!(
1693 docs.count("$.status", &Key::text("open"))
1694 .expect_err("refused")
1695 .code(),
1696 Code::Invalid
1697 );
1698 assert!(docs.index("$.status").is_none());
1699 assert!(!docs.drop_index("$.status"));
1700 }
1701
1702 #[test]
1703 fn declaring_the_same_index_twice_leaves_the_first_one_alone() {
1704 let mut docs = Docs::new();
1705 docs.create_index("$.status").expect("indexed");
1706 docs.put_bytes(b"order:1", &order(1, "open", 1))
1707 .expect("put");
1708 docs.create_index("$.status").expect("indexed again");
1709 assert_eq!(docs.indexes().len(), 1);
1710 assert_eq!(
1711 docs.index("$.status").expect("there").postings(),
1712 1,
1713 "a redeclaration did not double file anything"
1714 );
1715 assert!(docs.create_index("$.[").is_err(), "the path has to parse");
1716 }
1717
1718 #[test]
1719 fn clearing_a_collection_empties_its_indexes_and_keeps_them() {
1720 let mut docs = Docs::new();
1721 docs.create_index("$.status").expect("indexed");
1722 for i in 0..8i64 {
1723 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1724 .expect("put");
1725 }
1726 docs.clear();
1727 assert!(docs.is_empty());
1728 assert!(docs.index("$.status").expect("still declared").is_empty());
1729 assert_eq!(docs.count("$.status", &Key::text("open")).expect("i"), 0);
1730
1731 docs.put_bytes(b"order:9", &order(9, "open", 1))
1732 .expect("put");
1733 assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:9"]);
1734 }
1735
1736 #[test]
1737 fn two_indexes_intersect_as_the_sets_they_are() {
1738 let mut docs = Docs::new();
1739 docs.create_index("$.status").expect("indexed");
1740 docs.create_index("$.customer").expect("indexed");
1741 for i in 0..32i64 {
1742 let status = if i % 2 == 0 { "open" } else { "shut" };
1743 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i % 4, status, 1))
1744 .expect("put");
1745 }
1746
1747 let open = Key::text("open");
1751 let customer = Key::int(14);
1752 let small = docs.count("$.customer", &customer).expect("indexed");
1753 let large = docs.count("$.status", &open).expect("indexed");
1754 assert_eq!((small, large), (8, 16));
1755
1756 let small = docs.index("$.customer").expect("there").get(&customer);
1757 let large = docs.index("$.status").expect("there").get(&open);
1758 let (Some(small), Some(large)) = (small, large) else {
1759 panic!("both keys are filed");
1760 };
1761 let mut both = Vec::new();
1762 index::each_id(small, |id| {
1763 if large.contains(id) {
1764 both.push(String::from_utf8_lossy(id).into_owned());
1765 }
1766 });
1767 both.sort();
1768 assert_eq!(
1769 both,
1770 [
1771 "order:10", "order:14", "order:18", "order:2", "order:22", "order:26", "order:30",
1772 "order:6"
1773 ]
1774 );
1775 }
1776
1777 fn ranged(docs: &Docs, lo: Bound<&Key>, hi: Bound<&Key>) -> Vec<i64> {
1779 let mut out = Vec::new();
1780 let n = docs
1781 .range("$.customer", lo, hi, |_, d| {
1782 out.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1783 })
1784 .expect("ordered");
1785 assert_eq!(n, out.len());
1786
1787 let mut back = Vec::new();
1788 docs.range_rev("$.customer", lo, hi, |_, d| {
1789 back.push(d.get(b"customer").and_then(|v| v.as_int()).expect("there"));
1790 })
1791 .expect("ordered");
1792 back.reverse();
1793 assert_eq!(out, back, "backwards is forwards read the other way");
1794 assert_eq!(
1795 docs.count_range("$.customer", lo, hi).expect("ordered"),
1796 out.len()
1797 );
1798 out
1799 }
1800
1801 #[test]
1802 fn an_ordered_index_answers_a_range_in_order() {
1803 let mut docs = Docs::new();
1804 docs.create_ordered_index("$.customer").expect("ordered");
1805 let n = if cfg!(miri) { 24i64 } else { 64 };
1810 for i in 0..n {
1811 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1812 .expect("put");
1813 }
1814
1815 assert_eq!(
1816 ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1817 (0..n).map(|i| i * 7).collect::<Vec<i64>>()
1818 );
1819 let (lo, hi) = (Key::int(70), Key::int(105));
1820 assert_eq!(
1821 ranged(&docs, Bound::Included(&lo), Bound::Included(&hi)),
1822 [70, 77, 84, 91, 98, 105]
1823 );
1824 assert_eq!(
1825 ranged(&docs, Bound::Excluded(&lo), Bound::Excluded(&hi)),
1826 [77, 84, 91, 98]
1827 );
1828 assert_eq!(
1830 ranged(
1831 &docs,
1832 Bound::Included(&Key::int(71)),
1833 Bound::Excluded(&Key::int(90))
1834 ),
1835 [77, 84]
1836 );
1837 assert!(ranged(&docs, Bound::Included(&Key::int(442)), Bound::Unbounded).is_empty());
1838
1839 assert_eq!(docs.count("$.customer", &Key::int(70)).expect("i"), 1);
1841 assert_eq!(
1842 docs.index("$.customer").expect("there").kind(),
1843 IndexKind::Ordered
1844 );
1845 }
1846
1847 #[test]
1848 fn a_range_stays_right_through_writes_and_removals() {
1849 let n = if cfg!(miri) { 32i64 } else { 128 };
1853 let mut docs = Docs::new();
1854 for i in 0..n {
1855 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1856 .expect("put");
1857 }
1858 docs.create_ordered_index("$.customer").expect("ordered");
1861 assert_eq!(
1862 ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(),
1863 n as usize
1864 );
1865
1866 for i in (0..n).step_by(2) {
1869 assert!(docs.remove(format!("order:{i}").as_bytes()));
1870 }
1871 assert_eq!(
1872 ranged(&docs, Bound::Unbounded, Bound::Unbounded),
1873 (0..n)
1874 .filter(|i| i % 2 == 1)
1875 .map(|i| i * 7)
1876 .collect::<Vec<i64>>()
1877 );
1878
1879 docs.put_bytes(b"order:1", &order(200, "open", 1))
1881 .expect("put");
1882 let after = ranged(&docs, Bound::Unbounded, Bound::Unbounded);
1883 assert_eq!(after.first(), Some(&21), "seven is gone");
1884 assert_eq!(after.last(), Some(&1400), "and it came back at the top");
1885 }
1886
1887 #[test]
1888 fn an_equality_index_refuses_a_range_rather_than_answering_nothing() {
1889 let mut docs = Docs::new();
1890 docs.create_index("$.customer").expect("indexed");
1891 docs.put_bytes(b"order:1", &order(1, "open", 1))
1892 .expect("put");
1893 let err = docs
1894 .range("$.customer", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1895 .expect_err("refused");
1896 assert_eq!(err.code(), Code::Invalid);
1897 assert!(err.to_string().contains("equality"), "{err}");
1898 assert_eq!(
1899 docs.range("$.status", Bound::Unbounded, Bound::Unbounded, |_, _| ())
1900 .expect_err("refused")
1901 .code(),
1902 Code::Invalid
1903 );
1904 }
1905
1906 #[test]
1907 fn asking_for_an_order_on_an_equality_index_upgrades_it() {
1908 let mut docs = Docs::new();
1909 docs.create_index("$.customer").expect("indexed");
1910 for i in 0..8i64 {
1911 docs.put_bytes(format!("order:{i}").as_bytes(), &order(i, "open", 1))
1912 .expect("put");
1913 }
1914 assert_eq!(
1915 docs.index("$.customer").expect("there").kind(),
1916 IndexKind::Equality
1917 );
1918
1919 docs.create_ordered_index("$.customer").expect("upgraded");
1920 assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
1921 assert_eq!(ranged(&docs, Bound::Unbounded, Bound::Unbounded).len(), 8);
1922
1923 docs.create_index("$.customer").expect("already there");
1926 assert_eq!(
1927 docs.index("$.customer").expect("there").kind(),
1928 IndexKind::Ordered
1929 );
1930 assert_eq!(docs.indexes().len(), 1);
1931 }
1932
1933 fn tagged(title: &str, tags: &[&str]) -> Vec<u8> {
1935 let mut b = Builder::new();
1936 b.begin_object().expect("open");
1937 b.key(b"title").expect("key");
1938 b.text(title).expect("value");
1939 b.key(b"tags").expect("key");
1940 b.begin_array().expect("open");
1941 for tag in tags {
1942 b.text(tag).expect("value");
1943 }
1944 b.end_array().expect("close");
1945 b.end_object().expect("close");
1946 b.finish().expect("finished").to_vec()
1947 }
1948
1949 #[test]
1950 fn an_array_index_files_a_document_under_every_element() {
1951 let mut docs = Docs::new();
1952 docs.create_array_index("$.tags").expect("indexed");
1953 docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1954 .expect("put");
1955 docs.put_bytes(b"b", &tagged("two", &["blue", "green"]))
1956 .expect("put");
1957 docs.put_bytes(b"c", &tagged("three", &[])).expect("put");
1958
1959 assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
1960 assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1961 assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["b"]);
1962 assert!(found(&docs, "$.tags", &Key::text("puce")).is_empty());
1963 assert_eq!(
1964 docs.index("$.tags").expect("there").len(),
1965 3,
1966 "three distinct tags over two documents"
1967 );
1968 }
1969
1970 #[test]
1971 fn an_array_index_takes_every_element_back_out_again() {
1972 let mut docs = Docs::new();
1973 docs.create_array_index("$.tags").expect("indexed");
1974 docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
1975 .expect("put");
1976 docs.put_bytes(b"b", &tagged("two", &["blue"]))
1977 .expect("put");
1978
1979 docs.put_bytes(b"a", &tagged("one", &["blue", "green"]))
1981 .expect("put");
1982 assert!(found(&docs, "$.tags", &Key::text("red")).is_empty());
1983 assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["a", "b"]);
1984 assert_eq!(found(&docs, "$.tags", &Key::text("green")), ["a"]);
1985
1986 assert!(docs.remove(b"a"));
1987 assert_eq!(found(&docs, "$.tags", &Key::text("blue")), ["b"]);
1988 assert!(found(&docs, "$.tags", &Key::text("green")).is_empty());
1989 assert_eq!(
1990 docs.index("$.tags").expect("there").len(),
1991 1,
1992 "a tag nobody has left is not a key any more"
1993 );
1994 }
1995
1996 #[test]
1997 fn an_array_index_treats_one_value_as_a_list_of_one() {
1998 let mut docs = Docs::new();
1999 docs.create_array_index("$.status").expect("indexed");
2000 docs.put_bytes(b"order:1", &order(1, "open", 1))
2001 .expect("put");
2002 assert_eq!(found(&docs, "$.status", &Key::text("open")), ["order:1"]);
2003 }
2004
2005 #[test]
2006 fn the_same_element_twice_is_one_posting() {
2007 let mut docs = Docs::new();
2008 docs.create_array_index("$.tags").expect("indexed");
2009 docs.put_bytes(b"a", &tagged("one", &["red", "red", "red"]))
2010 .expect("put");
2011 assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
2012 assert_eq!(docs.index("$.tags").expect("there").postings(), 1);
2013
2014 assert!(docs.remove(b"a"));
2016 assert_eq!(docs.index("$.tags").expect("there").postings(), 0);
2017 assert!(docs.index("$.tags").expect("there").is_empty());
2018 }
2019
2020 #[test]
2021 fn a_text_index_files_a_document_under_every_word() {
2022 let mut docs = Docs::new();
2023 docs.create_text_index("$.title").expect("indexed");
2024 docs.put_bytes(b"a", &tagged("A red bicycle", &[]))
2025 .expect("put");
2026 docs.put_bytes(b"b", &tagged("The red car, and a bicycle!", &[]))
2027 .expect("put");
2028
2029 assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a", "b"]);
2030 assert_eq!(found(&docs, "$.title", &word("car")), ["b"]);
2031 assert_eq!(
2032 found(&docs, "$.title", &word("RED")),
2033 ["a", "b"],
2034 "a search folds case the same way the write did"
2035 );
2036 assert!(found(&docs, "$.title", &word("lorry")).is_empty());
2037 }
2038
2039 #[test]
2040 fn a_text_index_follows_the_words_through_a_rewrite() {
2041 let mut docs = Docs::new();
2042 docs.create_text_index("$.title").expect("indexed");
2043 docs.put_bytes(b"a", &tagged("a red bicycle", &[]))
2044 .expect("put");
2045 docs.put_bytes(b"a", &tagged("a blue bicycle", &[]))
2046 .expect("put");
2047 assert!(found(&docs, "$.title", &word("red")).is_empty());
2048 assert_eq!(found(&docs, "$.title", &word("blue")), ["a"]);
2049 assert_eq!(found(&docs, "$.title", &word("bicycle")), ["a"]);
2050
2051 assert!(docs.remove(b"a"));
2052 assert!(docs.index("$.title").expect("there").is_empty());
2053 }
2054
2055 #[test]
2056 fn a_text_index_declared_after_the_documents_finds_them() {
2057 let mut docs = Docs::new();
2058 for i in 0..16i64 {
2059 let title = if i % 2 == 0 {
2060 "a red one"
2061 } else {
2062 "a blue one"
2063 };
2064 docs.put_bytes(format!("t:{i}").as_bytes(), &tagged(title, &[]))
2065 .expect("put");
2066 }
2067 docs.create_text_index("$.title").expect("indexed");
2068 assert_eq!(docs.count("$.title", &word("red")).expect("i"), 8);
2069 assert_eq!(docs.count("$.title", &word("one")).expect("i"), 16);
2070 assert_eq!(
2071 docs.index("$.title").expect("there").len(),
2072 4,
2073 "a, red, blue and one"
2074 );
2075 }
2076
2077 #[test]
2078 fn changing_what_an_index_is_asked_rebuilds_it() {
2079 let mut docs = Docs::new();
2080 docs.create_index("$.tags").expect("indexed");
2081 docs.put_bytes(b"a", &tagged("one", &["red", "blue"]))
2082 .expect("put");
2083 assert!(
2084 found(&docs, "$.tags", &Key::text("red")).is_empty(),
2085 "an equality index over an array files nothing"
2086 );
2087
2088 docs.create_array_index("$.tags").expect("rebuilt");
2089 assert_eq!(docs.indexes().len(), 1, "it replaced rather than added");
2090 assert_eq!(found(&docs, "$.tags", &Key::text("red")), ["a"]);
2091
2092 docs.create_array_index("$.tags").expect("already there");
2093 assert_eq!(docs.indexes().len(), 1);
2094 }
2095
2096 fn word(w: &str) -> Key {
2098 Key::word(w).expect("one word")
2099 }
2100
2101 #[test]
2102 #[cfg_attr(miri, ignore = "a full key table is the claim and it is 65536 names")]
2103 fn a_collection_whose_key_table_is_full_stores_the_rest_with_names() {
2104 let mut docs = Docs::new();
2107 for i in 0..crate::KEYS_MAX {
2108 let name = format!("filler{i}");
2109 assert!(docs.keys.intern(name.as_bytes()).is_some());
2110 }
2111 assert!(docs.keys().is_full());
2112
2113 docs.put_bytes(b"order:1", &order(1, "open", 1))
2114 .expect("put");
2115 let d = docs.get(b"order:1").expect("stored");
2116 assert!(!d.value().is_interned(), "there were no ids left to use");
2117 assert_eq!(d.get(b"status").and_then(|v| v.as_text()), Some("open"));
2118 assert_eq!(
2119 d.path("$.lines[0].sku")
2120 .expect("a path")
2121 .and_then(|v| v.as_text()),
2122 Some("sku-0")
2123 );
2124 }
2125
2126 fn item(lang: &str, v: &[f32]) -> Vec<u8> {
2131 let mut b = Builder::new();
2132 b.begin_object().expect("open");
2133 b.key(b"lang").expect("key");
2134 b.text(lang).expect("value");
2135 b.key(b"embedding").expect("key");
2136 b.begin_array().expect("open");
2137 for x in v {
2138 b.float(f64::from(*x)).expect("value");
2139 }
2140 b.end_array().expect("close");
2141 b.end_object().expect("close");
2142 b.finish().expect("finished").to_vec()
2143 }
2144
2145 fn bare(lang: &str) -> Vec<u8> {
2147 let mut b = Builder::new();
2148 b.begin_object().expect("open");
2149 b.key(b"lang").expect("key");
2150 b.text(lang).expect("value");
2151 b.end_object().expect("close");
2152 b.finish().expect("finished").to_vec()
2153 }
2154
2155 fn spread(n: u64) -> [f32; 8] {
2158 let mut s = n.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1;
2159 let mut v = [0.0f32; 8];
2160 for x in &mut v {
2161 s ^= s << 13;
2162 s ^= s >> 7;
2163 s ^= s << 17;
2164 *x = (s >> 40) as f32 / 4096.0 - 1.0;
2165 }
2166 v
2167 }
2168
2169 #[test]
2170 fn a_document_and_its_embedding_are_one_write() {
2171 let mut docs = Docs::new();
2172 docs.create_vector_index("$.embedding", 3).expect("index");
2173 for (id, v) in [
2174 ("a", [1.0, 0.0, 0.0]),
2175 ("b", [0.0, 1.0, 0.0]),
2176 ("c", [0.0, 0.0, 1.0]),
2177 ] {
2178 docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2179 }
2180 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 3);
2181
2182 let mut got = Vec::new();
2185 let n = docs
2186 .nearest("$.embedding", &[0.9, 0.1, 0.0], 3, |id, doc, d| {
2187 let lang = doc
2188 .get(b"lang")
2189 .and_then(|v| v.as_text())
2190 .map(str::to_owned);
2191 got.push((id.to_vec(), lang, d));
2192 })
2193 .expect("nearest");
2194 assert_eq!(n, 3);
2195 assert_eq!(got[0].0, b"a".to_vec());
2196 assert_eq!(got[0].1.as_deref(), Some("en"));
2197 assert!(got[0].2 <= got[1].2 && got[1].2 <= got[2].2);
2198
2199 assert!(
2201 docs.nearest("$.lang", &[1.0, 0.0, 0.0], 1, |_, _, _| {})
2202 .is_err()
2203 );
2204 }
2205
2206 #[test]
2207 fn an_embedding_of_the_wrong_shape_fails_the_write_and_stores_nothing() {
2208 let mut docs = Docs::new();
2209 docs.create_vector_index("$.embedding", 3).expect("index");
2210 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2211 .expect("put");
2212
2213 for wrong in [vec![1.0, 0.0], vec![1.0, 0.0, 0.0, 0.0]] {
2214 assert!(docs.put_bytes(b"b", &item("en", &wrong)).is_err());
2215 }
2216 assert!(docs.get(b"b").is_none(), "the write left nothing behind");
2217 assert_eq!(docs.len(), 1);
2218 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2219 }
2220
2221 #[test]
2222 fn a_document_with_nothing_at_the_path_is_not_in_the_index() {
2223 let mut docs = Docs::new();
2224 docs.create_vector_index("$.embedding", 3).expect("index");
2225 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2226 .expect("put");
2227 docs.put_bytes(b"b", &bare("en")).expect("put");
2228 assert_eq!(docs.len(), 2);
2229 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 1);
2230
2231 docs.put_bytes(b"a", &bare("en")).expect("put");
2234 assert!(
2235 docs.vector_index("$.embedding")
2236 .expect("declared")
2237 .is_empty()
2238 );
2239 assert!(docs.get(b"a").is_some());
2240
2241 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2243 .expect("put");
2244 assert!(docs.remove(b"a"));
2245 assert!(
2246 docs.vector_index("$.embedding")
2247 .expect("declared")
2248 .is_empty()
2249 );
2250 }
2251
2252 #[test]
2253 fn a_filter_finds_the_nearest_match_and_not_the_nearest_that_matches() {
2254 let mut docs = Docs::new();
2255 docs.create_index("$.lang").expect("index");
2256 docs.create_vector_index("$.embedding", 8).expect("index");
2257
2258 let count = if cfg!(miri) { 30u64 } else { 400 };
2263 for n in 0..count {
2264 let id = format!("en:{n}");
2265 docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2266 .expect("put");
2267 }
2268 docs.put_bytes(b"fr", &item("fr", &spread(9_999)))
2269 .expect("put");
2270
2271 let q = spread(3);
2272 let mut top = Vec::new();
2273 docs.nearest("$.embedding", &q, 20, |id, _, _| top.push(id.to_vec()))
2274 .expect("nearest");
2275 assert_eq!(top[0], b"en:3".to_vec());
2276 assert!(
2277 !top.iter().any(|id| id == b"fr"),
2278 "searching and then filtering would have answered nothing"
2279 );
2280
2281 let french = [("$.lang", Key::text("fr"))];
2283 let mut found = Vec::new();
2284 docs.nearest_where("$.embedding", &q, 5, &french, |id, _, _| {
2285 found.push(id.to_vec())
2286 })
2287 .expect("nearest");
2288 assert_eq!(found, [b"fr".to_vec()]);
2289
2290 let nothing = [("$.topic", Key::text("finance"))];
2292 assert!(
2293 docs.nearest_where("$.embedding", &q, 5, ¬hing, |_, _, _| {})
2294 .is_err()
2295 );
2296 }
2297
2298 #[test]
2299 fn declaring_either_index_last_gives_the_same_answers() {
2300 let q = spread(11);
2301 let french = [("$.lang", Key::text("fr"))];
2302
2303 let (count, every) = if cfg!(miri) { (20u64, 5) } else { (200, 50) };
2310 let mut late = Docs::new();
2311 late.create_vector_index("$.embedding", 8).expect("index");
2312 for n in 0..count {
2313 let lang = if n % every == 0 { "fr" } else { "en" };
2314 let id = format!("{n}");
2315 late.put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2316 .expect("put");
2317 }
2318 late.create_index("$.lang").expect("index");
2319
2320 let mut early = Docs::new();
2322 early.create_index("$.lang").expect("index");
2323 for n in 0..count {
2324 let lang = if n % every == 0 { "fr" } else { "en" };
2325 let id = format!("{n}");
2326 early
2327 .put_bytes(id.as_bytes(), &item(lang, &spread(n)))
2328 .expect("put");
2329 }
2330 early.create_vector_index("$.embedding", 8).expect("index");
2331
2332 let mut a = Vec::new();
2333 late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2334 a.push(id.to_vec())
2335 })
2336 .expect("nearest");
2337 let mut b = Vec::new();
2338 early
2339 .nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2340 b.push(id.to_vec())
2341 })
2342 .expect("nearest");
2343 assert_eq!(a.len(), 4, "there are four French documents to find");
2344 assert_eq!(a, b);
2345
2346 assert!(late.drop_index("$.lang"));
2348 late.create_index("$.lang").expect("index");
2349 let mut again = Vec::new();
2350 late.nearest_where("$.embedding", &q, 4, &french, |id, _, _| {
2351 again.push(id.to_vec());
2352 })
2353 .expect("nearest");
2354 assert_eq!(again, a);
2355 }
2356
2357 #[test]
2358 fn nearest_to_leaves_the_document_itself_out() {
2359 let mut docs = Docs::new();
2360 docs.create_vector_index("$.embedding", 3).expect("index");
2361 for (id, v) in [
2362 ("a", [1.0, 0.0, 0.0]),
2363 ("b", [0.9, 0.1, 0.0]),
2364 ("c", [0.0, 0.0, 1.0]),
2365 ] {
2366 docs.put_bytes(id.as_bytes(), &item("en", &v)).expect("put");
2367 }
2368
2369 let mut like = Vec::new();
2370 docs.nearest_to("$.embedding", b"a", 2, |id, _, _| like.push(id.to_vec()))
2371 .expect("nearest");
2372 assert_eq!(like, [b"b".to_vec(), b"c".to_vec()]);
2373
2374 docs.put_bytes(b"d", &bare("en")).expect("put");
2376 let mut none = 0;
2377 assert_eq!(
2378 docs.nearest_to("$.embedding", b"d", 2, |_, _, _| none += 1)
2379 .expect("nearest"),
2380 0
2381 );
2382 }
2383
2384 #[test]
2385 fn declaring_the_same_vector_index_again_rebuilds_nothing() {
2386 let mut docs = Docs::new();
2387 for n in 0..8u64 {
2388 let id = format!("{n}");
2389 docs.put_bytes(id.as_bytes(), &item("en", &spread(n)))
2390 .expect("put");
2391 }
2392 docs.create_vector_index("$.embedding", 8).expect("index");
2393 assert_eq!(docs.vector_index("$.embedding").expect("declared").len(), 8);
2394
2395 docs.create_vector_index("$.embedding", 8).expect("again");
2397 assert_eq!(docs.vector_indexes().len(), 1);
2398
2399 assert!(docs.create_vector_index("$.embedding", 4).is_err());
2403 let still = docs.vector_index("$.embedding").expect("still declared");
2404 assert_eq!(still.dim(), 8);
2405 assert_eq!(still.len(), 8);
2406
2407 assert!(docs.drop_vector_index("$.embedding"));
2408 assert!(!docs.drop_vector_index("$.embedding"));
2409 assert!(docs.vector_indexes().is_empty());
2410 assert_eq!(docs.len(), 8, "the documents are untouched");
2411 }
2412
2413 #[test]
2414 fn clearing_a_collection_empties_the_vector_index_and_keeps_it_declared() {
2415 let mut docs = Docs::new();
2416 docs.create_vector_index("$.embedding", 3).expect("index");
2417 docs.put_bytes(b"a", &item("en", &[1.0, 0.0, 0.0]))
2418 .expect("put");
2419 let full = docs.memory_bytes();
2420
2421 docs.clear();
2422 assert!(docs.is_empty());
2423 assert!(
2424 docs.vector_index("$.embedding")
2425 .expect("declared")
2426 .is_empty()
2427 );
2428 assert!(docs.memory_bytes() < full);
2429
2430 docs.put_bytes(b"b", &item("en", &[0.0, 1.0, 0.0]))
2431 .expect("put");
2432 let mut got = Vec::new();
2433 docs.nearest("$.embedding", &[0.0, 1.0, 0.0], 1, |id, _, _| {
2434 got.push(id.to_vec())
2435 })
2436 .expect("nearest");
2437 assert_eq!(got, [b"b".to_vec()]);
2438 }
2439}