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        let expected_dim = self
827            .schema
828            .get_field_entry(field)
829            .and_then(|entry| entry.dense_vector_config.as_ref())
830            .map(|config| config.dim)
831            .ok_or_else(|| crate::Error::Schema("DenseVector field missing config".to_string()))?;
832        if dim != expected_dim {
833            return Err(crate::Error::Schema(format!(
834                "Dense vector dimension mismatch: schema expects {}, got {}",
835                expected_dim, dim
836            )));
837        }
838
839        let builder = self
840            .dense_vectors
841            .entry(field.0)
842            .or_insert_with(|| DenseVectorBuilder::new(dim));
843
844        // Verify dimension consistency
845        if builder.dim != dim && builder.len() > 0 {
846            return Err(crate::Error::Schema(format!(
847                "Dense vector dimension mismatch: expected {}, got {}",
848                builder.dim, dim
849            )));
850        }
851
852        builder.add(doc_id, ordinal, vector);
853
854        self.estimated_memory += std::mem::size_of_val(vector) + size_of::<(DocId, u16)>();
855
856        Ok(())
857    }
858
859    /// Index a binary dense vector field with ordinal tracking
860    fn index_binary_dense_vector_field(
861        &mut self,
862        field: Field,
863        doc_id: DocId,
864        ordinal: u16,
865        bytes: &[u8],
866    ) -> Result<()> {
867        let dim_bits = self
868            .schema
869            .get_field_entry(field)
870            .and_then(|e| e.binary_dense_vector_config.as_ref())
871            .map(|c| c.dim)
872            .ok_or_else(|| {
873                crate::Error::Schema("BinaryDenseVector field missing config".to_string())
874            })?;
875
876        let expected_byte_len = dim_bits.div_ceil(8);
877        if bytes.len() != expected_byte_len {
878            return Err(crate::Error::Schema(format!(
879                "Binary vector byte length mismatch: expected {} (dim={}), got {}",
880                expected_byte_len,
881                dim_bits,
882                bytes.len()
883            )));
884        }
885
886        let builder = self
887            .binary_dense_vectors
888            .entry(field.0)
889            .or_insert_with(|| BinaryDenseVectorBuilder::new(dim_bits));
890
891        builder.add(doc_id, ordinal, bytes);
892        self.estimated_memory += bytes.len() + size_of::<(DocId, u16)>();
893
894        Ok(())
895    }
896
897    /// Index a sparse vector field using dedicated sparse posting lists
898    ///
899    /// Collects (doc_id, ordinal, weight) postings per dimension. During commit, these are
900    /// converted to BlockSparsePostingList with proper quantization from SparseVectorConfig.
901    ///
902    /// Weights below the configured `weight_threshold` are not indexed. When
903    /// `doc_mass` is configured, only the top-|weight| entries covering that
904    /// fraction of the vector's total |weight| mass are kept (the excessive
905    /// tail of SPLADE-style vectors is cropped).
906    fn index_sparse_vector_field(
907        &mut self,
908        field: Field,
909        doc_id: DocId,
910        ordinal: u16,
911        entries: &[(u32, f32)],
912    ) -> Result<()> {
913        let (weight_threshold, doc_mass, min_terms) = self
914            .schema
915            .get_field_entry(field)
916            .and_then(|entry| entry.sparse_vector_config.as_ref())
917            .map(|config| (config.weight_threshold, config.doc_mass, config.min_terms))
918            .unwrap_or((0.0, None, 0));
919
920        let builder = self
921            .sparse_vectors
922            .entry(field.0)
923            .or_insert_with(SparseVectorBuilder::new);
924
925        builder.inc_vector_count();
926
927        // Document-side mass cropping: determine the per-vector weight cutoff
928        // below which entries fall outside the doc_mass fraction of total mass.
929        // Short vectors (<= min_terms entries) are never cropped.
930        let mass_cutoff = match doc_mass {
931            Some(mass) if mass < 1.0 && entries.len() > min_terms => {
932                let mut weights: Vec<f32> = entries
933                    .iter()
934                    .map(|&(_, w)| w.abs())
935                    .filter(|w| *w >= weight_threshold)
936                    .collect();
937                weights.sort_unstable_by(|a, b| b.total_cmp(a));
938                let total: f64 = weights.iter().map(|&w| w as f64).sum();
939                let target = total * mass as f64;
940                let mut cumulative = 0.0f64;
941                let mut cutoff = 0.0f32;
942                for &w in &weights {
943                    if cumulative >= target {
944                        break;
945                    }
946                    cumulative += w as f64;
947                    cutoff = w;
948                }
949                cutoff
950            }
951            _ => 0.0,
952        };
953
954        for &(dim_id, weight) in entries {
955            // Skip weights below threshold or outside the doc_mass prefix
956            if weight.abs() < weight_threshold || weight.abs() < mass_cutoff {
957                continue;
958            }
959
960            let is_new_dim = !builder.postings.contains_key(&dim_id);
961            builder.add(dim_id, doc_id, ordinal, weight);
962            self.estimated_memory += size_of::<(DocId, u16, f32)>();
963            if is_new_dim {
964                // HashMap entry overhead + Vec header
965                self.estimated_memory += size_of::<u32>() + size_of::<Vec<(DocId, u16, f32)>>() + 8; // 8 = hashmap control byte + padding
966            }
967        }
968
969        Ok(())
970    }
971
972    /// Write document to streaming store (reuses internal buffer to avoid per-doc allocation)
973    fn write_document_to_store(&mut self, doc: &Document) -> Result<()> {
974        use byteorder::{LittleEndian, WriteBytesExt};
975
976        super::store::serialize_document_into(doc, &self.schema, &mut self.doc_serialize_buffer)?;
977
978        #[cfg(feature = "native")]
979        {
980            self.store_file
981                .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
982            self.store_file.write_all(&self.doc_serialize_buffer)?;
983        }
984        #[cfg(not(feature = "native"))]
985        {
986            self.store_buffer
987                .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
988            self.store_buffer.write_all(&self.doc_serialize_buffer)?;
989        }
990
991        Ok(())
992    }
993
994    /// Build the final segment
995    ///
996    /// Streams all data directly to disk via StreamingWriter to avoid buffering
997    /// entire serialized outputs in memory. Each phase consumes and drops its
998    /// source data before the next phase begins.
999    pub async fn build<D: Directory + DirectoryWriter>(
1000        mut self,
1001        dir: &D,
1002        segment_id: SegmentId,
1003        trained: Option<&super::TrainedVectorStructures>,
1004    ) -> Result<SegmentMeta> {
1005        // Flush any buffered data
1006        #[cfg(feature = "native")]
1007        self.store_file.flush()?;
1008
1009        let files = SegmentFiles::new(segment_id.0);
1010
1011        // Phase 1: Stream positions directly to disk (consumes position_index)
1012        let position_index = std::mem::take(&mut self.position_index);
1013        let position_offsets = if !position_index.is_empty() {
1014            let mut pos_writer = dir.streaming_writer(&files.positions).await?;
1015            let offsets = postings::build_positions_streaming(
1016                position_index,
1017                &self.term_interner,
1018                &mut *pos_writer,
1019            )?;
1020            pos_writer.finish()?;
1021            offsets
1022        } else {
1023            FxHashMap::default()
1024        };
1025
1026        // Phase 2: 4-way parallel build — postings, store, dense vectors, sparse vectors
1027        // These are fully independent: different source data, different output files.
1028        let inverted_index = std::mem::take(&mut self.inverted_index);
1029        let term_interner = std::mem::replace(&mut self.term_interner, Rodeo::new());
1030        #[cfg(feature = "native")]
1031        let store_path = self.store_path.clone();
1032        #[cfg(feature = "native")]
1033        let num_compression_threads = self.config.num_compression_threads;
1034        let compression_level = self.config.compression_level;
1035        let dense_vectors = std::mem::take(&mut self.dense_vectors);
1036        let binary_dense_vectors = std::mem::take(&mut self.binary_dense_vectors);
1037        let mut sparse_vectors = std::mem::take(&mut self.sparse_vectors);
1038        let schema = &self.schema;
1039
1040        // Pre-create all streaming writers (async) before entering sync rayon scope
1041        // Wrapped in OffsetWriter to track bytes written per phase.
1042        let mut term_dict_writer =
1043            super::OffsetWriter::new(dir.streaming_writer(&files.term_dict).await?);
1044        let mut postings_writer =
1045            super::OffsetWriter::new(dir.streaming_writer(&files.postings).await?);
1046        let mut store_writer = super::OffsetWriter::new(dir.streaming_writer(&files.store).await?);
1047        let mut vectors_writer = if !dense_vectors.is_empty() || !binary_dense_vectors.is_empty() {
1048            Some(super::OffsetWriter::new(
1049                dir.streaming_writer(&files.vectors).await?,
1050            ))
1051        } else {
1052            None
1053        };
1054        let mut sparse_writer = if !sparse_vectors.is_empty() {
1055            Some(super::OffsetWriter::new(
1056                dir.streaming_writer(&files.sparse).await?,
1057            ))
1058        } else {
1059            None
1060        };
1061        let mut fast_fields = std::mem::take(&mut self.fast_fields);
1062        let num_docs = self.next_doc_id;
1063        let mut fast_writer = if !fast_fields.is_empty() {
1064            Some(super::OffsetWriter::new(
1065                dir.streaming_writer(&files.fast).await?,
1066            ))
1067        } else {
1068            None
1069        };
1070
1071        #[cfg(feature = "native")]
1072        {
1073            if let Some(ref mut f) = self.posting_spill_file {
1074                f.flush()?;
1075            }
1076            let posting_spill_index = std::mem::take(&mut self.posting_spill_index);
1077            let mut spill_reader_opt = if !posting_spill_index.is_empty() {
1078                let spill_file = std::fs::File::open(&self.posting_spill_path)?;
1079                Some((std::io::BufReader::new(spill_file), posting_spill_index))
1080            } else {
1081                None
1082            };
1083
1084            let ((postings_result, store_result), ((vectors_result, sparse_result), fast_result)) =
1085                rayon::join(
1086                    || {
1087                        rayon::join(
1088                            || {
1089                                let spill_arg = spill_reader_opt.as_mut().map(|(r, idx)| {
1090                                    (
1091                                        r as &mut std::io::BufReader<std::fs::File>,
1092                                        idx as &postings::SpillIndex,
1093                                    )
1094                                });
1095                                postings::build_postings_streaming(
1096                                    inverted_index,
1097                                    term_interner,
1098                                    &position_offsets,
1099                                    &mut term_dict_writer,
1100                                    &mut postings_writer,
1101                                    spill_arg,
1102                                )
1103                            },
1104                            || {
1105                                store::build_store_streaming(
1106                                    &store_path,
1107                                    num_compression_threads,
1108                                    compression_level,
1109                                    &mut store_writer,
1110                                    num_docs,
1111                                )
1112                            },
1113                        )
1114                    },
1115                    || {
1116                        rayon::join(
1117                            || {
1118                                rayon::join(
1119                                    || -> Result<()> {
1120                                        if let Some(ref mut w) = vectors_writer {
1121                                            dense::build_vectors_streaming(
1122                                                dense_vectors,
1123                                                binary_dense_vectors,
1124                                                schema,
1125                                                trained,
1126                                                w,
1127                                            )?;
1128                                        }
1129                                        Ok(())
1130                                    },
1131                                    || -> Result<()> {
1132                                        if let Some(ref mut w) = sparse_writer {
1133                                            sparse::build_sparse_streaming(
1134                                                &mut sparse_vectors,
1135                                                schema,
1136                                                w,
1137                                            )?;
1138                                        }
1139                                        Ok(())
1140                                    },
1141                                )
1142                            },
1143                            || -> Result<()> {
1144                                if let Some(ref mut w) = fast_writer {
1145                                    build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
1146                                }
1147                                Ok(())
1148                            },
1149                        )
1150                    },
1151                );
1152            postings_result?;
1153            store_result?;
1154            vectors_result?;
1155            sparse_result?;
1156            fast_result?;
1157        }
1158
1159        #[cfg(not(feature = "native"))]
1160        {
1161            postings::build_postings_streaming(
1162                inverted_index,
1163                term_interner,
1164                &position_offsets,
1165                &mut term_dict_writer,
1166                &mut postings_writer,
1167            )?;
1168            store::build_store_streaming_from_buffer(
1169                &self.store_buffer,
1170                compression_level,
1171                &mut store_writer,
1172                num_docs,
1173            )?;
1174            if let Some(ref mut w) = vectors_writer {
1175                dense::build_vectors_streaming(
1176                    dense_vectors,
1177                    binary_dense_vectors,
1178                    schema,
1179                    trained,
1180                    w,
1181                )?;
1182            }
1183            if let Some(ref mut w) = sparse_writer {
1184                sparse::build_sparse_streaming(&mut sparse_vectors, schema, w)?;
1185            }
1186            if let Some(ref mut w) = fast_writer {
1187                build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
1188            }
1189        }
1190
1191        let term_dict_bytes = term_dict_writer.offset() as usize;
1192        let postings_bytes = postings_writer.offset() as usize;
1193        let store_bytes = store_writer.offset() as usize;
1194        let vectors_bytes = vectors_writer.as_ref().map_or(0, |w| w.offset() as usize);
1195        let sparse_bytes = sparse_writer.as_ref().map_or(0, |w| w.offset() as usize);
1196        let fast_bytes = fast_writer.as_ref().map_or(0, |w| w.offset() as usize);
1197
1198        term_dict_writer.finish()?;
1199        postings_writer.finish()?;
1200        store_writer.finish()?;
1201        if let Some(w) = vectors_writer {
1202            w.finish()?;
1203        }
1204        if let Some(w) = sparse_writer {
1205            w.finish()?;
1206        }
1207        if let Some(w) = fast_writer {
1208            w.finish()?;
1209        }
1210        drop(position_offsets);
1211        drop(sparse_vectors);
1212
1213        log::info!(
1214            "[segment_build] {} docs: term_dict={}, postings={}, store={}, vectors={}, sparse={}, fast={}",
1215            num_docs,
1216            super::format_bytes(term_dict_bytes),
1217            super::format_bytes(postings_bytes),
1218            super::format_bytes(store_bytes),
1219            super::format_bytes(vectors_bytes),
1220            super::format_bytes(sparse_bytes),
1221            super::format_bytes(fast_bytes),
1222        );
1223
1224        let meta = SegmentMeta {
1225            id: segment_id.0,
1226            num_docs: self.next_doc_id,
1227            field_stats: self.field_stats.clone(),
1228        };
1229
1230        dir.write(&files.meta, &meta.serialize()?).await?;
1231
1232        // Cleanup temp files
1233        #[cfg(feature = "native")]
1234        {
1235            let _ = std::fs::remove_file(&self.store_path);
1236        }
1237
1238        Ok(meta)
1239    }
1240}
1241
1242/// Serialize all fast-field columns to a `.fast` file.
1243fn build_fast_fields_streaming(
1244    fast_fields: &mut FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
1245    num_docs: u32,
1246    writer: &mut dyn Write,
1247) -> Result<()> {
1248    use crate::structures::fast_field::{FastFieldTocEntry, write_fast_field_toc_and_footer};
1249
1250    if fast_fields.is_empty() {
1251        return Ok(());
1252    }
1253
1254    // Sort fields by id for deterministic output
1255    let mut field_ids: Vec<u32> = fast_fields.keys().copied().collect();
1256    field_ids.sort_unstable();
1257
1258    let mut toc_entries: Vec<FastFieldTocEntry> = Vec::with_capacity(field_ids.len());
1259    let mut current_offset = 0u64;
1260
1261    for &field_id in &field_ids {
1262        let ff = fast_fields.get_mut(&field_id).unwrap();
1263        ff.pad_to(num_docs);
1264
1265        let (mut toc, bytes_written) = ff.serialize(writer, current_offset)?;
1266        toc.field_id = field_id;
1267        current_offset += bytes_written;
1268        toc_entries.push(toc);
1269    }
1270
1271    // Write TOC + footer
1272    let toc_offset = current_offset;
1273    write_fast_field_toc_and_footer(writer, toc_offset, &toc_entries)?;
1274
1275    Ok(())
1276}
1277
1278#[cfg(feature = "native")]
1279impl Drop for SegmentBuilder {
1280    fn drop(&mut self) {
1281        let _ = std::fs::remove_file(&self.store_path);
1282        if self.posting_spill_file.is_some() {
1283            let _ = std::fs::remove_file(&self.posting_spill_path);
1284        }
1285    }
1286}