Skip to main content

zeppelin_embed/segment/
layout.rs

1//! Permanent little-endian segment layout declarations and region codecs.
2
3use 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
12/// Darwin-friendly region alignment; no row is padded.
13pub const REGION_ALIGNMENT: usize = 16 * 1024;
14/// Per-chunk checksum granularity.
15pub const CHECKSUM_CHUNK_BYTES: usize = 64 * 1024;
16/// Segment-specific bytes following the common file header.
17pub const SEGMENT_PREFIX_LEN: usize = 32;
18/// Permanent region-directory entry width.
19pub const REGION_ENTRY_LEN: usize = 32;
20/// Permanent vector-region geometry header width.
21pub const VECTOR_HEADER_LEN: usize = 32;
22
23/// Known and reserved region identifiers.
24#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
25#[repr(u16)]
26pub enum RegionKind {
27    /// Typed metadata arrays.
28    Columns = 1,
29    /// Alive bitmap and tombstone mask.
30    Alive = 2,
31    /// Structure-of-arrays packed vector codes.
32    VectorCodes = 3,
33    /// Structure-of-arrays factor records.
34    VectorFactors = 4,
35    /// Structure-of-arrays exact-rescore rows.
36    VectorRescore = 5,
37    /// Reserved for Task 13 postings.
38    Postings = 6,
39    /// Fixed-stride graph node blocks; originally reserved under the obsolete CSR name.
40    GraphNodeBlocks = 7,
41    /// Reserved for graph-neighbor-colocated codes.
42    GraphColocatedCodes = 8,
43    /// Reserved for a future sign-plane payload.
44    SignPlane = 9,
45    /// Reserved for optional clustered PDX blocks.
46    PdxClusteredBlocks = 10,
47    /// Per-64-KB xxh3 checksum table.
48    ChecksumTable = 11,
49    /// Dense `(doc_id:u128, revision:u64)` records for sealed store rows.
50    DocumentVersions = 12,
51    /// Dense length-delimited opaque stored metadata rows.
52    StoredMetadata = 13,
53    /// Dense optional UTF-8 source text rows.
54    StoredText = 14,
55    /// First id reserved for additive vector-space-N region triples.
56    VectorSpaceN = 4096,
57}
58
59impl RegionKind {
60    /// Returns the permanent numeric identifier.
61    #[must_use]
62    pub const fn id(self) -> u16 {
63        self as u16
64    }
65
66    /// Resolves a known or reserved identifier; unknown values remain skippable.
67    #[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/// One fixed-width region-directory entry.
111#[derive(Clone, Copy, Debug, Eq, PartialEq)]
112#[repr(C)]
113pub struct RegionEntry {
114    /// Known or forward-compatible region kind id.
115    pub kind: u16,
116    /// Additive region-family version.
117    pub version: u16,
118    /// Explicit zero padding reserved for future flags.
119    pub reserved: u32,
120    /// Absolute aligned file offset.
121    pub offset: u64,
122    /// Exact byte length.
123    pub length: u64,
124    /// Whole-region xxh3-64 checksum.
125    pub checksum: u64,
126}
127
128/// Permanent vector geometry and transform descriptor.
129#[derive(Clone, Copy, Debug, Eq, PartialEq)]
130#[repr(C)]
131pub struct VectorHeader {
132    /// Per-segment quantization or rescore scheme id.
133    pub scheme: u16,
134    /// Vector-space identifier; v1 requires zero.
135    pub vector_space_id: u16,
136    /// Logical dimension.
137    pub dims: u32,
138    /// Exact bytes between adjacent code/rescore rows.
139    pub row_stride_bytes: u32,
140    /// Exact bytes between factor records.
141    pub factor_stride_bytes: u16,
142    /// Transform kind; v1 requires identity zero.
143    pub transform_kind: u16,
144    /// Transform seed; v1 requires zero.
145    pub transform_seed: u64,
146    /// Dense segment-local row count.
147    pub row_count: u32,
148    /// Explicit must-be-zero padding.
149    pub reserved: u32,
150}
151
152impl VectorHeader {
153    /// Validates all v1 invariants before any payload cast.
154    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/// Permanent eight-byte affine Int8 factor record.
180#[derive(Clone, Copy, Debug, PartialEq)]
181#[repr(C)]
182pub struct Int8Factors {
183    /// Reconstruction step.
184    pub scale: f32,
185    /// Reconstructed value at code zero.
186    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
475/// Decodes a columns region. With `collection`, the result carries that
476/// schema, which must be an additive evolution of the segment's own.
477pub(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    // A segment sealed before an additive schema evolution lacks the added
557    // nullable columns; decoding it against the collection schema reads them
558    // as null. Any other difference is a violated contract.
559    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        // Final column: presence length, presence bitmap, full ID, zero null slot.
941        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}