1use 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
35pub(crate) const CONSOLIDATION_CHECKPOINT_FILE: &str = ".consolidate.checkpoint";
37pub const CONSOLIDATION_CHECKPOINT_MAGIC: [u8; 8] = *b"ZECONCP1";
39const CHECKPOINT_VERSION: u16 = 1;
40const CHECKPOINT_PREFIX_BYTES: usize = 16;
42const CHECKPOINT_TAIL_BYTES: usize = 16 + 16 + 8 + 4 + 4 + 8 + 8;
44const 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#[derive(Debug)]
51pub enum ConsolidateError {
52 Segment(SegmentError),
54 Store(StoreError),
56 Geometry(String),
58 DuplicateDocument {
60 doc_id: u128,
62 },
63 Lexical(String),
65 CheckpointIo {
67 path: PathBuf,
69 source: std::io::Error,
71 },
72 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#[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#[must_use]
152pub(crate) fn consolidation_merge_id(inputs: &[SegmentId]) -> SegmentId {
153 derived_segment_id(inputs, MERGE_ID_SEED)
154}
155
156#[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#[derive(Clone, Debug, Eq, PartialEq)]
181pub(crate) struct ConsolidationCheckpoint {
182 pub(crate) inputs: Vec<SegmentId>,
184 pub(crate) merge_id: SegmentId,
186 pub(crate) output_id: SegmentId,
188 pub(crate) rows_emitted: u64,
190 pub(crate) output_hash: u64,
192}
193
194pub(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 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
235pub(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
328pub 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
371pub(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
413pub(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
426pub(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
445pub(crate) struct MergedIntermediate {
447 pub(crate) live_rows: u64,
449 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
464pub(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
618pub(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 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}