1use crate::format::{FormatFamily, FormatRegistry};
4use crate::meta::{
5 AliveSet, Column, ColumnDefinition, ColumnId, ColumnInput, ColumnStore, ColumnStoreBuilder,
6 ColumnType, ColumnValue, Schema, TIMESTAMP_COLUMN,
7};
8use crate::quant::Bit4Factors;
9
10use super::{MetadataDecodeProvenance, SegmentError};
11
12pub const REGION_ALIGNMENT: usize = 16 * 1024;
14pub const CHECKSUM_CHUNK_BYTES: usize = 64 * 1024;
16pub const SEGMENT_PREFIX_LEN: usize = 32;
18pub const REGION_ENTRY_LEN: usize = 32;
20pub const VECTOR_HEADER_LEN: usize = 32;
22
23#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
25#[repr(u16)]
26pub enum RegionKind {
27 Columns = 1,
29 Alive = 2,
31 VectorCodes = 3,
33 VectorFactors = 4,
35 VectorRescore = 5,
37 Postings = 6,
39 GraphNodeBlocks = 7,
41 GraphColocatedCodes = 8,
43 SignPlane = 9,
45 PdxClusteredBlocks = 10,
47 ChecksumTable = 11,
49 DocumentVersions = 12,
51 StoredMetadata = 13,
53 StoredText = 14,
55 VectorSpaceN = 4096,
57}
58
59impl RegionKind {
60 #[must_use]
62 pub const fn id(self) -> u16 {
63 self as u16
64 }
65
66 #[must_use]
68 pub const fn from_id(id: u16) -> Option<Self> {
69 match id {
70 1 => Some(Self::Columns),
71 2 => Some(Self::Alive),
72 3 => Some(Self::VectorCodes),
73 4 => Some(Self::VectorFactors),
74 5 => Some(Self::VectorRescore),
75 6 => Some(Self::Postings),
76 7 => Some(Self::GraphNodeBlocks),
77 8 => Some(Self::GraphColocatedCodes),
78 9 => Some(Self::SignPlane),
79 10 => Some(Self::PdxClusteredBlocks),
80 11 => Some(Self::ChecksumTable),
81 12 => Some(Self::DocumentVersions),
82 13 => Some(Self::StoredMetadata),
83 14 => Some(Self::StoredText),
84 4096 => Some(Self::VectorSpaceN),
85 _ => None,
86 }
87 }
88
89 pub(crate) const fn family(self) -> Option<FormatFamily> {
90 match self {
91 Self::Columns => Some(FormatFamily::Columns),
92 Self::Alive => Some(FormatFamily::Alive),
93 Self::VectorCodes => Some(FormatFamily::VectorCodes),
94 Self::VectorFactors => Some(FormatFamily::VectorFactors),
95 Self::VectorRescore => Some(FormatFamily::VectorRescore),
96 Self::Postings => Some(FormatFamily::Postings),
97 Self::GraphNodeBlocks => Some(FormatFamily::GraphNodeBlocks),
98 Self::DocumentVersions => Some(FormatFamily::DocumentVersions),
99 Self::StoredMetadata => Some(FormatFamily::StoredMetadata),
100 Self::StoredText => Some(FormatFamily::StoredText),
101 Self::ChecksumTable => Some(FormatFamily::ChecksumTable),
102 Self::GraphColocatedCodes
103 | Self::SignPlane
104 | Self::PdxClusteredBlocks
105 | Self::VectorSpaceN => None,
106 }
107 }
108}
109
110#[derive(Clone, Copy, Debug, Eq, PartialEq)]
112#[repr(C)]
113pub struct RegionEntry {
114 pub kind: u16,
116 pub version: u16,
118 pub reserved: u32,
120 pub offset: u64,
122 pub length: u64,
124 pub checksum: u64,
126}
127
128#[derive(Clone, Copy, Debug, Eq, PartialEq)]
130#[repr(C)]
131pub struct VectorHeader {
132 pub scheme: u16,
134 pub vector_space_id: u16,
136 pub dims: u32,
138 pub row_stride_bytes: u32,
140 pub factor_stride_bytes: u16,
142 pub transform_kind: u16,
144 pub transform_seed: u64,
146 pub row_count: u32,
148 pub reserved: u32,
150}
151
152impl VectorHeader {
153 pub fn validate(self) -> Result<(), SegmentError> {
155 FormatRegistry::require_scheme(self.scheme)
156 .map_err(|error| SegmentError::Geometry(error.to_string()))?;
157 if self.vector_space_id != 0 {
158 return Err(SegmentError::Geometry(format!(
159 "v1 vector_space_id must be 0, got {}",
160 self.vector_space_id
161 )));
162 }
163 if self.transform_kind != 0 || self.transform_seed != 0 {
164 return Err(SegmentError::Geometry(format!(
165 "v1 transform must be identity/0, got kind {} seed {}",
166 self.transform_kind, self.transform_seed
167 )));
168 }
169 if self.reserved != 0 {
170 return Err(SegmentError::Geometry(format!(
171 "vector header reserved field must be 0, got {}",
172 self.reserved
173 )));
174 }
175 Ok(())
176 }
177}
178
179#[derive(Clone, Copy, Debug, PartialEq)]
181#[repr(C)]
182pub struct Int8Factors {
183 pub scale: f32,
185 pub offset: f32,
187}
188
189const _: [(); 12] = [(); std::mem::size_of::<Bit4Factors>()];
190const _: [(); 8] = [(); std::mem::size_of::<Int8Factors>()];
191const _: [(); REGION_ENTRY_LEN] = [(); std::mem::size_of::<RegionEntry>()];
192const _: [(); VECTOR_HEADER_LEN] = [(); std::mem::size_of::<VectorHeader>()];
193
194pub(crate) fn align_up(value: usize, alignment: usize) -> Result<usize, SegmentError> {
195 let adjustment = alignment.saturating_sub(1);
196 value
197 .checked_add(adjustment)
198 .map(|rounded| rounded / alignment * alignment)
199 .ok_or_else(|| SegmentError::Geometry("aligned offset overflow".to_owned()))
200}
201
202pub(crate) fn encode_vector_header(header: VectorHeader, output: &mut Vec<u8>) {
203 output.extend_from_slice(&header.scheme.to_le_bytes());
204 output.extend_from_slice(&header.vector_space_id.to_le_bytes());
205 output.extend_from_slice(&header.dims.to_le_bytes());
206 output.extend_from_slice(&header.row_stride_bytes.to_le_bytes());
207 output.extend_from_slice(&header.factor_stride_bytes.to_le_bytes());
208 output.extend_from_slice(&header.transform_kind.to_le_bytes());
209 output.extend_from_slice(&header.transform_seed.to_le_bytes());
210 output.extend_from_slice(&header.row_count.to_le_bytes());
211 output.extend_from_slice(&header.reserved.to_le_bytes());
212}
213
214pub(crate) fn decode_vector_header(bytes: &[u8]) -> Result<VectorHeader, SegmentError> {
215 let mut cursor = Cursor::new("vector-header", bytes);
216 let header = VectorHeader {
217 scheme: cursor.u16()?,
218 vector_space_id: cursor.u16()?,
219 dims: cursor.u32()?,
220 row_stride_bytes: cursor.u32()?,
221 factor_stride_bytes: cursor.u16()?,
222 transform_kind: cursor.u16()?,
223 transform_seed: cursor.u64()?,
224 row_count: cursor.u32()?,
225 reserved: cursor.u32()?,
226 };
227 header.validate()?;
228 Ok(header)
229}
230
231pub(crate) fn encode_alive(alive: &AliveSet) -> Vec<u8> {
232 let mut output = Vec::new();
233 output.extend_from_slice(&alive.row_count().to_le_bytes());
234 let byte_len = usize::try_from(alive.row_count())
235 .unwrap_or(usize::MAX)
236 .div_ceil(8);
237 let mut bits = vec![0_u8; byte_len];
238 for row in alive.iter_alive() {
239 let position = row as usize;
240 let byte = position / 8;
241 let bit = position % 8;
242 if let Some(value) = bits.get_mut(byte) {
243 *value |= 1_u8 << bit;
244 }
245 }
246 output.extend_from_slice(&(bits.len() as u32).to_le_bytes());
247 output.extend_from_slice(&bits);
248 output
249}
250
251pub(crate) fn decode_alive(bytes: &[u8]) -> Result<AliveSet, SegmentError> {
252 let mut cursor = Cursor::new("alive", bytes);
253 let row_count = cursor.u32().map_err(alive_cursor_error)?;
254 let declared_bytes = cursor.u32().map_err(alive_cursor_error)?;
255 let byte_len = usize::try_from(declared_bytes)
256 .map_err(|_| SegmentError::Alive("bitmap length exceeds usize".to_owned()))?;
257 let expected = (row_count as usize).div_ceil(8);
258 if byte_len != expected {
259 return Err(SegmentError::Alive(format!(
260 "bitmap length {byte_len}, expected {expected}"
261 )));
262 }
263 let byte_offset = cursor.position_u64().map_err(alive_cursor_error)?;
264 if cursor.remaining() < byte_len {
265 let observed_bytes = u32::try_from(cursor.remaining()).map_err(|_| {
266 SegmentError::Alive("truncated bitmap remainder exceeds u32".to_owned())
267 })?;
268 return Err(SegmentError::MetadataSemantic {
269 detail: format!(
270 "alive truncated at {}, need {byte_len}, total {}",
271 cursor.position,
272 bytes.len()
273 ),
274 provenance: MetadataDecodeProvenance::AliveBitmapTruncation {
275 row_count,
276 byte_offset,
277 declared_bytes,
278 observed_bytes,
279 },
280 });
281 }
282 let bits = cursor.take(byte_len).map_err(alive_cursor_error)?;
283 cursor.finish().map_err(SegmentError::Alive)?;
284 if !row_count.is_multiple_of(8) {
285 let used = row_count % 8;
286 if let Some(last) = bits.last() {
287 let mask = !((1_u8 << used) - 1);
288 if last & mask != 0 {
289 return Err(SegmentError::Alive(
290 "non-zero bitmap tail padding".to_owned(),
291 ));
292 }
293 }
294 }
295 let mut alive = AliveSet::new(row_count);
296 for row in 0..row_count {
297 let position = row as usize;
298 let is_alive = bits
299 .get(position / 8)
300 .is_some_and(|value| value & (1_u8 << (position % 8)) != 0);
301 if !is_alive {
302 alive
303 .tombstone(row)
304 .map_err(|error| SegmentError::Alive(error.to_string()))?;
305 }
306 }
307 Ok(alive)
308}
309
310fn alive_cursor_error(error: SegmentError) -> SegmentError {
311 match error {
312 SegmentError::Columns(detail) => SegmentError::Alive(detail),
313 other => other,
314 }
315}
316
317fn column_type_id(column_type: ColumnType) -> u16 {
318 match column_type {
319 ColumnType::Id128 => 7,
320 ColumnType::U64 => 1,
321 ColumnType::I64 => 2,
322 ColumnType::F64 => 3,
323 ColumnType::Bool => 4,
324 ColumnType::DictionaryString => 5,
325 ColumnType::RawString => 6,
326 }
327}
328
329fn column_type_from_id(id: u16) -> Result<ColumnType, SegmentError> {
330 match id {
331 7 => Ok(ColumnType::Id128),
332 1 => Ok(ColumnType::U64),
333 2 => Ok(ColumnType::I64),
334 3 => Ok(ColumnType::F64),
335 4 => Ok(ColumnType::Bool),
336 5 => Ok(ColumnType::DictionaryString),
337 6 => Ok(ColumnType::RawString),
338 _ => Err(SegmentError::Columns(format!("unknown column type {id}"))),
339 }
340}
341
342pub(crate) fn encode_columns(store: &ColumnStore) -> Result<Vec<u8>, SegmentError> {
343 let mut output = Vec::new();
344 output.extend_from_slice(&store.row_count().to_le_bytes());
345 output.extend_from_slice(&(store.schema().columns().len() as u32).to_le_bytes());
346 for definition in store.schema().columns() {
347 output.extend_from_slice(&definition.id().get().to_le_bytes());
348 output.extend_from_slice(&column_type_id(definition.column_type()).to_le_bytes());
349 output.extend_from_slice(&u16::from(definition.is_nullable()).to_le_bytes());
350 let name = definition.name().as_bytes();
351 let name_length = u32::try_from(name.len())
352 .map_err(|_| SegmentError::Columns("column name exceeds u32".to_owned()))?;
353 output.extend_from_slice(&name_length.to_le_bytes());
354 output.extend_from_slice(name);
355 }
356 for definition in store.schema().columns() {
357 let column = store.column(definition.id()).ok_or_else(|| {
358 SegmentError::Columns(format!("missing column {}", definition.id().get()))
359 })?;
360 let presence = encode_presence(column, store.row_count());
361 output.extend_from_slice(&(presence.len() as u32).to_le_bytes());
362 output.extend_from_slice(&presence);
363 encode_column_values(column, store.row_count(), &mut output)?;
364 }
365 Ok(output)
366}
367
368fn encode_presence(column: &Column, row_count: u32) -> Vec<u8> {
369 let mut bytes = vec![0_u8; (row_count as usize).div_ceil(8)];
370 for row in column.present().iter() {
371 let position = row as usize;
372 if let Some(value) = bytes.get_mut(position / 8) {
373 *value |= 1_u8 << (position % 8);
374 }
375 }
376 bytes
377}
378
379fn encode_column_values(
380 column: &Column,
381 row_count: u32,
382 output: &mut Vec<u8>,
383) -> Result<(), SegmentError> {
384 match column {
385 Column::Id128(values) => {
386 for row in 0..row_count {
387 output.extend_from_slice(
388 &values
389 .get(row)
390 .map_or(0, crate::ingest::DocId::get)
391 .to_le_bytes(),
392 );
393 }
394 }
395 Column::U64(values) => {
396 for row in 0..row_count {
397 output.extend_from_slice(&values.get(row).unwrap_or(0).to_le_bytes());
398 }
399 }
400 Column::I64(values) => {
401 for row in 0..row_count {
402 output.extend_from_slice(&values.get(row).unwrap_or(0).to_le_bytes());
403 }
404 }
405 Column::F64(values) => {
406 for row in 0..row_count {
407 output.extend_from_slice(&values.get(row).unwrap_or(0.0).to_bits().to_le_bytes());
408 }
409 }
410 Column::Bool(values) => {
411 for row in 0..row_count {
412 output.push(u8::from(values.get(row).unwrap_or(false)));
413 }
414 }
415 Column::DictionaryString(values) => {
416 let dictionary_length = u32::try_from(values.dictionary().len())
417 .map_err(|_| SegmentError::Columns("dictionary exceeds u32".to_owned()))?;
418 output.extend_from_slice(&dictionary_length.to_le_bytes());
419 for code in 0..dictionary_length {
420 let value = values.dictionary().get(code).ok_or_else(|| {
421 SegmentError::Columns(format!("missing dictionary code {code}"))
422 })?;
423 encode_string(value, output)?;
424 }
425 output.extend_from_slice(
426 &match values.codes().width() {
427 crate::meta::CodeWidth::U16 => 2_u16,
428 crate::meta::CodeWidth::U32 => 4_u16,
429 }
430 .to_le_bytes(),
431 );
432 output.extend_from_slice(&0_u16.to_le_bytes());
433 for row in 0..row_count {
434 let code = values.codes().get(row as usize).unwrap_or(0);
435 match values.codes().width() {
436 crate::meta::CodeWidth::U16 => {
437 output.extend_from_slice(&(code as u16).to_le_bytes())
438 }
439 crate::meta::CodeWidth::U32 => output.extend_from_slice(&code.to_le_bytes()),
440 }
441 }
442 }
443 Column::RawString(values) => {
444 for row in 0..row_count {
445 encode_string(values.get(row).unwrap_or(""), output)?;
446 }
447 }
448 }
449 Ok(())
450}
451
452fn encode_string(value: &str, output: &mut Vec<u8>) -> Result<(), SegmentError> {
453 let length = u32::try_from(value.len())
454 .map_err(|_| SegmentError::Columns("string exceeds u32".to_owned()))?;
455 output.extend_from_slice(&length.to_le_bytes());
456 output.extend_from_slice(value.as_bytes());
457 Ok(())
458}
459
460enum DecodedValues {
461 U64(Vec<u64>),
462 Id128(Vec<crate::ingest::DocId>),
463 I64(Vec<i64>),
464 F64(Vec<f64>),
465 Bool(Vec<bool>),
466 String(Vec<String>),
467}
468
469struct DecodedColumn {
470 definition: ColumnDefinition,
471 present: Vec<u8>,
472 values: DecodedValues,
473}
474
475pub(crate) fn decode_columns(
478 bytes: &[u8],
479 collection: Option<&Schema>,
480) -> Result<ColumnStore, SegmentError> {
481 let mut cursor = Cursor::new("columns", bytes);
482 let row_count = cursor.u32()?;
483 let column_count = cursor.usize_from_u32()?;
484 if column_count == 0 {
485 return Err(SegmentError::Columns(
486 "required timestamp column is absent".to_owned(),
487 ));
488 }
489 let mut definitions = Vec::with_capacity(column_count);
490 for _ in 0..column_count {
491 let id = ColumnId::new(cursor.u32()?);
492 let column_type = column_type_from_id(cursor.u16()?)?;
493 let nullable = match cursor.u16()? {
494 0 => false,
495 1 => true,
496 value => {
497 return Err(SegmentError::Columns(format!(
498 "invalid nullable flag {value}"
499 )));
500 }
501 };
502 let name_length = cursor.usize_from_u32()?;
503 let name = std::str::from_utf8(cursor.take(name_length)?)
504 .map_err(|error| SegmentError::Columns(format!("column name UTF-8: {error}")))?
505 .to_owned();
506 definitions.push(ColumnDefinition::new(id, name, column_type, nullable));
507 }
508 let timestamp = definitions.first().ok_or_else(|| {
509 SegmentError::Columns("required timestamp definition is absent".to_owned())
510 })?;
511 if timestamp.id() != TIMESTAMP_COLUMN
512 || timestamp.name() != "ts"
513 || timestamp.column_type() != ColumnType::I64
514 || timestamp.is_nullable()
515 {
516 return Err(SegmentError::Columns(
517 "timestamp definition is not canonical".to_owned(),
518 ));
519 }
520
521 let mut decoded = Vec::with_capacity(column_count);
522 for definition in definitions {
523 let presence_length = cursor.usize_from_u32()?;
524 let expected_presence = (row_count as usize).div_ceil(8);
525 if presence_length != expected_presence {
526 return Err(SegmentError::Columns(format!(
527 "column {} presence length {presence_length}, expected {expected_presence}",
528 definition.id().get()
529 )));
530 }
531 let presence_offset = cursor.position_u64()?;
532 let present = cursor.take(presence_length)?.to_vec();
533 validate_presence_tail(definition.id(), row_count, presence_offset, &present)?;
534 let values = decode_column_values(
535 definition.id(),
536 definition.column_type(),
537 row_count,
538 &present,
539 &mut cursor,
540 )?;
541 decoded.push(DecodedColumn {
542 definition,
543 present,
544 values,
545 });
546 }
547 cursor.finish().map_err(SegmentError::Columns)?;
548
549 let user_definitions = decoded
550 .iter()
551 .skip(1)
552 .map(|column| column.definition.clone())
553 .collect();
554 let schema =
555 Schema::new(user_definitions).map_err(|error| SegmentError::Columns(error.to_string()))?;
556 let schema = match collection {
560 None => schema,
561 Some(collection) => {
562 schema.additive_evolution(collection).map_err(|conflict| {
563 SegmentError::Columns(format!(
564 "segment columns conflict with the collection schema: {conflict}"
565 ))
566 })?;
567 collection.clone()
568 }
569 };
570 let mut builder = ColumnStoreBuilder::new(schema);
571 for row in 0..row_count {
572 let timestamp = decoded
573 .first()
574 .and_then(|column| match &column.values {
575 DecodedValues::I64(values) => values.get(row as usize).copied(),
576 _ => None,
577 })
578 .ok_or_else(|| SegmentError::Columns(format!("missing timestamp row {row}")))?;
579 let mut inputs = Vec::new();
580 for column in decoded.iter().skip(1) {
581 if !presence_contains(&column.present, row) {
582 continue;
583 }
584 let value = decoded_value(&column.values, row)?;
585 inputs.push(ColumnInput {
586 column: column.definition.id(),
587 value,
588 });
589 }
590 builder
591 .push_row(timestamp, &inputs)
592 .map_err(|error| SegmentError::Columns(error.to_string()))?;
593 }
594 builder
595 .finish()
596 .map_err(|error| SegmentError::Columns(error.to_string()))
597}
598
599fn decode_column_values(
600 column_id: ColumnId,
601 column_type: ColumnType,
602 row_count: u32,
603 present: &[u8],
604 cursor: &mut Cursor<'_>,
605) -> Result<DecodedValues, SegmentError> {
606 match column_type {
607 ColumnType::Id128 => (0..row_count)
608 .map(|_| cursor.u128().map(crate::ingest::DocId::new))
609 .collect::<Result<Vec<_>, _>>()
610 .map(DecodedValues::Id128),
611 ColumnType::U64 => (0..row_count)
612 .map(|_| cursor.u64())
613 .collect::<Result<Vec<_>, _>>()
614 .map(DecodedValues::U64),
615 ColumnType::I64 => (0..row_count)
616 .map(|_| cursor.i64())
617 .collect::<Result<Vec<_>, _>>()
618 .map(DecodedValues::I64),
619 ColumnType::F64 => (0..row_count)
620 .map(|_| cursor.u64().map(f64::from_bits))
621 .collect::<Result<Vec<_>, _>>()
622 .map(DecodedValues::F64),
623 ColumnType::Bool => (0..row_count)
624 .map(|_| match cursor.u8()? {
625 0 => Ok(false),
626 1 => Ok(true),
627 value => Err(SegmentError::Columns(format!(
628 "invalid Boolean byte {value}"
629 ))),
630 })
631 .collect::<Result<Vec<_>, _>>()
632 .map(DecodedValues::Bool),
633 ColumnType::DictionaryString => {
634 let dictionary_cardinality = cursor.u32()?;
635 let dictionary_length = usize::try_from(dictionary_cardinality).map_err(|_| {
636 SegmentError::Columns("dictionary cardinality exceeds usize".to_owned())
637 })?;
638 let remaining_entry_capacity = cursor.remaining() / std::mem::size_of::<u32>();
639 if cursor.remaining() >= std::mem::size_of::<u32>()
640 && dictionary_length > remaining_entry_capacity
641 {
642 return Err(SegmentError::Columns(format!(
643 "dictionary cardinality exceeds remaining entry capacity: declared {dictionary_cardinality}, capacity {remaining_entry_capacity}"
644 )));
645 }
646 let mut dictionary = Vec::with_capacity(dictionary_length);
647 for entry in 0..dictionary_length {
648 let declared_bytes = cursor.u32()?;
649 let length = usize::try_from(declared_bytes).map_err(|_| {
650 SegmentError::Columns(
651 "dictionary string storage length exceeds usize".to_owned(),
652 )
653 })?;
654 if cursor.remaining() < length && length > cursor.bytes.len() {
655 return Err(SegmentError::Columns(format!(
656 "dictionary string storage exceeds remaining Columns bytes at entry {entry}: declared {declared_bytes}, available {}",
657 cursor.remaining()
658 )));
659 }
660 let value = std::str::from_utf8(cursor.take(length)?)
661 .map(str::to_owned)
662 .map_err(|error| SegmentError::Columns(format!("columns UTF-8: {error}")))?;
663 dictionary.push(value);
664 }
665 let width = cursor.u16()?;
666 if cursor.u16()? != 0 {
667 return Err(SegmentError::Columns(
668 "dictionary reserved field is non-zero".to_owned(),
669 ));
670 }
671 let mut values = Vec::with_capacity(row_count as usize);
672 for row in 0..row_count {
673 let byte_offset = cursor.position_u64()?;
674 let code = match width {
675 2 => u32::from(cursor.u16()?),
676 4 => cursor.u32()?,
677 _ => {
678 return Err(SegmentError::Columns(format!(
679 "invalid dictionary width {width}"
680 )));
681 }
682 };
683 if presence_contains(present, row) {
684 let value = dictionary.get(code as usize).ok_or_else(|| {
685 SegmentError::MetadataSemantic {
686 detail: format!("dictionary code {code} out of range"),
687 provenance: MetadataDecodeProvenance::ColumnsDictionaryCode {
688 column_id: column_id.get(),
689 row,
690 byte_offset,
691 code,
692 dictionary_cardinality,
693 },
694 }
695 })?;
696 values.push(value.clone());
697 } else {
698 if code != 0 {
699 return Err(SegmentError::Columns(format!(
700 "dictionary null row {row} has non-zero code {code}"
701 )));
702 }
703 values.push(String::new());
704 }
705 }
706 Ok(DecodedValues::String(values))
707 }
708 ColumnType::RawString => (0..row_count)
709 .map(|row| decode_raw_string(column_id, row, cursor))
710 .collect::<Result<Vec<_>, _>>()
711 .map(DecodedValues::String),
712 }
713}
714
715fn decoded_value(values: &DecodedValues, row: u32) -> Result<ColumnValue<'_>, SegmentError> {
716 let position = row as usize;
717 match values {
718 DecodedValues::Id128(values) => values.get(position).copied().map(ColumnValue::Id128),
719 DecodedValues::U64(values) => values.get(position).copied().map(ColumnValue::U64),
720 DecodedValues::I64(values) => values.get(position).copied().map(ColumnValue::I64),
721 DecodedValues::F64(values) => values.get(position).copied().map(ColumnValue::F64),
722 DecodedValues::Bool(values) => values.get(position).copied().map(ColumnValue::Bool),
723 DecodedValues::String(values) => {
724 values.get(position).map(|value| ColumnValue::String(value))
725 }
726 }
727 .ok_or_else(|| SegmentError::Columns(format!("missing row {row}")))
728}
729
730fn decode_raw_string(
731 column_id: ColumnId,
732 row: u32,
733 cursor: &mut Cursor<'_>,
734) -> Result<String, SegmentError> {
735 let byte_offset = cursor.position_u64()?;
736 let declared_bytes = cursor.u32()?;
737 let length = usize::try_from(declared_bytes)
738 .map_err(|_| SegmentError::Columns("raw string length exceeds usize".to_owned()))?;
739 if cursor.remaining() < length {
740 let available_bytes = u64::try_from(cursor.remaining())
741 .map_err(|_| SegmentError::Columns("raw string remainder exceeds u64".to_owned()))?;
742 return Err(SegmentError::MetadataSemantic {
743 detail: format!(
744 "columns truncated at {}, need {length}, total {}",
745 cursor.position,
746 cursor.bytes.len()
747 ),
748 provenance: MetadataDecodeProvenance::ColumnsRawStringLength {
749 column_id: column_id.get(),
750 row,
751 byte_offset,
752 declared_bytes,
753 available_bytes,
754 },
755 });
756 }
757 std::str::from_utf8(cursor.take(length)?)
758 .map(str::to_owned)
759 .map_err(|error| SegmentError::Columns(format!("columns UTF-8: {error}")))
760}
761
762fn validate_presence_tail(
763 column_id: ColumnId,
764 row_count: u32,
765 byte_offset: u64,
766 bytes: &[u8],
767) -> Result<(), SegmentError> {
768 if !row_count.is_multiple_of(8) {
769 let used = row_count % 8;
770 let allowed_mask = (1_u8 << used) - 1;
771 if let Some(observed_byte) = bytes.last().copied()
772 && observed_byte & !allowed_mask != 0
773 {
774 let last_offset = u64::try_from(bytes.len().saturating_sub(1))
775 .ok()
776 .and_then(|delta| byte_offset.checked_add(delta))
777 .ok_or_else(|| SegmentError::Columns("presence offset overflow".to_owned()))?;
778 return Err(SegmentError::MetadataSemantic {
779 detail: "non-zero presence tail padding".to_owned(),
780 provenance: MetadataDecodeProvenance::ColumnsPresenceTail {
781 column_id: column_id.get(),
782 row_count,
783 byte_offset: last_offset,
784 observed_byte,
785 allowed_mask,
786 },
787 });
788 }
789 }
790 Ok(())
791}
792
793fn presence_contains(bytes: &[u8], row: u32) -> bool {
794 let position = row as usize;
795 bytes
796 .get(position / 8)
797 .is_some_and(|byte| byte & (1_u8 << (position % 8)) != 0)
798}
799
800pub(crate) struct Cursor<'a> {
801 artifact: &'static str,
802 bytes: &'a [u8],
803 position: usize,
804}
805
806impl<'a> Cursor<'a> {
807 pub(crate) const fn new(artifact: &'static str, bytes: &'a [u8]) -> Self {
808 Self {
809 artifact,
810 bytes,
811 position: 0,
812 }
813 }
814
815 pub(crate) fn take(&mut self, length: usize) -> Result<&'a [u8], SegmentError> {
816 let end = self
817 .position
818 .checked_add(length)
819 .ok_or_else(|| SegmentError::Columns(format!("{} offset overflow", self.artifact)))?;
820 let value = self.bytes.get(self.position..end).ok_or_else(|| {
821 SegmentError::Columns(format!(
822 "{} truncated at {}, need {length}, total {}",
823 self.artifact,
824 self.position,
825 self.bytes.len()
826 ))
827 })?;
828 self.position = end;
829 Ok(value)
830 }
831
832 pub(crate) fn position_u64(&self) -> Result<u64, SegmentError> {
833 u64::try_from(self.position)
834 .map_err(|_| SegmentError::Columns(format!("{} offset exceeds u64", self.artifact)))
835 }
836
837 pub(crate) fn remaining(&self) -> usize {
838 self.bytes.len().saturating_sub(self.position)
839 }
840
841 pub(crate) fn u8(&mut self) -> Result<u8, SegmentError> {
842 self.take(1)?
843 .first()
844 .copied()
845 .ok_or_else(|| SegmentError::Columns(format!("{} missing u8", self.artifact)))
846 }
847
848 pub(crate) fn u16(&mut self) -> Result<u16, SegmentError> {
849 let raw: [u8; 2] = self
850 .take(2)?
851 .try_into()
852 .map_err(|_| SegmentError::Columns(format!("{} invalid u16", self.artifact)))?;
853 Ok(u16::from_le_bytes(raw))
854 }
855
856 pub(crate) fn u32(&mut self) -> Result<u32, SegmentError> {
857 let raw: [u8; 4] = self
858 .take(4)?
859 .try_into()
860 .map_err(|_| SegmentError::Columns(format!("{} invalid u32", self.artifact)))?;
861 Ok(u32::from_le_bytes(raw))
862 }
863
864 pub(crate) fn u128(&mut self) -> Result<u128, SegmentError> {
865 let raw = self
866 .take(16)?
867 .try_into()
868 .map_err(|_| SegmentError::Columns(format!("{} invalid u128", self.artifact)))?;
869 Ok(u128::from_le_bytes(raw))
870 }
871
872 pub(crate) fn u64(&mut self) -> Result<u64, SegmentError> {
873 let raw: [u8; 8] = self
874 .take(8)?
875 .try_into()
876 .map_err(|_| SegmentError::Columns(format!("{} invalid u64", self.artifact)))?;
877 Ok(u64::from_le_bytes(raw))
878 }
879
880 pub(crate) fn i64(&mut self) -> Result<i64, SegmentError> {
881 let raw: [u8; 8] = self
882 .take(8)?
883 .try_into()
884 .map_err(|_| SegmentError::Columns(format!("{} invalid i64", self.artifact)))?;
885 Ok(i64::from_le_bytes(raw))
886 }
887
888 pub(crate) fn usize_from_u32(&mut self) -> Result<usize, SegmentError> {
889 usize::try_from(self.u32()?)
890 .map_err(|_| SegmentError::Columns(format!("{} u32 exceeds usize", self.artifact)))
891 }
892
893 pub(crate) fn finish(&self) -> Result<(), String> {
894 if self.position == self.bytes.len() {
895 Ok(())
896 } else {
897 Err(format!(
898 "{} has {} trailing bytes",
899 self.artifact,
900 self.bytes.len().saturating_sub(self.position)
901 ))
902 }
903 }
904}
905
906#[cfg(test)]
907#[allow(clippy::expect_used, clippy::panic)]
908mod collection_schema_tests {
909 use super::{
910 ColumnDefinition, ColumnId, ColumnInput, ColumnStoreBuilder, ColumnType, ColumnValue,
911 Schema, SegmentError, decode_columns, encode_columns,
912 };
913
914 const RANK: ColumnId = ColumnId::new(1);
915 const LANG: ColumnId = ColumnId::new(2);
916
917 #[test]
918 fn id128_segment_layout_and_null_presence_are_explicit() {
919 use crate::ingest::DocId;
920 let schema = Schema::new(vec![ColumnDefinition::new(
921 RANK,
922 "id",
923 ColumnType::Id128,
924 true,
925 )])
926 .expect("schema");
927 let mut builder = ColumnStoreBuilder::new(schema);
928 builder
929 .push_row(
930 0,
931 &[ColumnInput {
932 column: RANK,
933 value: ColumnValue::Id128(DocId::new(0xffeeddccbbaa99887766554433221100)),
934 }],
935 )
936 .expect("row");
937 builder.push_row(0, &[]).expect("null row");
938 let store = builder.finish().expect("columns");
939 let bytes = encode_columns(&store).expect("encode");
940 let tail: &[u8] = &[
942 1, 0, 0, 0, 1, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb,
943 0xcc, 0xdd, 0xee, 0xff, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
944 ];
945 assert!(bytes.ends_with(tail));
946 assert_eq!(super::column_type_id(ColumnType::Id128), 7);
947 assert_eq!(decode_columns(&bytes, None).expect("decode"), store);
948 let mut truncated = bytes;
949 truncated.pop();
950 assert!(decode_columns(&truncated, None).is_err());
951 }
952
953 fn rank_region() -> Vec<u8> {
954 let schema = Schema::new(vec![ColumnDefinition::new(
955 RANK,
956 "rank",
957 ColumnType::U64,
958 false,
959 )])
960 .expect("segment schema");
961 let mut builder = ColumnStoreBuilder::new(schema);
962 for value in [4_u64, 9] {
963 builder
964 .push_row(
965 0,
966 &[ColumnInput {
967 column: RANK,
968 value: ColumnValue::U64(value),
969 }],
970 )
971 .expect("push row");
972 }
973 encode_columns(&builder.finish().expect("finish columns")).expect("encode columns")
974 }
975
976 fn collection(columns: Vec<ColumnDefinition>) -> Schema {
977 Schema::new(columns).expect("collection schema")
978 }
979
980 #[test]
981 fn a_segment_decodes_against_an_additive_collection_schema_as_null() {
982 let target = collection(vec![
983 ColumnDefinition::new(RANK, "rank", ColumnType::U64, false),
984 ColumnDefinition::new(LANG, "lang", ColumnType::DictionaryString, true),
985 ]);
986 let columns = decode_columns(&rank_region(), Some(&target)).expect("decode widened");
987 assert_eq!(columns.schema(), &target);
988 let lang = columns.column(LANG).expect("added column is materialized");
989 assert!(lang.present().is_empty(), "every old row reads null");
990 assert_eq!(lang.len(), 2);
991 let own = decode_columns(&rank_region(), None).expect("decode own schema");
992 assert_eq!(own.schema().user_column_count(), 1);
993 }
994
995 #[test]
996 fn a_segment_column_that_conflicts_with_the_collection_schema_fails_loudly() {
997 let conflicts = [
998 collection(vec![ColumnDefinition::new(
999 RANK,
1000 "rank",
1001 ColumnType::I64,
1002 false,
1003 )]),
1004 collection(vec![ColumnDefinition::new(
1005 LANG,
1006 "lang",
1007 ColumnType::DictionaryString,
1008 true,
1009 )]),
1010 collection(vec![
1011 ColumnDefinition::new(RANK, "rank", ColumnType::U64, false),
1012 ColumnDefinition::new(LANG, "lang", ColumnType::DictionaryString, false),
1013 ]),
1014 ];
1015 for target in conflicts {
1016 match decode_columns(&rank_region(), Some(&target)) {
1017 Err(SegmentError::Columns(detail)) => {
1018 assert!(
1019 detail.contains("conflict with the collection schema"),
1020 "{detail}"
1021 );
1022 }
1023 other => panic!("expected a column conflict for {target:?}, got {other:?}"),
1024 }
1025 }
1026 }
1027}