Skip to main content

scirs2_io/columnar/
types.rs

1//! Column types and data structures for the columnar format.
2//!
3//! Defines the core types for column-oriented storage including
4//! column data variants, encoding strategies, and table structures.
5
6use std::collections::HashMap;
7use std::fmt;
8
9use crate::error::{IoError, Result};
10
11/// Magic bytes identifying the columnar format file
12pub const COLUMNAR_MAGIC: &[u8; 8] = b"SCIRCOL\x01";
13
14/// Current format version
15///
16/// Version 2 adds an optional per-column null/validity bitmap (see
17/// [`Column::null_mask`]) written directly after each column's encoded data.
18/// Version 1 files (no null section) are still readable: [`super::reader`]
19/// treats every column in a v1 file as having no null mask.
20pub const FORMAT_VERSION: u32 = 2;
21
22/// Column data type tag stored in the file header
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24#[repr(u8)]
25pub enum ColumnTypeTag {
26    /// 64-bit floating point
27    Float64 = 0,
28    /// 64-bit signed integer
29    Int64 = 1,
30    /// UTF-8 string
31    Str = 2,
32    /// Boolean
33    Bool = 3,
34}
35
36impl TryFrom<u8> for ColumnTypeTag {
37    type Error = IoError;
38
39    fn try_from(value: u8) -> std::result::Result<Self, Self::Error> {
40        match value {
41            0 => Ok(ColumnTypeTag::Float64),
42            1 => Ok(ColumnTypeTag::Int64),
43            2 => Ok(ColumnTypeTag::Str),
44            3 => Ok(ColumnTypeTag::Bool),
45            _ => Err(IoError::FormatError(format!(
46                "Unknown column type tag: {}",
47                value
48            ))),
49        }
50    }
51}
52
53/// Encoding strategy for a column
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55#[repr(u8)]
56pub enum EncodingType {
57    /// Plain (no encoding)
58    Plain = 0,
59    /// Run-length encoding
60    Rle = 1,
61    /// Dictionary encoding
62    Dictionary = 2,
63    /// Delta encoding (for sorted numeric columns)
64    Delta = 3,
65}
66
67impl TryFrom<u8> for EncodingType {
68    type Error = IoError;
69
70    fn try_from(value: u8) -> std::result::Result<Self, Self::Error> {
71        match value {
72            0 => Ok(EncodingType::Plain),
73            1 => Ok(EncodingType::Rle),
74            2 => Ok(EncodingType::Dictionary),
75            3 => Ok(EncodingType::Delta),
76            _ => Err(IoError::FormatError(format!(
77                "Unknown encoding type: {}",
78                value
79            ))),
80        }
81    }
82}
83
84/// A single column's data
85#[derive(Debug, Clone)]
86pub enum ColumnData {
87    /// 64-bit floating point values
88    Float64(Vec<f64>),
89    /// 64-bit signed integer values
90    Int64(Vec<i64>),
91    /// UTF-8 string values
92    Str(Vec<String>),
93    /// Boolean values
94    Bool(Vec<bool>),
95}
96
97impl ColumnData {
98    /// Returns the number of values in this column
99    pub fn len(&self) -> usize {
100        match self {
101            ColumnData::Float64(v) => v.len(),
102            ColumnData::Int64(v) => v.len(),
103            ColumnData::Str(v) => v.len(),
104            ColumnData::Bool(v) => v.len(),
105        }
106    }
107
108    /// Returns true if the column is empty
109    pub fn is_empty(&self) -> bool {
110        self.len() == 0
111    }
112
113    /// Returns the column type tag
114    pub fn type_tag(&self) -> ColumnTypeTag {
115        match self {
116            ColumnData::Float64(_) => ColumnTypeTag::Float64,
117            ColumnData::Int64(_) => ColumnTypeTag::Int64,
118            ColumnData::Str(_) => ColumnTypeTag::Str,
119            ColumnData::Bool(_) => ColumnTypeTag::Bool,
120        }
121    }
122
123    /// Try to get f64 data
124    pub fn as_f64(&self) -> Result<&[f64]> {
125        match self {
126            ColumnData::Float64(v) => Ok(v),
127            _ => Err(IoError::ConversionError(format!(
128                "Column is {:?}, not Float64",
129                self.type_tag()
130            ))),
131        }
132    }
133
134    /// Try to get i64 data
135    pub fn as_i64(&self) -> Result<&[i64]> {
136        match self {
137            ColumnData::Int64(v) => Ok(v),
138            _ => Err(IoError::ConversionError(format!(
139                "Column is {:?}, not Int64",
140                self.type_tag()
141            ))),
142        }
143    }
144
145    /// Try to get string data
146    pub fn as_str(&self) -> Result<&[String]> {
147        match self {
148            ColumnData::Str(v) => Ok(v),
149            _ => Err(IoError::ConversionError(format!(
150                "Column is {:?}, not Str",
151                self.type_tag()
152            ))),
153        }
154    }
155
156    /// Try to get bool data
157    pub fn as_bool(&self) -> Result<&[bool]> {
158        match self {
159            ColumnData::Bool(v) => Ok(v),
160            _ => Err(IoError::ConversionError(format!(
161                "Column is {:?}, not Bool",
162                self.type_tag()
163            ))),
164        }
165    }
166
167    /// Determine best encoding for this column's data
168    pub fn best_encoding(&self) -> EncodingType {
169        match self {
170            ColumnData::Float64(v) => {
171                if is_sorted_f64(v) {
172                    EncodingType::Delta
173                } else if has_runs_f64(v) {
174                    EncodingType::Rle
175                } else {
176                    EncodingType::Plain
177                }
178            }
179            ColumnData::Int64(v) => {
180                if is_sorted_i64(v) {
181                    EncodingType::Delta
182                } else if has_runs_i64(v) {
183                    EncodingType::Rle
184                } else {
185                    EncodingType::Plain
186                }
187            }
188            ColumnData::Str(v) => {
189                let unique_count = count_unique_strings(v);
190                if unique_count < v.len() / 2 {
191                    EncodingType::Dictionary
192                } else if has_runs_str(v) {
193                    EncodingType::Rle
194                } else {
195                    EncodingType::Plain
196                }
197            }
198            ColumnData::Bool(v) => {
199                if has_runs_bool(v) {
200                    EncodingType::Rle
201                } else {
202                    EncodingType::Plain
203                }
204            }
205        }
206    }
207}
208
209impl fmt::Display for ColumnData {
210    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211        match self {
212            ColumnData::Float64(v) => write!(f, "Float64[{}]", v.len()),
213            ColumnData::Int64(v) => write!(f, "Int64[{}]", v.len()),
214            ColumnData::Str(v) => write!(f, "Str[{}]", v.len()),
215            ColumnData::Bool(v) => write!(f, "Bool[{}]", v.len()),
216        }
217    }
218}
219
220/// A named column in a table
221#[derive(Debug, Clone)]
222pub struct Column {
223    /// Column name
224    pub name: String,
225    /// Column data
226    pub data: ColumnData,
227    /// Optional per-row null/validity bitmap: `null_mask[i] == true` means row
228    /// `i` is null/missing. `None` means the column has no null tracking (every
229    /// row is a valid, present value) -- the case for all columns created via
230    /// the plain `float64`/`int64`/`string`/`boolean` constructors.
231    ///
232    /// The underlying `ColumnData` vector still holds a concrete placeholder
233    /// value at null positions (there being no other way to represent "no
234    /// value" in a plain `Vec<T>`); the mask is what determines whether that
235    /// placeholder is a real value or should be treated as missing by
236    /// `ColumnStats` and other consumers.
237    pub null_mask: Option<Vec<bool>>,
238}
239
240impl Column {
241    /// Create a new column with f64 data (no null tracking)
242    pub fn float64(name: impl Into<String>, data: Vec<f64>) -> Self {
243        Column {
244            name: name.into(),
245            data: ColumnData::Float64(data),
246            null_mask: None,
247        }
248    }
249
250    /// Create a new column with i64 data (no null tracking)
251    pub fn int64(name: impl Into<String>, data: Vec<i64>) -> Self {
252        Column {
253            name: name.into(),
254            data: ColumnData::Int64(data),
255            null_mask: None,
256        }
257    }
258
259    /// Create a new column with string data (no null tracking)
260    pub fn string(name: impl Into<String>, data: Vec<String>) -> Self {
261        Column {
262            name: name.into(),
263            data: ColumnData::Str(data),
264            null_mask: None,
265        }
266    }
267
268    /// Create a new column with bool data (no null tracking)
269    pub fn boolean(name: impl Into<String>, data: Vec<bool>) -> Self {
270        Column {
271            name: name.into(),
272            data: ColumnData::Bool(data),
273            null_mask: None,
274        }
275    }
276
277    /// Create a new column with f64 data and an explicit null mask.
278    ///
279    /// `null_mask[i] == true` marks row `i` as null; the corresponding entry
280    /// in `data` is a placeholder and is excluded from `ColumnStats`.
281    ///
282    /// # Errors
283    ///
284    /// Returns an error if `null_mask.len() != data.len()`.
285    pub fn float64_with_nulls(
286        name: impl Into<String>,
287        data: Vec<f64>,
288        null_mask: Vec<bool>,
289    ) -> Result<Self> {
290        Self::with_nulls(name, ColumnData::Float64(data), null_mask)
291    }
292
293    /// Create a new column with i64 data and an explicit null mask (see
294    /// [`Column::float64_with_nulls`]).
295    pub fn int64_with_nulls(
296        name: impl Into<String>,
297        data: Vec<i64>,
298        null_mask: Vec<bool>,
299    ) -> Result<Self> {
300        Self::with_nulls(name, ColumnData::Int64(data), null_mask)
301    }
302
303    /// Create a new column with string data and an explicit null mask (see
304    /// [`Column::float64_with_nulls`]).
305    pub fn string_with_nulls(
306        name: impl Into<String>,
307        data: Vec<String>,
308        null_mask: Vec<bool>,
309    ) -> Result<Self> {
310        Self::with_nulls(name, ColumnData::Str(data), null_mask)
311    }
312
313    /// Create a new column with bool data and an explicit null mask (see
314    /// [`Column::float64_with_nulls`]).
315    pub fn boolean_with_nulls(
316        name: impl Into<String>,
317        data: Vec<bool>,
318        null_mask: Vec<bool>,
319    ) -> Result<Self> {
320        Self::with_nulls(name, ColumnData::Bool(data), null_mask)
321    }
322
323    fn with_nulls(name: impl Into<String>, data: ColumnData, null_mask: Vec<bool>) -> Result<Self> {
324        if null_mask.len() != data.len() {
325            return Err(IoError::FormatError(format!(
326                "null_mask has {} entries, expected {} (one per row)",
327                null_mask.len(),
328                data.len()
329            )));
330        }
331        Ok(Column {
332            name: name.into(),
333            data,
334            null_mask: Some(null_mask),
335        })
336    }
337
338    /// Returns the length (number of rows) of this column
339    pub fn len(&self) -> usize {
340        self.data.len()
341    }
342
343    /// Returns true if column is empty
344    pub fn is_empty(&self) -> bool {
345        self.data.is_empty()
346    }
347
348    /// Returns `true` if row `row` is null. Always `false` when the column
349    /// has no null mask.
350    pub fn is_null(&self, row: usize) -> bool {
351        self.null_mask
352            .as_ref()
353            .and_then(|mask| mask.get(row))
354            .copied()
355            .unwrap_or(false)
356    }
357
358    /// Returns the number of null rows in this column (0 when there is no
359    /// null mask).
360    pub fn null_count(&self) -> usize {
361        self.null_mask
362            .as_ref()
363            .map(|mask| mask.iter().filter(|&&is_null| is_null).count())
364            .unwrap_or(0)
365    }
366}
367
368/// A table containing multiple named columns
369#[derive(Debug, Clone)]
370pub struct ColumnarTable {
371    /// Columns in order
372    columns: Vec<Column>,
373    /// Name-to-index lookup
374    index: HashMap<String, usize>,
375}
376
377impl ColumnarTable {
378    /// Create a new empty table
379    pub fn new() -> Self {
380        ColumnarTable {
381            columns: Vec::new(),
382            index: HashMap::new(),
383        }
384    }
385
386    /// Create a table from a list of columns
387    pub fn from_columns(columns: Vec<Column>) -> Result<Self> {
388        // Validate all columns have same length
389        if !columns.is_empty() {
390            let expected_len = columns[0].len();
391            for col in &columns[1..] {
392                if col.len() != expected_len {
393                    return Err(IoError::FormatError(format!(
394                        "Column '{}' has {} rows, expected {}",
395                        col.name,
396                        col.len(),
397                        expected_len
398                    )));
399                }
400            }
401        }
402
403        let mut index = HashMap::new();
404        for (i, col) in columns.iter().enumerate() {
405            if index.contains_key(&col.name) {
406                return Err(IoError::FormatError(format!(
407                    "Duplicate column name: '{}'",
408                    col.name
409                )));
410            }
411            index.insert(col.name.clone(), i);
412        }
413
414        Ok(ColumnarTable { columns, index })
415    }
416
417    /// Add a column to the table
418    pub fn add_column(&mut self, column: Column) -> Result<()> {
419        if !self.columns.is_empty() && column.len() != self.num_rows() {
420            return Err(IoError::FormatError(format!(
421                "Column '{}' has {} rows, expected {}",
422                column.name,
423                column.len(),
424                self.num_rows()
425            )));
426        }
427        if self.index.contains_key(&column.name) {
428            return Err(IoError::FormatError(format!(
429                "Duplicate column name: '{}'",
430                column.name
431            )));
432        }
433        let idx = self.columns.len();
434        self.index.insert(column.name.clone(), idx);
435        self.columns.push(column);
436        Ok(())
437    }
438
439    /// Number of rows in the table
440    pub fn num_rows(&self) -> usize {
441        self.columns.first().map(|c| c.len()).unwrap_or(0)
442    }
443
444    /// Number of columns
445    pub fn num_columns(&self) -> usize {
446        self.columns.len()
447    }
448
449    /// Get column names in order
450    pub fn column_names(&self) -> Vec<&str> {
451        self.columns.iter().map(|c| c.name.as_str()).collect()
452    }
453
454    /// Get a column by name
455    pub fn column(&self, name: &str) -> Result<&Column> {
456        self.index
457            .get(name)
458            .map(|&idx| &self.columns[idx])
459            .ok_or_else(|| IoError::NotFound(format!("Column '{}' not found", name)))
460    }
461
462    /// Get a column by index
463    pub fn column_by_index(&self, idx: usize) -> Result<&Column> {
464        self.columns
465            .get(idx)
466            .ok_or_else(|| IoError::NotFound(format!("Column index {} out of range", idx)))
467    }
468
469    /// Get all columns as a slice
470    pub fn columns(&self) -> &[Column] {
471        &self.columns
472    }
473
474    /// Get f64 column data by name
475    pub fn get_f64(&self, name: &str) -> Result<&[f64]> {
476        self.column(name)?.data.as_f64()
477    }
478
479    /// Get i64 column data by name
480    pub fn get_i64(&self, name: &str) -> Result<&[i64]> {
481        self.column(name)?.data.as_i64()
482    }
483
484    /// Get string column data by name
485    pub fn get_str(&self, name: &str) -> Result<&[String]> {
486        self.column(name)?.data.as_str()
487    }
488
489    /// Get bool column data by name
490    pub fn get_bool(&self, name: &str) -> Result<&[bool]> {
491        self.column(name)?.data.as_bool()
492    }
493}
494
495impl Default for ColumnarTable {
496    fn default() -> Self {
497        Self::new()
498    }
499}
500
501// Helper functions for encoding detection
502
503fn is_sorted_f64(data: &[f64]) -> bool {
504    if data.len() < 2 {
505        return true;
506    }
507    data.windows(2).all(|w| w[0] <= w[1])
508}
509
510fn is_sorted_i64(data: &[i64]) -> bool {
511    if data.len() < 2 {
512        return true;
513    }
514    data.windows(2).all(|w| w[0] <= w[1])
515}
516
517fn has_runs_f64(data: &[f64]) -> bool {
518    if data.len() < 4 {
519        return false;
520    }
521    let mut run_count = 0;
522    let mut i = 0;
523    while i < data.len() {
524        let val = data[i];
525        let mut run_len = 1;
526        while i + run_len < data.len() && data[i + run_len] == val {
527            run_len += 1;
528        }
529        if run_len > 1 {
530            run_count += 1;
531        }
532        i += run_len;
533    }
534    // Beneficial if at least 20% of groups are runs
535    run_count * 5 >= data.len()
536}
537
538fn has_runs_i64(data: &[i64]) -> bool {
539    if data.len() < 4 {
540        return false;
541    }
542    let mut run_count = 0;
543    let mut i = 0;
544    while i < data.len() {
545        let val = data[i];
546        let mut run_len = 1;
547        while i + run_len < data.len() && data[i + run_len] == val {
548            run_len += 1;
549        }
550        if run_len > 1 {
551            run_count += 1;
552        }
553        i += run_len;
554    }
555    run_count * 5 >= data.len()
556}
557
558fn has_runs_str(data: &[String]) -> bool {
559    if data.len() < 4 {
560        return false;
561    }
562    let mut run_count = 0;
563    let mut i = 0;
564    while i < data.len() {
565        let val = &data[i];
566        let mut run_len = 1;
567        while i + run_len < data.len() && &data[i + run_len] == val {
568            run_len += 1;
569        }
570        if run_len > 1 {
571            run_count += 1;
572        }
573        i += run_len;
574    }
575    run_count * 5 >= data.len()
576}
577
578fn has_runs_bool(data: &[bool]) -> bool {
579    if data.len() < 4 {
580        return false;
581    }
582    let mut run_count = 0;
583    let mut i = 0;
584    while i < data.len() {
585        let val = data[i];
586        let mut run_len = 1;
587        while i + run_len < data.len() && data[i + run_len] == val {
588            run_len += 1;
589        }
590        if run_len > 1 {
591            run_count += 1;
592        }
593        i += run_len;
594    }
595    run_count * 5 >= data.len()
596}
597
598fn count_unique_strings(data: &[String]) -> usize {
599    let mut seen = std::collections::HashSet::new();
600    for s in data {
601        seen.insert(s.as_str());
602    }
603    seen.len()
604}