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
11pub(crate) mod bmp;
12mod config;
13mod dense;
14#[cfg(feature = "diagnostics")]
15mod diagnostics;
16mod postings;
17pub(crate) mod simhash;
18mod sparse;
19mod store;
20
21pub use config::{MemoryBreakdown, SegmentBuilderConfig, SegmentBuilderStats};
22
23use std::fs::{File, OpenOptions};
24use std::io::{BufWriter, Write};
25use std::mem::size_of;
26use std::path::PathBuf;
27
28use hashbrown::HashMap;
29use lasso::{Rodeo, Spur};
30use rustc_hash::FxHashMap;
31
32use super::types::{FieldStats, SegmentFiles, SegmentId, SegmentMeta};
33use std::sync::Arc;
34
35use crate::directories::{Directory, DirectoryWriter};
36use crate::dsl::{Document, Field, FieldType, FieldValue, Schema};
37use crate::tokenizer::BoxedTokenizer;
38use crate::{DocId, Result};
39
40use dense::DenseVectorBuilder;
41use postings::{CompactPosting, PositionPostingListBuilder, PostingListBuilder, TermKey};
42use sparse::SparseVectorBuilder;
43
44/// Size of the document store buffer before writing to disk
45const STORE_BUFFER_SIZE: usize = 16 * 1024 * 1024; // 16MB
46
47/// Memory overhead per new term in the inverted index:
48/// HashMap entry control byte + padding + TermKey + PostingListBuilder + Vec header
49const NEW_TERM_OVERHEAD: usize = size_of::<TermKey>() + size_of::<PostingListBuilder>() + 24;
50
51/// Memory overhead per newly interned string: Spur + arena pointers (2 × usize)
52const INTERN_OVERHEAD: usize = size_of::<Spur>() + 2 * size_of::<usize>();
53
54/// Memory overhead per new term in the position index
55const NEW_POS_TERM_OVERHEAD: usize =
56    size_of::<TermKey>() + size_of::<PositionPostingListBuilder>() + 24;
57
58/// Segment builder with optimized memory usage
59///
60/// Features:
61/// - Streams documents to disk immediately (no in-memory document storage)
62/// - Uses string interning for terms (reduced allocations)
63/// - Uses hashbrown HashMap (faster than BTreeMap)
64pub struct SegmentBuilder {
65    schema: Arc<Schema>,
66    config: SegmentBuilderConfig,
67    tokenizers: FxHashMap<Field, BoxedTokenizer>,
68
69    /// String interner for terms - O(1) lookup and deduplication
70    term_interner: Rodeo,
71
72    /// Inverted index: term key -> posting list
73    inverted_index: HashMap<TermKey, PostingListBuilder>,
74
75    /// Streaming document store writer
76    store_file: BufWriter<File>,
77    store_path: PathBuf,
78
79    /// Document count
80    next_doc_id: DocId,
81
82    /// Per-field statistics for BM25F
83    field_stats: FxHashMap<u32, FieldStats>,
84
85    /// Per-document field lengths stored compactly
86    /// Uses a flat Vec instead of Vec<HashMap> for better cache locality
87    /// Layout: [doc0_field0_len, doc0_field1_len, ..., doc1_field0_len, ...]
88    doc_field_lengths: Vec<u32>,
89    num_indexed_fields: usize,
90    field_to_slot: FxHashMap<u32, usize>,
91
92    /// Reusable buffer for per-document term frequency aggregation
93    /// Avoids allocating a new hashmap for each document
94    local_tf_buffer: FxHashMap<Spur, u32>,
95
96    /// Reusable buffer for per-document position tracking (when positions enabled)
97    /// Avoids allocating a new hashmap for each text field per document
98    local_positions: FxHashMap<Spur, Vec<u32>>,
99
100    /// Reusable buffer for tokenization to avoid per-token String allocations
101    token_buffer: String,
102
103    /// Reusable buffer for numeric field term encoding (avoids format!() alloc per call)
104    numeric_buffer: String,
105
106    /// Dense vector storage per field: field -> (doc_ids, vectors)
107    /// Vectors are stored as flat f32 arrays for efficient RaBitQ indexing
108    dense_vectors: FxHashMap<u32, DenseVectorBuilder>,
109
110    /// Sparse vector storage per field: field -> SparseVectorBuilder
111    /// Uses proper BlockSparsePostingList with configurable quantization
112    sparse_vectors: FxHashMap<u32, SparseVectorBuilder>,
113
114    /// Position index for fields with positions enabled
115    /// term key -> position posting list
116    position_index: HashMap<TermKey, PositionPostingListBuilder>,
117
118    /// Fields that have position tracking enabled, with their mode
119    position_enabled_fields: FxHashMap<u32, Option<crate::dsl::PositionMode>>,
120
121    /// Current element ordinal for multi-valued fields (reset per document)
122    current_element_ordinal: FxHashMap<u32, u32>,
123
124    /// Incrementally tracked memory estimate (avoids expensive stats() calls)
125    estimated_memory: usize,
126
127    /// Reusable buffer for document serialization (avoids per-document allocation)
128    doc_serialize_buffer: Vec<u8>,
129
130    /// Fast-field columnar writers per field_id (only for fields with fast=true)
131    fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
132
133    /// Per-ordinal SimHash for BMP block clustering: field_id → (doc_id, ordinal) → simhash
134    ordinal_simhashes: FxHashMap<u32, FxHashMap<(DocId, u16), u64>>,
135}
136
137impl SegmentBuilder {
138    /// Create a new segment builder
139    pub fn new(schema: Arc<Schema>, config: SegmentBuilderConfig) -> Result<Self> {
140        let segment_id = uuid::Uuid::new_v4();
141        let store_path = config
142            .temp_dir
143            .join(format!("hermes_store_{}.tmp", segment_id));
144
145        let store_file = BufWriter::with_capacity(
146            STORE_BUFFER_SIZE,
147            OpenOptions::new()
148                .create(true)
149                .write(true)
150                .truncate(true)
151                .open(&store_path)?,
152        );
153
154        // Count indexed fields, track positions, and auto-configure tokenizers
155        let registry = crate::tokenizer::TokenizerRegistry::new();
156        let mut num_indexed_fields = 0;
157        let mut field_to_slot = FxHashMap::default();
158        let mut position_enabled_fields = FxHashMap::default();
159        let mut tokenizers = FxHashMap::default();
160        for (field, entry) in schema.fields() {
161            if entry.indexed && matches!(entry.field_type, FieldType::Text) {
162                field_to_slot.insert(field.0, num_indexed_fields);
163                num_indexed_fields += 1;
164                if entry.positions.is_some() {
165                    position_enabled_fields.insert(field.0, entry.positions);
166                }
167                if let Some(ref tok_name) = entry.tokenizer
168                    && let Some(tokenizer) = registry.get(tok_name)
169                {
170                    tokenizers.insert(field, tokenizer);
171                }
172            }
173        }
174
175        // Initialize fast-field writers for fields with fast=true
176        use crate::structures::fast_field::{FastFieldColumnType, FastFieldWriter};
177        let mut fast_fields = FxHashMap::default();
178        for (field, entry) in schema.fields() {
179            if entry.fast {
180                let writer = if entry.multi {
181                    match entry.field_type {
182                        FieldType::U64 => {
183                            FastFieldWriter::new_numeric_multi(FastFieldColumnType::U64)
184                        }
185                        FieldType::I64 => {
186                            FastFieldWriter::new_numeric_multi(FastFieldColumnType::I64)
187                        }
188                        FieldType::F64 => {
189                            FastFieldWriter::new_numeric_multi(FastFieldColumnType::F64)
190                        }
191                        FieldType::Text => FastFieldWriter::new_text_multi(),
192                        _ => continue,
193                    }
194                } else {
195                    match entry.field_type {
196                        FieldType::U64 => FastFieldWriter::new_numeric(FastFieldColumnType::U64),
197                        FieldType::I64 => FastFieldWriter::new_numeric(FastFieldColumnType::I64),
198                        FieldType::F64 => FastFieldWriter::new_numeric(FastFieldColumnType::F64),
199                        FieldType::Text => FastFieldWriter::new_text(),
200                        _ => continue,
201                    }
202                };
203                fast_fields.insert(field.0, writer);
204            }
205        }
206
207        Ok(Self {
208            schema,
209            tokenizers,
210            term_interner: Rodeo::new(),
211            inverted_index: HashMap::with_capacity(config.posting_map_capacity),
212            store_file,
213            store_path,
214            next_doc_id: 0,
215            field_stats: FxHashMap::default(),
216            doc_field_lengths: Vec::new(),
217            num_indexed_fields,
218            field_to_slot,
219            local_tf_buffer: FxHashMap::default(),
220            local_positions: FxHashMap::default(),
221            token_buffer: String::with_capacity(64),
222            numeric_buffer: String::with_capacity(32),
223            config,
224            dense_vectors: FxHashMap::default(),
225            sparse_vectors: FxHashMap::default(),
226            position_index: HashMap::new(),
227            position_enabled_fields,
228            current_element_ordinal: FxHashMap::default(),
229            estimated_memory: 0,
230            doc_serialize_buffer: Vec::with_capacity(256),
231            fast_fields,
232            ordinal_simhashes: FxHashMap::default(),
233        })
234    }
235
236    pub fn set_tokenizer(&mut self, field: Field, tokenizer: BoxedTokenizer) {
237        self.tokenizers.insert(field, tokenizer);
238    }
239
240    /// Get the current element ordinal for a field and increment it.
241    /// Used for multi-valued fields (text, dense_vector, sparse_vector).
242    fn next_element_ordinal(&mut self, field_id: u32) -> u32 {
243        let ordinal = *self.current_element_ordinal.get(&field_id).unwrap_or(&0);
244        *self.current_element_ordinal.entry(field_id).or_insert(0) += 1;
245        ordinal
246    }
247
248    pub fn num_docs(&self) -> u32 {
249        self.next_doc_id
250    }
251
252    /// Fast O(1) memory estimate - updated incrementally during indexing
253    #[inline]
254    pub fn estimated_memory_bytes(&self) -> usize {
255        self.estimated_memory
256    }
257
258    /// Count total unique sparse dimensions across all fields
259    pub fn sparse_dim_count(&self) -> usize {
260        self.sparse_vectors.values().map(|b| b.postings.len()).sum()
261    }
262
263    /// Get current statistics for debugging performance (expensive - iterates all data)
264    pub fn stats(&self) -> SegmentBuilderStats {
265        use std::mem::size_of;
266
267        let postings_in_memory: usize =
268            self.inverted_index.values().map(|p| p.postings.len()).sum();
269
270        // Size constants computed from actual types
271        let compact_posting_size = size_of::<CompactPosting>();
272        let vec_overhead = size_of::<Vec<u8>>(); // Vec header: ptr + len + cap = 24 bytes on 64-bit
273        let term_key_size = size_of::<TermKey>();
274        let posting_builder_size = size_of::<PostingListBuilder>();
275        let spur_size = size_of::<lasso::Spur>();
276        let sparse_entry_size = size_of::<(DocId, u16, f32)>();
277
278        // hashbrown HashMap entry overhead: key + value + 1 byte control + padding
279        // Measured: ~(key_size + value_size + 8) per entry on average
280        let hashmap_entry_base_overhead = 8usize;
281
282        // FxHashMap uses same layout as hashbrown
283        let fxhashmap_entry_overhead = hashmap_entry_base_overhead;
284
285        // Postings memory
286        let postings_bytes: usize = self
287            .inverted_index
288            .values()
289            .map(|p| p.postings.capacity() * compact_posting_size + vec_overhead)
290            .sum();
291
292        // Inverted index overhead
293        let index_overhead_bytes = self.inverted_index.len()
294            * (term_key_size + posting_builder_size + hashmap_entry_base_overhead);
295
296        // Term interner: Rodeo stores strings + metadata
297        // Rodeo internal: string bytes + Spur + arena overhead (~2 pointers per string)
298        let interner_arena_overhead = 2 * size_of::<usize>();
299        let avg_term_len = 8; // Estimated average term length
300        let interner_bytes =
301            self.term_interner.len() * (avg_term_len + spur_size + interner_arena_overhead);
302
303        // Doc field lengths
304        let field_lengths_bytes =
305            self.doc_field_lengths.capacity() * size_of::<u32>() + vec_overhead;
306
307        // Dense vectors
308        let mut dense_vectors_bytes: usize = 0;
309        let mut dense_vector_count: usize = 0;
310        let doc_id_ordinal_size = size_of::<(DocId, u16)>();
311        for b in self.dense_vectors.values() {
312            dense_vectors_bytes += b.vectors.capacity() * size_of::<f32>()
313                + b.doc_ids.capacity() * doc_id_ordinal_size
314                + 2 * vec_overhead; // Two Vecs
315            dense_vector_count += b.doc_ids.len();
316        }
317
318        // Local buffers
319        let local_tf_entry_size = spur_size + size_of::<u32>() + fxhashmap_entry_overhead;
320        let local_tf_buffer_bytes = self.local_tf_buffer.capacity() * local_tf_entry_size;
321
322        // Sparse vectors
323        let mut sparse_vectors_bytes: usize = 0;
324        for builder in self.sparse_vectors.values() {
325            for postings in builder.postings.values() {
326                sparse_vectors_bytes += postings.capacity() * sparse_entry_size + vec_overhead;
327            }
328            // Inner FxHashMap overhead: u32 key + Vec value ptr + overhead
329            let inner_entry_size = size_of::<u32>() + vec_overhead + fxhashmap_entry_overhead;
330            sparse_vectors_bytes += builder.postings.len() * inner_entry_size;
331        }
332        // Outer FxHashMap overhead
333        let outer_sparse_entry_size =
334            size_of::<u32>() + size_of::<SparseVectorBuilder>() + fxhashmap_entry_overhead;
335        sparse_vectors_bytes += self.sparse_vectors.len() * outer_sparse_entry_size;
336
337        // Position index
338        let mut position_index_bytes: usize = 0;
339        for pos_builder in self.position_index.values() {
340            for (_, positions) in &pos_builder.postings {
341                position_index_bytes += positions.capacity() * size_of::<u32>() + vec_overhead;
342            }
343            // Vec<(DocId, Vec<u32>)> entry size
344            let pos_entry_size = size_of::<DocId>() + vec_overhead;
345            position_index_bytes += pos_builder.postings.capacity() * pos_entry_size;
346        }
347        // HashMap overhead for position_index
348        let pos_index_entry_size =
349            term_key_size + size_of::<PositionPostingListBuilder>() + hashmap_entry_base_overhead;
350        position_index_bytes += self.position_index.len() * pos_index_entry_size;
351
352        let estimated_memory_bytes = postings_bytes
353            + index_overhead_bytes
354            + interner_bytes
355            + field_lengths_bytes
356            + dense_vectors_bytes
357            + local_tf_buffer_bytes
358            + sparse_vectors_bytes
359            + position_index_bytes;
360
361        let memory_breakdown = MemoryBreakdown {
362            postings_bytes,
363            index_overhead_bytes,
364            interner_bytes,
365            field_lengths_bytes,
366            dense_vectors_bytes,
367            dense_vector_count,
368            sparse_vectors_bytes,
369            position_index_bytes,
370        };
371
372        SegmentBuilderStats {
373            num_docs: self.next_doc_id,
374            unique_terms: self.inverted_index.len(),
375            postings_in_memory,
376            interned_strings: self.term_interner.len(),
377            doc_field_lengths_size: self.doc_field_lengths.len(),
378            estimated_memory_bytes,
379            memory_breakdown,
380        }
381    }
382
383    /// Add a document - streams to disk immediately
384    pub fn add_document(&mut self, doc: Document) -> Result<DocId> {
385        let doc_id = self.next_doc_id;
386        self.next_doc_id += 1;
387
388        // Initialize field lengths for this document
389        let base_idx = self.doc_field_lengths.len();
390        self.doc_field_lengths
391            .resize(base_idx + self.num_indexed_fields, 0);
392        self.estimated_memory += self.num_indexed_fields * std::mem::size_of::<u32>();
393
394        // Reset element ordinals for this document (for multi-valued fields)
395        self.current_element_ordinal.clear();
396
397        for (field, value) in doc.field_values() {
398            let Some(entry) = self.schema.get_field_entry(*field) else {
399                continue;
400            };
401
402            // Dense vectors are written to .vectors when indexed || stored
403            // Other field types require indexed or fast
404            if !matches!(&entry.field_type, FieldType::DenseVector) && !entry.indexed && !entry.fast
405            {
406                continue;
407            }
408
409            match (&entry.field_type, value) {
410                (FieldType::Text, FieldValue::Text(text)) => {
411                    if entry.indexed {
412                        let element_ordinal = self.next_element_ordinal(field.0);
413                        let token_count =
414                            self.index_text_field(*field, doc_id, text, element_ordinal)?;
415
416                        let stats = self.field_stats.entry(field.0).or_default();
417                        stats.total_tokens += token_count as u64;
418                        if element_ordinal == 0 {
419                            stats.doc_count += 1;
420                        }
421
422                        if let Some(&slot) = self.field_to_slot.get(&field.0) {
423                            self.doc_field_lengths[base_idx + slot] = token_count;
424                        }
425                    }
426
427                    // Fast-field: store raw text for text ordinal column
428                    if let Some(ff) = self.fast_fields.get_mut(&field.0) {
429                        ff.add_text(doc_id, text);
430                    }
431                }
432                (FieldType::U64, FieldValue::U64(v)) => {
433                    if entry.indexed {
434                        self.index_numeric_field(*field, doc_id, *v)?;
435                    }
436                    if let Some(ff) = self.fast_fields.get_mut(&field.0) {
437                        ff.add_u64(doc_id, *v);
438                    }
439                }
440                (FieldType::I64, FieldValue::I64(v)) => {
441                    if entry.indexed {
442                        self.index_numeric_field(*field, doc_id, *v as u64)?;
443                    }
444                    if let Some(ff) = self.fast_fields.get_mut(&field.0) {
445                        ff.add_i64(doc_id, *v);
446                    }
447                }
448                (FieldType::F64, FieldValue::F64(v)) => {
449                    if entry.indexed {
450                        self.index_numeric_field(*field, doc_id, v.to_bits())?;
451                    }
452                    if let Some(ff) = self.fast_fields.get_mut(&field.0) {
453                        ff.add_f64(doc_id, *v);
454                    }
455                }
456                (FieldType::DenseVector, FieldValue::DenseVector(vec))
457                    if entry.indexed || entry.stored =>
458                {
459                    let ordinal = self.next_element_ordinal(field.0);
460                    self.index_dense_vector_field(*field, doc_id, ordinal as u16, vec)?;
461                }
462                (FieldType::SparseVector, FieldValue::SparseVector(entries)) => {
463                    let has_simhash = entry.simhash;
464                    let ordinal = self.next_element_ordinal(field.0);
465                    self.index_sparse_vector_field(*field, doc_id, ordinal as u16, entries)?;
466                    // V12: compute simhash for every ordinal (not just ordinal 0)
467                    if has_simhash {
468                        let h = simhash::simhash_from_sparse_vector(entries);
469                        let is_new_field = !self.ordinal_simhashes.contains_key(&field.0);
470                        self.ordinal_simhashes
471                            .entry(field.0)
472                            .or_default()
473                            .insert((doc_id, ordinal as u16), h);
474                        // Memory: (doc_id, ordinal) key (6B) + u64 hash (8B) + hashmap overhead (~8B)
475                        self.estimated_memory += size_of::<(DocId, u16)>() + size_of::<u64>() + 8;
476                        if is_new_field {
477                            self.estimated_memory +=
478                                size_of::<u32>() + size_of::<FxHashMap<(DocId, u16), u64>>() + 8;
479                        }
480                    }
481                }
482                _ => {}
483            }
484        }
485
486        // Stream document to disk immediately
487        self.write_document_to_store(&doc)?;
488
489        Ok(doc_id)
490    }
491
492    /// Index a text field using interned terms
493    ///
494    /// Uses a custom tokenizer when set for the field (via `set_tokenizer`),
495    /// otherwise falls back to an inline zero-allocation path (split_whitespace
496    /// + lowercase + strip non-alphanumeric).
497    ///
498    /// If position recording is enabled for this field, also records token positions
499    /// encoded as (element_ordinal << 20) | token_position.
500    fn index_text_field(
501        &mut self,
502        field: Field,
503        doc_id: DocId,
504        text: &str,
505        element_ordinal: u32,
506    ) -> Result<u32> {
507        use crate::dsl::PositionMode;
508
509        let field_id = field.0;
510        let position_mode = self
511            .position_enabled_fields
512            .get(&field_id)
513            .copied()
514            .flatten();
515
516        // Phase 1: Aggregate term frequencies within this document
517        // Also collect positions if enabled
518        // Reuse buffers to avoid allocations
519        self.local_tf_buffer.clear();
520        // Clear position Vecs in-place (keeps allocated capacity for reuse)
521        for v in self.local_positions.values_mut() {
522            v.clear();
523        }
524
525        let mut token_position = 0u32;
526
527        // Tokenize: use custom tokenizer if set, else inline zero-alloc path.
528        // The owned Vec<Token> is computed first so the immutable borrow of
529        // self.tokenizers ends before we mutate other fields.
530        let custom_tokens = self.tokenizers.get(&field).map(|t| t.tokenize(text));
531
532        if let Some(tokens) = custom_tokens {
533            // Custom tokenizer path
534            for token in &tokens {
535                let term_spur = if let Some(spur) = self.term_interner.get(&token.text) {
536                    spur
537                } else {
538                    let spur = self.term_interner.get_or_intern(&token.text);
539                    self.estimated_memory += token.text.len() + INTERN_OVERHEAD;
540                    spur
541                };
542                *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
543
544                if let Some(mode) = position_mode {
545                    let encoded_pos = match mode {
546                        PositionMode::Ordinal => element_ordinal << 20,
547                        PositionMode::TokenPosition => token.position,
548                        PositionMode::Full => (element_ordinal << 20) | token.position,
549                    };
550                    self.local_positions
551                        .entry(term_spur)
552                        .or_default()
553                        .push(encoded_pos);
554                }
555            }
556            token_position = tokens.len() as u32;
557        } else {
558            // Inline zero-allocation path: split_whitespace + lowercase + strip non-alphanumeric
559            for word in text.split_whitespace() {
560                self.token_buffer.clear();
561                for c in word.chars() {
562                    if c.is_alphanumeric() {
563                        for lc in c.to_lowercase() {
564                            self.token_buffer.push(lc);
565                        }
566                    }
567                }
568
569                if self.token_buffer.is_empty() {
570                    continue;
571                }
572
573                let term_spur = if let Some(spur) = self.term_interner.get(&self.token_buffer) {
574                    spur
575                } else {
576                    let spur = self.term_interner.get_or_intern(&self.token_buffer);
577                    self.estimated_memory += self.token_buffer.len() + INTERN_OVERHEAD;
578                    spur
579                };
580                *self.local_tf_buffer.entry(term_spur).or_insert(0) += 1;
581
582                if let Some(mode) = position_mode {
583                    let encoded_pos = match mode {
584                        PositionMode::Ordinal => element_ordinal << 20,
585                        PositionMode::TokenPosition => token_position,
586                        PositionMode::Full => (element_ordinal << 20) | token_position,
587                    };
588                    self.local_positions
589                        .entry(term_spur)
590                        .or_default()
591                        .push(encoded_pos);
592                }
593
594                token_position += 1;
595            }
596        }
597
598        // Phase 2: Insert aggregated terms into inverted index
599        // Now we only do one inverted_index lookup per unique term in doc
600        for (&term_spur, &tf) in &self.local_tf_buffer {
601            let term_key = TermKey {
602                field: field_id,
603                term: term_spur,
604            };
605
606            match self.inverted_index.entry(term_key) {
607                hashbrown::hash_map::Entry::Occupied(mut o) => {
608                    o.get_mut().add(doc_id, tf);
609                    self.estimated_memory += size_of::<CompactPosting>();
610                }
611                hashbrown::hash_map::Entry::Vacant(v) => {
612                    let mut posting = PostingListBuilder::new();
613                    posting.add(doc_id, tf);
614                    v.insert(posting);
615                    self.estimated_memory += size_of::<CompactPosting>() + NEW_TERM_OVERHEAD;
616                }
617            }
618
619            if position_mode.is_some()
620                && let Some(positions) = self.local_positions.get(&term_spur)
621            {
622                match self.position_index.entry(term_key) {
623                    hashbrown::hash_map::Entry::Occupied(mut o) => {
624                        for &pos in positions {
625                            o.get_mut().add_position(doc_id, pos);
626                        }
627                        self.estimated_memory += positions.len() * size_of::<u32>();
628                    }
629                    hashbrown::hash_map::Entry::Vacant(v) => {
630                        let mut pos_posting = PositionPostingListBuilder::new();
631                        for &pos in positions {
632                            pos_posting.add_position(doc_id, pos);
633                        }
634                        self.estimated_memory +=
635                            positions.len() * size_of::<u32>() + NEW_POS_TERM_OVERHEAD;
636                        v.insert(pos_posting);
637                    }
638                }
639            }
640        }
641
642        Ok(token_position)
643    }
644
645    fn index_numeric_field(&mut self, field: Field, doc_id: DocId, value: u64) -> Result<()> {
646        use std::fmt::Write;
647
648        self.numeric_buffer.clear();
649        write!(self.numeric_buffer, "__num_{}", value).unwrap();
650        let term_spur = if let Some(spur) = self.term_interner.get(&self.numeric_buffer) {
651            spur
652        } else {
653            let spur = self.term_interner.get_or_intern(&self.numeric_buffer);
654            self.estimated_memory += self.numeric_buffer.len() + INTERN_OVERHEAD;
655            spur
656        };
657
658        let term_key = TermKey {
659            field: field.0,
660            term: term_spur,
661        };
662
663        match self.inverted_index.entry(term_key) {
664            hashbrown::hash_map::Entry::Occupied(mut o) => {
665                o.get_mut().add(doc_id, 1);
666                self.estimated_memory += size_of::<CompactPosting>();
667            }
668            hashbrown::hash_map::Entry::Vacant(v) => {
669                let mut posting = PostingListBuilder::new();
670                posting.add(doc_id, 1);
671                v.insert(posting);
672                self.estimated_memory += size_of::<CompactPosting>() + NEW_TERM_OVERHEAD;
673            }
674        }
675
676        Ok(())
677    }
678
679    /// Index a dense vector field with ordinal tracking
680    fn index_dense_vector_field(
681        &mut self,
682        field: Field,
683        doc_id: DocId,
684        ordinal: u16,
685        vector: &[f32],
686    ) -> Result<()> {
687        let dim = vector.len();
688
689        let builder = self
690            .dense_vectors
691            .entry(field.0)
692            .or_insert_with(|| DenseVectorBuilder::new(dim));
693
694        // Verify dimension consistency
695        if builder.dim != dim && builder.len() > 0 {
696            return Err(crate::Error::Schema(format!(
697                "Dense vector dimension mismatch: expected {}, got {}",
698                builder.dim, dim
699            )));
700        }
701
702        builder.add(doc_id, ordinal, vector);
703
704        self.estimated_memory += std::mem::size_of_val(vector) + size_of::<(DocId, u16)>();
705
706        Ok(())
707    }
708
709    /// Index a sparse vector field using dedicated sparse posting lists
710    ///
711    /// Collects (doc_id, ordinal, weight) postings per dimension. During commit, these are
712    /// converted to BlockSparsePostingList with proper quantization from SparseVectorConfig.
713    ///
714    /// Weights below the configured `weight_threshold` are not indexed.
715    fn index_sparse_vector_field(
716        &mut self,
717        field: Field,
718        doc_id: DocId,
719        ordinal: u16,
720        entries: &[(u32, f32)],
721    ) -> Result<()> {
722        // Get weight threshold from field config (default 0.0 = no filtering)
723        let weight_threshold = self
724            .schema
725            .get_field_entry(field)
726            .and_then(|entry| entry.sparse_vector_config.as_ref())
727            .map(|config| config.weight_threshold)
728            .unwrap_or(0.0);
729
730        let builder = self
731            .sparse_vectors
732            .entry(field.0)
733            .or_insert_with(SparseVectorBuilder::new);
734
735        builder.inc_vector_count();
736
737        for &(dim_id, weight) in entries {
738            // Skip weights below threshold
739            if weight.abs() < weight_threshold {
740                continue;
741            }
742
743            let is_new_dim = !builder.postings.contains_key(&dim_id);
744            builder.add(dim_id, doc_id, ordinal, weight);
745            self.estimated_memory += size_of::<(DocId, u16, f32)>();
746            if is_new_dim {
747                // HashMap entry overhead + Vec header
748                self.estimated_memory += size_of::<u32>() + size_of::<Vec<(DocId, u16, f32)>>() + 8; // 8 = hashmap control byte + padding
749            }
750        }
751
752        Ok(())
753    }
754
755    /// Write document to streaming store (reuses internal buffer to avoid per-doc allocation)
756    fn write_document_to_store(&mut self, doc: &Document) -> Result<()> {
757        use byteorder::{LittleEndian, WriteBytesExt};
758
759        super::store::serialize_document_into(doc, &self.schema, &mut self.doc_serialize_buffer)?;
760
761        self.store_file
762            .write_u32::<LittleEndian>(self.doc_serialize_buffer.len() as u32)?;
763        self.store_file.write_all(&self.doc_serialize_buffer)?;
764
765        Ok(())
766    }
767
768    /// Build the final segment
769    ///
770    /// Streams all data directly to disk via StreamingWriter to avoid buffering
771    /// entire serialized outputs in memory. Each phase consumes and drops its
772    /// source data before the next phase begins.
773    pub async fn build<D: Directory + DirectoryWriter>(
774        mut self,
775        dir: &D,
776        segment_id: SegmentId,
777        trained: Option<&super::TrainedVectorStructures>,
778    ) -> Result<SegmentMeta> {
779        // Flush any buffered data
780        self.store_file.flush()?;
781
782        let files = SegmentFiles::new(segment_id.0);
783
784        // Phase 1: Stream positions directly to disk (consumes position_index)
785        let position_index = std::mem::take(&mut self.position_index);
786        let position_offsets = if !position_index.is_empty() {
787            let mut pos_writer = dir.streaming_writer(&files.positions).await?;
788            let offsets = postings::build_positions_streaming(
789                position_index,
790                &self.term_interner,
791                &mut *pos_writer,
792            )?;
793            pos_writer.finish()?;
794            offsets
795        } else {
796            FxHashMap::default()
797        };
798
799        // Phase 2: 4-way parallel build — postings, store, dense vectors, sparse vectors
800        // These are fully independent: different source data, different output files.
801        let inverted_index = std::mem::take(&mut self.inverted_index);
802        let term_interner = std::mem::replace(&mut self.term_interner, Rodeo::new());
803        let store_path = self.store_path.clone();
804        let num_compression_threads = self.config.num_compression_threads;
805        let compression_level = self.config.compression_level;
806        let dense_vectors = std::mem::take(&mut self.dense_vectors);
807        let mut sparse_vectors = std::mem::take(&mut self.sparse_vectors);
808        let ordinal_simhashes = std::mem::take(&mut self.ordinal_simhashes);
809        let schema = &self.schema;
810
811        // Pre-create all streaming writers (async) before entering sync rayon scope
812        // Wrapped in OffsetWriter to track bytes written per phase.
813        let mut term_dict_writer =
814            super::OffsetWriter::new(dir.streaming_writer(&files.term_dict).await?);
815        let mut postings_writer =
816            super::OffsetWriter::new(dir.streaming_writer(&files.postings).await?);
817        let mut store_writer = super::OffsetWriter::new(dir.streaming_writer(&files.store).await?);
818        let mut vectors_writer = if !dense_vectors.is_empty() {
819            Some(super::OffsetWriter::new(
820                dir.streaming_writer(&files.vectors).await?,
821            ))
822        } else {
823            None
824        };
825        let mut sparse_writer = if !sparse_vectors.is_empty() {
826            Some(super::OffsetWriter::new(
827                dir.streaming_writer(&files.sparse).await?,
828            ))
829        } else {
830            None
831        };
832        let mut fast_fields = std::mem::take(&mut self.fast_fields);
833        let num_docs = self.next_doc_id;
834        let mut fast_writer = if !fast_fields.is_empty() {
835            Some(super::OffsetWriter::new(
836                dir.streaming_writer(&files.fast).await?,
837            ))
838        } else {
839            None
840        };
841
842        let ((postings_result, store_result), ((vectors_result, sparse_result), fast_result)) =
843            rayon::join(
844                || {
845                    rayon::join(
846                        || {
847                            postings::build_postings_streaming(
848                                inverted_index,
849                                term_interner,
850                                &position_offsets,
851                                &mut term_dict_writer,
852                                &mut postings_writer,
853                            )
854                        },
855                        || {
856                            store::build_store_streaming(
857                                &store_path,
858                                num_compression_threads,
859                                compression_level,
860                                &mut store_writer,
861                                num_docs,
862                            )
863                        },
864                    )
865                },
866                || {
867                    rayon::join(
868                        || {
869                            rayon::join(
870                                || -> Result<()> {
871                                    if let Some(ref mut w) = vectors_writer {
872                                        dense::build_vectors_streaming(
873                                            dense_vectors,
874                                            schema,
875                                            trained,
876                                            w,
877                                        )?;
878                                    }
879                                    Ok(())
880                                },
881                                || -> Result<()> {
882                                    if let Some(ref mut w) = sparse_writer {
883                                        sparse::build_sparse_streaming(
884                                            &mut sparse_vectors,
885                                            schema,
886                                            &ordinal_simhashes,
887                                            w,
888                                        )?;
889                                    }
890                                    Ok(())
891                                },
892                            )
893                        },
894                        || -> Result<()> {
895                            if let Some(ref mut w) = fast_writer {
896                                build_fast_fields_streaming(&mut fast_fields, num_docs, w)?;
897                            }
898                            Ok(())
899                        },
900                    )
901                },
902            );
903        postings_result?;
904        store_result?;
905        vectors_result?;
906        sparse_result?;
907        fast_result?;
908
909        let term_dict_bytes = term_dict_writer.offset() as usize;
910        let postings_bytes = postings_writer.offset() as usize;
911        let store_bytes = store_writer.offset() as usize;
912        let vectors_bytes = vectors_writer.as_ref().map_or(0, |w| w.offset() as usize);
913        let sparse_bytes = sparse_writer.as_ref().map_or(0, |w| w.offset() as usize);
914        let fast_bytes = fast_writer.as_ref().map_or(0, |w| w.offset() as usize);
915
916        term_dict_writer.finish()?;
917        postings_writer.finish()?;
918        store_writer.finish()?;
919        if let Some(w) = vectors_writer {
920            w.finish()?;
921        }
922        if let Some(w) = sparse_writer {
923            w.finish()?;
924        }
925        if let Some(w) = fast_writer {
926            w.finish()?;
927        }
928        drop(position_offsets);
929        drop(sparse_vectors);
930
931        log::info!(
932            "[segment_build] {} docs: term_dict={}, postings={}, store={}, vectors={}, sparse={}, fast={}",
933            num_docs,
934            super::format_bytes(term_dict_bytes),
935            super::format_bytes(postings_bytes),
936            super::format_bytes(store_bytes),
937            super::format_bytes(vectors_bytes),
938            super::format_bytes(sparse_bytes),
939            super::format_bytes(fast_bytes),
940        );
941
942        let meta = SegmentMeta {
943            id: segment_id.0,
944            num_docs: self.next_doc_id,
945            field_stats: self.field_stats.clone(),
946        };
947
948        dir.write(&files.meta, &meta.serialize()?).await?;
949
950        // Cleanup temp files
951        let _ = std::fs::remove_file(&self.store_path);
952
953        Ok(meta)
954    }
955}
956
957/// Serialize all fast-field columns to a `.fast` file.
958fn build_fast_fields_streaming(
959    fast_fields: &mut FxHashMap<u32, crate::structures::fast_field::FastFieldWriter>,
960    num_docs: u32,
961    writer: &mut dyn Write,
962) -> Result<()> {
963    use crate::structures::fast_field::{FastFieldTocEntry, write_fast_field_toc_and_footer};
964
965    if fast_fields.is_empty() {
966        return Ok(());
967    }
968
969    // Sort fields by id for deterministic output
970    let mut field_ids: Vec<u32> = fast_fields.keys().copied().collect();
971    field_ids.sort_unstable();
972
973    let mut toc_entries: Vec<FastFieldTocEntry> = Vec::with_capacity(field_ids.len());
974    let mut current_offset = 0u64;
975
976    for &field_id in &field_ids {
977        let ff = fast_fields.get_mut(&field_id).unwrap();
978        ff.pad_to(num_docs);
979
980        let (mut toc, bytes_written) = ff.serialize(writer, current_offset)?;
981        toc.field_id = field_id;
982        current_offset += bytes_written;
983        toc_entries.push(toc);
984    }
985
986    // Write TOC + footer
987    let toc_offset = current_offset;
988    write_fast_field_toc_and_footer(writer, toc_offset, &toc_entries)?;
989
990    Ok(())
991}
992
993impl Drop for SegmentBuilder {
994    fn drop(&mut self) {
995        // Cleanup temp files on drop
996        let _ = std::fs::remove_file(&self.store_path);
997    }
998}