1use 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#[derive(Clone, Copy, Debug)]
27pub enum SegmentFactors<'a> {
28 F32,
30 Bit4(&'a [Bit4Factors]),
32 Int8(&'a [Int8Factors]),
34}
35
36#[derive(Clone, Copy, Debug)]
38pub struct SegmentBuild<'a> {
39 pub id: SegmentId,
41 pub scheme: u16,
43 pub dims: u32,
45 pub codes: &'a [u8],
47 pub factors: SegmentFactors<'a>,
49 pub rescore: &'a [f32],
51 pub columns: &'a ColumnStore,
53 pub alive: &'a AliveSet,
55}
56
57#[derive(Clone, Copy, Debug)]
59pub struct SegmentDocumentVersions<'a> {
60 pub doc_ids: &'a [DocId],
62 pub revisions: &'a [Revision],
64}
65
66#[derive(Clone, Copy, Debug)]
68pub struct SegmentStoredMetadata<'a> {
69 pub end_offsets: &'a [u64],
71 pub bytes: &'a [u8],
73}
74
75#[derive(Clone, Copy, Debug)]
77pub struct SegmentStoredText<'a> {
78 pub present: &'a [u8],
80 pub end_offsets: &'a [u64],
82 pub bytes: &'a [u8],
84}
85
86#[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#[derive(Clone, Copy, Debug)]
97pub struct SegmentPostings<'a> {
98 pub bytes: &'a [u8],
100}
101
102struct RegionBytes {
103 kind: RegionKind,
104 family: FormatFamily,
105 bytes: Vec<u8>,
106}
107
108#[derive(Clone, Copy, Debug)]
110pub(crate) struct SegmentCopiedRegion<'a> {
111 pub kind: u16,
113 pub version: u16,
115 pub bytes: &'a [u8],
117}
118
119pub fn encode_segment(build: SegmentBuild<'_>) -> Result<Vec<u8>, SegmentError> {
121 encode_segment_inner(build, None, None, None, None, None, &[])
122}
123
124pub 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
147pub 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, ®ions, 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(®ion.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(®ion.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(®ion.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(®ion.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
741pub 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
752pub 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
764pub 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
785pub 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
798pub 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
819pub 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
834pub 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
856pub(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
876pub 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
888pub 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
901pub(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
999pub(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 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}