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 fn next_vector_ordinal(&mut self, field_id: u32) -> Result<u16> {
348 let ordinal = self.next_element_ordinal(field_id);
349 u16::try_from(ordinal).map_err(|_| {
350 crate::Error::Document(format!(
351 "field {field_id} has more than {} vector values in one document",
352 u16::MAX as usize + 1
353 ))
354 })
355 }
356
357 pub fn num_docs(&self) -> u32 {
358 self.next_doc_id
359 }
360
361 #[inline]
363 pub fn estimated_memory_bytes(&self) -> usize {
364 self.estimated_memory
365 }
366
367 pub fn sparse_dim_count(&self) -> usize {
369 self.sparse_vectors.values().map(|b| b.postings.len()).sum()
370 }
371
372 pub fn stats(&self) -> SegmentBuilderStats {
374 use std::mem::size_of;
375
376 let postings_in_memory: usize =
377 self.inverted_index.values().map(|p| p.postings.len()).sum();
378
379 let compact_posting_size = size_of::<CompactPosting>();
381 let vec_overhead = size_of::<Vec<u8>>(); let term_key_size = size_of::<TermKey>();
383 let posting_builder_size = size_of::<PostingListBuilder>();
384 let spur_size = size_of::<Spur>();
385 let sparse_entry_size = size_of::<(DocId, u16, f32)>();
386
387 let hashmap_entry_base_overhead = 8usize;
390
391 let fxhashmap_entry_overhead = hashmap_entry_base_overhead;
393
394 let postings_bytes: usize = self
396 .inverted_index
397 .values()
398 .map(|p| p.postings.capacity() * compact_posting_size + vec_overhead)
399 .sum();
400
401 let index_overhead_bytes = self.inverted_index.len()
403 * (term_key_size + posting_builder_size + hashmap_entry_base_overhead);
404
405 let interner_arena_overhead = 2 * size_of::<usize>();
408 let avg_term_len = 8; let interner_bytes =
410 self.term_interner.len() * (avg_term_len + spur_size + interner_arena_overhead);
411
412 let field_lengths_bytes =
414 self.doc_field_lengths.capacity() * size_of::<u32>() + vec_overhead;
415
416 let mut dense_vectors_bytes: usize = 0;
418 let mut dense_vector_count: usize = 0;
419 let doc_id_ordinal_size = size_of::<(DocId, u16)>();
420 for b in self.dense_vectors.values() {
421 dense_vectors_bytes += b.vectors.capacity() * size_of::<f32>()
422 + b.doc_ids.capacity() * doc_id_ordinal_size
423 + 2 * vec_overhead; dense_vector_count += b.doc_ids.len();
425 }
426 for b in self.binary_dense_vectors.values() {
428 dense_vectors_bytes += b.vectors.capacity()
429 + b.doc_ids.capacity() * doc_id_ordinal_size
430 + 2 * vec_overhead;
431 dense_vector_count += b.doc_ids.len();
432 }
433
434 let local_tf_entry_size = spur_size + size_of::<u32>() + fxhashmap_entry_overhead;
436 let local_tf_buffer_bytes = self.local_tf_buffer.capacity() * local_tf_entry_size;
437
438 let mut sparse_vectors_bytes: usize = 0;
440 for builder in self.sparse_vectors.values() {
441 for postings in builder.postings.values() {
442 sparse_vectors_bytes += postings.capacity() * sparse_entry_size + vec_overhead;
443 }
444 let inner_entry_size = size_of::<u32>() + vec_overhead + fxhashmap_entry_overhead;
446 sparse_vectors_bytes += builder.postings.len() * inner_entry_size;
447 }
448 let outer_sparse_entry_size =
450 size_of::<u32>() + size_of::<SparseVectorBuilder>() + fxhashmap_entry_overhead;
451 sparse_vectors_bytes += self.sparse_vectors.len() * outer_sparse_entry_size;
452
453 let mut position_index_bytes: usize = 0;
455 for pos_builder in self.position_index.values() {
456 for (_, positions) in &pos_builder.postings {
457 position_index_bytes += positions.capacity() * size_of::<u32>() + vec_overhead;
458 }
459 let pos_entry_size = size_of::<DocId>() + vec_overhead;
461 position_index_bytes += pos_builder.postings.capacity() * pos_entry_size;
462 }
463 let pos_index_entry_size =
465 term_key_size + size_of::<PositionPostingListBuilder>() + hashmap_entry_base_overhead;
466 position_index_bytes += self.position_index.len() * pos_index_entry_size;
467
468 let estimated_memory_bytes = postings_bytes
469 + index_overhead_bytes
470 + interner_bytes
471 + field_lengths_bytes
472 + dense_vectors_bytes
473 + local_tf_buffer_bytes
474 + sparse_vectors_bytes
475 + position_index_bytes;
476
477 let memory_breakdown = MemoryBreakdown {
478 postings_bytes,
479 index_overhead_bytes,
480 interner_bytes,
481 field_lengths_bytes,
482 dense_vectors_bytes,
483 dense_vector_count,
484 sparse_vectors_bytes,
485 position_index_bytes,
486 };
487
488 SegmentBuilderStats {
489 num_docs: self.next_doc_id,
490 unique_terms: self.inverted_index.len(),
491 postings_in_memory,
492 interned_strings: self.term_interner.len(),
493 doc_field_lengths_size: self.doc_field_lengths.len(),
494 estimated_memory_bytes,
495 memory_breakdown,
496 }
497 }
498
499 pub fn add_document(&mut self, doc: Document) -> Result<DocId> {
501 let doc_id = self.next_doc_id;
502 self.next_doc_id += 1;
503
504 let base_idx = self.doc_field_lengths.len();
506 self.doc_field_lengths
507 .resize(base_idx + self.num_indexed_fields, 0);
508 self.estimated_memory += self.num_indexed_fields * std::mem::size_of::<u32>();
509
510 self.current_element_ordinal.clear();
512
513 for (field, value) in doc.field_values() {
514 let Some(entry) = self.schema.get_field_entry(*field) else {
515 continue;
516 };
517
518 if !matches!(
521 &entry.field_type,
522 FieldType::DenseVector | FieldType::BinaryDenseVector
523 ) && !entry.indexed
524 && !entry.fast
525 {
526 continue;
527 }
528
529 match (&entry.field_type, value) {
530 (FieldType::Text, FieldValue::Text(text)) => {
531 if entry.indexed {
532 let element_ordinal = self.next_element_ordinal(field.0);
533 let token_count =
534 self.index_text_field(*field, doc_id, text, element_ordinal)?;
535
536 let stats = self.field_stats.entry(field.0).or_default();
537 stats.total_tokens += token_count as u64;
538 if element_ordinal == 0 {
539 stats.doc_count += 1;
540 }
541
542 if let Some(&slot) = self.field_to_slot.get(&field.0) {
543 self.doc_field_lengths[base_idx + slot] = token_count;
544 }
545 }
546
547 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
549 ff.add_text(doc_id, text);
550 }
551 }
552 (FieldType::U64, FieldValue::U64(v)) => {
553 if entry.indexed {
554 self.index_numeric_field(*field, doc_id, *v)?;
555 }
556 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
557 ff.add_u64(doc_id, *v);
558 }
559 }
560 (FieldType::I64, FieldValue::I64(v)) => {
561 if entry.indexed {
562 self.index_numeric_field(*field, doc_id, *v as u64)?;
563 }
564 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
565 ff.add_i64(doc_id, *v);
566 }
567 }
568 (FieldType::F64, FieldValue::F64(v)) => {
569 if entry.indexed {
570 self.index_numeric_field(*field, doc_id, v.to_bits())?;
571 }
572 if let Some(ff) = self.fast_fields.get_mut(&field.0) {
573 ff.add_f64(doc_id, *v);
574 }
575 }
576 (FieldType::DenseVector, FieldValue::DenseVector(vec))
577 if entry.indexed || entry.stored =>
578 {
579 let ordinal = self.next_vector_ordinal(field.0)?;
580 self.index_dense_vector_field(*field, doc_id, ordinal, vec)?;
581 }
582 (FieldType::BinaryDenseVector, FieldValue::BinaryDenseVector(bytes))
583 if entry.indexed || entry.stored =>
584 {
585 let ordinal = self.next_vector_ordinal(field.0)?;
586 self.index_binary_dense_vector_field(*field, doc_id, ordinal, bytes)?;
587 }
588 (FieldType::SparseVector, FieldValue::SparseVector(entries)) => {
589 let ordinal = self.next_vector_ordinal(field.0)?;
590 self.index_sparse_vector_field(*field, doc_id, ordinal, entries)?;
591 }
592 _ => {}
593 }
594 }
595
596 self.write_document_to_store(&doc)?;
598
599 Ok(doc_id)
600 }
601
602 fn index_text_field(
611 &mut self,
612 field: Field,
613 doc_id: DocId,
614 text: &str,
615 element_ordinal: u32,
616 ) -> Result<u32> {
617 use crate::dsl::PositionMode;
618
619 let field_id = field.0;
620 let position_mode = self
621 .position_enabled_fields
622 .get(&field_id)
623 .copied()
624 .flatten();
625
626 self.local_tf_buffer.clear();
630 for v in self.local_positions.values_mut() {
632 v.clear();
633 }
634
635 let mut token_position = 0u32;
636
637 let custom_tokens = self.tokenizers.get(&field).map(|t| t.tokenize(text));
641
642 if let Some(tokens) = custom_tokens {
643 for token in &tokens {
645 let term_spur = if let Some(spur) = self.term_interner.get(&token.text) {
646 spur
647 } else {
648 let spur = self.term_interner.get_or_intern(&token.text);
649 self.estimated_memory += token.text.len() + INTERN_OVERHEAD;
650 spur
651 };
652 *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
653
654 if let Some(mode) = position_mode {
655 let encoded_pos = match mode {
656 PositionMode::Ordinal => element_ordinal << 20,
657 PositionMode::TokenPosition => token.position,
658 PositionMode::Full => (element_ordinal << 20) | token.position,
659 };
660 self.local_positions
661 .entry(term_spur)
662 .or_default()
663 .push(encoded_pos);
664 }
665 }
666 token_position = tokens.len() as u32;
667 } else {
668 for word in text.split_whitespace() {
670 self.token_buffer.clear();
671 for c in word.chars() {
672 if c.is_alphanumeric() {
673 for lc in c.to_lowercase() {
674 self.token_buffer.push(lc);
675 }
676 }
677 }
678
679 if self.token_buffer.is_empty() {
680 continue;
681 }
682
683 let term_spur = if let Some(spur) = self.term_interner.get(&self.token_buffer) {
684 spur
685 } else {
686 let spur = self.term_interner.get_or_intern(&self.token_buffer);
687 self.estimated_memory += self.token_buffer.len() + INTERN_OVERHEAD;
688 spur
689 };
690 *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
691
692 if let Some(mode) = position_mode {
693 let encoded_pos = match mode {
694 PositionMode::Ordinal => element_ordinal << 20,
695 PositionMode::TokenPosition => token_position,
696 PositionMode::Full => (element_ordinal << 20) | token_position,
697 };
698 self.local_positions
699 .entry(term_spur)
700 .or_default()
701 .push(encoded_pos);
702 }
703
704 token_position += 1;
705 }
706 }
707
708 for (&term_spur, &tf) in &self.local_tf_buffer {
711 let term_key = TermKey {
712 field: field_id,
713 term: term_spur,
714 };
715
716 match self.inverted_index.entry(term_key) {
717 hashbrown::hash_map::Entry::Occupied(mut o) => {
718 o.get_mut().add(doc_id, tf);
719 self.estimated_memory += size_of::<CompactPosting>();
720 #[cfg(feature = "native")]
722 if o.get().should_spill() {
723 use byteorder::{LittleEndian, WriteBytesExt};
724
725 let builder = o.get_mut();
726 let count = builder.postings.len() as u32;
727 let offset = self.posting_spill_offset;
728
729 let spill_file = if let Some(ref mut f) = self.posting_spill_file {
731 f
732 } else {
733 self.posting_spill_file = Some(BufWriter::with_capacity(
734 256 * 1024,
735 OpenOptions::new()
736 .create(true)
737 .write(true)
738 .truncate(true)
739 .open(&self.posting_spill_path)?,
740 ));
741 self.posting_spill_file.as_mut().unwrap()
742 };
743 for p in &builder.postings {
744 spill_file.write_u32::<LittleEndian>(p.doc_id)?;
745 spill_file.write_u16::<LittleEndian>(p.term_freq)?;
746 }
747 self.posting_spill_offset += count as u64 * 6;
748 self.posting_spill_index
749 .entry(term_key)
750 .or_default()
751 .push((offset, count));
752
753 let freed = builder.postings.len() * size_of::<CompactPosting>();
754 builder.spilled_count += count;
755 builder.postings.clear();
756 self.estimated_memory -= freed;
757 }
758 }
759 hashbrown::hash_map::Entry::Vacant(v) => {
760 let mut posting = PostingListBuilder::new();
761 posting.add(doc_id, tf);
762 v.insert(posting);
763 self.estimated_memory += size_of::<CompactPosting>() + NEW_TERM_OVERHEAD;
764 }
765 }
766
767 if position_mode.is_some()
768 && let Some(positions) = self.local_positions.get(&term_spur)
769 {
770 match self.position_index.entry(term_key) {
771 hashbrown::hash_map::Entry::Occupied(mut o) => {
772 for &pos in positions {
773 o.get_mut().add_position(doc_id, pos);
774 }
775 self.estimated_memory += positions.len() * size_of::<u32>();
776 }
777 hashbrown::hash_map::Entry::Vacant(v) => {
778 let mut pos_posting = PositionPostingListBuilder::new();
779 for &pos in positions {
780 pos_posting.add_position(doc_id, pos);
781 }
782 self.estimated_memory +=
783 positions.len() * size_of::<u32>() + NEW_POS_TERM_OVERHEAD;
784 v.insert(pos_posting);
785 }
786 }
787 }
788 }
789
790 Ok(token_position)
791 }
792
793 fn index_numeric_field(&mut self, field: Field, doc_id: DocId, value: u64) -> Result<()> {
794 use std::fmt::Write;
795
796 self.numeric_buffer.clear();
797 write!(self.numeric_buffer, "__num_{}", value).unwrap();
798 let term_spur = if let Some(spur) = self.term_interner.get(&self.numeric_buffer) {
799 spur
800 } else {
801 let spur = self.term_interner.get_or_intern(&self.numeric_buffer);
802 self.estimated_memory += self.numeric_buffer.len() + INTERN_OVERHEAD;
803 spur
804 };
805
806 let term_key = TermKey {
807 field: field.0,
808 term: term_spur,
809 };
810
811 match self.inverted_index.entry(term_key) {
812 hashbrown::hash_map::Entry::Occupied(mut o) => {
813 o.get_mut().add(doc_id, 1);
814 self.estimated_memory += size_of::<CompactPosting>();
815 }
816 hashbrown::hash_map::Entry::Vacant(v) => {
817 let mut posting = PostingListBuilder::new();
818 posting.add(doc_id, 1);
819 v.insert(posting);
820 self.estimated_memory += size_of::<CompactPosting>() + NEW_TERM_OVERHEAD;
821 }
822 }
823
824 Ok(())
825 }
826
827 fn index_dense_vector_field(
829 &mut self,
830 field: Field,
831 doc_id: DocId,
832 ordinal: u16,
833 vector: &[f32],
834 ) -> Result<()> {
835 let dim = vector.len();
836 let expected_dim = self
837 .schema
838 .get_field_entry(field)
839 .and_then(|entry| entry.dense_vector_config.as_ref())
840 .map(|config| config.dim)
841 .ok_or_else(|| crate::Error::Schema("DenseVector field missing config".to_string()))?;
842 if dim != expected_dim {
843 return Err(crate::Error::Schema(format!(
844 "Dense vector dimension mismatch: schema expects {}, got {}",
845 expected_dim, dim
846 )));
847 }
848 if let Some((index, value)) = vector
849 .iter()
850 .enumerate()
851 .find(|(_, value)| !value.is_finite())
852 {
853 return Err(crate::Error::Document(format!(
854 "dense vector contains non-finite value {value} at index {index}"
855 )));
856 }
857
858 let builder = self
859 .dense_vectors
860 .entry(field.0)
861 .or_insert_with(|| DenseVectorBuilder::new(dim));
862
863 if builder.dim != dim && builder.len() > 0 {
865 return Err(crate::Error::Schema(format!(
866 "Dense vector dimension mismatch: expected {}, got {}",
867 builder.dim, dim
868 )));
869 }
870
871 builder.add(doc_id, ordinal, vector);
872
873 self.estimated_memory += std::mem::size_of_val(vector) + size_of::<(DocId, u16)>();
874
875 Ok(())
876 }
877
878 fn index_binary_dense_vector_field(
880 &mut self,
881 field: Field,
882 doc_id: DocId,
883 ordinal: u16,
884 bytes: &[u8],
885 ) -> Result<()> {
886 let dim_bits = self
887 .schema
888 .get_field_entry(field)
889 .and_then(|e| e.binary_dense_vector_config.as_ref())
890 .map(|c| c.dim)
891 .ok_or_else(|| {
892 crate::Error::Schema("BinaryDenseVector field missing config".to_string())
893 })?;
894
895 let expected_byte_len = dim_bits.div_ceil(8);
896 if dim_bits == 0 || !dim_bits.is_multiple_of(8) {
897 return Err(crate::Error::Schema(format!(
898 "Binary vector dimension must be a positive multiple of 8, got {dim_bits}"
899 )));
900 }
901 if bytes.len() != expected_byte_len {
902 return Err(crate::Error::Schema(format!(
903 "Binary vector byte length mismatch: expected {} (dim={}), got {}",
904 expected_byte_len,
905 dim_bits,
906 bytes.len()
907 )));
908 }
909
910 let builder = self
911 .binary_dense_vectors
912 .entry(field.0)
913 .or_insert_with(|| BinaryDenseVectorBuilder::new(dim_bits));
914
915 builder.add(doc_id, ordinal, bytes);
916 self.estimated_memory += bytes.len() + size_of::<(DocId, u16)>();
917
918 Ok(())
919 }
920
921 fn index_sparse_vector_field(
931 &mut self,
932 field: Field,
933 doc_id: DocId,
934 ordinal: u16,
935 entries: &[(u32, f32)],
936 ) -> Result<()> {
937 if let Some((index, (_, weight))) = entries
938 .iter()
939 .enumerate()
940 .find(|(_, (_, weight))| !weight.is_finite())
941 {
942 return Err(crate::Error::Document(format!(
943 "sparse vector contains non-finite weight {weight} at index {index}"
944 )));
945 }
946 let (weight_threshold, doc_mass, min_terms) = self
947 .schema
948 .get_field_entry(field)
949 .and_then(|entry| entry.sparse_vector_config.as_ref())
950 .map(|config| (config.weight_threshold, config.doc_mass, config.min_terms))
951 .unwrap_or((0.0, None, 0));
952
953 let builder = self
954 .sparse_vectors
955 .entry(field.0)
956 .or_insert_with(SparseVectorBuilder::new);
957
958 builder.inc_vector_count();
959
960 let mass_cutoff = match doc_mass {
964 Some(mass) if mass < 1.0 && entries.len() > min_terms => {
965 let mut weights: Vec<f32> = entries
966 .iter()
967 .map(|&(_, w)| w.abs())
968 .filter(|w| *w >= weight_threshold)
969 .collect();
970 weights.sort_unstable_by(|a, b| b.total_cmp(a));
971 let total: f64 = weights.iter().map(|&w| w as f64).sum();
972 let target = total * mass as f64;
973 let mut cumulative = 0.0f64;
974 let mut cutoff = 0.0f32;
975 for &w in &weights {
976 if cumulative >= target {
977 break;
978 }
979 cumulative += w as f64;
980 cutoff = w;
981 }
982 cutoff
983 }
984 _ => 0.0,
985 };
986
987 for &(dim_id, weight) in entries {
988 if weight.abs() < weight_threshold || weight.abs() < mass_cutoff {
990 continue;
991 }
992
993 let is_new_dim = !builder.postings.contains_key(&dim_id);
994 builder.add(dim_id, doc_id, ordinal, weight);
995 self.estimated_memory += size_of::<(DocId, u16, f32)>();
996 if is_new_dim {
997 self.estimated_memory += size_of::<u32>() + size_of::<Vec<(DocId, u16, f32)>>() + 8; }
1000 }
1001
1002 Ok(())
1003 }
1004
1005 fn write_document_to_store(&mut self, doc: &Document) -> Result<()> {
1007 use byteorder::{LittleEndian, WriteBytesExt};
1008
1009 super::store::serialize_document_into(doc, &self.schema, &mut self.doc_serialize_buffer)?;
1010
1011 #[cfg(feature = "native")]
1012 {
1013 self.store_file
1014 .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
1015 self.store_file.write_all(&self.doc_serialize_buffer)?;
1016 }
1017 #[cfg(not(feature = "native"))]
1018 {
1019 self.store_buffer
1020 .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
1021 self.store_buffer.write_all(&self.doc_serialize_buffer)?;
1022 }
1023
1024 Ok(())
1025 }
1026
1027 pub async fn build<D: Directory + DirectoryWriter>(
1033 mut self,
1034 dir: &D,
1035 segment_id: SegmentId,
1036 trained: Option<&super::TrainedVectorStructures>,
1037 ) -> Result<SegmentMeta> {
1038 #[cfg(feature = "native")]
1040 self.store_file.flush()?;
1041
1042 let files = SegmentFiles::new(segment_id.0);
1043
1044 let position_index = std::mem::take(&mut self.position_index);
1046 let position_offsets = if !position_index.is_empty() {
1047 let mut pos_writer = dir.streaming_writer(&files.positions).await?;
1048 let offsets = postings::build_positions_streaming(
1049 position_index,
1050 &self.term_interner,
1051 &mut *pos_writer,
1052 )?;
1053 pos_writer.finish()?;
1054 offsets
1055 } else {
1056 FxHashMap::default()
1057 };
1058
1059 let inverted_index = std::mem::take(&mut self.inverted_index);
1062 let term_interner = std::mem::replace(&mut self.term_interner, Rodeo::new());
1063 #[cfg(feature = "native")]
1064 let store_path = self.store_path.clone();
1065 #[cfg(feature = "native")]
1066 let num_compression_threads = self.config.num_compression_threads;
1067 let compression_level = self.config.compression_level;
1068 let dense_vectors = std::mem::take(&mut self.dense_vectors);
1069 let binary_dense_vectors = std::mem::take(&mut self.binary_dense_vectors);
1070 let mut sparse_vectors = std::mem::take(&mut self.sparse_vectors);
1071 let schema = &self.schema;
1072
1073 let mut term_dict_writer =
1076 super::OffsetWriter::new(dir.streaming_writer(&files.term_dict).await?);
1077 let mut postings_writer =
1078 super::OffsetWriter::new(dir.streaming_writer(&files.postings).await?);
1079 let mut store_writer = super::OffsetWriter::new(dir.streaming_writer(&files.store).await?);
1080 let mut vectors_writer = if !dense_vectors.is_empty() || !binary_dense_vectors.is_empty() {
1081 Some(super::OffsetWriter::new(
1082 dir.streaming_writer(&files.vectors).await?,
1083 ))
1084 } else {
1085 None
1086 };
1087 let mut sparse_writer = if !sparse_vectors.is_empty() {
1088 Some(super::OffsetWriter::new(
1089 dir.streaming_writer(&files.sparse).await?,
1090 ))
1091 } else {
1092 None
1093 };
1094 let mut fast_fields = std::mem::take(&mut self.fast_fields);
1095 let num_docs = self.next_doc_id;
1096 let mut fast_writer = if !fast_fields.is_empty() {
1097 Some(super::OffsetWriter::new(
1098 dir.streaming_writer(&files.fast).await?,
1099 ))
1100 } else {
1101 None
1102 };
1103
1104 #[cfg(feature = "native")]
1105 {
1106 if let Some(ref mut f) = self.posting_spill_file {
1107 f.flush()?;
1108 }
1109 let posting_spill_index = std::mem::take(&mut self.posting_spill_index);
1110 let mut spill_reader_opt = if !posting_spill_index.is_empty() {
1111 let spill_file = std::fs::File::open(&self.posting_spill_path)?;
1112 Some((std::io::BufReader::new(spill_file), posting_spill_index))
1113 } else {
1114 None
1115 };
1116
1117 let ((postings_result, store_result), ((vectors_result, sparse_result), fast_result)) =
1118 rayon::join(
1119 || {
1120 rayon::join(
1121 || {
1122 let spill_arg = spill_reader_opt.as_mut().map(|(r, idx)| {
1123 (
1124 r as &mut std::io::BufReader<std::fs::File>,
1125 idx as &postings::SpillIndex,
1126 )
1127 });
1128 postings::build_postings_streaming(
1129 inverted_index,
1130 term_interner,
1131 &position_offsets,
1132 &mut term_dict_writer,
1133 &mut postings_writer,
1134 spill_arg,
1135 )
1136 },
1137 || {
1138 store::build_store_streaming(
1139 &store_path,
1140 num_compression_threads,
1141 compression_level,
1142 &mut store_writer,
1143 num_docs,
1144 )
1145 },
1146 )
1147 },
1148 || {
1149 rayon::join(
1150 || {
1151 rayon::join(
1152 || -> Result<()> {
1153 if let Some(ref mut w) = vectors_writer {
1154 dense::build_vectors_streaming(
1155 dense_vectors,
1156 binary_dense_vectors,
1157 schema,
1158 trained,
1159 w,
1160 )?;
1161 }
1162 Ok(())
1163 },
1164 || -> Result<()> {
1165 if let Some(ref mut w) = sparse_writer {
1166 sparse::build_sparse_streaming(
1167 &mut sparse_vectors,
1168 schema,
1169 w,
1170 )?;
1171 }
1172 Ok(())
1173 },
1174 )
1175 },
1176 || -> Result<()> {
1177 if let Some(ref mut w) = fast_writer {
1178 build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
1179 }
1180 Ok(())
1181 },
1182 )
1183 },
1184 );
1185 postings_result?;
1186 store_result?;
1187 vectors_result?;
1188 sparse_result?;
1189 fast_result?;
1190 }
1191
1192 #[cfg(not(feature = "native"))]
1193 {
1194 postings::build_postings_streaming(
1195 inverted_index,
1196 term_interner,
1197 &position_offsets,
1198 &mut term_dict_writer,
1199 &mut postings_writer,
1200 )?;
1201 store::build_store_streaming_from_buffer(
1202 &self.store_buffer,
1203 compression_level,
1204 &mut store_writer,
1205 num_docs,
1206 )?;
1207 if let Some(ref mut w) = vectors_writer {
1208 dense::build_vectors_streaming(
1209 dense_vectors,
1210 binary_dense_vectors,
1211 schema,
1212 trained,
1213 w,
1214 )?;
1215 }
1216 if let Some(ref mut w) = sparse_writer {
1217 sparse::build_sparse_streaming(&mut sparse_vectors, schema, w)?;
1218 }
1219 if let Some(ref mut w) = fast_writer {
1220 build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
1221 }
1222 }
1223
1224 let term_dict_bytes = term_dict_writer.offset() as usize;
1225 let postings_bytes = postings_writer.offset() as usize;
1226 let store_bytes = store_writer.offset() as usize;
1227 let vectors_bytes = vectors_writer.as_ref().map_or(0, |w| w.offset() as usize);
1228 let sparse_bytes = sparse_writer.as_ref().map_or(0, |w| w.offset() as usize);
1229 let fast_bytes = fast_writer.as_ref().map_or(0, |w| w.offset() as usize);
1230
1231 term_dict_writer.finish()?;
1232 postings_writer.finish()?;
1233 store_writer.finish()?;
1234 if let Some(w) = vectors_writer {
1235 w.finish()?;
1236 }
1237 if let Some(w) = sparse_writer {
1238 w.finish()?;
1239 }
1240 if let Some(w) = fast_writer {
1241 w.finish()?;
1242 }
1243 drop(position_offsets);
1244 drop(sparse_vectors);
1245
1246 log::info!(
1247 "[segment_build] {} docs: term_dict={}, postings={}, store={}, vectors={}, sparse={}, fast={}",
1248 num_docs,
1249 super::format_bytes(term_dict_bytes),
1250 super::format_bytes(postings_bytes),
1251 super::format_bytes(store_bytes),
1252 super::format_bytes(vectors_bytes),
1253 super::format_bytes(sparse_bytes),
1254 super::format_bytes(fast_bytes),
1255 );
1256
1257 let meta = SegmentMeta {
1258 id: segment_id.0,
1259 num_docs: self.next_doc_id,
1260 field_stats: self.field_stats.clone(),
1261 };
1262
1263 dir.write(&files.meta, &meta.serialize()?).await?;
1264
1265 #[cfg(feature = "native")]
1267 {
1268 let _ = std::fs::remove_file(&self.store_path);
1269 }
1270
1271 Ok(meta)
1272 }
1273}
1274
1275fn build_fast_fields_streaming(
1277 fast_fields: &mut FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
1278 num_docs: u32,
1279 writer: &mut dyn Write,
1280) -> Result<()> {
1281 use crate::structures::fast_field::{FastFieldTocEntry, write_fast_field_toc_and_footer};
1282
1283 if fast_fields.is_empty() {
1284 return Ok(());
1285 }
1286
1287 let mut field_ids: Vec<u32> = fast_fields.keys().copied().collect();
1289 field_ids.sort_unstable();
1290
1291 let mut toc_entries: Vec<FastFieldTocEntry> = Vec::with_capacity(field_ids.len());
1292 let mut current_offset = 0u64;
1293
1294 for &field_id in &field_ids {
1295 let ff = fast_fields.get_mut(&field_id).unwrap();
1296 ff.pad_to(num_docs);
1297
1298 let (mut toc, bytes_written) = ff.serialize(writer, current_offset)?;
1299 toc.field_id = field_id;
1300 current_offset += bytes_written;
1301 toc_entries.push(toc);
1302 }
1303
1304 let toc_offset = current_offset;
1306 write_fast_field_toc_and_footer(writer, toc_offset, &toc_entries)?;
1307
1308 Ok(())
1309}
1310
1311#[cfg(feature = "native")]
1312impl Drop for SegmentBuilder {
1313 fn drop(&mut self) {
1314 let _ = std::fs::remove_file(&self.store_path);
1315 if self.posting_spill_file.is_some() {
1316 let _ = std::fs::remove_file(&self.posting_spill_path);
1317 }
1318 }
1319}