Skip to main content

hermes_core/segment/builder/
mod.rs

1//! Streaming segment builder with optimized memory usage
2//!
3//! Key optimizations:
4//! - **String interning**: Terms are interned using `lasso` to avoid repeated allocations
5//! - **hashbrown HashMap**: O(1) average insertion instead of BTreeMap's O(log n)
6//! - **Streaming document store**: Documents written to disk immediately
7//! - **Zero-copy store build**: Pre-serialized doc bytes passed directly to compressor
8//! - **Parallel posting serialization**: Rayon parallel sort + serialize
9//! - **Inline posting fast path**: Small terms skip PostingList/BlockPostingList entirely
10
11#[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// String interning: lasso on native (fast arena), HashMap on WASM (no C deps)
39#[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    /// Simple string interner for WASM (replaces lasso::Rodeo).
50    /// Stores each string once in a Vec; HashMap maps &str → index.
51    pub struct Rodeo {
52        /// Canonical storage — each string lives here exactly once.
53        strings: Vec<Box<str>>,
54        /// Maps borrowed string slices (pointing into `strings`) to their index.
55        /// Safety: entries are never removed and Box<str> has a stable address.
56        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            // Safety: the Box<str> is stored in self.strings (append-only Vec)
78            // and never moved or freed while the Rodeo is alive.
79            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
110/// Size of the document store buffer before writing to disk
111const STORE_BUFFER_SIZE: usize = 16 * 1024 * 1024; // 16MB
112
113/// Memory overhead per new term in the inverted index:
114/// HashMap entry control byte + padding + TermKey + PostingListBuilder + Vec header
115const NEW_TERM_OVERHEAD: usize = size_of::<TermKey>() + size_of::<PostingListBuilder>() + 24;
116
117/// Memory overhead per newly interned string: Spur + arena pointers (2 × usize)
118const INTERN_OVERHEAD: usize = size_of::<Spur>() + 2 * size_of::<usize>();
119
120/// Memory overhead per new term in the position index
121const NEW_POS_TERM_OVERHEAD: usize =
122    size_of::<TermKey>() + size_of::<PositionPostingListBuilder>() + 24;
123
124/// Segment builder with optimized memory usage
125///
126/// Features:
127/// - Streams documents to disk immediately (no in-memory document storage)
128/// - Uses string interning for terms (reduced allocations)
129/// - Uses hashbrown HashMap (faster than BTreeMap)
130pub struct SegmentBuilder {
131    schema: Arc<Schema>,
132    config: SegmentBuilderConfig,
133    tokenizers: FxHashMap<Field, BoxedTokenizer>,
134
135    /// String interner for terms - O(1) lookup and deduplication
136    term_interner: Rodeo,
137
138    /// Inverted index: term key -> posting list
139    inverted_index: HashMap<TermKey, PostingListBuilder>,
140
141    /// Spill file for high-frequency posting lists (lazily created on first spill).
142    #[cfg(feature = "native")]
143    posting_spill_file: Option<BufWriter<File>>,
144    #[cfg(feature = "native")]
145    posting_spill_path: PathBuf,
146    /// Tracks spilled ranges per term key: (file_offset, posting_count).
147    #[cfg(feature = "native")]
148    posting_spill_index: HashMap<TermKey, Vec<(u64, u32)>>,
149    #[cfg(feature = "native")]
150    posting_spill_offset: u64,
151
152    /// Streaming document store writer (native: temp file on disk, WASM: in-memory buffer)
153    #[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    /// Document count
161    next_doc_id: DocId,
162
163    /// Per-field statistics for BM25F
164    field_stats: FxHashMap<u32, FieldStats>,
165
166    /// Per-document field lengths stored compactly
167    /// Uses a flat Vec instead of Vec<HashMap> for better cache locality
168    /// Layout: [doc0_field0_len, doc0_field1_len, ..., doc1_field0_len, ...]
169    doc_field_lengths: Vec<u32>,
170    num_indexed_fields: usize,
171    field_to_slot: FxHashMap<u32, usize>,
172
173    /// Reusable buffer for per-document term frequency aggregation
174    /// Avoids allocating a new hashmap for each document
175    local_tf_buffer: FxHashMap<Spur, u32>,
176
177    /// Reusable buffer for per-document position tracking (when positions enabled)
178    /// Avoids allocating a new hashmap for each text field per document
179    local_positions: FxHashMap<Spur, Vec<u32>>,
180
181    /// Reusable buffer for tokenization to avoid per-token String allocations
182    token_buffer: String,
183
184    /// Reusable buffer for numeric field term encoding (avoids format!() alloc per call)
185    numeric_buffer: String,
186
187    /// Dense vector storage per field: field -> (doc_ids, vectors)
188    /// Vectors are stored as flat f32 arrays for efficient RaBitQ indexing
189    dense_vectors: FxHashMap<u32, DenseVectorBuilder>,
190
191    /// Binary dense vector storage per field: field -> packed-bit vectors
192    binary_dense_vectors: FxHashMap<u32, BinaryDenseVectorBuilder>,
193
194    /// Sparse vector storage per field: field -> SparseVectorBuilder
195    /// Uses proper BlockSparsePostingList with configurable quantization
196    sparse_vectors: FxHashMap<u32, SparseVectorBuilder>,
197
198    /// Position index for fields with positions enabled
199    /// term key -> position posting list
200    position_index: HashMap<TermKey, PositionPostingListBuilder>,
201
202    /// Fields that have position tracking enabled, with their mode
203    position_enabled_fields: FxHashMap<u32, Option<crate::dsl::PositionMode>>,
204
205    /// Current element ordinal for multi-valued fields (reset per document)
206    current_element_ordinal: FxHashMap<u32, u32>,
207
208    /// Incrementally tracked memory estimate (avoids expensive stats() calls)
209    estimated_memory: usize,
210
211    /// Reusable buffer for document serialization (avoids per-document allocation)
212    doc_serialize_buffer: Vec<u8>,
213
214    /// Fast-field columnar writers per field_id (only for fields with fast=true)
215    fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
216}
217
218impl SegmentBuilder {
219    /// Create a new segment builder
220    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        // Count indexed fields, track positions, and auto-configure tokenizers
242        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        // Initialize fast-field writers for fields with fast=true
263        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    /// Get the current element ordinal for a field and increment it.
340    /// Used for multi-valued fields (text, dense_vector, sparse_vector).
341    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    /// Fast O(1) memory estimate - updated incrementally during indexing
352    #[inline]
353    pub fn estimated_memory_bytes(&self) -> usize {
354        self.estimated_memory
355    }
356
357    /// Count total unique sparse dimensions across all fields
358    pub fn sparse_dim_count(&self) -> usize {
359        self.sparse_vectors.values().map(|b| b.postings.len()).sum()
360    }
361
362    /// Get current statistics for debugging performance (expensive - iterates all data)
363    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        // Size constants computed from actual types
370        let compact_posting_size = size_of::<CompactPosting>();
371        let vec_overhead = size_of::<Vec<u8>>(); // Vec header: ptr + len + cap = 24 bytes on 64-bit
372        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        // hashbrown HashMap entry overhead: key + value + 1 byte control + padding
378        // Measured: ~(key_size + value_size + 8) per entry on average
379        let hashmap_entry_base_overhead = 8usize;
380
381        // FxHashMap uses same layout as hashbrown
382        let fxhashmap_entry_overhead = hashmap_entry_base_overhead;
383
384        // Postings memory
385        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        // Inverted index overhead
392        let index_overhead_bytes = self.inverted_index.len()
393            * (term_key_size + posting_builder_size + hashmap_entry_base_overhead);
394
395        // Term interner: Rodeo stores strings + metadata
396        // Rodeo internal: string bytes + Spur + arena overhead (~2 pointers per string)
397        let interner_arena_overhead = 2 * size_of::<usize>();
398        let avg_term_len = 8; // Estimated average term length
399        let interner_bytes =
400            self.term_interner.len() * (avg_term_len + spur_size + interner_arena_overhead);
401
402        // Doc field lengths
403        let field_lengths_bytes =
404            self.doc_field_lengths.capacity() * size_of::<u32>() + vec_overhead;
405
406        // Dense vectors
407        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; // Two Vecs
414            dense_vector_count += b.doc_ids.len();
415        }
416        // Binary dense vectors
417        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        // Local buffers
425        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        // Sparse vectors
429        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            // Inner FxHashMap overhead: u32 key + Vec value ptr + overhead
435            let inner_entry_size = size_of::<u32>() + vec_overhead + fxhashmap_entry_overhead;
436            sparse_vectors_bytes += builder.postings.len() * inner_entry_size;
437        }
438        // Outer FxHashMap overhead
439        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        // Position index
444        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            // Vec<(DocId, Vec<u32>)> entry size
450            let pos_entry_size = size_of::<DocId>() + vec_overhead;
451            position_index_bytes += pos_builder.postings.capacity() * pos_entry_size;
452        }
453        // HashMap overhead for position_index
454        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    /// Add a document - streams to disk immediately
490    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        // Initialize field lengths for this document
495        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        // Reset element ordinals for this document (for multi-valued fields)
501        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            // Dense/binary vectors are written to .vectors when indexed || stored
509            // Other field types require indexed or fast
510            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                    // Fast-field: store raw text for text ordinal column
538                    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        // Stream document to disk immediately
587        self.write_document_to_store(&doc)?;
588
589        Ok(doc_id)
590    }
591
592    /// Index a text field using interned terms
593    ///
594    /// Uses a custom tokenizer when set for the field (via `set_tokenizer`),
595    /// otherwise falls back to an inline zero-allocation path (split_whitespace
596    /// + lowercase + strip non-alphanumeric).
597    ///
598    /// If position recording is enabled for this field, also records token positions
599    /// encoded as (element_ordinal << 20) | token_position.
600    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        // Phase 1: Aggregate term frequencies within this document
617        // Also collect positions if enabled
618        // Reuse buffers to avoid allocations
619        self.local_tf_buffer.clear();
620        // Clear position Vecs in-place (keeps allocated capacity for reuse)
621        for v in self.local_positions.values_mut() {
622            v.clear();
623        }
624
625        let mut token_position = 0u32;
626
627        // Tokenize: use custom tokenizer if set, else inline zero-alloc path.
628        // The owned Vec<Token> is computed first so the immutable borrow of
629        // self.tokenizers ends before we mutate other fields.
630        let custom_tokens = self.tokenizers.get(&field).map(|t| t.tokenize(text));
631
632        if let Some(tokens) = custom_tokens {
633            // Custom tokenizer path
634            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            // Inline zero-allocation path: split_whitespace + lowercase + strip non-alphanumeric
659            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        // Phase 2: Insert aggregated terms into inverted index
699        // Now we only do one inverted_index lookup per unique term in doc
700        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                    // Spill large posting lists to disk to reduce peak memory
711                    #[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                        // Lazily create the spill file on first spill
720                        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    /// Index a dense vector field with ordinal tracking
818    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        // Verify dimension consistency
833        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    /// Index a binary dense vector field with ordinal tracking
848    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    /// Index a sparse vector field using dedicated sparse posting lists
886    ///
887    /// Collects (doc_id, ordinal, weight) postings per dimension. During commit, these are
888    /// converted to BlockSparsePostingList with proper quantization from SparseVectorConfig.
889    ///
890    /// Weights below the configured `weight_threshold` are not indexed. When
891    /// `doc_mass` is configured, only the top-|weight| entries covering that
892    /// fraction of the vector's total |weight| mass are kept (the excessive
893    /// tail of SPLADE-style vectors is cropped).
894    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        // Document-side mass cropping: determine the per-vector weight cutoff
916        // below which entries fall outside the doc_mass fraction of total mass.
917        // Short vectors (<= min_terms entries) are never cropped.
918        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            // Skip weights below threshold or outside the doc_mass prefix
944            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                // HashMap entry overhead + Vec header
953                self.estimated_memory += size_of::<u32>() + size_of::<Vec<(DocId, u16, f32)>>() + 8; // 8 = hashmap control byte + padding
954            }
955        }
956
957        Ok(())
958    }
959
960    /// Write document to streaming store (reuses internal buffer to avoid per-doc allocation)
961    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    /// Build the final segment
983    ///
984    /// Streams all data directly to disk via StreamingWriter to avoid buffering
985    /// entire serialized outputs in memory. Each phase consumes and drops its
986    /// source data before the next phase begins.
987    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        // Flush any buffered data
994        #[cfg(feature = "native")]
995        self.store_file.flush()?;
996
997        let files = SegmentFiles::new(segment_id.0);
998
999        // Phase 1: Stream positions directly to disk (consumes position_index)
1000        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        // Phase 2: 4-way parallel build — postings, store, dense vectors, sparse vectors
1015        // These are fully independent: different source data, different output files.
1016        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        // Pre-create all streaming writers (async) before entering sync rayon scope
1029        // Wrapped in OffsetWriter to track bytes written per phase.
1030        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        // Cleanup temp files
1221        #[cfg(feature = "native")]
1222        {
1223            let _ = std::fs::remove_file(&self.store_path);
1224        }
1225
1226        Ok(meta)
1227    }
1228}
1229
1230/// Serialize all fast-field columns to a `.fast` file.
1231fn 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    // Sort fields by id for deterministic output
1243    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    // Write TOC + footer
1260    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}