Skip to main content

zeppelin_embed/segment/
writer.rs

1//! In-memory segment assembly and atomic publication.
2
3use std::path::{Path, PathBuf};
4
5use xxhash_rust::xxh3::xxh3_64;
6
7use crate::format::frame::{
8    FILE_HEADER_LEN, FILE_MAGIC, FILE_TRAILER_LEN, FormatCheck, FormatError,
9};
10use crate::format::{FormatFamily, FormatRegistry};
11use crate::graph::block::{GraphNodeBlockBuild, decode_node_blocks, encode_node_blocks};
12use crate::ingest::{DocId, Revision};
13use crate::lifecycle::durability::{DurabilityPolicy, SyncRequirement};
14use crate::meta::{AliveSet, ColumnStore};
15use crate::quant::Bit4Factors;
16use crate::vfs::Vfs;
17
18use super::layout::{
19    CHECKSUM_CHUNK_BYTES, Int8Factors, REGION_ALIGNMENT, REGION_ENTRY_LEN, RegionEntry, RegionKind,
20    SEGMENT_PREFIX_LEN, VECTOR_HEADER_LEN, VectorHeader, align_up, encode_alive, encode_columns,
21    encode_vector_header,
22};
23use super::{SegmentError, SegmentId, SegmentMeta};
24
25/// Borrowed factor records for a supported per-segment quantization scheme.
26#[derive(Clone, Copy, Debug)]
27pub enum SegmentFactors<'a> {
28    /// Full-precision rows carry no separate factor records.
29    F32,
30    /// Permanent 12-byte Bit4 factor records.
31    Bit4(&'a [Bit4Factors]),
32    /// Permanent 8-byte affine Int8 factor records.
33    Int8(&'a [Int8Factors]),
34}
35
36/// Complete in-memory data sealed into one immutable segment.
37#[derive(Clone, Copy, Debug)]
38pub struct SegmentBuild<'a> {
39    /// Sortable immutable identity.
40    pub id: SegmentId,
41    /// Permanent per-segment quantization scheme id.
42    pub scheme: u16,
43    /// Logical vector dimension.
44    pub dims: u32,
45    /// Contiguous row-major packed codes with no row padding.
46    pub codes: &'a [u8],
47    /// Contiguous factor records matching `scheme`.
48    pub factors: SegmentFactors<'a>,
49    /// Contiguous row-major f32 exact-rescore source.
50    pub rescore: &'a [f32],
51    /// Typed metadata arrays aligned to row ids.
52    pub columns: &'a ColumnStore,
53    /// Alive/tombstone state aligned to row ids.
54    pub alive: &'a AliveSet,
55}
56
57/// Dense application identities aligned to one segment's local row ids.
58#[derive(Clone, Copy, Debug)]
59pub struct SegmentDocumentVersions<'a> {
60    /// Stable application document ids.
61    pub doc_ids: &'a [DocId],
62    /// Monotonic application revisions aligned to `doc_ids`.
63    pub revisions: &'a [Revision],
64}
65
66/// Dense opaque metadata bytes aligned to one segment's local row ids.
67#[derive(Clone, Copy, Debug)]
68pub struct SegmentStoredMetadata<'a> {
69    /// Cumulative exclusive end offset for every row.
70    pub end_offsets: &'a [u64],
71    /// Concatenated row payloads addressed by `end_offsets`.
72    pub bytes: &'a [u8],
73}
74
75/// Dense optional UTF-8 source text aligned to one segment's local row ids.
76#[derive(Clone, Copy, Debug)]
77pub struct SegmentStoredText<'a> {
78    /// One byte per dense row; zero is absent and one is present.
79    pub present: &'a [u8],
80    /// Cumulative exclusive end offset for every row.
81    pub end_offsets: &'a [u64],
82    /// Concatenated UTF-8 row payloads addressed by `end_offsets`.
83    pub bytes: &'a [u8],
84}
85
86/// Optional immutable payload regions aligned to document rows.
87#[derive(Clone, Copy, Debug)]
88pub(crate) struct SegmentPayloads<'a> {
89    pub(crate) documents: SegmentDocumentVersions<'a>,
90    pub(crate) metadata: Option<SegmentStoredMetadata<'a>>,
91    pub(crate) text: Option<SegmentStoredText<'a>>,
92    pub(crate) postings: Option<SegmentPostings<'a>>,
93}
94
95/// One complete encoded lexical segment for the Postings region.
96#[derive(Clone, Copy, Debug)]
97pub struct SegmentPostings<'a> {
98    /// Bytes produced by `SealedSegment::encode_region`.
99    pub bytes: &'a [u8],
100}
101
102struct RegionBytes {
103    kind: RegionKind,
104    family: FormatFamily,
105    bytes: Vec<u8>,
106}
107
108/// One source region copied byte-for-byte into a rewritten segment.
109#[derive(Clone, Copy, Debug)]
110pub(crate) struct SegmentCopiedRegion<'a> {
111    /// Permanent or forward-compatible source region kind.
112    pub kind: u16,
113    /// Exact additive region version from the source directory.
114    pub version: u16,
115    /// Original checksummed region payload.
116    pub bytes: &'a [u8],
117}
118
119/// Encodes one complete segment in RAM, validating every cross-region shape first.
120pub fn encode_segment(build: SegmentBuild<'_>) -> Result<Vec<u8>, SegmentError> {
121    encode_segment_inner(build, None, None, None, None, None, &[])
122}
123
124/// Encodes a complete segment with one fixed-stride graph node-block region.
125pub fn encode_segment_with_graph(
126    build: SegmentBuild<'_>,
127    graph: GraphNodeBlockBuild<'_>,
128) -> Result<Vec<u8>, SegmentError> {
129    if graph.layout.dims() != build.dims {
130        return Err(SegmentError::Geometry(format!(
131            "graph dimensions {}, segment dimensions {}",
132            graph.layout.dims(),
133            build.dims
134        )));
135    }
136    if graph.nodes.len() != build.columns.row_count() as usize {
137        return Err(SegmentError::Geometry(format!(
138            "graph nodes {}, segment rows {}",
139            graph.nodes.len(),
140            build.columns.row_count()
141        )));
142    }
143    let graph_bytes = encode_node_blocks(graph)?.into_bytes();
144    encode_segment_inner(build, Some(graph_bytes), None, None, None, None, &[])
145}
146
147/// Encodes a graph segment while preserving dense application document identities.
148pub fn encode_segment_with_graph_and_documents(
149    build: SegmentBuild<'_>,
150    graph: GraphNodeBlockBuild<'_>,
151    documents: SegmentDocumentVersions<'_>,
152) -> Result<Vec<u8>, SegmentError> {
153    if graph.layout.dims() != build.dims {
154        return Err(SegmentError::Geometry(format!(
155            "graph dimensions {}, segment dimensions {}",
156            graph.layout.dims(),
157            build.dims
158        )));
159    }
160    if graph.nodes.len() != build.columns.row_count() as usize {
161        return Err(SegmentError::Geometry(format!(
162            "graph nodes {}, segment rows {}",
163            graph.nodes.len(),
164            build.columns.row_count()
165        )));
166    }
167    let graph_bytes = encode_node_blocks(graph)?.into_bytes();
168    encode_segment_inner(
169        build,
170        Some(graph_bytes),
171        Some(documents),
172        None,
173        None,
174        None,
175        &[],
176    )
177}
178
179fn encode_segment_inner(
180    build: SegmentBuild<'_>,
181    graph_bytes: Option<Vec<u8>>,
182    document_versions: Option<SegmentDocumentVersions<'_>>,
183    stored_metadata: Option<SegmentStoredMetadata<'_>>,
184    stored_text: Option<SegmentStoredText<'_>>,
185    postings: Option<SegmentPostings<'_>>,
186    copied_regions: &[SegmentCopiedRegion<'_>],
187) -> Result<Vec<u8>, SegmentError> {
188    let row_count = build.columns.row_count();
189    if build.alive.row_count() != row_count {
190        return Err(SegmentError::Geometry(format!(
191            "alive rows {}, column rows {row_count}",
192            build.alive.row_count()
193        )));
194    }
195    FormatRegistry::require_scheme(build.scheme)
196        .map_err(|error| SegmentError::Geometry(error.to_string()))?;
197    let dims = usize::try_from(build.dims)
198        .map_err(|_| SegmentError::Geometry("dims exceeds usize".to_owned()))?;
199    let rows = row_count as usize;
200    let (row_stride, factor_stride, factor_bytes) = encode_factors(build, rows)?;
201    let expected_codes = row_stride
202        .checked_mul(rows)
203        .ok_or_else(|| SegmentError::Geometry("code length overflow".to_owned()))?;
204    if build.codes.len() != expected_codes {
205        return Err(SegmentError::Geometry(format!(
206            "code bytes {}, expected {expected_codes}",
207            build.codes.len()
208        )));
209    }
210    let expected_rescore = dims
211        .checked_mul(rows)
212        .ok_or_else(|| SegmentError::Geometry("rescore length overflow".to_owned()))?;
213    if build.rescore.len() != expected_rescore {
214        return Err(SegmentError::Geometry(format!(
215            "rescore values {}, expected {expected_rescore}",
216            build.rescore.len()
217        )));
218    }
219
220    let common_header = VectorHeader {
221        scheme: build.scheme,
222        vector_space_id: 0,
223        dims: build.dims,
224        row_stride_bytes: u32::try_from(row_stride)
225            .map_err(|_| SegmentError::Geometry("row stride exceeds u32".to_owned()))?,
226        factor_stride_bytes: factor_stride,
227        transform_kind: 0,
228        transform_seed: 0,
229        row_count,
230        reserved: 0,
231    };
232    common_header.validate()?;
233    let mut codes = Vec::with_capacity(VECTOR_HEADER_LEN.saturating_add(build.codes.len()));
234    encode_vector_header(common_header, &mut codes);
235    codes.extend_from_slice(build.codes);
236
237    let mut factors = Vec::with_capacity(VECTOR_HEADER_LEN.saturating_add(factor_bytes.len()));
238    encode_vector_header(common_header, &mut factors);
239    factors.extend_from_slice(&factor_bytes);
240
241    let rescore_stride = dims
242        .checked_mul(std::mem::size_of::<f32>())
243        .ok_or_else(|| SegmentError::Geometry("rescore stride overflow".to_owned()))?;
244    let rescore_header = VectorHeader {
245        scheme: 0,
246        vector_space_id: 0,
247        dims: build.dims,
248        row_stride_bytes: u32::try_from(rescore_stride)
249            .map_err(|_| SegmentError::Geometry("rescore stride exceeds u32".to_owned()))?,
250        factor_stride_bytes: 0,
251        transform_kind: 0,
252        transform_seed: 0,
253        row_count,
254        reserved: 0,
255    };
256    let mut rescore =
257        Vec::with_capacity(VECTOR_HEADER_LEN.saturating_add(build.rescore.len().saturating_mul(4)));
258    encode_vector_header(rescore_header, &mut rescore);
259    for value in build.rescore {
260        rescore.extend_from_slice(&value.to_bits().to_le_bytes());
261    }
262
263    let mut regions = vec![
264        RegionBytes {
265            kind: RegionKind::Columns,
266            family: FormatFamily::Columns,
267            bytes: encode_columns(build.columns)?,
268        },
269        RegionBytes {
270            kind: RegionKind::Alive,
271            family: FormatFamily::Alive,
272            bytes: encode_alive(build.alive),
273        },
274        RegionBytes {
275            kind: RegionKind::VectorFactors,
276            family: FormatFamily::VectorFactors,
277            bytes: factors,
278        },
279        RegionBytes {
280            kind: RegionKind::VectorCodes,
281            family: FormatFamily::VectorCodes,
282            bytes: codes,
283        },
284        RegionBytes {
285            kind: RegionKind::VectorRescore,
286            family: FormatFamily::VectorRescore,
287            bytes: rescore,
288        },
289    ];
290    if let Some(bytes) = graph_bytes {
291        regions.push(RegionBytes {
292            kind: RegionKind::GraphNodeBlocks,
293            family: FormatFamily::GraphNodeBlocks,
294            bytes,
295        });
296    }
297    if let Some(documents) = document_versions {
298        regions.push(RegionBytes {
299            kind: RegionKind::DocumentVersions,
300            family: FormatFamily::DocumentVersions,
301            bytes: encode_document_versions(documents, rows)?,
302        });
303    }
304    if let Some(metadata) = stored_metadata {
305        regions.push(RegionBytes {
306            kind: RegionKind::StoredMetadata,
307            family: FormatFamily::StoredMetadata,
308            bytes: encode_stored_metadata(metadata, rows)?,
309        });
310    }
311    if let Some(text) = stored_text {
312        regions.push(RegionBytes {
313            kind: RegionKind::StoredText,
314            family: FormatFamily::StoredText,
315            bytes: encode_stored_text(text, rows)?,
316        });
317    }
318    if let Some(postings) = postings {
319        let decoded = crate::fts::sealed::SealedSegment::decode_region(postings.bytes)?;
320        if decoded.row_count() != row_count {
321            return Err(SegmentError::Geometry(format!(
322                "postings rows {}, segment rows {row_count}",
323                decoded.row_count()
324            )));
325        }
326        regions.push(RegionBytes {
327            kind: RegionKind::Postings,
328            family: FormatFamily::Postings,
329            bytes: postings.bytes.to_vec(),
330        });
331    }
332    encode_regions(build, &regions, copied_regions)
333}
334
335fn encode_stored_metadata(
336    metadata: SegmentStoredMetadata<'_>,
337    rows: usize,
338) -> Result<Vec<u8>, SegmentError> {
339    if metadata.end_offsets.len() != rows {
340        return Err(SegmentError::Geometry(format!(
341            "stored-metadata offsets {}, expected {rows}",
342            metadata.end_offsets.len()
343        )));
344    }
345    let mut previous = 0_u64;
346    for end in metadata.end_offsets {
347        if *end < previous {
348            return Err(SegmentError::Geometry(
349                "stored-metadata offsets are not monotonic".to_owned(),
350            ));
351        }
352        previous = *end;
353    }
354    let byte_len = u64::try_from(metadata.bytes.len())
355        .map_err(|_| SegmentError::Geometry("stored-metadata bytes exceed u64".to_owned()))?;
356    if previous != byte_len {
357        return Err(SegmentError::Geometry(format!(
358            "stored-metadata final offset {previous}, bytes {byte_len}"
359        )));
360    }
361    let offset_count = rows.checked_add(1).ok_or_else(|| {
362        SegmentError::Geometry("stored-metadata offset count overflow".to_owned())
363    })?;
364    let capacity = 8_usize
365        .checked_add(offset_count.saturating_mul(8))
366        .and_then(|value| value.checked_add(metadata.bytes.len()))
367        .ok_or_else(|| SegmentError::Geometry("stored-metadata length overflow".to_owned()))?;
368    let row_count = u32::try_from(rows)
369        .map_err(|_| SegmentError::Geometry("stored-metadata rows exceed u32".to_owned()))?;
370    let mut encoded = Vec::with_capacity(capacity);
371    encoded.extend_from_slice(&row_count.to_le_bytes());
372    encoded.extend_from_slice(&0_u32.to_le_bytes());
373    encoded.extend_from_slice(&0_u64.to_le_bytes());
374    for end in metadata.end_offsets {
375        encoded.extend_from_slice(&end.to_le_bytes());
376    }
377    encoded.extend_from_slice(metadata.bytes);
378    Ok(encoded)
379}
380
381fn encode_stored_text(text: SegmentStoredText<'_>, rows: usize) -> Result<Vec<u8>, SegmentError> {
382    if text.present.len() != rows || text.end_offsets.len() != rows {
383        return Err(SegmentError::Geometry(format!(
384            "stored-text presence/offset rows {}/{}, expected {rows}",
385            text.present.len(),
386            text.end_offsets.len()
387        )));
388    }
389    let mut previous = 0_u64;
390    for (row, (present, end)) in text.present.iter().zip(text.end_offsets).enumerate() {
391        if *present > 1 {
392            return Err(SegmentError::Geometry(format!(
393                "stored-text presence row {row} is {present}, expected 0 or 1"
394            )));
395        }
396        if *end < previous {
397            return Err(SegmentError::Geometry(
398                "stored-text offsets are not monotonic".to_owned(),
399            ));
400        }
401        if *present == 0 && *end != previous {
402            return Err(SegmentError::Geometry(format!(
403                "stored-text absent row {row} owns bytes {previous}..{end}"
404            )));
405        }
406        let start = usize::try_from(previous)
407            .map_err(|_| SegmentError::Geometry("stored-text offset exceeds usize".to_owned()))?;
408        let finish = usize::try_from(*end)
409            .map_err(|_| SegmentError::Geometry("stored-text offset exceeds usize".to_owned()))?;
410        let row_bytes = text.bytes.get(start..finish).ok_or_else(|| {
411            SegmentError::Geometry(format!("stored-text row {row} exceeds payload"))
412        })?;
413        if *present == 1 {
414            std::str::from_utf8(row_bytes).map_err(|_| {
415                SegmentError::Geometry(format!("stored-text row {row} is not UTF-8"))
416            })?;
417        }
418        previous = *end;
419    }
420    let byte_len = u64::try_from(text.bytes.len())
421        .map_err(|_| SegmentError::Geometry("stored-text bytes exceed u64".to_owned()))?;
422    if previous != byte_len {
423        return Err(SegmentError::Geometry(format!(
424            "stored-text final offset {previous}, bytes {byte_len}"
425        )));
426    }
427    let bitmap_len = rows.div_ceil(8);
428    let capacity = 8_usize
429        .checked_add(bitmap_len)
430        .and_then(|value| value.checked_add(rows.saturating_mul(8)))
431        .and_then(|value| value.checked_add(text.bytes.len()))
432        .ok_or_else(|| SegmentError::Geometry("stored-text length overflow".to_owned()))?;
433    let row_count = u32::try_from(rows)
434        .map_err(|_| SegmentError::Geometry("stored-text rows exceed u32".to_owned()))?;
435    let mut encoded = Vec::with_capacity(capacity);
436    encoded.extend_from_slice(&row_count.to_le_bytes());
437    encoded.extend_from_slice(&0_u32.to_le_bytes());
438    encoded.resize(8 + bitmap_len, 0);
439    for (row, present) in text.present.iter().enumerate() {
440        if *present == 1 {
441            let byte = encoded
442                .get_mut(8 + row / 8)
443                .ok_or_else(|| SegmentError::Geometry("stored-text bitmap overflow".to_owned()))?;
444            *byte |= 1 << (row % 8);
445        }
446    }
447    for end in text.end_offsets {
448        encoded.extend_from_slice(&end.to_le_bytes());
449    }
450    encoded.extend_from_slice(text.bytes);
451    Ok(encoded)
452}
453
454fn encode_document_versions(
455    documents: SegmentDocumentVersions<'_>,
456    rows: usize,
457) -> Result<Vec<u8>, SegmentError> {
458    if documents.doc_ids.len() != rows || documents.revisions.len() != rows {
459        return Err(SegmentError::Geometry(format!(
460            "document-version rows {}/{}, expected {rows}",
461            documents.doc_ids.len(),
462            documents.revisions.len()
463        )));
464    }
465    let capacity = rows
466        .checked_mul(24)
467        .ok_or_else(|| SegmentError::Geometry("document-version length overflow".to_owned()))?;
468    let mut bytes = Vec::with_capacity(capacity);
469    for (doc_id, revision) in documents.doc_ids.iter().zip(documents.revisions) {
470        bytes.extend_from_slice(&doc_id.get().to_le_bytes());
471        bytes.extend_from_slice(&revision.get().to_le_bytes());
472    }
473    Ok(bytes)
474}
475
476fn encode_factors(
477    build: SegmentBuild<'_>,
478    rows: usize,
479) -> Result<(usize, u16, Vec<u8>), SegmentError> {
480    match build.factors {
481        SegmentFactors::F32 => {
482            if build.scheme != 0 {
483                return Err(SegmentError::Geometry(format!(
484                    "F32 factors require scheme 0, got {}",
485                    build.scheme
486                )));
487            }
488            let row_stride = (build.dims as usize)
489                .checked_mul(std::mem::size_of::<f32>())
490                .ok_or_else(|| SegmentError::Geometry("F32 row stride overflow".to_owned()))?;
491            Ok((row_stride, 0, Vec::new()))
492        }
493        SegmentFactors::Bit4(factors) => {
494            if build.scheme != 4 {
495                return Err(SegmentError::Geometry(format!(
496                    "Bit4 factors require scheme 4, got {}",
497                    build.scheme
498                )));
499            }
500            if factors.len() != rows {
501                return Err(SegmentError::Geometry(format!(
502                    "factor rows {}, expected {rows}",
503                    factors.len()
504                )));
505            }
506            let row_stride = (build.dims as usize).div_ceil(2);
507            let mut bytes = Vec::with_capacity(factors.len().saturating_mul(12));
508            for factor in factors {
509                for field in factor.persisted_fields() {
510                    bytes.extend_from_slice(&field.to_bits().to_le_bytes());
511                }
512            }
513            Ok((row_stride, 12, bytes))
514        }
515        SegmentFactors::Int8(factors) => {
516            if build.scheme != 2 {
517                return Err(SegmentError::Geometry(format!(
518                    "Int8 factors require scheme 2, got {}",
519                    build.scheme
520                )));
521            }
522            if factors.len() != rows {
523                return Err(SegmentError::Geometry(format!(
524                    "factor rows {}, expected {rows}",
525                    factors.len()
526                )));
527            }
528            let mut bytes = Vec::with_capacity(factors.len().saturating_mul(8));
529            for factor in factors {
530                bytes.extend_from_slice(&factor.scale.to_bits().to_le_bytes());
531                bytes.extend_from_slice(&factor.offset.to_bits().to_le_bytes());
532            }
533            Ok((build.dims as usize, 8, bytes))
534        }
535    }
536}
537
538fn encode_regions(
539    build: SegmentBuild<'_>,
540    regions: &[RegionBytes],
541    copied_regions: &[SegmentCopiedRegion<'_>],
542) -> Result<Vec<u8>, SegmentError> {
543    let region_count = regions
544        .len()
545        .saturating_add(copied_regions.len())
546        .saturating_add(1);
547    let header_length = FILE_HEADER_LEN
548        .checked_add(SEGMENT_PREFIX_LEN)
549        .and_then(|value| value.checked_add(region_count.saturating_mul(REGION_ENTRY_LEN)))
550        .and_then(|value| value.checked_add(8))
551        .ok_or_else(|| SegmentError::Geometry("header length overflow".to_owned()))?;
552    let mut next_offset = align_up(header_length, REGION_ALIGNMENT)?;
553    let mut entries = Vec::with_capacity(region_count);
554    for region in regions.iter() {
555        let length = u64::try_from(region.bytes.len())
556            .map_err(|_| SegmentError::Geometry("region length exceeds u64".to_owned()))?;
557        entries.push(RegionEntry {
558            kind: region.kind.id(),
559            version: current_version(region.family)?,
560            reserved: 0,
561            offset: next_offset as u64,
562            length,
563            checksum: xxh3_64(&region.bytes),
564        });
565        next_offset = next_offset
566            .checked_add(region.bytes.len())
567            .ok_or_else(|| SegmentError::Geometry("region end overflow".to_owned()))?;
568        next_offset = align_up(next_offset, REGION_ALIGNMENT)?;
569    }
570    for region in copied_regions {
571        if region.kind == RegionKind::ChecksumTable.id()
572            || entries.iter().any(|entry| entry.kind == region.kind)
573        {
574            return Err(SegmentError::Geometry(format!(
575                "copied region kind {} conflicts with the rewritten segment",
576                region.kind
577            )));
578        }
579        let length = u64::try_from(region.bytes.len())
580            .map_err(|_| SegmentError::Geometry("region length exceeds u64".to_owned()))?;
581        entries.push(RegionEntry {
582            kind: region.kind,
583            version: region.version,
584            reserved: 0,
585            offset: next_offset as u64,
586            length,
587            checksum: xxh3_64(region.bytes),
588        });
589        next_offset = next_offset
590            .checked_add(region.bytes.len())
591            .ok_or_else(|| SegmentError::Geometry("region end overflow".to_owned()))?;
592        next_offset = align_up(next_offset, REGION_ALIGNMENT)?;
593    }
594
595    let checksum_table = encode_checksum_table(regions, copied_regions);
596    entries.push(RegionEntry {
597        kind: RegionKind::ChecksumTable.id(),
598        version: current_version(FormatFamily::ChecksumTable)?,
599        reserved: 0,
600        offset: next_offset as u64,
601        length: checksum_table.len() as u64,
602        checksum: xxh3_64(&checksum_table),
603    });
604    let file_without_trailer = next_offset
605        .checked_add(checksum_table.len())
606        .ok_or_else(|| SegmentError::Geometry("file length overflow".to_owned()))?;
607    let file_length = file_without_trailer
608        .checked_add(FILE_TRAILER_LEN)
609        .ok_or_else(|| SegmentError::Geometry("file trailer overflow".to_owned()))?;
610
611    let mut output = Vec::with_capacity(file_length);
612    output.extend_from_slice(&FILE_MAGIC);
613    output.extend_from_slice(&FormatFamily::Segment.id().to_le_bytes());
614    output.extend_from_slice(&1_u16.to_le_bytes());
615    output.extend_from_slice(&0_u32.to_le_bytes());
616    output.extend_from_slice(&(header_length as u64).to_le_bytes());
617    output.extend_from_slice(&(file_length as u64).to_le_bytes());
618    output.extend_from_slice(build.id.as_bytes());
619    output.extend_from_slice(&build.columns.row_count().to_le_bytes());
620    output.extend_from_slice(&(region_count as u16).to_le_bytes());
621    output.extend_from_slice(&0_u16.to_le_bytes());
622    output.extend_from_slice(&build.scheme.to_le_bytes());
623    output.extend_from_slice(&0_u16.to_le_bytes());
624    output.extend_from_slice(&build.dims.to_le_bytes());
625    for entry in &entries {
626        output.extend_from_slice(&entry.kind.to_le_bytes());
627        output.extend_from_slice(&entry.version.to_le_bytes());
628        output.extend_from_slice(&entry.reserved.to_le_bytes());
629        output.extend_from_slice(&entry.offset.to_le_bytes());
630        output.extend_from_slice(&entry.length.to_le_bytes());
631        output.extend_from_slice(&entry.checksum.to_le_bytes());
632    }
633    let header_checksum = xxh3_64(&output);
634    output.extend_from_slice(&header_checksum.to_le_bytes());
635    if output.len() != header_length {
636        return Err(SegmentError::Geometry(format!(
637            "encoded header length {}, declared {header_length}",
638            output.len()
639        )));
640    }
641    let first_offset = entries
642        .first()
643        .and_then(|entry| usize::try_from(entry.offset).ok())
644        .ok_or_else(|| SegmentError::Geometry("missing first region offset".to_owned()))?;
645    output.resize(first_offset, 0);
646    for (region, entry) in regions.iter().zip(&entries) {
647        let offset = usize::try_from(entry.offset)
648            .map_err(|_| SegmentError::Geometry("region offset exceeds usize".to_owned()))?;
649        if output.len() != offset {
650            return Err(SegmentError::Geometry(format!(
651                "region {:?} offset {offset}, current length {}",
652                region.kind,
653                output.len()
654            )));
655        }
656        output.extend_from_slice(&region.bytes);
657        output.resize(align_up(output.len(), REGION_ALIGNMENT)?, 0);
658    }
659    for (region, entry) in copied_regions
660        .iter()
661        .zip(entries.iter().skip(regions.len()))
662    {
663        let offset = usize::try_from(entry.offset)
664            .map_err(|_| SegmentError::Geometry("region offset exceeds usize".to_owned()))?;
665        if output.len() != offset {
666            return Err(SegmentError::Geometry(format!(
667                "copied region {} offset {offset}, current length {}",
668                region.kind,
669                output.len()
670            )));
671        }
672        output.extend_from_slice(region.bytes);
673        output.resize(align_up(output.len(), REGION_ALIGNMENT)?, 0);
674    }
675    let checksum_entry = entries.last().ok_or_else(|| {
676        SegmentError::Geometry("missing checksum-table directory entry".to_owned())
677    })?;
678    let checksum_offset = usize::try_from(checksum_entry.offset)
679        .map_err(|_| SegmentError::Geometry("checksum offset exceeds usize".to_owned()))?;
680    if output.len() != checksum_offset {
681        return Err(SegmentError::Geometry(format!(
682            "checksum offset {checksum_offset}, current length {}",
683            output.len()
684        )));
685    }
686    output.extend_from_slice(&checksum_table);
687    let whole_file_checksum = xxh3_64(&output);
688    output.extend_from_slice(&whole_file_checksum.to_le_bytes());
689    if output.len() != file_length {
690        return Err(SegmentError::Geometry(format!(
691            "encoded file length {}, declared {file_length}",
692            output.len()
693        )));
694    }
695    Ok(output)
696}
697
698fn encode_checksum_table(
699    regions: &[RegionBytes],
700    copied_regions: &[SegmentCopiedRegion<'_>],
701) -> Vec<u8> {
702    let entry_count = regions
703        .iter()
704        .map(|region| region.bytes.len().div_ceil(CHECKSUM_CHUNK_BYTES))
705        .chain(
706            copied_regions
707                .iter()
708                .map(|region| region.bytes.len().div_ceil(CHECKSUM_CHUNK_BYTES)),
709        )
710        .sum::<usize>();
711    let mut output = Vec::with_capacity(8_usize.saturating_add(entry_count.saturating_mul(16)));
712    output.extend_from_slice(&(entry_count as u32).to_le_bytes());
713    output.extend_from_slice(&(CHECKSUM_CHUNK_BYTES as u32).to_le_bytes());
714    for region in regions {
715        for (chunk_index, chunk) in region.bytes.chunks(CHECKSUM_CHUNK_BYTES).enumerate() {
716            output.extend_from_slice(&region.kind.id().to_le_bytes());
717            output.extend_from_slice(&0_u16.to_le_bytes());
718            output.extend_from_slice(&(chunk_index as u32).to_le_bytes());
719            output.extend_from_slice(&xxh3_64(chunk).to_le_bytes());
720        }
721    }
722    for region in copied_regions {
723        for (chunk_index, chunk) in region.bytes.chunks(CHECKSUM_CHUNK_BYTES).enumerate() {
724            output.extend_from_slice(&region.kind.to_le_bytes());
725            output.extend_from_slice(&0_u16.to_le_bytes());
726            output.extend_from_slice(&(chunk_index as u32).to_le_bytes());
727            output.extend_from_slice(&xxh3_64(chunk).to_le_bytes());
728        }
729    }
730    output
731}
732
733fn current_version(family: FormatFamily) -> Result<u16, SegmentError> {
734    FormatRegistry::families()
735        .iter()
736        .find(|spec| spec.family == family)
737        .map(|spec| spec.current_version)
738        .ok_or_else(|| SegmentError::Geometry(format!("family {} is not registered", family.id())))
739}
740
741/// Writes a temp file, synchronizes it, renames it into place, then syncs the directory.
742pub fn write_segment(
743    vfs: &dyn Vfs,
744    directory: &Path,
745    build: SegmentBuild<'_>,
746    policy: DurabilityPolicy,
747) -> Result<SegmentMeta, SegmentError> {
748    let bytes = encode_segment(build)?;
749    publish_segment(vfs, directory, build, policy, &bytes)
750}
751
752/// Writes and publishes a segment with dense application document identities.
753pub fn write_segment_with_documents(
754    vfs: &dyn Vfs,
755    directory: &Path,
756    build: SegmentBuild<'_>,
757    documents: SegmentDocumentVersions<'_>,
758    policy: DurabilityPolicy,
759) -> Result<SegmentMeta, SegmentError> {
760    let bytes = encode_segment_inner(build, None, Some(documents), None, None, None, &[])?;
761    publish_segment(vfs, directory, build, policy, &bytes)
762}
763
764/// Writes and publishes document identities plus opaque stored metadata.
765pub fn write_segment_with_documents_and_metadata(
766    vfs: &dyn Vfs,
767    directory: &Path,
768    build: SegmentBuild<'_>,
769    documents: SegmentDocumentVersions<'_>,
770    metadata: SegmentStoredMetadata<'_>,
771    policy: DurabilityPolicy,
772) -> Result<SegmentMeta, SegmentError> {
773    let bytes = encode_segment_inner(
774        build,
775        None,
776        Some(documents),
777        Some(metadata),
778        None,
779        None,
780        &[],
781    )?;
782    publish_segment(vfs, directory, build, policy, &bytes)
783}
784
785/// Writes document identities plus optional lossless UTF-8 source rows.
786pub fn write_segment_with_documents_and_text(
787    vfs: &dyn Vfs,
788    directory: &Path,
789    build: SegmentBuild<'_>,
790    documents: SegmentDocumentVersions<'_>,
791    text: SegmentStoredText<'_>,
792    policy: DurabilityPolicy,
793) -> Result<SegmentMeta, SegmentError> {
794    let bytes = encode_segment_inner(build, None, Some(documents), None, Some(text), None, &[])?;
795    publish_segment(vfs, directory, build, policy, &bytes)
796}
797
798/// Writes document identities and a sealed lexical Postings region.
799pub fn write_segment_with_documents_and_postings(
800    vfs: &dyn Vfs,
801    directory: &Path,
802    build: SegmentBuild<'_>,
803    documents: SegmentDocumentVersions<'_>,
804    postings: SegmentPostings<'_>,
805    policy: DurabilityPolicy,
806) -> Result<SegmentMeta, SegmentError> {
807    let bytes = encode_segment_inner(
808        build,
809        None,
810        Some(documents),
811        None,
812        None,
813        Some(postings),
814        &[],
815    )?;
816    publish_segment(vfs, directory, build, policy, &bytes)
817}
818
819/// Writes a sealed lexical region without application document identities.
820///
821/// This remains readable for format compatibility, while store-level lexical
822/// and hybrid joins reject the missing identity as a typed query error.
823pub fn write_segment_with_postings(
824    vfs: &dyn Vfs,
825    directory: &Path,
826    build: SegmentBuild<'_>,
827    postings: SegmentPostings<'_>,
828    policy: DurabilityPolicy,
829) -> Result<SegmentMeta, SegmentError> {
830    let bytes = encode_segment_inner(build, None, None, None, None, Some(postings), &[])?;
831    publish_segment(vfs, directory, build, policy, &bytes)
832}
833
834/// Writes document identities, opaque metadata, and a sealed lexical region.
835pub fn write_segment_with_documents_metadata_and_postings(
836    vfs: &dyn Vfs,
837    directory: &Path,
838    build: SegmentBuild<'_>,
839    documents: SegmentDocumentVersions<'_>,
840    metadata: SegmentStoredMetadata<'_>,
841    postings: SegmentPostings<'_>,
842    policy: DurabilityPolicy,
843) -> Result<SegmentMeta, SegmentError> {
844    let bytes = encode_segment_inner(
845        build,
846        None,
847        Some(documents),
848        Some(metadata),
849        None,
850        Some(postings),
851        &[],
852    )?;
853    publish_segment(vfs, directory, build, policy, &bytes)
854}
855
856/// Writes document identities and any optional payload regions in one segment.
857pub(crate) fn write_segment_with_documents_payloads(
858    vfs: &dyn Vfs,
859    directory: &Path,
860    build: SegmentBuild<'_>,
861    payloads: SegmentPayloads<'_>,
862    policy: DurabilityPolicy,
863) -> Result<SegmentMeta, SegmentError> {
864    let bytes = encode_segment_inner(
865        build,
866        None,
867        Some(payloads.documents),
868        payloads.metadata,
869        payloads.text,
870        payloads.postings,
871        &[],
872    )?;
873    publish_segment(vfs, directory, build, policy, &bytes)
874}
875
876/// Writes and atomically publishes a segment containing fixed-stride graph blocks.
877pub fn write_segment_with_graph(
878    vfs: &dyn Vfs,
879    directory: &Path,
880    build: SegmentBuild<'_>,
881    graph: GraphNodeBlockBuild<'_>,
882    policy: DurabilityPolicy,
883) -> Result<SegmentMeta, SegmentError> {
884    let bytes = encode_segment_with_graph(build, graph)?;
885    publish_segment(vfs, directory, build, policy, &bytes)
886}
887
888/// Writes and atomically publishes graph blocks plus dense document identities.
889pub fn write_segment_with_graph_and_documents(
890    vfs: &dyn Vfs,
891    directory: &Path,
892    build: SegmentBuild<'_>,
893    graph: GraphNodeBlockBuild<'_>,
894    documents: SegmentDocumentVersions<'_>,
895    policy: DurabilityPolicy,
896) -> Result<SegmentMeta, SegmentError> {
897    let bytes = encode_segment_with_graph_and_documents(build, graph, documents)?;
898    publish_segment(vfs, directory, build, policy, &bytes)
899}
900
901/// Writes authenticated graph bytes while copying opaque source regions.
902pub(crate) fn write_segment_with_encoded_graph_and_copied_regions(
903    vfs: &dyn Vfs,
904    directory: &Path,
905    build: SegmentBuild<'_>,
906    graph_bytes: &[u8],
907    documents: Option<SegmentDocumentVersions<'_>>,
908    copied_regions: &[SegmentCopiedRegion<'_>],
909    policy: DurabilityPolicy,
910) -> Result<SegmentMeta, SegmentError> {
911    let graph = decode_node_blocks(graph_bytes)?;
912    if graph.layout().dims() != build.dims {
913        return Err(SegmentError::Geometry(format!(
914            "graph dimensions {}, segment dimensions {}",
915            graph.layout().dims(),
916            build.dims
917        )));
918    }
919    if graph.node_count() != build.columns.row_count() {
920        return Err(SegmentError::Geometry(format!(
921            "graph nodes {}, segment rows {}",
922            graph.node_count(),
923            build.columns.row_count()
924        )));
925    }
926    let bytes = encode_segment_inner(
927        build,
928        Some(graph_bytes.to_vec()),
929        documents,
930        None,
931        None,
932        None,
933        copied_regions,
934    )?;
935    publish_segment(vfs, directory, build, policy, &bytes)
936}
937
938fn publish_segment(
939    vfs: &dyn Vfs,
940    directory: &Path,
941    build: SegmentBuild<'_>,
942    policy: DurabilityPolicy,
943    bytes: &[u8],
944) -> Result<SegmentMeta, SegmentError> {
945    let final_path = directory.join(build.id.file_name());
946    let temporary_path = temporary_path(directory, build.id);
947    vfs.write(&temporary_path, bytes)
948        .map_err(|error| SegmentError::io(&temporary_path, error))?;
949    let observed = vfs
950        .read(&temporary_path)
951        .map_err(|error| SegmentError::io(&temporary_path, error))?;
952    if observed != bytes {
953        let (check, detail) = if observed.len() != bytes.len() {
954            (
955                FormatCheck::FileLength,
956                format!(
957                    "temporary segment write retained {} bytes, expected {}",
958                    observed.len(),
959                    bytes.len()
960                ),
961            )
962        } else {
963            (
964                FormatCheck::FileChecksum,
965                "temporary segment bytes differ from the encoded image".to_owned(),
966            )
967        };
968        return Err(FormatError::new(temporary_path.display().to_string(), check, detail).into());
969    }
970    match policy.data_file_sync() {
971        SyncRequirement::Skip => {}
972        SyncRequirement::Sync(kind) => vfs
973            .sync(&temporary_path, kind)
974            .map_err(|error| SegmentError::io(&temporary_path, error))?,
975    }
976    vfs.rename(&temporary_path, &final_path)
977        .map_err(|error| SegmentError::io(&final_path, error))?;
978    match policy.directory_sync() {
979        SyncRequirement::Skip => {}
980        SyncRequirement::Sync(kind) => vfs
981            .sync(directory, kind)
982            .map_err(|error| SegmentError::io(directory, error))?,
983    }
984    Ok(SegmentMeta {
985        id: build.id,
986        row_count: build.columns.row_count(),
987        scheme: build.scheme,
988        dims: build.dims,
989        file_size: bytes.len() as u64,
990        epoch_id: None,
991        clustering_key_range: super::ClusteringKeyRange::Unstamped,
992    })
993}
994
995fn temporary_path(directory: &Path, id: SegmentId) -> PathBuf {
996    directory.join(format!(".{}.tmp", id.file_name()))
997}
998
999/// Rebuilds only postings; all other source regions remain byte-identical.
1000pub(crate) fn reindex_segment(
1001    vfs: &dyn Vfs,
1002    directory: &Path,
1003    input: &super::reader::SegmentReader,
1004    id: SegmentId,
1005    analyzer: &crate::fts::tokenizer::Analyzer,
1006    policy: DurabilityPolicy,
1007) -> Result<SegmentMeta, SegmentError> {
1008    use crate::fts::{
1009        index::{Document, SegmentIndex},
1010        sealed::SealedSegment,
1011    };
1012    let text = input.stored_text()?;
1013    if text.is_none() && input.postings()?.is_some() {
1014        return Err(SegmentError::Geometry(
1015            "cannot reindex postings without stored text".to_owned(),
1016        ));
1017    }
1018    let mut index = SegmentIndex::new();
1019    for row in 0..input.meta().row_count as usize {
1020        let source = match text.as_ref() {
1021            Some(text) => text
1022                .row(row)
1023                .ok_or_else(|| SegmentError::Geometry("missing stored text row".to_owned()))?,
1024            None => None,
1025        };
1026        let document = source.map_or_else(Document::new, Document::with_text);
1027        index
1028            .push_document(analyzer, &document)
1029            .map_err(|error| SegmentError::Geometry(error.to_string()))?;
1030    }
1031    let postings = SealedSegment::seal(&index)
1032        .map_err(|error| SegmentError::Geometry(error.to_string()))?
1033        .encode_region()
1034        .map_err(|error| SegmentError::Geometry(error.to_string()))?;
1035    let mut copied = Vec::new();
1036    for entry in input.directory() {
1037        if entry.kind != RegionKind::Postings.id() && entry.kind != RegionKind::ChecksumTable.id() {
1038            copied.push(SegmentCopiedRegion {
1039                kind: entry.kind,
1040                version: entry.version,
1041                bytes: input.region_by_id(entry.kind)?,
1042            });
1043        }
1044    }
1045    copied.push(SegmentCopiedRegion {
1046        kind: RegionKind::Postings.id(),
1047        version: FormatFamily::Postings.current_version(),
1048        bytes: &postings,
1049    });
1050    copied.sort_by_key(|region| region.kind);
1051    let columns = input.columns()?;
1052    let alive = input.alive()?;
1053    // encode_regions uses only identity/geometry; vector bytes are copied above.
1054    let build = SegmentBuild {
1055        id,
1056        scheme: input.meta().scheme,
1057        dims: input.meta().dims,
1058        codes: &[],
1059        factors: SegmentFactors::F32,
1060        rescore: &[],
1061        columns: &columns,
1062        alive: &alive,
1063    };
1064    let bytes = encode_regions(build, &[], &copied)?;
1065    let mut meta = publish_segment(vfs, directory, build, policy, &bytes)?;
1066    meta.epoch_id = input.meta().epoch_id;
1067    meta.clustering_key_range = input.meta().clustering_key_range;
1068    Ok(meta)
1069}