Skip to main content

zeppelin_embed/graph/
consolidate.rs

1//! Copy-on-write consolidation of sealed graph segments (Task 19-M8).
2//!
3//! Consolidation merges every published sealed graph segment's live rows
4//! into one new sealed segment with new dense row ids, rebuilds every
5//! derived region (postings are retokenized from carried stored text, the
6//! graph is rebuilt over the union), and atomically publishes one N-to-1
7//! manifest replacement. Sealed inputs are never mutated in place; the
8//! manifest commit precedes every unlink, exactly like purge replacement.
9//!
10//! The only persisted artifact this module owns is the private
11//! `.consolidate.checkpoint` resumability file. It is never read by
12//! queries and is refused-and-cleared on any validation failure.
13
14use std::collections::HashSet;
15use std::path::{Path, PathBuf};
16
17use xxhash_rust::xxh3::{xxh3_64, xxh3_64_with_seed};
18
19use crate::fts::index::{Document as LexicalDocument, SegmentIndex};
20use crate::fts::sealed::SealedSegment;
21use crate::fts::tokenizer::Analyzer;
22use crate::ingest::DocId;
23use crate::lifecycle::StoreError;
24use crate::lifecycle::durability::{DurabilityPolicy, SyncRequirement};
25use crate::meta::{AliveSet, ColumnStoreBuilder};
26use crate::segment::layout::RegionKind;
27use crate::segment::reader::SegmentReader;
28use crate::segment::writer::{
29    SegmentBuild, SegmentDocumentVersions, SegmentPayloads, SegmentPostings, SegmentStoredMetadata,
30    SegmentStoredText, write_segment, write_segment_with_documents_payloads,
31};
32use crate::segment::{ClusteringKeyRange, SegmentError, SegmentId};
33use crate::vfs::Vfs;
34
35/// Fixed private name of the store's single consolidation checkpoint.
36pub(crate) const CONSOLIDATION_CHECKPOINT_FILE: &str = ".consolidate.checkpoint";
37/// Magic prefix of every consolidation checkpoint.
38pub const CONSOLIDATION_CHECKPOINT_MAGIC: [u8; 8] = *b"ZECONCP1";
39const CHECKPOINT_VERSION: u16 = 1;
40/// magic + version + reserved + input_count.
41const CHECKPOINT_PREFIX_BYTES: usize = 16;
42/// merge id + output id + rows_emitted + cursor pair + output hash + checksum.
43const CHECKPOINT_TAIL_BYTES: usize = 16 + 16 + 8 + 4 + 4 + 8 + 8;
44/// Manifest transitions never carry more sealed segments than this.
45const CHECKPOINT_MAX_INPUTS: u32 = 65_536;
46const MERGE_ID_SEED: u64 = 0x0019_4d38_0000_0001;
47const OUTPUT_ID_SEED: u64 = 0x0019_4d38_0000_0002;
48
49/// Typed failure of the consolidation merge pass or its checkpoint.
50#[derive(Debug)]
51pub enum ConsolidateError {
52    /// Reading or writing a segment failed its validated contract.
53    Segment(SegmentError),
54    /// Store lifecycle, accounting, or filesystem failure.
55    Store(StoreError),
56    /// Checked geometry over merged rows was violated.
57    Geometry(String),
58    /// Two live rows across the merge inputs share one document id.
59    DuplicateDocument {
60        /// Application document id present in more than one live input row.
61        doc_id: u128,
62    },
63    /// Rebuilding the merged lexical region failed.
64    Lexical(String),
65    /// Reading, replacing, or removing the consolidation checkpoint failed.
66    CheckpointIo {
67        /// Checkpoint or checkpoint-temporary path.
68        path: PathBuf,
69        /// Underlying filesystem failure.
70        source: std::io::Error,
71    },
72    /// Checkpoint bytes failed their closed contract.
73    CheckpointCorrupt(String),
74}
75
76impl std::fmt::Display for ConsolidateError {
77    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        match self {
79            Self::Segment(error) => write!(formatter, "consolidation segment: {error}"),
80            Self::Store(error) => write!(formatter, "consolidation store: {error}"),
81            Self::Geometry(detail) => write!(formatter, "consolidation geometry: {detail}"),
82            Self::DuplicateDocument { doc_id } => write!(
83                formatter,
84                "consolidation found document {doc_id} live in more than one input segment"
85            ),
86            Self::Lexical(detail) => write!(formatter, "consolidation lexical rebuild: {detail}"),
87            Self::CheckpointIo { path, source } => write!(
88                formatter,
89                "consolidation checkpoint I/O {}: {source}",
90                path.display()
91            ),
92            Self::CheckpointCorrupt(detail) => {
93                write!(formatter, "consolidation checkpoint is corrupt: {detail}")
94            }
95        }
96    }
97}
98
99impl std::error::Error for ConsolidateError {
100    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
101        match self {
102            Self::Segment(error) => Some(error),
103            Self::Store(error) => Some(error),
104            Self::CheckpointIo { source, .. } => Some(source),
105            Self::Geometry(_)
106            | Self::DuplicateDocument { .. }
107            | Self::Lexical(_)
108            | Self::CheckpointCorrupt(_) => None,
109        }
110    }
111}
112
113impl From<SegmentError> for ConsolidateError {
114    fn from(error: SegmentError) -> Self {
115        Self::Segment(error)
116    }
117}
118
119impl From<StoreError> for ConsolidateError {
120    fn from(error: StoreError) -> Self {
121        Self::Store(error)
122    }
123}
124
125/// Whether the merge pass can carry one source region kind.
126///
127/// Everything the merge re-encodes, rebuilds, or replaces is supported;
128/// reserved accelerator regions and forward-compatible unknown kinds defer
129/// consolidation because their bytes cannot be re-derived for new row ids.
130#[must_use]
131pub(crate) const fn consolidation_supports_region(kind: u16) -> bool {
132    matches!(
133        RegionKind::from_id(kind),
134        Some(
135            RegionKind::Columns
136                | RegionKind::Alive
137                | RegionKind::VectorCodes
138                | RegionKind::VectorFactors
139                | RegionKind::VectorRescore
140                | RegionKind::DocumentVersions
141                | RegionKind::GraphNodeBlocks
142                | RegionKind::ChecksumTable
143                | RegionKind::Postings
144                | RegionKind::StoredMetadata
145                | RegionKind::StoredText,
146        )
147    )
148}
149
150/// Deterministic unpublished merge-intermediate id for one input set.
151#[must_use]
152pub(crate) fn consolidation_merge_id(inputs: &[SegmentId]) -> SegmentId {
153    derived_segment_id(inputs, MERGE_ID_SEED)
154}
155
156/// Deterministic published output id for one input set.
157#[must_use]
158pub(crate) fn consolidation_output_id(inputs: &[SegmentId]) -> SegmentId {
159    derived_segment_id(inputs, OUTPUT_ID_SEED)
160}
161
162fn derived_segment_id(inputs: &[SegmentId], seed: u64) -> SegmentId {
163    let mut input_bytes = Vec::with_capacity(inputs.len().saturating_mul(16));
164    for id in inputs {
165        input_bytes.extend_from_slice(id.as_bytes());
166    }
167    let first = xxh3_64_with_seed(&input_bytes, seed).to_be_bytes();
168    let second = xxh3_64_with_seed(&input_bytes, !seed).to_be_bytes();
169    let mut bytes = [0_u8; 16];
170    if let Some(prefix) = bytes.get_mut(..8) {
171        prefix.copy_from_slice(&first);
172    }
173    if let Some(suffix) = bytes.get_mut(8..) {
174        suffix.copy_from_slice(&second);
175    }
176    SegmentId::from_bytes(bytes)
177}
178
179/// Durable progress marker between the merge pass and the graph build.
180#[derive(Clone, Debug, Eq, PartialEq)]
181pub(crate) struct ConsolidationCheckpoint {
182    /// Sealed input segment ids, strictly ascending.
183    pub(crate) inputs: Vec<SegmentId>,
184    /// Unpublished merge-intermediate segment id.
185    pub(crate) merge_id: SegmentId,
186    /// Output segment id the eventual publish will use.
187    pub(crate) output_id: SegmentId,
188    /// Live rows emitted into the intermediate segment.
189    pub(crate) rows_emitted: u64,
190    /// Xxh3-64 over the complete intermediate segment file bytes.
191    pub(crate) output_hash: u64,
192}
193
194/// Encodes one checkpoint in the frozen `ZECONCP1` little-endian layout.
195pub(crate) fn encode_consolidation_checkpoint(
196    checkpoint: &ConsolidationCheckpoint,
197) -> Result<Vec<u8>, ConsolidateError> {
198    let input_count = u32::try_from(checkpoint.inputs.len())
199        .map_err(|_| ConsolidateError::Geometry("checkpoint input count exceeds u32".to_owned()))?;
200    if input_count == 0 || input_count > CHECKPOINT_MAX_INPUTS {
201        return Err(ConsolidateError::Geometry(format!(
202            "checkpoint input count {input_count} is outside 1..={CHECKPOINT_MAX_INPUTS}"
203        )));
204    }
205    let total = CHECKPOINT_PREFIX_BYTES
206        .checked_add(checkpoint.inputs.len().saturating_mul(16))
207        .and_then(|bytes| bytes.checked_add(CHECKPOINT_TAIL_BYTES))
208        .ok_or_else(|| ConsolidateError::Geometry("checkpoint length overflow".to_owned()))?;
209    let mut bytes = Vec::with_capacity(total);
210    bytes.extend_from_slice(&CONSOLIDATION_CHECKPOINT_MAGIC);
211    bytes.extend_from_slice(&CHECKPOINT_VERSION.to_le_bytes());
212    bytes.extend_from_slice(&0_u16.to_le_bytes());
213    bytes.extend_from_slice(&input_count.to_le_bytes());
214    for id in &checkpoint.inputs {
215        bytes.extend_from_slice(id.as_bytes());
216    }
217    bytes.extend_from_slice(checkpoint.merge_id.as_bytes());
218    bytes.extend_from_slice(checkpoint.output_id.as_bytes());
219    bytes.extend_from_slice(&checkpoint.rows_emitted.to_le_bytes());
220    // Version 1 admits the merge pass whole, so the persisted cursor is
221    // always terminal: segment cursor at input_count, row cursor at zero.
222    bytes.extend_from_slice(&input_count.to_le_bytes());
223    bytes.extend_from_slice(&0_u32.to_le_bytes());
224    bytes.extend_from_slice(&checkpoint.output_hash.to_le_bytes());
225    bytes.extend_from_slice(&xxh3_64(&bytes).to_le_bytes());
226    if bytes.len() != total {
227        return Err(ConsolidateError::Geometry(format!(
228            "checkpoint encoded {} bytes, expected {total}",
229            bytes.len()
230        )));
231    }
232    Ok(bytes)
233}
234
235/// Decodes and fully validates one `ZECONCP1` checkpoint.
236pub(crate) fn decode_consolidation_checkpoint(
237    bytes: &[u8],
238) -> Result<ConsolidationCheckpoint, ConsolidateError> {
239    let corrupt = |detail: String| ConsolidateError::CheckpointCorrupt(detail);
240    let minimum = CHECKPOINT_PREFIX_BYTES
241        .checked_add(16)
242        .and_then(|prefix| prefix.checked_add(CHECKPOINT_TAIL_BYTES))
243        .ok_or_else(|| corrupt("minimum length overflow".to_owned()))?;
244    if bytes.len() < minimum {
245        return Err(corrupt(format!(
246            "truncated: need at least {minimum} bytes, got {}",
247            bytes.len()
248        )));
249    }
250    if bytes.get(..8) != Some(CONSOLIDATION_CHECKPOINT_MAGIC.as_slice()) {
251        return Err(corrupt("magic does not match ZECONCP1".to_owned()));
252    }
253    let checksum_start = bytes
254        .len()
255        .checked_sub(8)
256        .ok_or_else(|| corrupt("checksum underflow".to_owned()))?;
257    let stored_checksum = read_u64(bytes, checksum_start)?;
258    let checksummed = bytes
259        .get(..checksum_start)
260        .ok_or_else(|| corrupt("checksummed prefix is unavailable".to_owned()))?;
261    let actual_checksum = xxh3_64(checksummed);
262    if stored_checksum != actual_checksum {
263        return Err(corrupt(format!(
264            "xxh3-64 expected {stored_checksum:#018x}, computed {actual_checksum:#018x}"
265        )));
266    }
267    let version = read_u16(bytes, 8)?;
268    if version != CHECKPOINT_VERSION {
269        return Err(corrupt(format!(
270            "version {version}, expected {CHECKPOINT_VERSION}"
271        )));
272    }
273    if read_u16(bytes, 10)? != 0 {
274        return Err(corrupt("reserved bytes are non-zero".to_owned()));
275    }
276    let input_count = read_u32(bytes, 12)?;
277    if input_count == 0 || input_count > CHECKPOINT_MAX_INPUTS {
278        return Err(corrupt(format!(
279            "input count {input_count} is outside 1..={CHECKPOINT_MAX_INPUTS}"
280        )));
281    }
282    let expected = CHECKPOINT_PREFIX_BYTES
283        .checked_add((input_count as usize).saturating_mul(16))
284        .and_then(|bytes| bytes.checked_add(CHECKPOINT_TAIL_BYTES))
285        .ok_or_else(|| corrupt("declared length overflow".to_owned()))?;
286    if bytes.len() != expected {
287        return Err(corrupt(format!(
288            "length {} bytes, declared layout needs {expected}",
289            bytes.len()
290        )));
291    }
292    let mut inputs = Vec::with_capacity(input_count as usize);
293    let mut previous: Option<[u8; 16]> = None;
294    for index in 0..input_count as usize {
295        let offset = CHECKPOINT_PREFIX_BYTES
296            .checked_add(index.saturating_mul(16))
297            .ok_or_else(|| corrupt("input offset overflow".to_owned()))?;
298        let id = read_id(bytes, offset)?;
299        if previous.is_some_and(|earlier| earlier >= id) {
300            return Err(corrupt("input ids are not strictly ascending".to_owned()));
301        }
302        previous = Some(id);
303        inputs.push(SegmentId::from_bytes(id));
304    }
305    let tail = CHECKPOINT_PREFIX_BYTES
306        .checked_add((input_count as usize).saturating_mul(16))
307        .ok_or_else(|| corrupt("tail offset overflow".to_owned()))?;
308    let merge_id = SegmentId::from_bytes(read_id(bytes, tail)?);
309    let output_id = SegmentId::from_bytes(read_id(bytes, tail.saturating_add(16))?);
310    let rows_emitted = read_u64(bytes, tail.saturating_add(32))?;
311    let cursor_segment = read_u32(bytes, tail.saturating_add(40))?;
312    let cursor_row = read_u32(bytes, tail.saturating_add(44))?;
313    if cursor_segment != input_count || cursor_row != 0 {
314        return Err(corrupt(format!(
315            "cursor ({cursor_segment}, {cursor_row}) is not terminal for {input_count} inputs"
316        )));
317    }
318    let output_hash = read_u64(bytes, tail.saturating_add(48))?;
319    Ok(ConsolidationCheckpoint {
320        inputs,
321        merge_id,
322        output_id,
323        rows_emitted,
324        output_hash,
325    })
326}
327
328/// Validates arbitrary consolidation-checkpoint bytes without a store.
329///
330/// This is the shared parser seam for the dependency-free fuzz target.
331pub fn validate_consolidation_checkpoint(bytes: &[u8]) -> Result<(), String> {
332    match decode_consolidation_checkpoint(bytes) {
333        Ok(_) => Ok(()),
334        Err(error) => Err(error.to_string()),
335    }
336}
337
338fn read_u16(bytes: &[u8], offset: usize) -> Result<u16, ConsolidateError> {
339    bytes
340        .get(offset..offset.saturating_add(2))
341        .and_then(|value| value.try_into().ok())
342        .map(u16::from_le_bytes)
343        .ok_or_else(|| ConsolidateError::CheckpointCorrupt(format!("u16 at {offset} is truncated")))
344}
345
346fn read_u32(bytes: &[u8], offset: usize) -> Result<u32, ConsolidateError> {
347    bytes
348        .get(offset..offset.saturating_add(4))
349        .and_then(|value| value.try_into().ok())
350        .map(u32::from_le_bytes)
351        .ok_or_else(|| ConsolidateError::CheckpointCorrupt(format!("u32 at {offset} is truncated")))
352}
353
354fn read_u64(bytes: &[u8], offset: usize) -> Result<u64, ConsolidateError> {
355    bytes
356        .get(offset..offset.saturating_add(8))
357        .and_then(|value| value.try_into().ok())
358        .map(u64::from_le_bytes)
359        .ok_or_else(|| ConsolidateError::CheckpointCorrupt(format!("u64 at {offset} is truncated")))
360}
361
362fn read_id(bytes: &[u8], offset: usize) -> Result<[u8; 16], ConsolidateError> {
363    bytes
364        .get(offset..offset.saturating_add(16))
365        .and_then(|value| value.try_into().ok())
366        .ok_or_else(|| {
367            ConsolidateError::CheckpointCorrupt(format!("segment id at {offset} is truncated"))
368        })
369}
370
371/// Writes the checkpoint through the temp-sync-rename-dirsync discipline.
372pub(crate) fn write_consolidation_checkpoint(
373    vfs: &dyn Vfs,
374    directory: &Path,
375    checkpoint: &ConsolidationCheckpoint,
376    policy: DurabilityPolicy,
377) -> Result<(), ConsolidateError> {
378    let bytes = encode_consolidation_checkpoint(checkpoint)?;
379    let path = directory.join(CONSOLIDATION_CHECKPOINT_FILE);
380    let temporary = directory.join(format!("{CONSOLIDATION_CHECKPOINT_FILE}.tmp"));
381    if let Err(source) = vfs.write(&temporary, &bytes) {
382        let _ = vfs.delete(&temporary);
383        return Err(ConsolidateError::CheckpointIo {
384            path: temporary,
385            source,
386        });
387    }
388    if let SyncRequirement::Sync(kind) = policy.data_file_sync()
389        && let Err(source) = vfs.sync(&temporary, kind)
390    {
391        let _ = vfs.delete(&temporary);
392        return Err(ConsolidateError::CheckpointIo {
393            path: temporary,
394            source,
395        });
396    }
397    if let Err(source) = vfs.rename(&temporary, &path) {
398        let _ = vfs.delete(&temporary);
399        return Err(ConsolidateError::CheckpointIo { path, source });
400    }
401    match policy.directory_sync() {
402        SyncRequirement::Skip => Ok(()),
403        SyncRequirement::Sync(kind) => {
404            vfs.sync(directory, kind)
405                .map_err(|source| ConsolidateError::CheckpointIo {
406                    path: directory.to_path_buf(),
407                    source,
408                })
409        }
410    }
411}
412
413/// Reads the checkpoint; absent is `Ok(None)`, invalid bytes are typed corrupt.
414pub(crate) fn read_consolidation_checkpoint(
415    vfs: &dyn Vfs,
416    directory: &Path,
417) -> Result<Option<ConsolidationCheckpoint>, ConsolidateError> {
418    let path = directory.join(CONSOLIDATION_CHECKPOINT_FILE);
419    match vfs.read(&path) {
420        Ok(bytes) => decode_consolidation_checkpoint(&bytes).map(Some),
421        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
422        Err(source) => Err(ConsolidateError::CheckpointIo { path, source }),
423    }
424}
425
426/// Removes the checkpoint and its temporary; both tolerate absence.
427pub(crate) fn remove_consolidation_checkpoint(
428    vfs: &dyn Vfs,
429    directory: &Path,
430) -> Result<(), ConsolidateError> {
431    for name in [
432        format!("{CONSOLIDATION_CHECKPOINT_FILE}.tmp"),
433        CONSOLIDATION_CHECKPOINT_FILE.to_owned(),
434    ] {
435        let path = directory.join(name);
436        match vfs.delete(&path) {
437            Ok(()) => {}
438            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
439            Err(source) => return Err(ConsolidateError::CheckpointIo { path, source }),
440        }
441    }
442    Ok(())
443}
444
445/// Result of one completed merge pass, before any graph exists.
446pub(crate) struct MergedIntermediate {
447    /// Live rows copied out of the inputs.
448    pub(crate) live_rows: u64,
449    /// Xxh3-64 over the complete intermediate file bytes.
450    pub(crate) file_hash: u64,
451}
452
453struct MergeGather {
454    codes: Vec<u8>,
455    factors: Vec<crate::quant::Bit4Factors>,
456    rescore: Vec<f32>,
457    doc_ids: Vec<DocId>,
458    revisions: Vec<crate::ingest::Revision>,
459    metadata: Option<(Vec<u64>, Vec<u8>)>,
460    text: Option<(Vec<u8>, Vec<u64>, Vec<u8>)>,
461    lexical: Option<SegmentIndex>,
462}
463
464/// Merges every input's live rows into one unpublished sealed segment.
465///
466/// Inputs must be ascending by segment id; rows are emitted in input order
467/// then ascending source row id, so the output is deterministic for one
468/// input set. Returns `Ok(None)` when no input holds a live row.
469pub(crate) fn merge_segments(
470    vfs: &dyn Vfs,
471    directory: &Path,
472    inputs: &[&SegmentReader],
473    merge_id: SegmentId,
474    analyzer: &Analyzer,
475    policy: DurabilityPolicy,
476) -> Result<Option<MergedIntermediate>, ConsolidateError> {
477    validate_merge_inputs(inputs)?;
478    let dims = merge_dims(inputs)?;
479    let carry_documents = region_present(inputs, RegionKind::DocumentVersions)?;
480    let carry_metadata = inputs
481        .iter()
482        .any(|reader| has_region(reader, RegionKind::StoredMetadata));
483    let rebuild_lexical = inputs
484        .iter()
485        .any(|reader| has_region(reader, RegionKind::Postings));
486    let survivors = live_rows(inputs)?;
487    let total_rows: usize = survivors.iter().map(Vec::len).sum();
488    if total_rows == 0 {
489        return Ok(None);
490    }
491    let row_count = u32::try_from(total_rows).map_err(|_| {
492        ConsolidateError::Geometry("merged rows exceed the u32 row space".to_owned())
493    })?;
494    let mut gather = MergeGather {
495        codes: Vec::new(),
496        factors: Vec::new(),
497        rescore: Vec::new(),
498        doc_ids: Vec::new(),
499        revisions: Vec::new(),
500        metadata: carry_metadata.then(|| (Vec::new(), Vec::new())),
501        text: None,
502        lexical: rebuild_lexical.then(SegmentIndex::new),
503    };
504    let merge_schema = inputs
505        .first()
506        .ok_or_else(|| ConsolidateError::Geometry("merge requires at least one input".to_owned()))?
507        .columns()?
508        .schema()
509        .clone();
510    let mut columns = ColumnStoreBuilder::new(merge_schema.clone());
511    let carry_text = inputs
512        .iter()
513        .any(|reader| has_region(reader, RegionKind::StoredText));
514    if carry_text {
515        gather.text = Some((Vec::new(), Vec::new(), Vec::new()));
516    }
517    for (reader, rows) in inputs.iter().zip(&survivors) {
518        gather_input(
519            reader,
520            rows,
521            dims,
522            carry_documents,
523            analyzer,
524            &merge_schema,
525            &mut gather,
526            &mut columns,
527        )?;
528    }
529    if carry_documents {
530        let mut seen = HashSet::with_capacity(gather.doc_ids.len());
531        for doc_id in &gather.doc_ids {
532            if !seen.insert(*doc_id) {
533                return Err(ConsolidateError::DuplicateDocument {
534                    doc_id: doc_id.get(),
535                });
536            }
537        }
538    }
539    let columns = columns
540        .finish()
541        .map_err(|source| ConsolidateError::Segment(SegmentError::Columns(source.to_string())))?;
542    let alive = AliveSet::new(row_count);
543    let scheme = merge_scheme(inputs)?;
544    let build = SegmentBuild {
545        id: merge_id,
546        scheme,
547        dims,
548        codes: &gather.codes,
549        factors: crate::segment::writer::SegmentFactors::Bit4(&gather.factors),
550        rescore: &gather.rescore,
551        columns: &columns,
552        alive: &alive,
553    };
554    let text_has_rows = gather
555        .text
556        .as_ref()
557        .is_some_and(|(present, _, _)| present.contains(&1));
558    let postings_bytes = match gather.lexical.as_ref() {
559        Some(lexical) if text_has_rows => Some(
560            SealedSegment::seal(lexical)
561                .map_err(|error| ConsolidateError::Lexical(error.to_string()))?
562                .encode_region()
563                .map_err(|error| ConsolidateError::Lexical(error.to_string()))?,
564        ),
565        _ => None,
566    };
567    let _meta = if carry_documents {
568        let documents = SegmentDocumentVersions {
569            doc_ids: &gather.doc_ids,
570            revisions: &gather.revisions,
571        };
572        let metadata = gather
573            .metadata
574            .as_ref()
575            .map(|(end_offsets, bytes)| SegmentStoredMetadata { end_offsets, bytes });
576        let text = gather
577            .text
578            .as_ref()
579            .filter(|(present, _, _)| present.contains(&1))
580            .map(|(present, end_offsets, bytes)| SegmentStoredText {
581                present,
582                end_offsets,
583                bytes,
584            });
585        let postings = postings_bytes
586            .as_deref()
587            .map(|bytes| SegmentPostings { bytes });
588        write_segment_with_documents_payloads(
589            vfs,
590            directory,
591            build,
592            SegmentPayloads {
593                documents,
594                metadata,
595                text,
596                postings,
597            },
598            policy,
599        )?
600    } else {
601        if gather.metadata.is_some() || text_has_rows || postings_bytes.is_some() {
602            return Err(ConsolidateError::Geometry(
603                "merged payload regions require document identities".to_owned(),
604            ));
605        }
606        write_segment(vfs, directory, build, policy)?
607    };
608    let path = directory.join(merge_id.file_name());
609    let file_bytes = vfs
610        .read(&path)
611        .map_err(|source| ConsolidateError::Segment(SegmentError::io(&path, source)))?;
612    Ok(Some(MergedIntermediate {
613        live_rows: u64::from(row_count),
614        file_hash: xxh3_64(&file_bytes),
615    }))
616}
617
618/// Recomputes the clustering-key union over merged live rows (10-C rules).
619pub(crate) fn merged_clustering(
620    columns: &crate::meta::ColumnStore,
621    alive: &AliveSet,
622) -> Result<ClusteringKeyRange, ConsolidateError> {
623    crate::ingest::purge_support::clustering_range(columns, alive).map_err(ConsolidateError::Store)
624}
625
626fn validate_merge_inputs(inputs: &[&SegmentReader]) -> Result<(), ConsolidateError> {
627    if inputs.len() < 2 {
628        return Err(ConsolidateError::Geometry(format!(
629            "consolidation requires at least two inputs, got {}",
630            inputs.len()
631        )));
632    }
633    let mut previous: Option<&SegmentId> = None;
634    for reader in inputs {
635        let id = &reader.meta().id;
636        if previous.is_some_and(|earlier| earlier.as_bytes() >= id.as_bytes()) {
637            return Err(ConsolidateError::Geometry(
638                "consolidation inputs are not strictly ascending by segment id".to_owned(),
639            ));
640        }
641        previous = Some(id);
642        for entry in reader.directory() {
643            if !consolidation_supports_region(entry.kind) {
644                return Err(ConsolidateError::Geometry(format!(
645                    "consolidation cannot carry source region kind {}",
646                    entry.kind
647                )));
648            }
649        }
650        if has_region(reader, RegionKind::Postings) && !has_region(reader, RegionKind::StoredText) {
651            return Err(ConsolidateError::Lexical(format!(
652                "segment {} carries postings without the stored text needed to rebuild them",
653                reader.meta().id
654            )));
655        }
656    }
657    Ok(())
658}
659
660fn merge_dims(inputs: &[&SegmentReader]) -> Result<u32, ConsolidateError> {
661    let mut dims = None;
662    for reader in inputs {
663        let candidate = reader.meta().dims;
664        if dims.is_some_and(|expected| expected != candidate) {
665            return Err(ConsolidateError::Geometry(
666                "consolidation inputs disagree on vector dimensions".to_owned(),
667            ));
668        }
669        dims = Some(candidate);
670    }
671    dims.ok_or_else(|| ConsolidateError::Geometry("consolidation has no inputs".to_owned()))
672}
673
674fn merge_scheme(inputs: &[&SegmentReader]) -> Result<u16, ConsolidateError> {
675    for reader in inputs {
676        if reader.meta().scheme != 4 {
677            return Err(ConsolidateError::Geometry(format!(
678                "consolidation requires Bit4 scheme 4, segment {} has {}",
679                reader.meta().id,
680                reader.meta().scheme
681            )));
682        }
683    }
684    Ok(4)
685}
686
687fn region_present(inputs: &[&SegmentReader], kind: RegionKind) -> Result<bool, ConsolidateError> {
688    let mut present = None;
689    for reader in inputs {
690        let has = has_region(reader, kind);
691        if present.is_some_and(|expected| expected != has) {
692            return Err(ConsolidateError::Geometry(format!(
693                "consolidation inputs disagree on region kind {} presence",
694                kind.id()
695            )));
696        }
697        present = Some(has);
698    }
699    Ok(present.unwrap_or(false))
700}
701
702fn has_region(reader: &SegmentReader, kind: RegionKind) -> bool {
703    reader
704        .directory()
705        .iter()
706        .any(|entry| entry.kind == kind.id())
707}
708
709fn live_rows(inputs: &[&SegmentReader]) -> Result<Vec<Vec<usize>>, ConsolidateError> {
710    let mut survivors = Vec::with_capacity(inputs.len());
711    for reader in inputs {
712        let alive = reader.alive()?;
713        survivors.push(alive.iter_alive().map(|row| row as usize).collect());
714    }
715    Ok(survivors)
716}
717
718#[allow(clippy::too_many_arguments)]
719fn gather_input(
720    reader: &SegmentReader,
721    rows: &[usize],
722    dims: u32,
723    carry_documents: bool,
724    analyzer: &Analyzer,
725    merge_schema: &crate::meta::Schema,
726    gather: &mut MergeGather,
727    columns: &mut ColumnStoreBuilder,
728) -> Result<(), ConsolidateError> {
729    let dims = dims as usize;
730    let (codes, factors) = crate::ingest::purge_support::gather_survivor_codes(reader, rows, dims)
731        .map_err(ConsolidateError::Store)?;
732    let crate::ingest::purge_support::OwnedFactors::Bit4(factors) = factors else {
733        return Err(ConsolidateError::Geometry(
734            "consolidation gathered non-Bit4 factors".to_owned(),
735        ));
736    };
737    gather.codes.extend_from_slice(&codes);
738    gather.factors.extend_from_slice(&factors);
739    let rescore = crate::ingest::purge_support::gather_survivor_rescore(reader, rows, dims)
740        .map_err(ConsolidateError::Store)?;
741    gather.rescore.extend_from_slice(&rescore);
742    let source_columns = reader.columns()?;
743    if source_columns.schema() != merge_schema {
744        return Err(ConsolidateError::Geometry(format!(
745            "segment {} schema differs from the merge schema",
746            reader.meta().id
747        )));
748    }
749    crate::ingest::purge_support::append_survivor_columns(columns, &source_columns, rows)
750        .map_err(ConsolidateError::Store)?;
751    if carry_documents {
752        let (doc_ids, revisions) =
753            crate::ingest::purge_support::gather_survivor_documents(reader, rows)
754                .map_err(ConsolidateError::Store)?;
755        gather.doc_ids.extend_from_slice(&doc_ids);
756        gather.revisions.extend_from_slice(&revisions);
757    }
758    if let Some((end_offsets, bytes)) = gather.metadata.as_mut() {
759        let source = reader.stored_metadata()?;
760        for row in rows {
761            if let Some(source) = source {
762                bytes.extend_from_slice(source.row(*row).ok_or_else(|| {
763                    ConsolidateError::Segment(SegmentError::Geometry(format!(
764                        "stored metadata row {row} is missing"
765                    )))
766                })?);
767            }
768            end_offsets.push(u64::try_from(bytes.len()).map_err(|_| {
769                ConsolidateError::Geometry("stored metadata offset exceeds u64".to_owned())
770            })?);
771        }
772    }
773    let source_text = reader.stored_text()?;
774    if let Some((present, end_offsets, bytes)) = gather.text.as_mut() {
775        for row in rows {
776            let text = match source_text {
777                Some(view) => view.row(*row).ok_or_else(|| {
778                    ConsolidateError::Segment(SegmentError::Geometry(format!(
779                        "stored text row {row} is missing"
780                    )))
781                })?,
782                None => None,
783            };
784            match text {
785                Some(text) => {
786                    present.push(1);
787                    bytes.extend_from_slice(text.as_bytes());
788                }
789                None => present.push(0),
790            }
791            end_offsets.push(u64::try_from(bytes.len()).map_err(|_| {
792                ConsolidateError::Geometry("stored text offset exceeds u64".to_owned())
793            })?);
794        }
795    }
796    if let Some(lexical) = gather.lexical.as_mut() {
797        for row in rows {
798            let text = match source_text {
799                Some(view) => view.row(*row).ok_or_else(|| {
800                    ConsolidateError::Segment(SegmentError::Geometry(format!(
801                        "stored text row {row} is missing for the lexical rebuild"
802                    )))
803                })?,
804                None => None,
805            };
806            let document = text.map_or_else(LexicalDocument::new, LexicalDocument::with_text);
807            lexical
808                .push_document(analyzer, &document)
809                .map_err(|error| ConsolidateError::Lexical(error.to_string()))?;
810        }
811    }
812    Ok(())
813}
814
815#[cfg(test)]
816#[allow(
817    clippy::expect_used,
818    clippy::unwrap_used,
819    clippy::indexing_slicing,
820    clippy::panic
821)]
822mod tests {
823    use super::{
824        CONSOLIDATION_CHECKPOINT_FILE, ConsolidateError, ConsolidationCheckpoint,
825        consolidation_merge_id, consolidation_output_id, consolidation_supports_region,
826        decode_consolidation_checkpoint, encode_consolidation_checkpoint, merge_segments,
827        read_consolidation_checkpoint, remove_consolidation_checkpoint,
828        validate_consolidation_checkpoint, write_consolidation_checkpoint,
829    };
830    use crate::fts::tokenizer::{Analyzer, TokenizerConfig};
831    use crate::graph::block::{GraphNodeBlockBuild, GraphNodeBlockInput, GraphNodeLayout};
832    use crate::lifecycle::durability::{CommitTier, DurabilityMode, DurabilityPolicy};
833    use crate::meta::{AliveSet, ColumnStoreBuilder, Schema};
834    use crate::quant::quantize_bit4;
835    use crate::segment::SegmentId;
836    use crate::segment::layout::RegionKind;
837    use crate::segment::reader::SegmentReader;
838    use crate::segment::writer::{SegmentBuild, SegmentFactors, write_segment_with_graph};
839    use crate::vfs::StdVfs;
840
841    fn checkpoint_fixture() -> ConsolidationCheckpoint {
842        let inputs = vec![
843            SegmentId::new(1, [1; 10]),
844            SegmentId::new(2, [2; 10]),
845            SegmentId::new(3, [3; 10]),
846        ];
847        ConsolidationCheckpoint {
848            merge_id: consolidation_merge_id(&inputs),
849            output_id: consolidation_output_id(&inputs),
850            inputs,
851            rows_emitted: 384,
852            output_hash: 0x0102_0304_0506_0708,
853        }
854    }
855
856    #[test]
857    fn consolidation_checkpoint_round_trips_and_ids_are_deterministic() {
858        let checkpoint = checkpoint_fixture();
859        let bytes = encode_consolidation_checkpoint(&checkpoint).expect("encode checkpoint");
860        let decoded = decode_consolidation_checkpoint(&bytes).expect("decode checkpoint");
861        assert_eq!(decoded, checkpoint);
862        assert!(validate_consolidation_checkpoint(&bytes).is_ok());
863
864        assert_eq!(
865            consolidation_merge_id(&checkpoint.inputs),
866            checkpoint.merge_id
867        );
868        assert_ne!(checkpoint.merge_id, checkpoint.output_id);
869        let mut reordered = checkpoint.inputs.clone();
870        reordered.swap(0, 2);
871        assert_ne!(
872            consolidation_output_id(&reordered),
873            consolidation_output_id(&checkpoint.inputs),
874            "id derivation must bind the exact input order"
875        );
876    }
877
878    #[test]
879    fn consolidation_checkpoint_rejects_every_corruption_class() {
880        let checkpoint = checkpoint_fixture();
881        let valid = encode_consolidation_checkpoint(&checkpoint).expect("encode checkpoint");
882
883        assert!(validate_consolidation_checkpoint(&[]).is_err());
884        assert!(validate_consolidation_checkpoint(&valid[..valid.len() - 1]).is_err());
885
886        let mut magic = valid.clone();
887        magic[0] ^= 0xff;
888        assert!(validate_consolidation_checkpoint(&magic).is_err());
889
890        let mut version = valid.clone();
891        version[8] = 9;
892        assert!(
893            validate_consolidation_checkpoint(&reseal(version)).is_err(),
894            "future version must be refused"
895        );
896
897        let mut reserved = valid.clone();
898        reserved[10] = 1;
899        assert!(validate_consolidation_checkpoint(&reseal(reserved)).is_err());
900
901        let mut checksum = valid.clone();
902        let last = checksum.len() - 1;
903        checksum[last] ^= 0xff;
904        assert!(validate_consolidation_checkpoint(&checksum).is_err());
905
906        let mut count = valid.clone();
907        count[12..16].copy_from_slice(&0_u32.to_le_bytes());
908        assert!(validate_consolidation_checkpoint(&reseal(count)).is_err());
909
910        let mut unsorted = valid.clone();
911        let (left, right) = (16, 32);
912        for offset in 0..16 {
913            unsorted.swap(left + offset, right + offset);
914        }
915        assert!(validate_consolidation_checkpoint(&reseal(unsorted)).is_err());
916
917        let mut cursor = valid.clone();
918        let tail = 16 + 3 * 16;
919        let cursor_offset = tail + 16 + 16 + 8;
920        cursor[cursor_offset..cursor_offset + 4].copy_from_slice(&1_u32.to_le_bytes());
921        assert!(
922            validate_consolidation_checkpoint(&reseal(cursor)).is_err(),
923            "a non-terminal cursor cannot be resumed by version 1"
924        );
925    }
926
927    fn reseal(mut bytes: Vec<u8>) -> Vec<u8> {
928        let start = bytes.len() - 8;
929        let checksum = xxhash_rust::xxh3::xxh3_64(&bytes[..start]).to_le_bytes();
930        bytes[start..].copy_from_slice(&checksum);
931        bytes
932    }
933
934    #[test]
935    fn consolidation_checkpoint_file_round_trips_and_clears() {
936        let directory = tempfile::tempdir().expect("checkpoint directory");
937        let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
938            .expect("derived policy");
939        let checkpoint = checkpoint_fixture();
940
941        assert!(
942            read_consolidation_checkpoint(&StdVfs, directory.path())
943                .expect("absent checkpoint reads clean")
944                .is_none()
945        );
946        write_consolidation_checkpoint(&StdVfs, directory.path(), &checkpoint, policy)
947            .expect("write checkpoint");
948        let read = read_consolidation_checkpoint(&StdVfs, directory.path())
949            .expect("read checkpoint")
950            .expect("checkpoint present");
951        assert_eq!(read, checkpoint);
952
953        let path = directory.path().join(CONSOLIDATION_CHECKPOINT_FILE);
954        let mut bytes = std::fs::read(&path).expect("checkpoint bytes");
955        bytes[20] ^= 0xff;
956        std::fs::write(&path, bytes).expect("corrupt checkpoint");
957        assert!(read_consolidation_checkpoint(&StdVfs, directory.path()).is_err());
958
959        remove_consolidation_checkpoint(&StdVfs, directory.path()).expect("clear checkpoint");
960        assert!(!path.exists());
961        remove_consolidation_checkpoint(&StdVfs, directory.path())
962            .expect("second clear is idempotent");
963    }
964
965    const MERGE_DIMS: usize = 128;
966
967    fn graph_segment(
968        directory: &std::path::Path,
969        id: SegmentId,
970        rows: usize,
971        base: f32,
972        tombstones: &[u32],
973    ) -> SegmentId {
974        let vectors = (0..rows)
975            .flat_map(|row| std::iter::repeat_n(base + row as f32, MERGE_DIMS))
976            .collect::<Vec<_>>();
977        let stride = MERGE_DIMS.div_ceil(2);
978        let mut codes = vec![0_u8; rows * stride];
979        let mut factors = Vec::with_capacity(rows);
980        for (vector, encoded) in vectors
981            .chunks_exact(MERGE_DIMS)
982            .zip(codes.chunks_exact_mut(stride))
983        {
984            factors.push(quantize_bit4(vector, encoded).expect("finite merge fixture row"));
985        }
986        let schema = Schema::new(Vec::new()).expect("timestamp-only schema");
987        let mut columns = ColumnStoreBuilder::new(schema);
988        for row in 0..rows {
989            columns
990                .push_row(row as i64, &[])
991                .expect("merge fixture column row");
992        }
993        let columns = columns.finish().expect("merge fixture columns");
994        let mut alive = AliveSet::new(rows as u32);
995        for row in tombstones {
996            alive.tombstone(*row).expect("merge fixture tombstone");
997        }
998        let neighbors = (0..rows)
999            .map(|row| {
1000                u32::try_from((row + 1) % rows)
1001                    .ok()
1002                    .into_iter()
1003                    .collect::<Vec<_>>()
1004            })
1005            .collect::<Vec<_>>();
1006        let nodes = codes
1007            .chunks_exact(stride)
1008            .zip(&factors)
1009            .zip(&neighbors)
1010            .map(|((codes, factors), neighbors)| GraphNodeBlockInput {
1011                codes,
1012                factors: *factors,
1013                flags: 1,
1014                neighbors,
1015            })
1016            .collect::<Vec<_>>();
1017        write_segment_with_graph(
1018            &StdVfs,
1019            directory,
1020            SegmentBuild {
1021                id,
1022                scheme: 4,
1023                dims: MERGE_DIMS as u32,
1024                codes: &codes,
1025                factors: SegmentFactors::Bit4(&factors),
1026                rescore: &vectors,
1027                columns: &columns,
1028                alive: &alive,
1029            },
1030            GraphNodeBlockBuild {
1031                layout: GraphNodeLayout::new(MERGE_DIMS as u32, 128, 1)
1032                    .expect("merge fixture layout"),
1033                nodes: &nodes,
1034            },
1035            DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
1036                .expect("derived policy"),
1037        )
1038        .expect("write merge fixture segment");
1039        id
1040    }
1041
1042    #[test]
1043    fn merge_segments_without_documents_carries_live_rows_and_rejects_bad_inputs() {
1044        let directory = tempfile::tempdir().expect("merge fixture directory");
1045        let policy = DurabilityPolicy::new(DurabilityMode::Derived, CommitTier::None)
1046            .expect("derived policy");
1047        let analyzer = Analyzer::new(TokenizerConfig::text_default()).expect("default analyzer");
1048        let first_id = graph_segment(directory.path(), SegmentId::new(1, [1; 10]), 4, 0.0, &[1]);
1049        let second_id = graph_segment(directory.path(), SegmentId::new(2, [2; 10]), 3, 10.0, &[]);
1050        let first = SegmentReader::open(
1051            &StdVfs,
1052            &directory.path().join(first_id.file_name()),
1053            first_id,
1054        )
1055        .expect("open first merge input");
1056        let second = SegmentReader::open(
1057            &StdVfs,
1058            &directory.path().join(second_id.file_name()),
1059            second_id,
1060        )
1061        .expect("open second merge input");
1062
1063        let merge_id = consolidation_merge_id(&[first_id, second_id]);
1064        let merged = merge_segments(
1065            &StdVfs,
1066            directory.path(),
1067            &[&first, &second],
1068            merge_id,
1069            &analyzer,
1070            policy,
1071        )
1072        .expect("merge without documents")
1073        .expect("live rows were carried");
1074        assert_eq!(merged.live_rows, 3 + 3, "one tombstoned row is dropped");
1075        let reader = SegmentReader::open(
1076            &StdVfs,
1077            &directory.path().join(merge_id.file_name()),
1078            merge_id,
1079        )
1080        .expect("open merged intermediate");
1081        assert_eq!(reader.meta().row_count, 6);
1082        assert_eq!(reader.alive().expect("merged alive").live_count(), 6);
1083        assert!(
1084            reader
1085                .document_version(0)
1086                .expect("merged document row")
1087                .is_none(),
1088            "inputs without document identities merge without them"
1089        );
1090
1091        // Fewer than two inputs, and unsorted inputs, are typed geometry errors.
1092        assert!(matches!(
1093            merge_segments(
1094                &StdVfs,
1095                directory.path(),
1096                &[&first],
1097                merge_id,
1098                &analyzer,
1099                policy
1100            ),
1101            Err(ConsolidateError::Geometry(_))
1102        ));
1103        assert!(matches!(
1104            merge_segments(
1105                &StdVfs,
1106                directory.path(),
1107                &[&second, &first],
1108                merge_id,
1109                &analyzer,
1110                policy
1111            ),
1112            Err(ConsolidateError::Geometry(_))
1113        ));
1114    }
1115
1116    #[test]
1117    fn consolidation_region_support_is_the_frozen_merge_set() {
1118        for kind in [
1119            RegionKind::Columns,
1120            RegionKind::Alive,
1121            RegionKind::VectorCodes,
1122            RegionKind::VectorFactors,
1123            RegionKind::VectorRescore,
1124            RegionKind::DocumentVersions,
1125            RegionKind::GraphNodeBlocks,
1126            RegionKind::ChecksumTable,
1127            RegionKind::Postings,
1128            RegionKind::StoredMetadata,
1129            RegionKind::StoredText,
1130        ] {
1131            assert!(consolidation_supports_region(kind.id()), "{kind:?}");
1132        }
1133        for kind in [
1134            RegionKind::GraphColocatedCodes,
1135            RegionKind::SignPlane,
1136            RegionKind::PdxClusteredBlocks,
1137            RegionKind::VectorSpaceN,
1138        ] {
1139            assert!(!consolidation_supports_region(kind.id()), "{kind:?}");
1140        }
1141        assert!(!consolidation_supports_region(65_000));
1142    }
1143}