1#[cfg_attr(not(feature = "native"), allow(dead_code))]
12pub(crate) mod bmp;
13mod config;
14mod dense;
15#[cfg(feature = "diagnostics")]
16mod diagnostics;
17#[cfg_attr(not(feature = "native"), allow(dead_code))]
18pub(crate) mod graph_bisection;
19pub use graph_bisection::BpBudget;
20mod postings;
21mod sparse;
22mod store;
23
24pub use config::{MemoryBreakdown, SegmentBuilderConfig, SegmentBuilderStats};
25
26#[cfg(feature = "native")]
27use std::fs::{File, OpenOptions};
28#[cfg(feature = "native")]
29use std::io::BufWriter;
30use std::io::Write;
31use std::mem::size_of;
32#[cfg(feature = "native")]
33use std::path::PathBuf;
34
35use hashbrown::HashMap;
36use rustc_hash::FxHashMap;
37
38#[cfg(feature = "native")]
40use lasso::{Rodeo, Spur};
41
42#[cfg(not(feature = "native"))]
43pub(crate) mod simple_interner {
44 use hashbrown::HashMap;
45
46 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
47 pub struct Spur(u32);
48
49 pub struct Rodeo {
52 strings: Vec<Box<str>>,
54 map: HashMap<&'static str, u32>,
57 }
58
59 impl Rodeo {
60 pub fn new() -> Self {
61 Self {
62 strings: Vec::new(),
63 map: HashMap::new(),
64 }
65 }
66
67 pub fn get(&self, key: &str) -> Option<Spur> {
68 self.map.get(key).map(|&id| Spur(id))
69 }
70
71 pub fn get_or_intern(&mut self, key: &str) -> Spur {
72 if let Some(&id) = self.map.get(key) {
73 return Spur(id);
74 }
75 let id = self.strings.len() as u32;
76 let boxed: Box<str> = key.into();
77 let static_ref: &'static str = unsafe { &*(boxed.as_ref() as *const str) };
80 self.strings.push(boxed);
81 self.map.insert(static_ref, id);
82 Spur(id)
83 }
84
85 pub fn resolve(&self, spur: &Spur) -> &str {
86 &self.strings[spur.0 as usize]
87 }
88
89 pub fn len(&self) -> usize {
90 self.strings.len()
91 }
92 }
93}
94
95#[cfg(not(feature = "native"))]
96use simple_interner::{Rodeo, Spur};
97
98use super::types::{FieldStats, SegmentFiles, SegmentId, SegmentMeta};
99use std::sync::Arc;
100
101use crate::directories::{Directory, DirectoryWriter};
102use crate::dsl::{Document, Field, FieldType, FieldValue, Schema};
103use crate::tokenizer::BoxedTokenizer;
104use crate::{DocId, Result};
105
106use dense::{BinaryDenseVectorBuilder, DenseVectorBuilder};
107use postings::{CompactPosting, PositionPostingListBuilder, PostingListBuilder, TermKey};
108use sparse::SparseVectorBuilder;
109
110const STORE_BUFFER_SIZE: usize = 16 * 1024 * 1024; const NEW_TERM_OVERHEAD: usize = size_of::<TermKey>() + size_of::<PostingListBuilder>() + 24;
116
117const INTERN_OVERHEAD: usize = size_of::<Spur>() + 2 * size_of::<usize>();
119
120const NEW_POS_TERM_OVERHEAD: usize =
122 size_of::<TermKey>() + size_of::<PositionPostingListBuilder>() + 24;
123
124pub struct SegmentBuilder {
131 schema: Arc<Schema>,
132 config: SegmentBuilderConfig,
133 tokenizers: FxHashMap<Field, BoxedTokenizer>,
134
135 term_interner: Rodeo,
137
138 inverted_index: HashMap<TermKey, PostingListBuilder>,
140
141 #[cfg(feature = "native")]
143 posting_spill_file: Option<BufWriter<File>>,
144 #[cfg(feature = "native")]
145 posting_spill_path: PathBuf,
146 #[cfg(feature = "native")]
148 posting_spill_index: HashMap<TermKey, Vec<(u64, u32)>>,
149 #[cfg(feature = "native")]
150 posting_spill_offset: u64,
151
152 #[cfg(feature = "native")]
154 store_file: BufWriter<File>,
155 #[cfg(feature = "native")]
156 store_path: PathBuf,
157 #[cfg(not(feature = "native"))]
158 store_buffer: Vec<u8>,
159
160 next_doc_id: DocId,
162
163 field_stats: FxHashMap<u32, FieldStats>,
165
166 doc_field_lengths: Vec<u32>,
170 num_indexed_fields: usize,
171 field_to_slot: FxHashMap<u32, usize>,
172
173 local_tf_buffer: FxHashMap<Spur, u32>,
176
177 local_positions: FxHashMap<Spur, Vec<u32>>,
180
181 token_buffer: String,
183
184 numeric_buffer: String,
186
187 dense_vectors: FxHashMap<u32, DenseVectorBuilder>,
190
191 binary_dense_vectors: FxHashMap<u32, BinaryDenseVectorBuilder>,
193
194 sparse_vectors: FxHashMap<u32, SparseVectorBuilder>,
197
198 position_index: HashMap<TermKey, PositionPostingListBuilder>,
201
202 position_enabled_fields: FxHashMap<u32, Option<crate::dsl::PositionMode>>,
204
205 current_element_ordinal: FxHashMap<u32, u32>,
207
208 estimated_memory: usize,
210
211 doc_serialize_buffer: Vec<u8>,
213
214 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
216}
217
218impl SegmentBuilder {
219 pub fn new(schema: Arc<Schema>, config: SegmentBuilderConfig) -> Result<Self> {
221 #[cfg(feature = "native")]
222 let (store_file, store_path, spill_path) = {
223 let segment_id = uuid::Uuid::new_v4();
224 let store_path = config
225 .temp_dir
226 .join(format!("hermes_store_{}.tmp", segment_id));
227 let store_file = BufWriter::with_capacity(
228 STORE_BUFFER_SIZE,
229 OpenOptions::new()
230 .create(true)
231 .write(true)
232 .truncate(true)
233 .open(&store_path)?,
234 );
235 let spill_path = config
236 .temp_dir
237 .join(format!("hermes_spill_{}.tmp", segment_id));
238 (store_file, store_path, spill_path)
239 };
240
241 let registry = crate::tokenizer::TokenizerRegistry::new();
243 let mut num_indexed_fields = 0;
244 let mut field_to_slot = FxHashMap::default();
245 let mut position_enabled_fields = FxHashMap::default();
246 let mut tokenizers = FxHashMap::default();
247 for (field, entry) in schema.fields() {
248 if entry.indexed && matches!(entry.field_type, FieldType::Text) {
249 field_to_slot.insert(field.0, num_indexed_fields);
250 num_indexed_fields += 1;
251 if entry.positions.is_some() {
252 position_enabled_fields.insert(field.0, entry.positions);
253 }
254 if let Some(ref tok_name) = entry.tokenizer
255 && let Some(tokenizer) = registry.get(tok_name)
256 {
257 tokenizers.insert(field, tokenizer);
258 }
259 }
260 }
261
262 use crate::structures::fast_field::{FastFieldColumnType, FastFieldWriter};
264 let mut fast_fields = FxHashMap::default();
265 for (field, entry) in schema.fields() {
266 if entry.fast {
267 let writer = if entry.multi {
268 match entry.field_type {
269 FieldType::U64 => {
270 FastFieldWriter::new_numeric_multi(FastFieldColumnType::U64)
271 }
272 FieldType::I64 => {
273 FastFieldWriter::new_numeric_multi(FastFieldColumnType::I64)
274 }
275 FieldType::F64 => {
276 FastFieldWriter::new_numeric_multi(FastFieldColumnType::F64)
277 }
278 FieldType::Text => FastFieldWriter::new_text_multi(),
279 _ => continue,
280 }
281 } else {
282 match entry.field_type {
283 FieldType::U64 => FastFieldWriter::new_numeric(FastFieldColumnType::U64),
284 FieldType::I64 => FastFieldWriter::new_numeric(FastFieldColumnType::I64),
285 FieldType::F64 => FastFieldWriter::new_numeric(FastFieldColumnType::F64),
286 FieldType::Text => FastFieldWriter::new_text(),
287 _ => continue,
288 }
289 };
290 fast_fields.insert(field.0, writer);
291 }
292 }
293
294 Ok(Self {
295 schema,
296 tokenizers,
297 term_interner: Rodeo::new(),
298 inverted_index: HashMap::with_capacity(config.posting_map_capacity),
299 #[cfg(feature = "native")]
300 posting_spill_file: None,
301 #[cfg(feature = "native")]
302 posting_spill_path: spill_path,
303 #[cfg(feature = "native")]
304 posting_spill_index: HashMap::new(),
305 #[cfg(feature = "native")]
306 posting_spill_offset: 0,
307 #[cfg(feature = "native")]
308 store_file,
309 #[cfg(feature = "native")]
310 store_path,
311 #[cfg(not(feature = "native"))]
312 store_buffer: Vec::with_capacity(STORE_BUFFER_SIZE),
313 next_doc_id: 0,
314 field_stats: FxHashMap::default(),
315 doc_field_lengths: Vec::new(),
316 num_indexed_fields,
317 field_to_slot,
318 local_tf_buffer: FxHashMap::default(),
319 local_positions: FxHashMap::default(),
320 token_buffer: String::with_capacity(64),
321 numeric_buffer: String::with_capacity(32),
322 config,
323 dense_vectors: FxHashMap::default(),
324 binary_dense_vectors: FxHashMap::default(),
325 sparse_vectors: FxHashMap::default(),
326 position_index: HashMap::new(),
327 position_enabled_fields,
328 current_element_ordinal: FxHashMap::default(),
329 estimated_memory: 0,
330 doc_serialize_buffer: Vec::with_capacity(256),
331 fast_fields,
332 })
333 }
334
335 pub fn set_tokenizer(&mut self, field: Field, tokenizer: BoxedTokenizer) {
336 self.tokenizers.insert(field, tokenizer);
337 }
338
339 fn next_element_ordinal(&mut self, field_id: u32) -> u32 {
342 let ordinal = *self.current_element_ordinal.get(&field_id).unwrap_or(&0);
343 *self.current_element_ordinal.entry(field_id).or_insert(0) += 1;
344 ordinal
345 }
346
347 pub fn num_docs(&self) -> u32 {
348 self.next_doc_id
349 }
350
351 #[inline]
353 pub fn estimated_memory_bytes(&self) -> usize {
354 self.estimated_memory
355 }
356
357 pub fn sparse_dim_count(&self) -> usize {
359 self.sparse_vectors.values().map(|b| b.postings.len()).sum()
360 }
361
362 pub fn stats(&self) -> SegmentBuilderStats {
364 use std::mem::size_of;
365
366 let postings_in_memory: usize =
367 self.inverted_index.values().map(|p| p.postings.len()).sum();
368
369 let compact_posting_size = size_of::<CompactPosting>();
371 let vec_overhead = size_of::<Vec<u8>>(); let term_key_size = size_of::<TermKey>();
373 let posting_builder_size = size_of::<PostingListBuilder>();
374 let spur_size = size_of::<Spur>();
375 let sparse_entry_size = size_of::<(DocId, u16, f32)>();
376
377 let hashmap_entry_base_overhead = 8usize;
380
381 let fxhashmap_entry_overhead = hashmap_entry_base_overhead;
383
384 let postings_bytes: usize = self
386 .inverted_index
387 .values()
388 .map(|p| p.postings.capacity() * compact_posting_size + vec_overhead)
389 .sum();
390
391 let index_overhead_bytes = self.inverted_index.len()
393 * (term_key_size + posting_builder_size + hashmap_entry_base_overhead);
394
395 let interner_arena_overhead = 2 * size_of::<usize>();
398 let avg_term_len = 8; let interner_bytes =
400 self.term_interner.len() * (avg_term_len + spur_size + interner_arena_overhead);
401
402 let field_lengths_bytes =
404 self.doc_field_lengths.capacity() * size_of::<u32>() + vec_overhead;
405
406 let mut dense_vectors_bytes: usize = 0;
408 let mut dense_vector_count: usize = 0;
409 let doc_id_ordinal_size = size_of::<(DocId, u16)>();
410 for b in self.dense_vectors.values() {
411 dense_vectors_bytes += b.vectors.capacity() * size_of::<f32>()
412 + b.doc_ids.capacity() * doc_id_ordinal_size
413 + 2 * vec_overhead; dense_vector_count += b.doc_ids.len();
415 }
416 for b in self.binary_dense_vectors.values() {
418 dense_vectors_bytes += b.vectors.capacity()
419 + b.doc_ids.capacity() * doc_id_ordinal_size
420 + 2 * vec_overhead;
421 dense_vector_count += b.doc_ids.len();
422 }
423
424 let local_tf_entry_size = spur_size + size_of::<u32>() + fxhashmap_entry_overhead;
426 let local_tf_buffer_bytes = self.local_tf_buffer.capacity() * local_tf_entry_size;
427
428 let mut sparse_vectors_bytes: usize = 0;
430 for builder in self.sparse_vectors.values() {
431 for postings in builder.postings.values() {
432 sparse_vectors_bytes += postings.capacity() * sparse_entry_size + vec_overhead;
433 }
434 let inner_entry_size = size_of::<u32>() + vec_overhead + fxhashmap_entry_overhead;
436 sparse_vectors_bytes += builder.postings.len() * inner_entry_size;
437 }
438 let outer_sparse_entry_size =
440 size_of::<u32>() + size_of::<SparseVectorBuilder>() + fxhashmap_entry_overhead;
441 sparse_vectors_bytes += self.sparse_vectors.len() * outer_sparse_entry_size;
442
443 let mut position_index_bytes: usize = 0;
445 for pos_builder in self.position_index.values() {
446 for (_, positions) in &pos_builder.postings {
447 position_index_bytes += positions.capacity() * size_of::<u32>() + vec_overhead;
448 }
449 let pos_entry_size = size_of::<DocId>() + vec_overhead;
451 position_index_bytes += pos_builder.postings.capacity() * pos_entry_size;
452 }
453 let pos_index_entry_size =
455 term_key_size + size_of::<PositionPostingListBuilder>() + hashmap_entry_base_overhead;
456 position_index_bytes += self.position_index.len() * pos_index_entry_size;
457
458 let estimated_memory_bytes = postings_bytes
459 + index_overhead_bytes
460 + interner_bytes
461 + field_lengths_bytes
462 + dense_vectors_bytes
463 + local_tf_buffer_bytes
464 + sparse_vectors_bytes
465 + position_index_bytes;
466
467 let memory_breakdown = MemoryBreakdown {
468 postings_bytes,
469 index_overhead_bytes,
470 interner_bytes,
471 field_lengths_bytes,
472 dense_vectors_bytes,
473 dense_vector_count,
474 sparse_vectors_bytes,
475 position_index_bytes,
476 };
477
478 SegmentBuilderStats {
479 num_docs: self.next_doc_id,
480 unique_terms: self.inverted_index.len(),
481 postings_in_memory,
482 interned_strings: self.term_interner.len(),
483 doc_field_lengths_size: self.doc_field_lengths.len(),
484 estimated_memory_bytes,
485 memory_breakdown,
486 }
487 }
488
489 pub fn add_document(&mut self, doc: Document) -> Result<DocId> {
491 let doc_id = self.next_doc_id;
492 self.next_doc_id += 1;
493
494 let base_idx = self.doc_field_lengths.len();
496 self.doc_field_lengths
497 .resize(base_idx + self.num_indexed_fields, 0);
498 self.estimated_memory += self.num_indexed_fields * std::mem::size_of::<u32>();
499
500 self.current_element_ordinal.clear();
502
503 for (field, value) in doc.field_values() {
504 let Some(entry) = self.schema.get_field_entry(*field) else {
505 continue;
506 };
507
508 if !matches!(
511 &entry.field_type,
512 FieldType::DenseVector | FieldType::BinaryDenseVector
513 ) && !entry.indexed
514 && !entry.fast
515 {
516 continue;
517 }
518
519 match (&entry.field_type, value) {
520 (FieldType::Text, FieldValue::Text(text)) => {
521 if entry.indexed {
522 let element_ordinal = self.next_element_ordinal(field.0);
523 let token_count =
524 self.index_text_field(*field, doc_id, text, element_ordinal)?;
525
526 let stats = self.field_stats.entry(field.0).or_default();
527 stats.total_tokens += token_count as u64;
528 if element_ordinal == 0 {
529 stats.doc_count += 1;
530 }
531
532 if let Some(&slot) = self.field_to_slot.get(&field.0) {
533 self.doc_field_lengths[base_idx + slot] = token_count;
534 }
535 }
536
537 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
539 ff.add_text(doc_id, text);
540 }
541 }
542 (FieldType::U64, FieldValue::U64(v)) => {
543 if entry.indexed {
544 self.index_numeric_field(*field, doc_id, *v)?;
545 }
546 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
547 ff.add_u64(doc_id, *v);
548 }
549 }
550 (FieldType::I64, FieldValue::I64(v)) => {
551 if entry.indexed {
552 self.index_numeric_field(*field, doc_id, *v as u64)?;
553 }
554 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
555 ff.add_i64(doc_id, *v);
556 }
557 }
558 (FieldType::F64, FieldValue::F64(v)) => {
559 if entry.indexed {
560 self.index_numeric_field(*field, doc_id, v.to_bits())?;
561 }
562 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
563 ff.add_f64(doc_id, *v);
564 }
565 }
566 (FieldType::DenseVector, FieldValue::DenseVector(vec))
567 if entry.indexed || entry.stored =>
568 {
569 let ordinal = self.next_element_ordinal(field.0);
570 self.index_dense_vector_field(*field, doc_id, ordinal as u16, vec)?;
571 }
572 (FieldType::BinaryDenseVector, FieldValue::BinaryDenseVector(bytes))
573 if entry.indexed || entry.stored =>
574 {
575 let ordinal = self.next_element_ordinal(field.0);
576 self.index_binary_dense_vector_field(*field, doc_id, ordinal as u16, bytes)?;
577 }
578 (FieldType::SparseVector, FieldValue::SparseVector(entries)) => {
579 let ordinal = self.next_element_ordinal(field.0);
580 self.index_sparse_vector_field(*field, doc_id, ordinal as u16, entries)?;
581 }
582 _ => {}
583 }
584 }
585
586 self.write_document_to_store(&doc)?;
588
589 Ok(doc_id)
590 }
591
592 fn index_text_field(
601 &mut self,
602 field: Field,
603 doc_id: DocId,
604 text: &str,
605 element_ordinal: u32,
606 ) -> Result<u32> {
607 use crate::dsl::PositionMode;
608
609 let field_id = field.0;
610 let position_mode = self
611 .position_enabled_fields
612 .get(&field_id)
613 .copied()
614 .flatten();
615
616 self.local_tf_buffer.clear();
620 for v in self.local_positions.values_mut() {
622 v.clear();
623 }
624
625 let mut token_position = 0u32;
626
627 let custom_tokens = self.tokenizers.get(&field).map(|t| t.tokenize(text));
631
632 if let Some(tokens) = custom_tokens {
633 for token in &tokens {
635 let term_spur = if let Some(spur) = self.term_interner.get(&token.text) {
636 spur
637 } else {
638 let spur = self.term_interner.get_or_intern(&token.text);
639 self.estimated_memory += token.text.len() + INTERN_OVERHEAD;
640 spur
641 };
642 *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
643
644 if let Some(mode) = position_mode {
645 let encoded_pos = match mode {
646 PositionMode::Ordinal => element_ordinal << 20,
647 PositionMode::TokenPosition => token.position,
648 PositionMode::Full => (element_ordinal << 20) | token.position,
649 };
650 self.local_positions
651 .entry(term_spur)
652 .or_default()
653 .push(encoded_pos);
654 }
655 }
656 token_position = tokens.len() as u32;
657 } else {
658 for word in text.split_whitespace() {
660 self.token_buffer.clear();
661 for c in word.chars() {
662 if c.is_alphanumeric() {
663 for lc in c.to_lowercase() {
664 self.token_buffer.push(lc);
665 }
666 }
667 }
668
669 if self.token_buffer.is_empty() {
670 continue;
671 }
672
673 let term_spur = if let Some(spur) = self.term_interner.get(&self.token_buffer) {
674 spur
675 } else {
676 let spur = self.term_interner.get_or_intern(&self.token_buffer);
677 self.estimated_memory += self.token_buffer.len() + INTERN_OVERHEAD;
678 spur
679 };
680 *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
681
682 if let Some(mode) = position_mode {
683 let encoded_pos = match mode {
684 PositionMode::Ordinal => element_ordinal << 20,
685 PositionMode::TokenPosition => token_position,
686 PositionMode::Full => (element_ordinal << 20) | token_position,
687 };
688 self.local_positions
689 .entry(term_spur)
690 .or_default()
691 .push(encoded_pos);
692 }
693
694 token_position += 1;
695 }
696 }
697
698 for (&term_spur, &tf) in &self.local_tf_buffer {
701 let term_key = TermKey {
702 field: field_id,
703 term: term_spur,
704 };
705
706 match self.inverted_index.entry(term_key) {
707 hashbrown::hash_map::Entry::Occupied(mut o) => {
708 o.get_mut().add(doc_id, tf);
709 self.estimated_memory += size_of::<CompactPosting>();
710 #[cfg(feature = "native")]
712 if o.get().should_spill() {
713 use byteorder::{LittleEndian, WriteBytesExt};
714
715 let builder = o.get_mut();
716 let count = builder.postings.len() as u32;
717 let offset = self.posting_spill_offset;
718
719 let spill_file = if let Some(ref mut f) = self.posting_spill_file {
721 f
722 } else {
723 self.posting_spill_file = Some(BufWriter::with_capacity(
724 256 * 1024,
725 OpenOptions::new()
726 .create(true)
727 .write(true)
728 .truncate(true)
729 .open(&self.posting_spill_path)?,
730 ));
731 self.posting_spill_file.as_mut().unwrap()
732 };
733 for p in &builder.postings {
734 spill_file.write_u32::<LittleEndian>(p.doc_id)?;
735 spill_file.write_u16::<LittleEndian>(p.term_freq)?;
736 }
737 self.posting_spill_offset += count as u64 * 6;
738 self.posting_spill_index
739 .entry(term_key)
740 .or_default()
741 .push((offset, count));
742
743 let freed = builder.postings.len() * size_of::<CompactPosting>();
744 builder.spilled_count += count;
745 builder.postings.clear();
746 self.estimated_memory -= freed;
747 }
748 }
749 hashbrown::hash_map::Entry::Vacant(v) => {
750 let mut posting = PostingListBuilder::new();
751 posting.add(doc_id, tf);
752 v.insert(posting);
753 self.estimated_memory += size_of::<CompactPosting>() + NEW_TERM_OVERHEAD;
754 }
755 }
756
757 if position_mode.is_some()
758 && let Some(positions) = self.local_positions.get(&term_spur)
759 {
760 match self.position_index.entry(term_key) {
761 hashbrown::hash_map::Entry::Occupied(mut o) => {
762 for &pos in positions {
763 o.get_mut().add_position(doc_id, pos);
764 }
765 self.estimated_memory += positions.len() * size_of::<u32>();
766 }
767 hashbrown::hash_map::Entry::Vacant(v) => {
768 let mut pos_posting = PositionPostingListBuilder::new();
769 for &pos in positions {
770 pos_posting.add_position(doc_id, pos);
771 }
772 self.estimated_memory +=
773 positions.len() * size_of::<u32>() + NEW_POS_TERM_OVERHEAD;
774 v.insert(pos_posting);
775 }
776 }
777 }
778 }
779
780 Ok(token_position)
781 }
782
783 fn index_numeric_field(&mut self, field: Field, doc_id: DocId, value: u64) -> Result<()> {
784 use std::fmt::Write;
785
786 self.numeric_buffer.clear();
787 write!(self.numeric_buffer, "__num_{}", value).unwrap();
788 let term_spur = if let Some(spur) = self.term_interner.get(&self.numeric_buffer) {
789 spur
790 } else {
791 let spur = self.term_interner.get_or_intern(&self.numeric_buffer);
792 self.estimated_memory += self.numeric_buffer.len() + INTERN_OVERHEAD;
793 spur
794 };
795
796 let term_key = TermKey {
797 field: field.0,
798 term: term_spur,
799 };
800
801 match self.inverted_index.entry(term_key) {
802 hashbrown::hash_map::Entry::Occupied(mut o) => {
803 o.get_mut().add(doc_id, 1);
804 self.estimated_memory += size_of::<CompactPosting>();
805 }
806 hashbrown::hash_map::Entry::Vacant(v) => {
807 let mut posting = PostingListBuilder::new();
808 posting.add(doc_id, 1);
809 v.insert(posting);
810 self.estimated_memory += size_of::<CompactPosting>() + NEW_TERM_OVERHEAD;
811 }
812 }
813
814 Ok(())
815 }
816
817 fn index_dense_vector_field(
819 &mut self,
820 field: Field,
821 doc_id: DocId,
822 ordinal: u16,
823 vector: &[f32],
824 ) -> Result<()> {
825 let dim = vector.len();
826
827 let builder = self
828 .dense_vectors
829 .entry(field.0)
830 .or_insert_with(|| DenseVectorBuilder::new(dim));
831
832 if builder.dim != dim && builder.len() > 0 {
834 return Err(crate::Error::Schema(format!(
835 "Dense vector dimension mismatch: expected {}, got {}",
836 builder.dim, dim
837 )));
838 }
839
840 builder.add(doc_id, ordinal, vector);
841
842 self.estimated_memory += std::mem::size_of_val(vector) + size_of::<(DocId, u16)>();
843
844 Ok(())
845 }
846
847 fn index_binary_dense_vector_field(
849 &mut self,
850 field: Field,
851 doc_id: DocId,
852 ordinal: u16,
853 bytes: &[u8],
854 ) -> Result<()> {
855 let dim_bits = self
856 .schema
857 .get_field_entry(field)
858 .and_then(|e| e.binary_dense_vector_config.as_ref())
859 .map(|c| c.dim)
860 .ok_or_else(|| {
861 crate::Error::Schema("BinaryDenseVector field missing config".to_string())
862 })?;
863
864 let expected_byte_len = dim_bits.div_ceil(8);
865 if bytes.len() != expected_byte_len {
866 return Err(crate::Error::Schema(format!(
867 "Binary vector byte length mismatch: expected {} (dim={}), got {}",
868 expected_byte_len,
869 dim_bits,
870 bytes.len()
871 )));
872 }
873
874 let builder = self
875 .binary_dense_vectors
876 .entry(field.0)
877 .or_insert_with(|| BinaryDenseVectorBuilder::new(dim_bits));
878
879 builder.add(doc_id, ordinal, bytes);
880 self.estimated_memory += bytes.len() + size_of::<(DocId, u16)>();
881
882 Ok(())
883 }
884
885 fn index_sparse_vector_field(
895 &mut self,
896 field: Field,
897 doc_id: DocId,
898 ordinal: u16,
899 entries: &[(u32, f32)],
900 ) -> Result<()> {
901 let (weight_threshold, doc_mass, min_terms) = self
902 .schema
903 .get_field_entry(field)
904 .and_then(|entry| entry.sparse_vector_config.as_ref())
905 .map(|config| (config.weight_threshold, config.doc_mass, config.min_terms))
906 .unwrap_or((0.0, None, 0));
907
908 let builder = self
909 .sparse_vectors
910 .entry(field.0)
911 .or_insert_with(SparseVectorBuilder::new);
912
913 builder.inc_vector_count();
914
915 let mass_cutoff = match doc_mass {
919 Some(mass) if mass < 1.0 && entries.len() > min_terms => {
920 let mut weights: Vec<f32> = entries
921 .iter()
922 .map(|&(_, w)| w.abs())
923 .filter(|w| *w >= weight_threshold)
924 .collect();
925 weights.sort_unstable_by(|a, b| b.total_cmp(a));
926 let total: f64 = weights.iter().map(|&w| w as f64).sum();
927 let target = total * mass as f64;
928 let mut cumulative = 0.0f64;
929 let mut cutoff = 0.0f32;
930 for &w in &weights {
931 if cumulative >= target {
932 break;
933 }
934 cumulative += w as f64;
935 cutoff = w;
936 }
937 cutoff
938 }
939 _ => 0.0,
940 };
941
942 for &(dim_id, weight) in entries {
943 if weight.abs() < weight_threshold || weight.abs() < mass_cutoff {
945 continue;
946 }
947
948 let is_new_dim = !builder.postings.contains_key(&dim_id);
949 builder.add(dim_id, doc_id, ordinal, weight);
950 self.estimated_memory += size_of::<(DocId, u16, f32)>();
951 if is_new_dim {
952 self.estimated_memory += size_of::<u32>() + size_of::<Vec<(DocId, u16, f32)>>() + 8; }
955 }
956
957 Ok(())
958 }
959
960 fn write_document_to_store(&mut self, doc: &Document) -> Result<()> {
962 use byteorder::{LittleEndian, WriteBytesExt};
963
964 super::store::serialize_document_into(doc, &self.schema, &mut self.doc_serialize_buffer)?;
965
966 #[cfg(feature = "native")]
967 {
968 self.store_file
969 .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
970 self.store_file.write_all(&self.doc_serialize_buffer)?;
971 }
972 #[cfg(not(feature = "native"))]
973 {
974 self.store_buffer
975 .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
976 self.store_buffer.write_all(&self.doc_serialize_buffer)?;
977 }
978
979 Ok(())
980 }
981
982 pub async fn build<D: Directory + DirectoryWriter>(
988 mut self,
989 dir: &D,
990 segment_id: SegmentId,
991 trained: Option<&super::TrainedVectorStructures>,
992 ) -> Result<SegmentMeta> {
993 #[cfg(feature = "native")]
995 self.store_file.flush()?;
996
997 let files = SegmentFiles::new(segment_id.0);
998
999 let position_index = std::mem::take(&mut self.position_index);
1001 let position_offsets = if !position_index.is_empty() {
1002 let mut pos_writer = dir.streaming_writer(&files.positions).await?;
1003 let offsets = postings::build_positions_streaming(
1004 position_index,
1005 &self.term_interner,
1006 &mut *pos_writer,
1007 )?;
1008 pos_writer.finish()?;
1009 offsets
1010 } else {
1011 FxHashMap::default()
1012 };
1013
1014 let inverted_index = std::mem::take(&mut self.inverted_index);
1017 let term_interner = std::mem::replace(&mut self.term_interner, Rodeo::new());
1018 #[cfg(feature = "native")]
1019 let store_path = self.store_path.clone();
1020 #[cfg(feature = "native")]
1021 let num_compression_threads = self.config.num_compression_threads;
1022 let compression_level = self.config.compression_level;
1023 let dense_vectors = std::mem::take(&mut self.dense_vectors);
1024 let binary_dense_vectors = std::mem::take(&mut self.binary_dense_vectors);
1025 let mut sparse_vectors = std::mem::take(&mut self.sparse_vectors);
1026 let schema = &self.schema;
1027
1028 let mut term_dict_writer =
1031 super::OffsetWriter::new(dir.streaming_writer(&files.term_dict).await?);
1032 let mut postings_writer =
1033 super::OffsetWriter::new(dir.streaming_writer(&files.postings).await?);
1034 let mut store_writer = super::OffsetWriter::new(dir.streaming_writer(&files.store).await?);
1035 let mut vectors_writer = if !dense_vectors.is_empty() || !binary_dense_vectors.is_empty() {
1036 Some(super::OffsetWriter::new(
1037 dir.streaming_writer(&files.vectors).await?,
1038 ))
1039 } else {
1040 None
1041 };
1042 let mut sparse_writer = if !sparse_vectors.is_empty() {
1043 Some(super::OffsetWriter::new(
1044 dir.streaming_writer(&files.sparse).await?,
1045 ))
1046 } else {
1047 None
1048 };
1049 let mut fast_fields = std::mem::take(&mut self.fast_fields);
1050 let num_docs = self.next_doc_id;
1051 let mut fast_writer = if !fast_fields.is_empty() {
1052 Some(super::OffsetWriter::new(
1053 dir.streaming_writer(&files.fast).await?,
1054 ))
1055 } else {
1056 None
1057 };
1058
1059 #[cfg(feature = "native")]
1060 {
1061 if let Some(ref mut f) = self.posting_spill_file {
1062 f.flush()?;
1063 }
1064 let posting_spill_index = std::mem::take(&mut self.posting_spill_index);
1065 let mut spill_reader_opt = if !posting_spill_index.is_empty() {
1066 let spill_file = std::fs::File::open(&self.posting_spill_path)?;
1067 Some((std::io::BufReader::new(spill_file), posting_spill_index))
1068 } else {
1069 None
1070 };
1071
1072 let ((postings_result, store_result), ((vectors_result, sparse_result), fast_result)) =
1073 rayon::join(
1074 || {
1075 rayon::join(
1076 || {
1077 let spill_arg = spill_reader_opt.as_mut().map(|(r, idx)| {
1078 (
1079 r as &mut std::io::BufReader<std::fs::File>,
1080 idx as &postings::SpillIndex,
1081 )
1082 });
1083 postings::build_postings_streaming(
1084 inverted_index,
1085 term_interner,
1086 &position_offsets,
1087 &mut term_dict_writer,
1088 &mut postings_writer,
1089 spill_arg,
1090 )
1091 },
1092 || {
1093 store::build_store_streaming(
1094 &store_path,
1095 num_compression_threads,
1096 compression_level,
1097 &mut store_writer,
1098 num_docs,
1099 )
1100 },
1101 )
1102 },
1103 || {
1104 rayon::join(
1105 || {
1106 rayon::join(
1107 || -> Result<()> {
1108 if let Some(ref mut w) = vectors_writer {
1109 dense::build_vectors_streaming(
1110 dense_vectors,
1111 binary_dense_vectors,
1112 schema,
1113 trained,
1114 w,
1115 )?;
1116 }
1117 Ok(())
1118 },
1119 || -> Result<()> {
1120 if let Some(ref mut w) = sparse_writer {
1121 sparse::build_sparse_streaming(
1122 &mut sparse_vectors,
1123 schema,
1124 w,
1125 )?;
1126 }
1127 Ok(())
1128 },
1129 )
1130 },
1131 || -> Result<()> {
1132 if let Some(ref mut w) = fast_writer {
1133 build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
1134 }
1135 Ok(())
1136 },
1137 )
1138 },
1139 );
1140 postings_result?;
1141 store_result?;
1142 vectors_result?;
1143 sparse_result?;
1144 fast_result?;
1145 }
1146
1147 #[cfg(not(feature = "native"))]
1148 {
1149 postings::build_postings_streaming(
1150 inverted_index,
1151 term_interner,
1152 &position_offsets,
1153 &mut term_dict_writer,
1154 &mut postings_writer,
1155 )?;
1156 store::build_store_streaming_from_buffer(
1157 &self.store_buffer,
1158 compression_level,
1159 &mut store_writer,
1160 num_docs,
1161 )?;
1162 if let Some(ref mut w) = vectors_writer {
1163 dense::build_vectors_streaming(
1164 dense_vectors,
1165 binary_dense_vectors,
1166 schema,
1167 trained,
1168 w,
1169 )?;
1170 }
1171 if let Some(ref mut w) = sparse_writer {
1172 sparse::build_sparse_streaming(&mut sparse_vectors, schema, w)?;
1173 }
1174 if let Some(ref mut w) = fast_writer {
1175 build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
1176 }
1177 }
1178
1179 let term_dict_bytes = term_dict_writer.offset() as usize;
1180 let postings_bytes = postings_writer.offset() as usize;
1181 let store_bytes = store_writer.offset() as usize;
1182 let vectors_bytes = vectors_writer.as_ref().map_or(0, |w| w.offset() as usize);
1183 let sparse_bytes = sparse_writer.as_ref().map_or(0, |w| w.offset() as usize);
1184 let fast_bytes = fast_writer.as_ref().map_or(0, |w| w.offset() as usize);
1185
1186 term_dict_writer.finish()?;
1187 postings_writer.finish()?;
1188 store_writer.finish()?;
1189 if let Some(w) = vectors_writer {
1190 w.finish()?;
1191 }
1192 if let Some(w) = sparse_writer {
1193 w.finish()?;
1194 }
1195 if let Some(w) = fast_writer {
1196 w.finish()?;
1197 }
1198 drop(position_offsets);
1199 drop(sparse_vectors);
1200
1201 log::info!(
1202 "[segment_build] {} docs: term_dict={}, postings={}, store={}, vectors={}, sparse={}, fast={}",
1203 num_docs,
1204 super::format_bytes(term_dict_bytes),
1205 super::format_bytes(postings_bytes),
1206 super::format_bytes(store_bytes),
1207 super::format_bytes(vectors_bytes),
1208 super::format_bytes(sparse_bytes),
1209 super::format_bytes(fast_bytes),
1210 );
1211
1212 let meta = SegmentMeta {
1213 id: segment_id.0,
1214 num_docs: self.next_doc_id,
1215 field_stats: self.field_stats.clone(),
1216 };
1217
1218 dir.write(&files.meta, &meta.serialize()?).await?;
1219
1220 #[cfg(feature = "native")]
1222 {
1223 let _ = std::fs::remove_file(&self.store_path);
1224 }
1225
1226 Ok(meta)
1227 }
1228}
1229
1230fn build_fast_fields_streaming(
1232 fast_fields: &mut FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
1233 num_docs: u32,
1234 writer: &mut dyn Write,
1235) -> Result<()> {
1236 use crate::structures::fast_field::{FastFieldTocEntry, write_fast_field_toc_and_footer};
1237
1238 if fast_fields.is_empty() {
1239 return Ok(());
1240 }
1241
1242 let mut field_ids: Vec<u32> = fast_fields.keys().copied().collect();
1244 field_ids.sort_unstable();
1245
1246 let mut toc_entries: Vec<FastFieldTocEntry> = Vec::with_capacity(field_ids.len());
1247 let mut current_offset = 0u64;
1248
1249 for &field_id in &field_ids {
1250 let ff = fast_fields.get_mut(&field_id).unwrap();
1251 ff.pad_to(num_docs);
1252
1253 let (mut toc, bytes_written) = ff.serialize(writer, current_offset)?;
1254 toc.field_id = field_id;
1255 current_offset += bytes_written;
1256 toc_entries.push(toc);
1257 }
1258
1259 let toc_offset = current_offset;
1261 write_fast_field_toc_and_footer(writer, toc_offset, &toc_entries)?;
1262
1263 Ok(())
1264}
1265
1266#[cfg(feature = "native")]
1267impl Drop for SegmentBuilder {
1268 fn drop(&mut self) {
1269 let _ = std::fs::remove_file(&self.store_path);
1270 if self.posting_spill_file.is_some() {
1271 let _ = std::fs::remove_file(&self.posting_spill_path);
1272 }
1273 }
1274}