1#[cfg(feature = "native")]
12use std::fs::{File, OpenOptions};
13#[cfg(feature = "native")]
14use std::io::{BufWriter, Write};
15#[cfg(feature = "native")]
16use std::path::PathBuf;
17#[cfg(feature = "native")]
18use std::sync::Arc;
19
20#[cfg(feature = "native")]
21use hashbrown::HashMap;
22#[cfg(feature = "native")]
23use lasso::{Rodeo, Spur};
24#[cfg(feature = "native")]
25use rustc_hash::FxHashMap;
26
27#[cfg(feature = "native")]
28use super::store::StoreWriter;
29#[cfg(feature = "native")]
30use super::types::{FieldStats, SegmentFiles, SegmentId, SegmentMeta};
31#[cfg(feature = "native")]
32use crate::compression::CompressionLevel;
33#[cfg(feature = "native")]
34use crate::directories::{Directory, DirectoryWriter};
35#[cfg(feature = "native")]
36use crate::dsl::{Document, Field, FieldType, FieldValue, Schema};
37#[cfg(feature = "native")]
38use crate::structures::{PostingList, SSTableWriter, TermInfo};
39#[cfg(feature = "native")]
40use crate::tokenizer::{BoxedTokenizer, LowercaseTokenizer};
41#[cfg(feature = "native")]
42use crate::wand::WandData;
43#[cfg(feature = "native")]
44use crate::{DocId, Result};
45
46#[cfg(feature = "native")]
48const POSTING_FLUSH_THRESHOLD: usize = 100_000;
49
50#[cfg(feature = "native")]
52const SPILL_BUFFER_SIZE: usize = 16 * 1024 * 1024; #[cfg(feature = "native")]
57const NUM_INDEX_SHARDS: usize = 64;
58
59#[cfg(feature = "native")]
61#[derive(Clone, Copy, PartialEq, Eq, Hash)]
62struct TermKey {
63 field: u32,
64 term: Spur,
65}
66
67#[cfg(feature = "native")]
68impl TermKey {
69 #[inline]
71 fn shard(&self) -> usize {
72 (self.term.into_inner().get() as usize) & (NUM_INDEX_SHARDS - 1)
75 }
76}
77
78#[cfg(feature = "native")]
81struct ShardedInvertedIndex {
82 shards: Vec<HashMap<TermKey, SpillablePostingList>>,
83}
84
85#[cfg(feature = "native")]
86impl ShardedInvertedIndex {
87 fn new(capacity_per_shard: usize) -> Self {
88 let mut shards = Vec::with_capacity(NUM_INDEX_SHARDS);
89 for _ in 0..NUM_INDEX_SHARDS {
90 shards.push(HashMap::with_capacity(capacity_per_shard));
91 }
92 Self { shards }
93 }
94
95 #[inline]
96 fn get_mut(&mut self, key: &TermKey) -> Option<&mut SpillablePostingList> {
97 self.shards[key.shard()].get_mut(key)
98 }
99
100 #[inline]
102 fn get_or_insert(&mut self, key: TermKey) -> &mut SpillablePostingList {
103 self.shards[key.shard()]
104 .entry(key)
105 .or_insert_with(SpillablePostingList::new)
106 }
107
108 fn len(&self) -> usize {
109 self.shards.iter().map(|s| s.len()).sum()
110 }
111
112 fn total_postings_in_memory(&self) -> usize {
114 self.shards
115 .iter()
116 .flat_map(|s| s.values())
117 .map(|p| p.memory.len())
118 .sum()
119 }
120
121 fn shard_stats(&self) -> (usize, usize, usize) {
123 let sizes: Vec<usize> = self.shards.iter().map(|s| s.len()).collect();
124 let min = *sizes.iter().min().unwrap_or(&0);
125 let max = *sizes.iter().max().unwrap_or(&0);
126 let avg = if sizes.is_empty() {
127 0
128 } else {
129 sizes.iter().sum::<usize>() / sizes.len()
130 };
131 (min, max, avg)
132 }
133}
134
135#[cfg(feature = "native")]
137struct ShardedIndexIter<'a> {
138 shards: std::slice::Iter<'a, HashMap<TermKey, SpillablePostingList>>,
139 current: Option<hashbrown::hash_map::Iter<'a, TermKey, SpillablePostingList>>,
140}
141
142#[cfg(feature = "native")]
143impl<'a> Iterator for ShardedIndexIter<'a> {
144 type Item = (&'a TermKey, &'a SpillablePostingList);
145
146 fn next(&mut self) -> Option<Self::Item> {
147 loop {
148 if let Some(ref mut current) = self.current
149 && let Some(item) = current.next()
150 {
151 return Some(item);
152 }
153 match self.shards.next() {
155 Some(shard) => self.current = Some(shard.iter()),
156 None => return None,
157 }
158 }
159 }
160}
161
162#[cfg(feature = "native")]
163impl<'a> IntoIterator for &'a ShardedInvertedIndex {
164 type Item = (&'a TermKey, &'a SpillablePostingList);
165 type IntoIter = ShardedIndexIter<'a>;
166
167 fn into_iter(self) -> Self::IntoIter {
168 ShardedIndexIter {
169 shards: self.shards.iter(),
170 current: None,
171 }
172 }
173}
174
175#[cfg(feature = "native")]
177#[derive(Clone, Copy)]
178struct CompactPosting {
179 doc_id: DocId,
180 term_freq: u16, }
182
183#[cfg(feature = "native")]
185struct SpillablePostingList {
186 memory: Vec<CompactPosting>,
188 spill_offset: i64,
190 spill_count: u32,
192}
193
194#[cfg(feature = "native")]
195impl SpillablePostingList {
196 fn new() -> Self {
197 Self {
198 memory: Vec::new(),
199 spill_offset: -1,
200 spill_count: 0,
201 }
202 }
203
204 #[allow(dead_code)]
205 fn with_capacity(capacity: usize) -> Self {
206 Self {
207 memory: Vec::with_capacity(capacity),
208 spill_offset: -1,
209 spill_count: 0,
210 }
211 }
212
213 #[inline]
214 fn add(&mut self, doc_id: DocId, term_freq: u32) {
215 if let Some(last) = self.memory.last_mut()
217 && last.doc_id == doc_id
218 {
219 last.term_freq = last.term_freq.saturating_add(term_freq as u16);
220 return;
221 }
222 self.memory.push(CompactPosting {
223 doc_id,
224 term_freq: term_freq.min(u16::MAX as u32) as u16,
225 });
226 }
227
228 fn total_count(&self) -> usize {
229 self.memory.len() + self.spill_count as usize
230 }
231
232 fn needs_spill(&self) -> bool {
233 self.memory.len() >= POSTING_FLUSH_THRESHOLD
234 }
235}
236
237#[cfg(feature = "native")]
239#[derive(Debug, Clone)]
240pub struct SegmentBuilderStats {
241 pub num_docs: u32,
243 pub unique_terms: usize,
245 pub postings_in_memory: usize,
247 pub interned_strings: usize,
249 pub doc_field_lengths_size: usize,
251 pub shard_min: usize,
253 pub shard_max: usize,
254 pub shard_avg: usize,
255 pub spill_bytes: u64,
257}
258
259#[cfg(feature = "native")]
261#[derive(Clone)]
262pub struct SegmentBuilderConfig {
263 pub temp_dir: PathBuf,
265 pub compression_level: CompressionLevel,
267 pub num_compression_threads: usize,
269 pub interner_capacity: usize,
271 pub posting_map_capacity: usize,
273}
274
275#[cfg(feature = "native")]
276impl Default for SegmentBuilderConfig {
277 fn default() -> Self {
278 Self {
279 temp_dir: std::env::temp_dir(),
280 compression_level: CompressionLevel(7),
281 num_compression_threads: num_cpus::get(),
282 interner_capacity: 1_000_000,
283 posting_map_capacity: 500_000,
284 }
285 }
286}
287
288#[cfg(feature = "native")]
296pub struct SegmentBuilder {
297 schema: Schema,
298 config: SegmentBuilderConfig,
299 tokenizers: FxHashMap<Field, BoxedTokenizer>,
300
301 term_interner: Rodeo,
303
304 inverted_index: ShardedInvertedIndex,
307
308 store_file: BufWriter<File>,
310 store_path: PathBuf,
311
312 spill_file: Option<BufWriter<File>>,
314 spill_path: PathBuf,
315 spill_offset: u64,
316
317 next_doc_id: DocId,
319
320 field_stats: FxHashMap<u32, FieldStats>,
322
323 doc_field_lengths: Vec<u32>,
327 num_indexed_fields: usize,
328 field_to_slot: FxHashMap<u32, usize>,
329
330 wand_data: Option<Arc<WandData>>,
332
333 local_tf_buffer: FxHashMap<Spur, u32>,
336
337 terms_to_spill_buffer: Vec<TermKey>,
339}
340
341#[cfg(feature = "native")]
342impl SegmentBuilder {
343 pub fn new(schema: Schema, config: SegmentBuilderConfig) -> Result<Self> {
345 let segment_id = uuid::Uuid::new_v4();
346 let store_path = config
347 .temp_dir
348 .join(format!("hermes_store_{}.tmp", segment_id));
349 let spill_path = config
350 .temp_dir
351 .join(format!("hermes_spill_{}.tmp", segment_id));
352
353 let store_file = BufWriter::with_capacity(
354 SPILL_BUFFER_SIZE,
355 OpenOptions::new()
356 .create(true)
357 .write(true)
358 .truncate(true)
359 .open(&store_path)?,
360 );
361
362 let mut num_indexed_fields = 0;
364 let mut field_to_slot = FxHashMap::default();
365 for (field, entry) in schema.fields() {
366 if entry.indexed && matches!(entry.field_type, FieldType::Text) {
367 field_to_slot.insert(field.0, num_indexed_fields);
368 num_indexed_fields += 1;
369 }
370 }
371
372 Ok(Self {
373 schema,
374 tokenizers: FxHashMap::default(),
375 term_interner: Rodeo::new(),
376 inverted_index: ShardedInvertedIndex::new(
377 config.posting_map_capacity / NUM_INDEX_SHARDS,
378 ),
379 store_file,
380 store_path,
381 spill_file: None,
382 spill_path,
383 spill_offset: 0,
384 next_doc_id: 0,
385 field_stats: FxHashMap::default(),
386 doc_field_lengths: Vec::new(),
387 num_indexed_fields,
388 field_to_slot,
389 wand_data: None,
390 local_tf_buffer: FxHashMap::default(),
391 terms_to_spill_buffer: Vec::new(),
392 config,
393 })
394 }
395
396 pub fn with_wand_data(
398 schema: Schema,
399 config: SegmentBuilderConfig,
400 wand_data: Arc<WandData>,
401 ) -> Result<Self> {
402 let mut builder = Self::new(schema, config)?;
403 builder.wand_data = Some(wand_data);
404 Ok(builder)
405 }
406
407 pub fn set_tokenizer(&mut self, field: Field, tokenizer: BoxedTokenizer) {
408 self.tokenizers.insert(field, tokenizer);
409 }
410
411 pub fn num_docs(&self) -> u32 {
412 self.next_doc_id
413 }
414
415 pub fn stats(&self) -> SegmentBuilderStats {
417 let (shard_min, shard_max, shard_avg) = self.inverted_index.shard_stats();
418 SegmentBuilderStats {
419 num_docs: self.next_doc_id,
420 unique_terms: self.inverted_index.len(),
421 postings_in_memory: self.inverted_index.total_postings_in_memory(),
422 interned_strings: self.term_interner.len(),
423 doc_field_lengths_size: self.doc_field_lengths.len(),
424 shard_min,
425 shard_max,
426 shard_avg,
427 spill_bytes: self.spill_offset,
428 }
429 }
430
431 pub fn add_document(&mut self, doc: Document) -> Result<DocId> {
433 let doc_id = self.next_doc_id;
434 self.next_doc_id += 1;
435
436 let base_idx = self.doc_field_lengths.len();
438 self.doc_field_lengths
439 .resize(base_idx + self.num_indexed_fields, 0);
440
441 for (field, value) in doc.field_values() {
442 let entry = self.schema.get_field_entry(*field);
443 if entry.is_none() || !entry.unwrap().indexed {
444 continue;
445 }
446
447 let entry = entry.unwrap();
448 match (&entry.field_type, value) {
449 (FieldType::Text, FieldValue::Text(text)) => {
450 let token_count = self.index_text_field(*field, doc_id, text)?;
451
452 let stats = self.field_stats.entry(field.0).or_default();
454 stats.total_tokens += token_count as u64;
455 stats.doc_count += 1;
456
457 if let Some(&slot) = self.field_to_slot.get(&field.0) {
459 self.doc_field_lengths[base_idx + slot] = token_count;
460 }
461 }
462 (FieldType::U64, FieldValue::U64(v)) => {
463 self.index_numeric_field(*field, doc_id, *v)?;
464 }
465 (FieldType::I64, FieldValue::I64(v)) => {
466 self.index_numeric_field(*field, doc_id, *v as u64)?;
467 }
468 (FieldType::F64, FieldValue::F64(v)) => {
469 self.index_numeric_field(*field, doc_id, v.to_bits())?;
470 }
471 _ => {}
472 }
473 }
474
475 self.write_document_to_store(&doc)?;
477
478 Ok(doc_id)
479 }
480
481 fn index_text_field(&mut self, field: Field, doc_id: DocId, text: &str) -> Result<u32> {
487 let default_tokenizer = LowercaseTokenizer;
488 let tokenizer: &dyn crate::tokenizer::TokenizerClone = self
489 .tokenizers
490 .get(&field)
491 .map(|t| t.as_ref())
492 .unwrap_or(&default_tokenizer);
493
494 let tokens = tokenizer.tokenize(text);
495 let token_count = tokens.len() as u32;
496
497 self.local_tf_buffer.clear();
500
501 for token in tokens {
502 let term_spur = self.term_interner.get_or_intern(&token.text);
504 *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
505 }
506
507 let field_id = field.0;
511 self.terms_to_spill_buffer.clear();
512
513 for (&term_spur, &tf) in &self.local_tf_buffer {
514 let term_key = TermKey {
515 field: field_id,
516 term: term_spur,
517 };
518
519 let posting = self.inverted_index.get_or_insert(term_key);
520 posting.add(doc_id, tf);
521
522 if posting.needs_spill() {
524 self.terms_to_spill_buffer.push(term_key);
525 }
526 }
527
528 for i in 0..self.terms_to_spill_buffer.len() {
530 let term_key = self.terms_to_spill_buffer[i];
531 self.spill_posting_list(term_key)?;
532 }
533
534 Ok(token_count)
535 }
536
537 fn index_numeric_field(&mut self, field: Field, doc_id: DocId, value: u64) -> Result<()> {
538 let term_str = format!("__num_{}", value);
540 let term_spur = self.term_interner.get_or_intern(&term_str);
541
542 let term_key = TermKey {
543 field: field.0,
544 term: term_spur,
545 };
546
547 let posting = self.inverted_index.get_or_insert(term_key);
548 posting.add(doc_id, 1);
549
550 Ok(())
551 }
552
553 fn write_document_to_store(&mut self, doc: &Document) -> Result<()> {
555 use byteorder::{LittleEndian, WriteBytesExt};
556
557 let doc_bytes = super::store::serialize_document(doc, &self.schema)?;
558
559 self.store_file
560 .write_u32::<LittleEndian>(doc_bytes.len() as u32)?;
561 self.store_file.write_all(&doc_bytes)?;
562
563 Ok(())
564 }
565
566 fn spill_posting_list(&mut self, term_key: TermKey) -> Result<()> {
568 use byteorder::{LittleEndian, WriteBytesExt};
569
570 let posting = self.inverted_index.get_mut(&term_key).unwrap();
571
572 if self.spill_file.is_none() {
574 let file = OpenOptions::new()
575 .create(true)
576 .write(true)
577 .read(true)
578 .truncate(true)
579 .open(&self.spill_path)?;
580 self.spill_file = Some(BufWriter::with_capacity(SPILL_BUFFER_SIZE, file));
581 }
582
583 let spill_file = self.spill_file.as_mut().unwrap();
584
585 if posting.spill_offset < 0 {
587 posting.spill_offset = self.spill_offset as i64;
588 }
589
590 for p in &posting.memory {
592 spill_file.write_u32::<LittleEndian>(p.doc_id)?;
593 spill_file.write_u16::<LittleEndian>(p.term_freq)?;
594 self.spill_offset += 6; }
596
597 posting.spill_count += posting.memory.len() as u32;
598 posting.memory.clear();
599 posting.memory.shrink_to(POSTING_FLUSH_THRESHOLD / 4); Ok(())
602 }
603
604 pub async fn build<D: Directory + DirectoryWriter>(
606 mut self,
607 dir: &D,
608 segment_id: SegmentId,
609 ) -> Result<SegmentMeta> {
610 self.store_file.flush()?;
612 if let Some(ref mut spill) = self.spill_file {
613 spill.flush()?;
614 }
615
616 let files = SegmentFiles::new(segment_id.0);
617
618 let (term_dict_data, postings_data) = self.build_postings()?;
620
621 let store_data = self.build_store_from_stream()?;
623
624 dir.write(&files.term_dict, &term_dict_data).await?;
626 dir.write(&files.postings, &postings_data).await?;
627 dir.write(&files.store, &store_data).await?;
628
629 let meta = SegmentMeta {
630 id: segment_id.0,
631 num_docs: self.next_doc_id,
632 field_stats: self.field_stats.clone(),
633 };
634
635 dir.write(&files.meta, &meta.serialize()?).await?;
636
637 let _ = std::fs::remove_file(&self.store_path);
639 let _ = std::fs::remove_file(&self.spill_path);
640
641 Ok(meta)
642 }
643
644 fn build_postings(&mut self) -> Result<(Vec<u8>, Vec<u8>)> {
646 use std::collections::BTreeMap;
647
648 let mut sorted_terms: BTreeMap<Vec<u8>, &SpillablePostingList> = BTreeMap::new();
651
652 for (term_key, posting_list) in &self.inverted_index {
653 let term_str = self.term_interner.resolve(&term_key.term);
654 let mut key = Vec::with_capacity(4 + term_str.len());
655 key.extend_from_slice(&term_key.field.to_le_bytes());
656 key.extend_from_slice(term_str.as_bytes());
657 sorted_terms.insert(key, posting_list);
658 }
659
660 let mut term_dict = Vec::new();
661 let mut postings = Vec::new();
662 let mut writer = SSTableWriter::<TermInfo>::new(&mut term_dict);
663
664 let spill_mmap = if self.spill_file.is_some() {
666 drop(self.spill_file.take()); let file = File::open(&self.spill_path)?;
668 Some(unsafe { memmap2::Mmap::map(&file)? })
669 } else {
670 None
671 };
672
673 for (key, spill_posting) in sorted_terms {
674 let mut full_postings = PostingList::with_capacity(spill_posting.total_count());
676
677 if spill_posting.spill_offset >= 0
679 && let Some(ref mmap) = spill_mmap
680 {
681 let mut offset = spill_posting.spill_offset as usize;
682 for _ in 0..spill_posting.spill_count {
683 let doc_id = u32::from_le_bytes([
684 mmap[offset],
685 mmap[offset + 1],
686 mmap[offset + 2],
687 mmap[offset + 3],
688 ]);
689 let term_freq = u16::from_le_bytes([mmap[offset + 4], mmap[offset + 5]]);
690 full_postings.push(doc_id, term_freq as u32);
691 offset += 6;
692 }
693 }
694
695 for p in &spill_posting.memory {
697 full_postings.push(p.doc_id, p.term_freq as u32);
698 }
699
700 let doc_ids: Vec<u32> = full_postings.iter().map(|p| p.doc_id).collect();
702 let term_freqs: Vec<u32> = full_postings.iter().map(|p| p.term_freq).collect();
703
704 let term_info = if let Some(inline) = TermInfo::try_inline(&doc_ids, &term_freqs) {
705 inline
706 } else {
707 let posting_offset = postings.len() as u64;
708 let block_list =
709 crate::structures::BlockPostingList::from_posting_list(&full_postings)?;
710 block_list.serialize(&mut postings)?;
711 TermInfo::external(
712 posting_offset,
713 (postings.len() as u64 - posting_offset) as u32,
714 full_postings.doc_count(),
715 )
716 };
717
718 writer.insert(&key, &term_info)?;
719 }
720
721 writer.finish()?;
722 Ok((term_dict, postings))
723 }
724
725 fn build_store_from_stream(&mut self) -> Result<Vec<u8>> {
727 drop(std::mem::replace(
729 &mut self.store_file,
730 BufWriter::new(File::create("/dev/null")?),
731 ));
732
733 let file = File::open(&self.store_path)?;
734 let mmap = unsafe { memmap2::Mmap::map(&file)? };
735
736 let mut store_data = Vec::new();
738 let mut store_writer =
739 StoreWriter::with_compression_level(&mut store_data, self.config.compression_level);
740
741 let mut offset = 0usize;
742 while offset < mmap.len() {
743 if offset + 4 > mmap.len() {
744 break;
745 }
746
747 let doc_len = u32::from_le_bytes([
748 mmap[offset],
749 mmap[offset + 1],
750 mmap[offset + 2],
751 mmap[offset + 3],
752 ]) as usize;
753 offset += 4;
754
755 if offset + doc_len > mmap.len() {
756 break;
757 }
758
759 let doc_bytes = &mmap[offset..offset + doc_len];
760 offset += doc_len;
761
762 if let Ok(doc) = super::store::deserialize_document(doc_bytes, &self.schema) {
764 store_writer.store(&doc, &self.schema)?;
765 }
766 }
767
768 store_writer.finish()?;
769 Ok(store_data)
770 }
771}
772
773#[cfg(feature = "native")]
774impl Drop for SegmentBuilder {
775 fn drop(&mut self) {
776 let _ = std::fs::remove_file(&self.store_path);
778 let _ = std::fs::remove_file(&self.spill_path);
779 }
780}