zeppelin_embed/meta/
columns.rs

1//! Typed, contiguous column arrays and row-aligned builders.
2
3use super::bitmap::DocBitmap;
4use super::dict::{
5    DictionaryBuilder, DictionaryCodes, DictionaryError, StringDictionary, StringStorage,
6};
7use super::{ColumnId, ColumnType, Schema, TIMESTAMP_COLUMN};
8
9/// A borrowed, typed value supplied for one metadata row.
10#[derive(Clone, Copy, Debug, PartialEq)]
11pub enum ColumnValue<'a> {
12    /// An unsigned integer.
13    U64(u64),
14    /// A signed integer.
15    I64(i64),
16    /// A floating-point value.
17    F64(f64),
18    /// A Boolean value.
19    Bool(bool),
20    /// A UTF-8 value for either string physical representation.
21    String(&'a str),
22}
23
24impl ColumnValue<'_> {
25    fn compatible_with(self, column_type: ColumnType) -> bool {
26        matches!(
27            (self, column_type),
28            (Self::U64(_), ColumnType::U64)
29                | (Self::I64(_), ColumnType::I64)
30                | (Self::F64(_), ColumnType::F64)
31                | (Self::Bool(_), ColumnType::Bool)
32                | (
33                    Self::String(_),
34                    ColumnType::DictionaryString | ColumnType::RawString
35                )
36        )
37    }
38
39    fn value_type(self) -> ColumnType {
40        match self {
41            Self::U64(_) => ColumnType::U64,
42            Self::I64(_) => ColumnType::I64,
43            Self::F64(_) => ColumnType::F64,
44            Self::Bool(_) => ColumnType::Bool,
45            Self::String(_) => ColumnType::RawString,
46        }
47    }
48}
49
50/// One column/value pair supplied while appending a row.
51#[derive(Clone, Copy, Debug, PartialEq)]
52pub struct ColumnInput<'a> {
53    /// The schema column receiving the value.
54    pub column: ColumnId,
55    /// The typed value to append.
56    pub value: ColumnValue<'a>,
57}
58
59/// A typed failure that leaves the logical row count unchanged.
60#[derive(Clone, Debug, Eq, PartialEq)]
61pub enum BuildError {
62    /// An input referred to an identifier absent from the schema.
63    UnknownColumn(ColumnId),
64    /// The same column appeared twice in one row.
65    DuplicateColumn(ColumnId),
66    /// A non-nullable column was omitted.
67    MissingRequiredColumn(ColumnId),
68    /// The timestamp must use the dedicated append argument.
69    TimestampProvidedAsInput,
70    /// A value did not match its declared physical type.
71    TypeMismatch {
72        /// The mismatched column.
73        column: ColumnId,
74        /// The declared physical type.
75        expected: ColumnType,
76        /// The supplied value type.
77        actual: ColumnType,
78    },
79    /// The segment reached the supported u32 document-id space.
80    TooManyRows,
81    /// Dictionary or string side-array construction failed.
82    Dictionary(DictionaryError),
83}
84
85/// Narrow limits used only to prove naturally huge metadata builder guards.
86#[cfg(any(test, feature = "test-support"))]
87#[derive(Clone, Copy, Debug, Eq, PartialEq)]
88#[doc(hidden)]
89pub struct MetadataBuildTestLimits {
90    pub max_rows: u32,
91    pub max_dictionary_entries: u64,
92    pub max_string_bytes: u64,
93}
94
95#[cfg(any(test, feature = "test-support"))]
96impl MetadataBuildTestLimits {
97    #[must_use]
98    pub const fn new(max_rows: u32, max_dictionary_entries: u64, max_string_bytes: u64) -> Self {
99        Self {
100            max_rows,
101            max_dictionary_entries,
102            max_string_bytes,
103        }
104    }
105}
106
107impl std::fmt::Display for BuildError {
108    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        match self {
110            Self::UnknownColumn(column) => write!(formatter, "unknown column {}", column.get()),
111            Self::DuplicateColumn(column) => {
112                write!(formatter, "duplicate column {}", column.get())
113            }
114            Self::MissingRequiredColumn(column) => {
115                write!(formatter, "missing required column {}", column.get())
116            }
117            Self::TimestampProvidedAsInput => {
118                formatter.write_str("timestamp must use the dedicated ts argument")
119            }
120            Self::TypeMismatch {
121                column,
122                expected,
123                actual,
124            } => write!(
125                formatter,
126                "column {} expects {expected:?}, received {actual:?}",
127                column.get()
128            ),
129            Self::TooManyRows => formatter.write_str("segment row count exceeds u32"),
130            Self::Dictionary(error) => error.fmt(formatter),
131        }
132    }
133}
134
135impl std::error::Error for BuildError {}
136
137impl From<DictionaryError> for BuildError {
138    fn from(error: DictionaryError) -> Self {
139        Self::Dictionary(error)
140    }
141}
142
143/// A nullable numeric array plus its explicit presence bitmap.
144#[derive(Clone, Debug, PartialEq)]
145#[repr(C)]
146pub struct NumericColumn<T> {
147    values: Vec<T>,
148    present: DocBitmap,
149}
150
151impl<T: Copy> NumericColumn<T> {
152    /// Returns the physical row count including null placeholders.
153    #[must_use]
154    pub fn len(&self) -> usize {
155        self.values.len()
156    }
157
158    /// Returns whether the physical array is empty.
159    #[must_use]
160    pub fn is_empty(&self) -> bool {
161        self.values.is_empty()
162    }
163
164    /// Reads a non-null value, returning `None` for null or out-of-range rows.
165    #[must_use]
166    pub fn get(&self, row: u32) -> Option<T> {
167        if !self.present.contains(row) {
168            return None;
169        }
170        let position = usize::try_from(row).ok()?;
171        self.values.get(position).copied()
172    }
173
174    /// Returns the presence bitmap used by `exists` and `null` predicates.
175    #[must_use]
176    pub fn present(&self) -> &DocBitmap {
177        &self.present
178    }
179
180    pub(crate) fn values(&self) -> &[T] {
181        &self.values
182    }
183}
184
185/// A nullable Boolean array stored as contiguous zero/one bytes.
186#[derive(Clone, Debug, Eq, PartialEq)]
187#[repr(C)]
188pub struct BoolColumn {
189    values: Vec<u8>,
190    present: DocBitmap,
191}
192
193impl BoolColumn {
194    /// Returns the physical row count including null placeholders.
195    #[must_use]
196    pub fn len(&self) -> usize {
197        self.values.len()
198    }
199
200    /// Returns whether the physical array is empty.
201    #[must_use]
202    pub fn is_empty(&self) -> bool {
203        self.values.is_empty()
204    }
205
206    /// Reads a non-null value, returning `None` for null or out-of-range rows.
207    #[must_use]
208    pub fn get(&self, row: u32) -> Option<bool> {
209        if !self.present.contains(row) {
210            return None;
211        }
212        let position = usize::try_from(row).ok()?;
213        self.values.get(position).map(|value| *value != 0)
214    }
215
216    /// Returns the presence bitmap used by `exists` and `null` predicates.
217    #[must_use]
218    pub fn present(&self) -> &DocBitmap {
219        &self.present
220    }
221
222    pub(crate) fn values(&self) -> &[u8] {
223        &self.values
224    }
225}
226
227/// A nullable, low-cardinality dictionary string column.
228#[derive(Clone, Debug, Eq, PartialEq)]
229#[repr(C)]
230pub struct DictionaryColumn {
231    codes: DictionaryCodes,
232    dictionary: StringDictionary,
233    present: DocBitmap,
234}
235
236impl DictionaryColumn {
237    /// Returns the physical row count including null placeholders.
238    #[must_use]
239    pub fn len(&self) -> usize {
240        self.codes.len()
241    }
242
243    /// Returns whether the physical array is empty.
244    #[must_use]
245    pub fn is_empty(&self) -> bool {
246        self.codes.is_empty()
247    }
248
249    /// Decodes a non-null row value.
250    #[must_use]
251    pub fn get(&self, row: u32) -> Option<&str> {
252        if !self.present.contains(row) {
253            return None;
254        }
255        let position = usize::try_from(row).ok()?;
256        let code = self.codes.get(position)?;
257        self.dictionary.get(code)
258    }
259
260    /// Returns the selected physical code array.
261    #[must_use]
262    pub fn codes(&self) -> &DictionaryCodes {
263        &self.codes
264    }
265
266    /// Returns the immutable string dictionary.
267    #[must_use]
268    pub fn dictionary(&self) -> &StringDictionary {
269        &self.dictionary
270    }
271
272    /// Returns the presence bitmap used by `exists` and `null` predicates.
273    #[must_use]
274    pub fn present(&self) -> &DocBitmap {
275        &self.present
276    }
277}
278
279/// A nullable raw UTF-8 column backed by contiguous offsets and bytes.
280#[derive(Clone, Debug, Eq, PartialEq)]
281#[repr(C)]
282pub struct RawStringColumn {
283    values: StringStorage,
284    present: DocBitmap,
285}
286
287impl RawStringColumn {
288    /// Returns the physical row count including null placeholders.
289    #[must_use]
290    pub fn len(&self) -> usize {
291        self.values.len()
292    }
293
294    /// Returns whether the physical array is empty.
295    #[must_use]
296    pub fn is_empty(&self) -> bool {
297        self.values.len() == 0
298    }
299
300    /// Reads a non-null row value.
301    #[must_use]
302    pub fn get(&self, row: u32) -> Option<&str> {
303        if !self.present.contains(row) {
304            return None;
305        }
306        let position = usize::try_from(row).ok()?;
307        self.values.get(position)
308    }
309
310    /// Returns the presence bitmap used by `exists` and `null` predicates.
311    #[must_use]
312    pub fn present(&self) -> &DocBitmap {
313        &self.present
314    }
315
316    pub(crate) fn physical_get(&self, position: usize) -> Option<&str> {
317        self.values.get(position)
318    }
319}
320
321/// One immutable typed physical column.
322#[derive(Clone, Debug, PartialEq)]
323pub enum Column {
324    /// An unsigned-integer array.
325    U64(NumericColumn<u64>),
326    /// A signed-integer array, including `ts`.
327    I64(NumericColumn<i64>),
328    /// A floating-point array.
329    F64(NumericColumn<f64>),
330    /// A Boolean byte array.
331    Bool(BoolColumn),
332    /// A dictionary string array.
333    DictionaryString(DictionaryColumn),
334    /// A raw UTF-8 array.
335    RawString(RawStringColumn),
336}
337
338impl Column {
339    /// Returns the physical column type.
340    #[must_use]
341    pub const fn column_type(&self) -> ColumnType {
342        match self {
343            Self::U64(_) => ColumnType::U64,
344            Self::I64(_) => ColumnType::I64,
345            Self::F64(_) => ColumnType::F64,
346            Self::Bool(_) => ColumnType::Bool,
347            Self::DictionaryString(_) => ColumnType::DictionaryString,
348            Self::RawString(_) => ColumnType::RawString,
349        }
350    }
351
352    /// Returns the physical row count including null placeholders.
353    #[must_use]
354    pub fn len(&self) -> usize {
355        match self {
356            Self::U64(column) => column.len(),
357            Self::I64(column) => column.len(),
358            Self::F64(column) => column.len(),
359            Self::Bool(column) => column.len(),
360            Self::DictionaryString(column) => column.len(),
361            Self::RawString(column) => column.len(),
362        }
363    }
364
365    /// Returns whether the column contains no physical rows.
366    #[must_use]
367    pub fn is_empty(&self) -> bool {
368        self.len() == 0
369    }
370
371    /// Returns the non-null row bitmap.
372    #[must_use]
373    pub fn present(&self) -> &DocBitmap {
374        match self {
375            Self::U64(column) => column.present(),
376            Self::I64(column) => column.present(),
377            Self::F64(column) => column.present(),
378            Self::Bool(column) => column.present(),
379            Self::DictionaryString(column) => column.present(),
380            Self::RawString(column) => column.present(),
381        }
382    }
383}
384
385/// Immutable metadata columns aligned to one segment's document-id space.
386#[derive(Clone, Debug, PartialEq)]
387pub struct ColumnStore {
388    schema: Schema,
389    columns: Vec<Column>,
390    row_count: u32,
391}
392
393impl ColumnStore {
394    pub(crate) fn compact_for_cache(&mut self) {
395        for column in &mut self.columns {
396            match column {
397                Column::U64(column) => column.present.compact_for_cache(),
398                Column::I64(column) => column.present.compact_for_cache(),
399                Column::F64(column) => column.present.compact_for_cache(),
400                Column::Bool(column) => column.present.compact_for_cache(),
401                Column::DictionaryString(column) => column.present.compact_for_cache(),
402                Column::RawString(column) => column.present.compact_for_cache(),
403            }
404        }
405    }
406
407    pub(crate) fn resident_bytes(&self) -> Option<usize> {
408        let direct = self.schema.resident_bytes()?.checked_add(
409            self.columns
410                .capacity()
411                .checked_mul(std::mem::size_of::<Column>())?,
412        )?;
413        self.columns.iter().try_fold(direct, |total, column| {
414            total.checked_add(column_resident_bytes(column)?)
415        })
416    }
417
418    /// Returns the schema that defines these physical arrays.
419    #[must_use]
420    pub fn schema(&self) -> &Schema {
421        &self.schema
422    }
423
424    /// Returns the shared segment row count.
425    #[must_use]
426    pub const fn row_count(&self) -> u32 {
427        self.row_count
428    }
429
430    /// Finds an immutable physical column by typed identifier.
431    #[must_use]
432    pub fn column(&self, id: ColumnId) -> Option<&Column> {
433        self.schema
434            .columns()
435            .iter()
436            .zip(&self.columns)
437            .find_map(|(definition, column)| (definition.id() == id).then_some(column))
438    }
439
440    /// Reads one required timestamp.
441    #[must_use]
442    pub fn timestamp(&self, row: u32) -> Option<i64> {
443        match self.column(TIMESTAMP_COLUMN)? {
444            Column::I64(column) => column.get(row),
445            _ => None,
446        }
447    }
448
449    /// Returns the complete required timestamp side-array.
450    #[must_use]
451    pub fn timestamps(&self) -> &[i64] {
452        match self.column(TIMESTAMP_COLUMN) {
453            Some(Column::I64(column)) => column.values(),
454            _ => &[],
455        }
456    }
457}
458
459fn column_resident_bytes(column: &Column) -> Option<usize> {
460    match column {
461        Column::U64(column) => column
462            .values
463            .capacity()
464            .checked_mul(std::mem::size_of::<u64>())?
465            .checked_add(column.present.resident_bytes()?),
466        Column::I64(column) => column
467            .values
468            .capacity()
469            .checked_mul(std::mem::size_of::<i64>())?
470            .checked_add(column.present.resident_bytes()?),
471        Column::F64(column) => column
472            .values
473            .capacity()
474            .checked_mul(std::mem::size_of::<f64>())?
475            .checked_add(column.present.resident_bytes()?),
476        Column::Bool(column) => column
477            .values
478            .capacity()
479            .checked_add(column.present.resident_bytes()?),
480        Column::DictionaryString(column) => column
481            .codes
482            .resident_bytes()?
483            .checked_add(column.dictionary.resident_bytes()?)?
484            .checked_add(column.present.resident_bytes()?),
485        Column::RawString(column) => column
486            .values
487            .resident_bytes()?
488            .checked_add(column.present.resident_bytes()?),
489    }
490}
491
492enum BuilderColumn {
493    U64(Vec<u64>, DocBitmap),
494    I64(Vec<i64>, DocBitmap),
495    F64(Vec<f64>, DocBitmap),
496    Bool(Vec<u8>, DocBitmap),
497    Dictionary {
498        codes: Vec<u32>,
499        dictionary: DictionaryBuilder,
500        present: DocBitmap,
501    },
502    RawString(StringStorage, DocBitmap),
503}
504
505impl BuilderColumn {
506    fn new(column_type: ColumnType) -> Self {
507        match column_type {
508            ColumnType::U64 => Self::U64(Vec::new(), DocBitmap::new()),
509            ColumnType::I64 => Self::I64(Vec::new(), DocBitmap::new()),
510            ColumnType::F64 => Self::F64(Vec::new(), DocBitmap::new()),
511            ColumnType::Bool => Self::Bool(Vec::new(), DocBitmap::new()),
512            ColumnType::DictionaryString => Self::Dictionary {
513                codes: Vec::new(),
514                dictionary: DictionaryBuilder::new(),
515                present: DocBitmap::new(),
516            },
517            ColumnType::RawString => Self::RawString(StringStorage::new(), DocBitmap::new()),
518        }
519    }
520
521    fn preflight(&self, value: Option<ColumnValue<'_>>) -> Result<(), DictionaryError> {
522        match (self, value) {
523            (Self::Dictionary { dictionary, .. }, Some(ColumnValue::String(value))) => {
524                dictionary.can_intern(value)
525            }
526            (Self::RawString(values, _), Some(ColumnValue::String(value))) => {
527                values.can_push(value)
528            }
529            (Self::RawString(values, _), None) => values.can_push(""),
530            _ => Ok(()),
531        }
532    }
533
534    #[cfg(any(test, feature = "test-support"))]
535    fn check_test_limits(
536        &self,
537        value: Option<ColumnValue<'_>>,
538        limits: MetadataBuildTestLimits,
539    ) -> Result<(), DictionaryError> {
540        match (self, value) {
541            (Self::Dictionary { dictionary, .. }, Some(ColumnValue::String(value))) => {
542                if !dictionary.contains(value)
543                    && dictionary.cardinality() >= limits.max_dictionary_entries
544                {
545                    return Err(DictionaryError::CardinalityOverflow);
546                }
547                let bytes = u64::try_from(dictionary.string_bytes())
548                    .unwrap_or(u64::MAX)
549                    .checked_add(u64::try_from(value.len()).unwrap_or(u64::MAX))
550                    .ok_or(DictionaryError::StringStorageOverflow)?;
551                if !dictionary.contains(value) && bytes > limits.max_string_bytes {
552                    return Err(DictionaryError::StringStorageOverflow);
553                }
554                Ok(())
555            }
556            (Self::RawString(values, _), Some(ColumnValue::String(value))) => {
557                let bytes = u64::try_from(values.byte_len())
558                    .unwrap_or(u64::MAX)
559                    .checked_add(u64::try_from(value.len()).unwrap_or(u64::MAX))
560                    .ok_or(DictionaryError::StringStorageOverflow)?;
561                if bytes > limits.max_string_bytes {
562                    Err(DictionaryError::StringStorageOverflow)
563                } else {
564                    Ok(())
565                }
566            }
567            (Self::RawString(values, _), None) => {
568                if u64::try_from(values.byte_len()).unwrap_or(u64::MAX) > limits.max_string_bytes {
569                    Err(DictionaryError::StringStorageOverflow)
570                } else {
571                    Ok(())
572                }
573            }
574            _ => Ok(()),
575        }
576    }
577
578    fn push(
579        &mut self,
580        row: u32,
581        column: ColumnId,
582        value: Option<ColumnValue<'_>>,
583    ) -> Result<(), BuildError> {
584        let expected = self.column_type();
585        match (self, value) {
586            (Self::U64(values, present), Some(ColumnValue::U64(value))) => {
587                values.push(value);
588                present.insert(row);
589            }
590            (Self::U64(values, _), None) => values.push(0),
591            (Self::I64(values, present), Some(ColumnValue::I64(value))) => {
592                values.push(value);
593                present.insert(row);
594            }
595            (Self::I64(values, _), None) => values.push(0),
596            (Self::F64(values, present), Some(ColumnValue::F64(value))) => {
597                values.push(value);
598                present.insert(row);
599            }
600            (Self::F64(values, _), None) => values.push(0.0),
601            (Self::Bool(values, present), Some(ColumnValue::Bool(value))) => {
602                values.push(u8::from(value));
603                present.insert(row);
604            }
605            (Self::Bool(values, _), None) => values.push(0),
606            (
607                Self::Dictionary {
608                    codes,
609                    dictionary,
610                    present,
611                },
612                Some(ColumnValue::String(value)),
613            ) => {
614                codes.push(dictionary.intern(value)?);
615                present.insert(row);
616            }
617            (Self::Dictionary { codes, .. }, None) => codes.push(0),
618            (Self::RawString(values, present), Some(ColumnValue::String(value))) => {
619                values.push(value)?;
620                present.insert(row);
621            }
622            (Self::RawString(values, _), None) => values.push("")?,
623            (_, Some(value)) => {
624                return Err(BuildError::TypeMismatch {
625                    column,
626                    expected,
627                    actual: value.value_type(),
628                });
629            }
630        }
631        Ok(())
632    }
633
634    fn column_type(&self) -> ColumnType {
635        match self {
636            Self::U64(_, _) => ColumnType::U64,
637            Self::I64(_, _) => ColumnType::I64,
638            Self::F64(_, _) => ColumnType::F64,
639            Self::Bool(_, _) => ColumnType::Bool,
640            Self::Dictionary { .. } => ColumnType::DictionaryString,
641            Self::RawString(_, _) => ColumnType::RawString,
642        }
643    }
644
645    fn finish(self) -> Result<Column, DictionaryError> {
646        match self {
647            Self::U64(values, present) => Ok(Column::U64(NumericColumn { values, present })),
648            Self::I64(values, present) => Ok(Column::I64(NumericColumn { values, present })),
649            Self::F64(values, present) => Ok(Column::F64(NumericColumn { values, present })),
650            Self::Bool(values, present) => Ok(Column::Bool(BoolColumn { values, present })),
651            Self::Dictionary {
652                codes,
653                dictionary,
654                present,
655            } => {
656                let cardinality = dictionary.cardinality();
657                Ok(Column::DictionaryString(DictionaryColumn {
658                    codes: DictionaryCodes::from_u32(codes, cardinality)?,
659                    dictionary: dictionary.finish(),
660                    present,
661                }))
662            }
663            Self::RawString(values, present) => {
664                Ok(Column::RawString(RawStringColumn { values, present }))
665            }
666        }
667    }
668}
669
670/// Mutable row builder that keeps every physical column exactly aligned.
671pub struct ColumnStoreBuilder {
672    schema: Schema,
673    columns: Vec<BuilderColumn>,
674    row_count: u32,
675    #[cfg(any(test, feature = "test-support"))]
676    test_limits: Option<MetadataBuildTestLimits>,
677}
678
679impl ColumnStoreBuilder {
680    /// Creates empty physical arrays for a validated schema.
681    #[must_use]
682    pub fn new(schema: Schema) -> Self {
683        let columns = schema
684            .columns()
685            .iter()
686            .map(|definition| BuilderColumn::new(definition.column_type()))
687            .collect();
688        Self {
689            schema,
690            columns,
691            row_count: 0,
692            #[cfg(any(test, feature = "test-support"))]
693            test_limits: None,
694        }
695    }
696
697    /// Creates a builder with lower limits at the production guard sites.
698    #[cfg(any(test, feature = "test-support"))]
699    #[doc(hidden)]
700    #[must_use]
701    pub fn new_with_test_limits(schema: Schema, limits: MetadataBuildTestLimits) -> Self {
702        let mut builder = Self::new(schema);
703        builder.test_limits = Some(limits);
704        builder
705    }
706
707    /// Appends one timestamp and its typed user-column values atomically.
708    pub fn push_row(
709        &mut self,
710        timestamp: i64,
711        inputs: &[ColumnInput<'_>],
712    ) -> Result<u32, BuildError> {
713        if self.row_count == u32::MAX {
714            return Err(BuildError::TooManyRows);
715        }
716        #[cfg(any(test, feature = "test-support"))]
717        if self
718            .test_limits
719            .is_some_and(|limits| self.row_count >= limits.max_rows)
720        {
721            return Err(BuildError::TooManyRows);
722        }
723        self.validate_inputs(inputs)?;
724
725        for (definition, builder) in self.schema.columns().iter().zip(&self.columns) {
726            let value = if definition.id() == TIMESTAMP_COLUMN {
727                Some(ColumnValue::I64(timestamp))
728            } else {
729                find_input(inputs, definition.id()).map(|input| input.value)
730            };
731            #[cfg(any(test, feature = "test-support"))]
732            if let Some(limits) = self.test_limits {
733                builder.check_test_limits(value, limits)?;
734            }
735            builder.preflight(value)?;
736        }
737
738        let row = self.row_count;
739        for (definition, builder) in self.schema.columns().iter().zip(&mut self.columns) {
740            let value = if definition.id() == TIMESTAMP_COLUMN {
741                Some(ColumnValue::I64(timestamp))
742            } else {
743                find_input(inputs, definition.id()).map(|input| input.value)
744            };
745            builder.push(row, definition.id(), value)?;
746        }
747        self.row_count = self.row_count.saturating_add(1);
748        Ok(row)
749    }
750
751    /// Seals the aligned arrays into an immutable column store.
752    pub fn finish(self) -> Result<ColumnStore, BuildError> {
753        let columns = self
754            .columns
755            .into_iter()
756            .map(BuilderColumn::finish)
757            .collect::<Result<Vec<_>, _>>()?;
758        Ok(ColumnStore {
759            schema: self.schema,
760            columns,
761            row_count: self.row_count,
762        })
763    }
764
765    fn validate_inputs(&self, inputs: &[ColumnInput<'_>]) -> Result<(), BuildError> {
766        for (position, input) in inputs.iter().enumerate() {
767            if input.column == TIMESTAMP_COLUMN {
768                return Err(BuildError::TimestampProvidedAsInput);
769            }
770            if inputs
771                .iter()
772                .take(position)
773                .any(|candidate| candidate.column == input.column)
774            {
775                return Err(BuildError::DuplicateColumn(input.column));
776            }
777            let definition = self
778                .schema
779                .column(input.column)
780                .ok_or(BuildError::UnknownColumn(input.column))?;
781            if !input.value.compatible_with(definition.column_type()) {
782                return Err(BuildError::TypeMismatch {
783                    column: input.column,
784                    expected: definition.column_type(),
785                    actual: input.value.value_type(),
786                });
787            }
788        }
789
790        for definition in self
791            .schema
792            .columns()
793            .iter()
794            .filter(|definition| definition.id() != TIMESTAMP_COLUMN)
795        {
796            if !definition.is_nullable() && find_input(inputs, definition.id()).is_none() {
797                return Err(BuildError::MissingRequiredColumn(definition.id()));
798            }
799        }
800        Ok(())
801    }
802}
803
804fn find_input<'a, 'value>(
805    inputs: &'a [ColumnInput<'value>],
806    column: ColumnId,
807) -> Option<&'a ColumnInput<'value>> {
808    inputs.iter().find(|input| input.column == column)
809}
810
811#[allow(
812    clippy::expect_used,
813    clippy::indexing_slicing,
814    clippy::panic,
815    clippy::unwrap_used
816)]
817#[cfg(test)]
818mod tests {
819    use super::*;
820    use crate::meta::{CodeWidth, ColumnDefinition};
821    use proptest::prelude::*;
822    use proptest::test_runner::{Config, RngSeed, TestRunner};
823    use rand::RngCore;
824
825    #[derive(Clone, Debug)]
826    struct ArbitraryRow {
827        timestamp: i64,
828        unsigned: Option<u64>,
829        signed: Option<i64>,
830        float: Option<f64>,
831        boolean: Option<bool>,
832        dictionary: Option<String>,
833        raw: Option<String>,
834    }
835
836    impl Arbitrary for ArbitraryRow {
837        type Parameters = ();
838        type Strategy = BoxedStrategy<Self>;
839
840        fn arbitrary_with((): Self::Parameters) -> Self::Strategy {
841            (
842                any::<i64>(),
843                proptest::option::of(any::<u64>()),
844                proptest::option::of(any::<i64>()),
845                proptest::option::of(any::<i32>().prop_map(f64::from)),
846                proptest::option::of(any::<bool>()),
847                proptest::option::of("[a-z]{0,12}"),
848                proptest::option::of("[ -~]{0,24}"),
849            )
850                .prop_map(
851                    |(timestamp, unsigned, signed, float, boolean, dictionary, raw)| Self {
852                        timestamp,
853                        unsigned,
854                        signed,
855                        float,
856                        boolean,
857                        dictionary,
858                        raw,
859                    },
860                )
861                .boxed()
862        }
863    }
864
865    fn complete_schema(nullable: bool) -> Schema {
866        Schema::new(vec![
867            ColumnDefinition::new(ColumnId::new(1), "u", ColumnType::U64, nullable),
868            ColumnDefinition::new(ColumnId::new(2), "i", ColumnType::I64, nullable),
869            ColumnDefinition::new(ColumnId::new(3), "f", ColumnType::F64, nullable),
870            ColumnDefinition::new(ColumnId::new(4), "b", ColumnType::Bool, nullable),
871            ColumnDefinition::new(
872                ColumnId::new(5),
873                "dictionary",
874                ColumnType::DictionaryString,
875                nullable,
876            ),
877            ColumnDefinition::new(ColumnId::new(6), "raw", ColumnType::RawString, nullable),
878        ])
879        .expect("schema is valid")
880    }
881
882    fn build_arbitrary_store(rows: &[ArbitraryRow]) -> ColumnStore {
883        let mut builder = ColumnStoreBuilder::new(complete_schema(true));
884        for row in rows {
885            let mut inputs = Vec::new();
886            if let Some(value) = row.unsigned {
887                inputs.push(ColumnInput {
888                    column: ColumnId::new(1),
889                    value: ColumnValue::U64(value),
890                });
891            }
892            if let Some(value) = row.signed {
893                inputs.push(ColumnInput {
894                    column: ColumnId::new(2),
895                    value: ColumnValue::I64(value),
896                });
897            }
898            if let Some(value) = row.float {
899                inputs.push(ColumnInput {
900                    column: ColumnId::new(3),
901                    value: ColumnValue::F64(value),
902                });
903            }
904            if let Some(value) = row.boolean {
905                inputs.push(ColumnInput {
906                    column: ColumnId::new(4),
907                    value: ColumnValue::Bool(value),
908                });
909            }
910            if let Some(value) = row.dictionary.as_deref() {
911                inputs.push(ColumnInput {
912                    column: ColumnId::new(5),
913                    value: ColumnValue::String(value),
914                });
915            }
916            if let Some(value) = row.raw.as_deref() {
917                inputs.push(ColumnInput {
918                    column: ColumnId::new(6),
919                    value: ColumnValue::String(value),
920                });
921            }
922            builder
923                .push_row(row.timestamp, &inputs)
924                .expect("arbitrary row is schema-compatible");
925        }
926        builder.finish().expect("bounded store finishes")
927    }
928
929    #[test]
930    fn prop_column_store_builders_and_accessors_roundtrip() {
931        let name = "meta::columns::tests::prop_column_store_builders_and_accessors_roundtrip";
932        let mut seeded = crate::test_support::seeded_rng(name);
933        let config = Config {
934            rng_seed: RngSeed::Fixed(seeded.next_u64()),
935            cases: 256,
936            ..Config::default()
937        };
938        let mut runner = TestRunner::new(config);
939        let rows = prop::collection::vec(any::<ArbitraryRow>(), 0..=32);
940        let result = runner.run(&rows, |rows| {
941            let store = build_arbitrary_store(&rows);
942            prop_assert_eq!(store.row_count(), rows.len() as u32);
943            prop_assert_eq!(store.schema(), &complete_schema(true));
944            prop_assert_eq!(
945                store.timestamps(),
946                rows.iter().map(|row| row.timestamp).collect::<Vec<_>>()
947            );
948
949            for (offset, expected_type) in [
950                ColumnType::I64,
951                ColumnType::U64,
952                ColumnType::I64,
953                ColumnType::F64,
954                ColumnType::Bool,
955                ColumnType::DictionaryString,
956                ColumnType::RawString,
957            ]
958            .into_iter()
959            .enumerate()
960            {
961                let column = store
962                    .column(ColumnId::new(offset as u32))
963                    .expect("defined column");
964                prop_assert_eq!(column.column_type(), expected_type);
965                prop_assert_eq!(column.len(), rows.len());
966                prop_assert_eq!(column.is_empty(), rows.is_empty());
967            }
968
969            for (position, row) in rows.iter().enumerate() {
970                let row_id = position as u32;
971                prop_assert_eq!(store.timestamp(row_id), Some(row.timestamp));
972                match store.column(ColumnId::new(1)) {
973                    Some(Column::U64(column)) => prop_assert_eq!(column.get(row_id), row.unsigned),
974                    actual => prop_assert!(false, "unexpected u64 column: {actual:?}"),
975                }
976                match store.column(ColumnId::new(2)) {
977                    Some(Column::I64(column)) => prop_assert_eq!(column.get(row_id), row.signed),
978                    actual => prop_assert!(false, "unexpected i64 column: {actual:?}"),
979                }
980                match store.column(ColumnId::new(3)) {
981                    Some(Column::F64(column)) => prop_assert_eq!(column.get(row_id), row.float),
982                    actual => prop_assert!(false, "unexpected f64 column: {actual:?}"),
983                }
984                match store.column(ColumnId::new(4)) {
985                    Some(Column::Bool(column)) => {
986                        prop_assert_eq!(column.get(row_id), row.boolean)
987                    }
988                    actual => prop_assert!(false, "unexpected bool column: {actual:?}"),
989                }
990                match store.column(ColumnId::new(5)) {
991                    Some(Column::DictionaryString(column)) => {
992                        prop_assert_eq!(column.get(row_id), row.dictionary.as_deref());
993                    }
994                    actual => prop_assert!(false, "unexpected dictionary column: {actual:?}"),
995                }
996                match store.column(ColumnId::new(6)) {
997                    Some(Column::RawString(column)) => {
998                        prop_assert_eq!(column.get(row_id), row.raw.as_deref());
999                    }
1000                    actual => prop_assert!(false, "unexpected raw column: {actual:?}"),
1001                }
1002            }
1003            let outside = rows.len() as u32;
1004            prop_assert_eq!(store.timestamp(outside), None);
1005            prop_assert_eq!(store.column(ColumnId::new(99)), None);
1006            Ok(())
1007        });
1008        assert!(result.is_ok(), "property result: {result:?}");
1009    }
1010
1011    #[test]
1012    fn builder_errors_are_typed_and_atomic_at_capacity_edges() {
1013        let mut builder = ColumnStoreBuilder::new(complete_schema(false));
1014        let error = builder
1015            .push_row(
1016                7,
1017                &[ColumnInput {
1018                    column: ColumnId::new(1),
1019                    value: ColumnValue::Bool(true),
1020                }],
1021            )
1022            .expect_err("wrong physical type must fail");
1023        assert_eq!(
1024            error,
1025            BuildError::TypeMismatch {
1026                column: ColumnId::new(1),
1027                expected: ColumnType::U64,
1028                actual: ColumnType::Bool,
1029            }
1030        );
1031        assert_eq!(builder.row_count, 0);
1032
1033        assert_eq!(
1034            builder.push_row(7, &[]),
1035            Err(BuildError::MissingRequiredColumn(ColumnId::new(1)))
1036        );
1037        assert_eq!(builder.row_count, 0, "length mismatch must be atomic");
1038
1039        assert_eq!(
1040            builder.push_row(
1041                7,
1042                &[ColumnInput {
1043                    column: ColumnId::new(99),
1044                    value: ColumnValue::U64(1),
1045                }]
1046            ),
1047            Err(BuildError::UnknownColumn(ColumnId::new(99)))
1048        );
1049        assert_eq!(
1050            builder.push_row(
1051                7,
1052                &[
1053                    ColumnInput {
1054                        column: ColumnId::new(1),
1055                        value: ColumnValue::U64(1),
1056                    },
1057                    ColumnInput {
1058                        column: ColumnId::new(1),
1059                        value: ColumnValue::U64(2),
1060                    },
1061                ]
1062            ),
1063            Err(BuildError::DuplicateColumn(ColumnId::new(1)))
1064        );
1065        assert_eq!(
1066            builder.push_row(
1067                7,
1068                &[ColumnInput {
1069                    column: TIMESTAMP_COLUMN,
1070                    value: ColumnValue::I64(7),
1071                }]
1072            ),
1073            Err(BuildError::TimestampProvidedAsInput)
1074        );
1075        assert_eq!(builder.row_count, 0);
1076
1077        let mut at_capacity = ColumnStoreBuilder::new(complete_schema(true));
1078        at_capacity.row_count = u32::MAX;
1079        assert_eq!(at_capacity.push_row(8, &[]), Err(BuildError::TooManyRows));
1080        assert_eq!(at_capacity.row_count, u32::MAX);
1081
1082        let mut nullable = ColumnStoreBuilder::new(complete_schema(true));
1083        assert_eq!(nullable.push_row(11, &[]), Ok(0));
1084        let store = nullable.finish().expect("null row finishes");
1085        assert_eq!(store.timestamps(), &[11]);
1086        for id in 1..=6 {
1087            let column = store.column(ColumnId::new(id)).expect("nullable column");
1088            assert_eq!(column.len(), 1);
1089            assert_eq!(column.present().cardinality(), 0);
1090        }
1091        match store.column(ColumnId::new(1)) {
1092            Some(Column::U64(column)) => assert_eq!(column.get(0), None),
1093            actual => panic!("unexpected column: {actual:?}"),
1094        }
1095        match store.column(ColumnId::new(4)) {
1096            Some(Column::Bool(column)) => assert_eq!(column.get(0), None),
1097            actual => panic!("unexpected column: {actual:?}"),
1098        }
1099        match store.column(ColumnId::new(5)) {
1100            Some(Column::DictionaryString(column)) => assert_eq!(column.get(0), None),
1101            actual => panic!("unexpected column: {actual:?}"),
1102        }
1103        match store.column(ColumnId::new(6)) {
1104            Some(Column::RawString(column)) => assert_eq!(column.get(0), None),
1105            actual => panic!("unexpected column: {actual:?}"),
1106        }
1107    }
1108
1109    fn dictionary_store(cardinality: u32) -> ColumnStore {
1110        let schema = Schema::new(vec![ColumnDefinition::new(
1111            ColumnId::new(1),
1112            "dictionary",
1113            ColumnType::DictionaryString,
1114            false,
1115        )])
1116        .expect("schema is valid");
1117        let mut builder = ColumnStoreBuilder::new(schema);
1118        for value in 0..cardinality {
1119            let text = value.to_string();
1120            builder
1121                .push_row(
1122                    i64::from(value),
1123                    &[ColumnInput {
1124                        column: ColumnId::new(1),
1125                        value: ColumnValue::String(&text),
1126                    }],
1127                )
1128                .expect("bounded dictionary row");
1129        }
1130        builder.finish().expect("bounded dictionary finishes")
1131    }
1132
1133    #[test]
1134    fn column_builder_widens_dictionary_at_u16_boundary() {
1135        let narrow = dictionary_store(u32::from(u16::MAX));
1136        match narrow.column(ColumnId::new(1)) {
1137            Some(Column::DictionaryString(column)) => {
1138                assert_eq!(column.codes().width(), CodeWidth::U16);
1139                assert_eq!(column.get(u32::from(u16::MAX) - 1), Some("65534"));
1140            }
1141            actual => panic!("unexpected dictionary column: {actual:?}"),
1142        }
1143
1144        let wide = dictionary_store(u32::from(u16::MAX).saturating_add(1));
1145        match wide.column(ColumnId::new(1)) {
1146            Some(Column::DictionaryString(column)) => {
1147                assert_eq!(column.codes().width(), CodeWidth::U32);
1148                assert_eq!(column.get(u32::from(u16::MAX)), Some("65535"));
1149            }
1150            actual => panic!("unexpected dictionary column: {actual:?}"),
1151        }
1152    }
1153}
zeppelin_embed/meta/columns.rs

zeppelin_embed/meta/
columns.rs
