1use std::collections::{BTreeMap, HashMap, HashSet};
2use std::io::{self, Read, Write};
3use std::net::Ipv6Addr;
4
5use columnar::MonotonicallyMappableToU128;
6use common::{read_u32_vint_no_advance, serialize_vint_u32, BinarySerializable, DateTime, VInt};
7use serde_json::Map;
8pub use CompactDoc as TantivyDocument;
9
10use super::{ReferenceValue, ReferenceValueLeaf, Value};
11use crate::schema::document::{
12 DeserializeError, Document, DocumentDeserialize, DocumentDeserializer,
13};
14use crate::schema::field_type::ValueParsingError;
15use crate::schema::{Facet, Field, NamedFieldDocument, OwnedValue, Schema};
16use crate::tokenizer::PreTokenizedString;
17
18#[repr(C, packed)]
19#[derive(Debug, Clone)]
20struct FieldValueAddr {
22 pub field: u16,
23 pub value_addr: ValueAddr,
24}
25
26#[derive(Debug, Clone)]
27pub struct CompactDoc {
29 pub node_data: Vec<u8>,
33 field_values: Vec<FieldValueAddr>,
35}
36
37impl Default for CompactDoc {
38 fn default() -> Self {
39 Self::new()
40 }
41}
42
43impl CompactDoc {
44 pub fn with_capacity(bytes: usize) -> CompactDoc {
47 CompactDoc {
48 node_data: Vec::with_capacity(bytes),
49 field_values: Vec::with_capacity(4),
50 }
51 }
52
53 pub fn new() -> CompactDoc {
55 CompactDoc::with_capacity(1024)
56 }
57
58 pub fn shrink_to_fit(&mut self) {
60 self.node_data.shrink_to_fit();
61 self.field_values.shrink_to_fit();
62 }
63
64 pub fn len(&self) -> usize {
66 self.field_values.len()
67 }
68
69 pub fn add_facet<F>(&mut self, field: Field, path: F)
71 where Facet: From<F> {
72 let facet = Facet::from(path);
73 self.add_leaf_field_value(field, ReferenceValueLeaf::Facet(facet.encoded_str()));
74 }
75
76 pub fn add_text<S: AsRef<str>>(&mut self, field: Field, text: S) {
78 self.add_leaf_field_value(field, ReferenceValueLeaf::Str(text.as_ref()));
79 }
80
81 pub fn add_pre_tokenized_text(&mut self, field: Field, pre_tokenized_text: PreTokenizedString) {
83 self.add_leaf_field_value(field, pre_tokenized_text);
84 }
85
86 pub fn add_u64(&mut self, field: Field, value: u64) {
88 self.add_leaf_field_value(field, value);
89 }
90
91 pub fn add_ip_addr(&mut self, field: Field, value: Ipv6Addr) {
93 self.add_leaf_field_value(field, value);
94 }
95
96 pub fn add_i64(&mut self, field: Field, value: i64) {
98 self.add_leaf_field_value(field, value);
99 }
100
101 pub fn add_f64(&mut self, field: Field, value: f64) {
103 self.add_leaf_field_value(field, value);
104 }
105
106 pub fn add_bool(&mut self, field: Field, value: bool) {
108 self.add_leaf_field_value(field, value);
109 }
110
111 pub fn add_date(&mut self, field: Field, value: DateTime) {
113 self.add_leaf_field_value(field, value);
114 }
115
116 pub fn add_bytes(&mut self, field: Field, value: &[u8]) {
118 self.add_leaf_field_value(field, value);
119 }
120
121 pub fn add_object(&mut self, field: Field, object: BTreeMap<String, OwnedValue>) {
123 self.add_field_value(field, &OwnedValue::from(object));
124 }
125
126 pub fn add_field_value<'a, V: Value<'a>>(&mut self, field: Field, value: V) {
131 let field_value = FieldValueAddr {
132 field: field
133 .field_id()
134 .try_into()
135 .expect("support only up to u16::MAX field ids"),
136 value_addr: self.add_value(value),
137 };
138 self.field_values.push(field_value);
139 }
140
141 pub fn add_leaf_field_value<'a, T: Into<ReferenceValueLeaf<'a>>>(
144 &mut self,
145 field: Field,
146 typed_val: T,
147 ) {
148 let value = typed_val.into();
149 let field_value = FieldValueAddr {
150 field: field
151 .field_id()
152 .try_into()
153 .expect("support only up to u16::MAX field ids"),
154 value_addr: self.add_value_leaf(value),
155 };
156 self.field_values.push(field_value);
157 }
158
159 pub fn field_values(&self) -> impl Iterator<Item = (Field, CompactDocValue<'_>)> {
161 self.field_values.iter().map(|field_val| {
162 let field = Field::from_field_id(field_val.field as u32);
163 let val = self.get_compact_doc_value(field_val.value_addr);
164 (field, val)
165 })
166 }
167
168 pub fn get_all(&self, field: Field) -> impl Iterator<Item = CompactDocValue<'_>> + '_ {
170 self.field_values
171 .iter()
172 .filter(move |field_value| Field::from_field_id(field_value.field as u32) == field)
173 .map(|val| self.get_compact_doc_value(val.value_addr))
174 }
175
176 pub fn get_first(&self, field: Field) -> Option<CompactDocValue<'_>> {
178 self.get_all(field).next()
179 }
180
181 pub fn convert_named_doc(
183 schema: &Schema,
184 named_doc: NamedFieldDocument,
185 ) -> Result<Self, DocParsingError> {
186 let mut document = Self::new();
187 for (field_name, values) in named_doc.0 {
188 if let Ok(field) = schema.get_field(&field_name) {
189 for value in values {
190 document.add_field_value(field, &value);
191 }
192 }
193 }
194 Ok(document)
195 }
196
197 pub fn parse_json(schema: &Schema, doc_json: &str) -> Result<Self, DocParsingError> {
199 let json_obj: Map<String, serde_json::Value> =
200 serde_json::from_str(doc_json).map_err(|_| DocParsingError::invalid_json(doc_json))?;
201 Self::from_json_object(schema, json_obj)
202 }
203
204 pub fn from_json_object(
206 schema: &Schema,
207 json_obj: Map<String, serde_json::Value>,
208 ) -> Result<Self, DocParsingError> {
209 let mut doc = Self::default();
210 for (field_name, json_value) in json_obj {
211 if let Ok(field) = schema.get_field(&field_name) {
212 let field_entry = schema.get_field_entry(field);
213 let field_type = field_entry.field_type();
214 match json_value {
215 serde_json::Value::Array(json_items) => {
216 for json_item in json_items {
217 let value = field_type
218 .value_from_json(json_item)
219 .map_err(|e| DocParsingError::ValueError(field_name.clone(), e))?;
220 doc.add_field_value(field, &value);
221 }
222 }
223 _ => {
224 let value = field_type
225 .value_from_json(json_value)
226 .map_err(|e| DocParsingError::ValueError(field_name.clone(), e))?;
227 doc.add_field_value(field, &value);
228 }
229 }
230 }
231 }
232 Ok(doc)
233 }
234
235 fn add_value_leaf(&mut self, leaf: ReferenceValueLeaf) -> ValueAddr {
236 let type_id = ValueType::from(&leaf);
237 let val_addr = match leaf {
240 ReferenceValueLeaf::Null => 0,
241 ReferenceValueLeaf::Str(bytes) => {
242 write_bytes_into(&mut self.node_data, bytes.as_bytes())
243 }
244 ReferenceValueLeaf::Facet(bytes) => {
245 write_bytes_into(&mut self.node_data, bytes.as_bytes())
246 }
247 ReferenceValueLeaf::Bytes(bytes) => write_bytes_into(&mut self.node_data, bytes),
248 ReferenceValueLeaf::U64(num) => write_into(&mut self.node_data, num),
249 ReferenceValueLeaf::I64(num) => write_into(&mut self.node_data, num),
250 ReferenceValueLeaf::F64(num) => write_into(&mut self.node_data, num),
251 ReferenceValueLeaf::Bool(b) => b as u32,
252 ReferenceValueLeaf::Date(date) => {
253 write_into(&mut self.node_data, date.into_timestamp_nanos())
254 }
255 ReferenceValueLeaf::IpAddr(num) => write_into(&mut self.node_data, num.to_u128()),
256 ReferenceValueLeaf::PreTokStr(pre_tok) => write_into(&mut self.node_data, *pre_tok),
257 };
258 ValueAddr { type_id, val_addr }
259 }
260 fn add_value<'a, V: Value<'a>>(&mut self, value: V) -> ValueAddr {
262 let value = value.as_value();
263 let type_id = ValueType::from(&value);
264 match value {
265 ReferenceValue::Leaf(leaf) => self.add_value_leaf(leaf),
266 ReferenceValue::Array(elements) => {
267 let mut addresses = Vec::new();
271 for elem in elements {
272 let value_addr = self.add_value(elem);
273 write_into(&mut addresses, value_addr);
274 }
275 ValueAddr {
276 type_id,
277 val_addr: write_bytes_into(&mut self.node_data, &addresses),
278 }
279 }
280 ReferenceValue::Object(entries) => {
281 let mut addresses = Vec::new();
283 for (key, value) in entries {
284 let key_addr = self.add_value_leaf(ReferenceValueLeaf::Str(key));
285 let value_addr = self.add_value(value);
286 write_into(&mut addresses, key_addr);
287 write_into(&mut addresses, value_addr);
288 }
289 ValueAddr {
290 type_id,
291 val_addr: write_bytes_into(&mut self.node_data, &addresses),
292 }
293 }
294 }
295 }
296
297 fn get_compact_doc_value(&self, value_addr: ValueAddr) -> CompactDocValue<'_> {
299 CompactDocValue {
300 container: self,
301 value_addr,
302 }
303 }
304
305 fn extract_bytes(&self, addr: Addr) -> &[u8] {
307 binary_deserialize_bytes(self.get_slice(addr))
308 }
309
310 fn extract_str(&self, addr: Addr) -> &str {
312 let data = self.extract_bytes(addr);
313 unsafe { std::str::from_utf8_unchecked(data) }
315 }
316
317 fn read_from<T: BinarySerializable>(&self, addr: Addr) -> io::Result<T> {
319 let data_slice = &self.node_data[addr as usize..];
320 let mut cursor = std::io::Cursor::new(data_slice);
321 T::deserialize(&mut cursor)
322 }
323
324 fn get_slice(&self, addr: Addr) -> &[u8] {
326 &self.node_data[addr as usize..]
327 }
328}
329
330fn binary_deserialize_bytes(data: &[u8]) -> &[u8] {
332 let (len, bytes_read) = read_u32_vint_no_advance(data);
333 &data[bytes_read..bytes_read + len as usize]
334}
335
336fn write_bytes_into(vec: &mut Vec<u8>, data: &[u8]) -> u32 {
340 let pos = vec.len() as u32;
341 let mut buf = [0u8; 8];
342 let len_vint_bytes = serialize_vint_u32(data.len() as u32, &mut buf);
343 vec.extend_from_slice(len_vint_bytes);
344 vec.extend_from_slice(data);
345 pos
346}
347
348fn write_into<T: BinarySerializable>(vec: &mut Vec<u8>, value: T) -> u32 {
350 let pos = vec.len() as u32;
351 value.serialize(vec).unwrap();
352 pos
353}
354
355impl PartialEq for CompactDoc {
356 fn eq(&self, other: &Self) -> bool {
357 let convert_to_comparable_map = |doc: &CompactDoc| {
359 let mut field_value_set: HashMap<Field, HashSet<String>> = Default::default();
360 for field_value in doc.field_values.iter() {
361 let value: OwnedValue = doc.get_compact_doc_value(field_value.value_addr).into();
362 let value = serde_json::to_string(&value).unwrap();
363 field_value_set
364 .entry(Field::from_field_id(field_value.field as u32))
365 .or_default()
366 .insert(value);
367 }
368 field_value_set
369 };
370 let self_field_values: HashMap<Field, HashSet<String>> = convert_to_comparable_map(self);
371 let other_field_values: HashMap<Field, HashSet<String>> = convert_to_comparable_map(other);
372 self_field_values.eq(&other_field_values)
373 }
374}
375
376impl Eq for CompactDoc {}
377
378impl DocumentDeserialize for CompactDoc {
379 fn deserialize<'de, D>(mut deserializer: D) -> Result<Self, DeserializeError>
380 where D: DocumentDeserializer<'de> {
381 let mut doc = CompactDoc::default();
382 while let Some((field, value)) = deserializer.next_field::<OwnedValue>()? {
385 doc.add_field_value(field, &value);
386 }
387 Ok(doc)
388 }
389}
390
391#[derive(Debug, Clone, Copy)]
393pub struct CompactDocValue<'a> {
394 container: &'a CompactDoc,
395 value_addr: ValueAddr,
396}
397impl PartialEq for CompactDocValue<'_> {
398 fn eq(&self, other: &Self) -> bool {
399 let value1: OwnedValue = (*self).into();
400 let value2: OwnedValue = (*other).into();
401 value1 == value2
402 }
403}
404impl From<CompactDocValue<'_>> for OwnedValue {
405 fn from(value: CompactDocValue) -> Self {
406 value.as_value().into()
407 }
408}
409impl<'a> Value<'a> for CompactDocValue<'a> {
410 type ArrayIter = CompactDocArrayIter<'a>;
411
412 type ObjectIter = CompactDocObjectIter<'a>;
413
414 fn as_value(&self) -> ReferenceValue<'a, Self> {
415 self.get_ref_value().unwrap()
416 }
417}
418impl<'a> CompactDocValue<'a> {
419 fn get_ref_value(&self) -> io::Result<ReferenceValue<'a, CompactDocValue<'a>>> {
420 let addr = self.value_addr.val_addr;
421 match self.value_addr.type_id {
422 ValueType::Null => Ok(ReferenceValueLeaf::Null.into()),
423 ValueType::Str => {
424 let str_ref = self.container.extract_str(addr);
425 Ok(ReferenceValueLeaf::Str(str_ref).into())
426 }
427 ValueType::Facet => {
428 let str_ref = self.container.extract_str(addr);
429 Ok(ReferenceValueLeaf::Facet(str_ref).into())
430 }
431 ValueType::Bytes => {
432 let data = self.container.extract_bytes(addr);
433 Ok(ReferenceValueLeaf::Bytes(data).into())
434 }
435 ValueType::U64 => self
436 .container
437 .read_from::<u64>(addr)
438 .map(ReferenceValueLeaf::U64)
439 .map(Into::into),
440 ValueType::I64 => self
441 .container
442 .read_from::<i64>(addr)
443 .map(ReferenceValueLeaf::I64)
444 .map(Into::into),
445 ValueType::F64 => self
446 .container
447 .read_from::<f64>(addr)
448 .map(ReferenceValueLeaf::F64)
449 .map(Into::into),
450 ValueType::Bool => Ok(ReferenceValueLeaf::Bool(addr != 0).into()),
451 ValueType::Date => self
452 .container
453 .read_from::<i64>(addr)
454 .map(|ts| ReferenceValueLeaf::Date(DateTime::from_timestamp_nanos(ts)))
455 .map(Into::into),
456 ValueType::IpAddr => self
457 .container
458 .read_from::<u128>(addr)
459 .map(|num| ReferenceValueLeaf::IpAddr(Ipv6Addr::from_u128(num)))
460 .map(Into::into),
461 ValueType::PreTokStr => self
462 .container
463 .read_from::<PreTokenizedString>(addr)
464 .map(Into::into)
465 .map(ReferenceValueLeaf::PreTokStr)
466 .map(Into::into),
467 ValueType::Object => Ok(ReferenceValue::Object(CompactDocObjectIter::new(
468 self.container,
469 addr,
470 )?)),
471 ValueType::Array => Ok(ReferenceValue::Array(CompactDocArrayIter::new(
472 self.container,
473 addr,
474 )?)),
475 }
476 }
477}
478
479type Addr = u32;
481
482#[derive(Clone, Copy, Default)]
483#[repr(C, packed)]
484struct ValueAddr {
486 type_id: ValueType,
487 val_addr: Addr,
489}
490impl BinarySerializable for ValueAddr {
491 fn serialize<W: Write + ?Sized>(&self, writer: &mut W) -> io::Result<()> {
492 self.type_id.serialize(writer)?;
493 VInt(self.val_addr as u64).serialize(writer)
494 }
495
496 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
497 let type_id = ValueType::deserialize(reader)?;
498 let val_addr = VInt::deserialize(reader)?.0 as u32;
499 Ok(ValueAddr { type_id, val_addr })
500 }
501}
502impl std::fmt::Debug for ValueAddr {
503 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
504 let val_addr = self.val_addr;
505 f.write_fmt(format_args!("{:?} at {:?}", self.type_id, val_addr))
506 }
507}
508
509#[derive(Default, Clone, Copy, Debug, PartialEq)]
516#[repr(u8)]
517pub enum ValueType {
518 #[default]
520 Null = 0,
521 Str = 1,
523 U64 = 2,
525 I64 = 3,
527 F64 = 4,
529 Date = 5,
531 Facet = 6,
533 Bytes = 7,
535 IpAddr = 8,
537 Bool = 9,
539 PreTokStr = 10,
541 Object = 11,
543 Array = 12,
545}
546
547impl BinarySerializable for ValueType {
548 fn serialize<W: Write + ?Sized>(&self, writer: &mut W) -> io::Result<()> {
549 (*self as u8).serialize(writer)?;
550 Ok(())
551 }
552
553 fn deserialize<R: Read>(reader: &mut R) -> io::Result<Self> {
554 let num = u8::deserialize(reader)?;
555 let type_id = if (0..=12).contains(&num) {
556 unsafe { std::mem::transmute::<u8, ValueType>(num) }
557 } else {
558 return Err(io::Error::new(
559 io::ErrorKind::InvalidData,
560 format!("Invalid value type id: {num}"),
561 ));
562 };
563 Ok(type_id)
564 }
565}
566
567impl<'a, V: Value<'a>> From<&ReferenceValue<'a, V>> for ValueType {
568 fn from(value: &ReferenceValue<'a, V>) -> Self {
569 match value {
570 ReferenceValue::Leaf(leaf) => leaf.into(),
571 ReferenceValue::Array(_) => ValueType::Array,
572 ReferenceValue::Object(_) => ValueType::Object,
573 }
574 }
575}
576impl<'a> From<&ReferenceValueLeaf<'a>> for ValueType {
577 fn from(value: &ReferenceValueLeaf<'a>) -> Self {
578 match value {
579 ReferenceValueLeaf::Null => ValueType::Null,
580 ReferenceValueLeaf::Str(_) => ValueType::Str,
581 ReferenceValueLeaf::U64(_) => ValueType::U64,
582 ReferenceValueLeaf::I64(_) => ValueType::I64,
583 ReferenceValueLeaf::F64(_) => ValueType::F64,
584 ReferenceValueLeaf::Bool(_) => ValueType::Bool,
585 ReferenceValueLeaf::Date(_) => ValueType::Date,
586 ReferenceValueLeaf::IpAddr(_) => ValueType::IpAddr,
587 ReferenceValueLeaf::PreTokStr(_) => ValueType::PreTokStr,
588 ReferenceValueLeaf::Facet(_) => ValueType::Facet,
589 ReferenceValueLeaf::Bytes(_) => ValueType::Bytes,
590 }
591 }
592}
593
594#[derive(Debug, Clone)]
595pub struct CompactDocObjectIter<'a> {
597 container: &'a CompactDoc,
598 node_addresses_slice: &'a [u8],
599}
600
601impl<'a> CompactDocObjectIter<'a> {
602 fn new(container: &'a CompactDoc, addr: Addr) -> io::Result<Self> {
603 let node_addresses_slice = container.extract_bytes(addr);
605 Ok(Self {
606 container,
607 node_addresses_slice,
608 })
609 }
610}
611
612impl<'a> Iterator for CompactDocObjectIter<'a> {
613 type Item = (&'a str, CompactDocValue<'a>);
614
615 fn next(&mut self) -> Option<Self::Item> {
616 if self.node_addresses_slice.is_empty() {
617 return None;
618 }
619 let key_addr = ValueAddr::deserialize(&mut self.node_addresses_slice).ok()?;
620 let key = self.container.extract_str(key_addr.val_addr);
621 let value = ValueAddr::deserialize(&mut self.node_addresses_slice).ok()?;
622 let value = CompactDocValue {
623 container: self.container,
624 value_addr: value,
625 };
626 Some((key, value))
627 }
628}
629
630#[derive(Debug, Clone)]
631pub struct CompactDocArrayIter<'a> {
633 container: &'a CompactDoc,
634 node_addresses_slice: &'a [u8],
635}
636
637impl<'a> CompactDocArrayIter<'a> {
638 fn new(container: &'a CompactDoc, addr: Addr) -> io::Result<Self> {
639 let node_addresses_slice = container.extract_bytes(addr);
641 Ok(Self {
642 container,
643 node_addresses_slice,
644 })
645 }
646}
647
648impl<'a> Iterator for CompactDocArrayIter<'a> {
649 type Item = CompactDocValue<'a>;
650
651 fn next(&mut self) -> Option<Self::Item> {
652 if self.node_addresses_slice.is_empty() {
653 return None;
654 }
655 let value = ValueAddr::deserialize(&mut self.node_addresses_slice).ok()?;
656 let value = CompactDocValue {
657 container: self.container,
658 value_addr: value,
659 };
660 Some(value)
661 }
662}
663
664impl Document for CompactDoc {
665 type Value<'a> = CompactDocValue<'a>;
666 type FieldsValuesIter<'a> = FieldValueIterRef<'a>;
667
668 fn iter_fields_and_values(&self) -> Self::FieldsValuesIter<'_> {
669 FieldValueIterRef {
670 slice: self.field_values.iter(),
671 container: self,
672 }
673 }
674}
675
676pub struct FieldValueIterRef<'a> {
678 slice: std::slice::Iter<'a, FieldValueAddr>,
679 container: &'a CompactDoc,
680}
681
682impl<'a> Iterator for FieldValueIterRef<'a> {
683 type Item = (Field, CompactDocValue<'a>);
684
685 fn next(&mut self) -> Option<Self::Item> {
686 self.slice.next().map(|field_value| {
687 (
688 Field::from_field_id(field_value.field as u32),
689 CompactDocValue::<'a> {
690 container: self.container,
691 value_addr: field_value.value_addr,
692 },
693 )
694 })
695 }
696}
697
698#[derive(Debug, Error, PartialEq)]
701pub enum DocParsingError {
702 #[error("The provided string is not valid JSON")]
704 InvalidJson(String),
705 #[error("The field '{0:?}' could not be parsed: {1:?}")]
707 ValueError(String, ValueParsingError),
708}
709
710impl DocParsingError {
711 fn invalid_json(invalid_json: &str) -> Self {
713 let sample = invalid_json.chars().take(20).collect();
714 DocParsingError::InvalidJson(sample)
715 }
716}
717
718#[cfg(test)]
719mod tests {
720 use crate::schema::*;
721
722 #[test]
723 fn test_doc() {
724 let mut schema_builder = Schema::builder();
725 let text_field = schema_builder.add_text_field("title", TEXT);
726 let mut doc = TantivyDocument::default();
727 doc.add_text(text_field, "My title");
728 assert_eq!(doc.field_values().count(), 1);
729
730 let schema = schema_builder.build();
731 let _val = doc.get_first(text_field).unwrap();
732 let _json = doc.to_named_doc(&schema);
733 }
734
735 #[test]
736 fn test_json_value() {
737 let json_str = r#"{
738 "toto": "titi",
739 "float": -0.2,
740 "bool": true,
741 "unsigned": 1,
742 "signed": -2,
743 "complexobject": {
744 "field.with.dot": 1
745 },
746 "date": "1985-04-12T23:20:50.52Z",
747 "my_arr": [2, 3, {"my_key": "two tokens"}, 4, {"nested_array": [2, 5, 6, [7, 8, {"a": [{"d": {"e":[99]}}, 9000]}, 9, 10], [5, 5]]}]
748 }"#;
749 let json_val: std::collections::BTreeMap<String, OwnedValue> =
750 serde_json::from_str(json_str).unwrap();
751
752 let mut schema_builder = Schema::builder();
753 let json_field = schema_builder.add_json_field("json", TEXT);
754 let mut doc = TantivyDocument::default();
755 doc.add_object(json_field, json_val);
756
757 let schema = schema_builder.build();
758 let json = doc.to_json(&schema);
759 let actual_json: serde_json::Value = serde_json::from_str(&json).unwrap();
760 let expected_json: serde_json::Value = serde_json::from_str(json_str).unwrap();
761 assert_eq!(actual_json["json"][0], expected_json);
762 }
763
764 }