scirs2_io/formats/
orc_lite.rs

1//! ORC-lite columnar format
2//!
3//! A simplified implementation of the Apache ORC columnar storage format,
4//! supporting stripes, a file footer, and the ORC integer RLE v2 codec.
5//!
6//! Format layout:
7//! ```text
8//! [MAGIC "ORCLITE\0"] [stripe_0] ... [stripe_N] [FOOTER] [footer_length: u32 LE]
9//! ```
10//!
11//! Each stripe:
12//! ```text
13//! [STRIPE_MAGIC "STRP"] [num_rows: u32 LE] [num_cols: u32 LE]
14//!   [col_0_header] ... [col_N_header]
15//!   [col_0_data]   ... [col_N_data]
16//! ```
17//!
18//! Column header: `[name_len: u16 LE][name bytes][encoding: u8][data_len: u32 LE]`
19
20use std::collections::HashMap;
21use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
22use std::path::Path;
23
24use crate::error::{IoError, Result};
25
26// ──────────────────────────────────────────────────────────────────────────────
27// Constants
28// ──────────────────────────────────────────────────────────────────────────────
29
30const FILE_MAGIC: &[u8; 8] = b"ORCLITE\0";
31const STRIPE_MAGIC: &[u8; 4] = b"STRP";
32const FOOTER_MAGIC: &[u8; 4] = b"FOOT";
33const FILE_MAGIC_LEN: usize = 8;
34
35// ──────────────────────────────────────────────────────────────────────────────
36// Column encoding
37// ──────────────────────────────────────────────────────────────────────────────
38
39/// Column data encoding strategy.
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41#[repr(u8)]
42pub enum ColumnEncoding {
43    /// Raw values stored directly (8 bytes per element for f64, 8 for i64).
44    Direct = 0,
45    /// Dictionary encoding: values replaced by u16 indices into a dictionary.
46    Dictionary = 1,
47    /// ORC integer RLE v2 (delta + base encoding for i64 columns).
48    RleV2 = 2,
49    /// UTF-8 length-prefixed strings (variable width).
50    DirectString = 3,
51}
52
53impl ColumnEncoding {
54    fn from_u8(v: u8) -> Result<Self> {
55        match v {
56            0 => Ok(ColumnEncoding::Direct),
57            1 => Ok(ColumnEncoding::Dictionary),
58            2 => Ok(ColumnEncoding::RleV2),
59            3 => Ok(ColumnEncoding::DirectString),
60            other => Err(IoError::FormatError(format!(
61                "Unknown ORC-lite column encoding byte: {}",
62                other
63            ))),
64        }
65    }
66}
67
68// ──────────────────────────────────────────────────────────────────────────────
69// Column data variants
70// ──────────────────────────────────────────────────────────────────────────────
71
72/// In-memory column data.
73#[derive(Debug, Clone, PartialEq)]
74pub enum OrcColumnData {
75    Int64(Vec<i64>),
76    Float64(Vec<f64>),
77    Boolean(Vec<bool>),
78    String(Vec<String>),
79    /// Nullable integer (None = null).
80    NullableInt64(Vec<Option<i64>>),
81}
82
83impl OrcColumnData {
84    /// Number of values in the column.
85    pub fn len(&self) -> usize {
86        match self {
87            OrcColumnData::Int64(v) => v.len(),
88            OrcColumnData::Float64(v) => v.len(),
89            OrcColumnData::Boolean(v) => v.len(),
90            OrcColumnData::String(v) => v.len(),
91            OrcColumnData::NullableInt64(v) => v.len(),
92        }
93    }
94
95    pub fn is_empty(&self) -> bool {
96        self.len() == 0
97    }
98
99    fn type_tag(&self) -> u8 {
100        match self {
101            OrcColumnData::Int64(_) => 0,
102            OrcColumnData::Float64(_) => 1,
103            OrcColumnData::Boolean(_) => 2,
104            OrcColumnData::String(_) => 3,
105            OrcColumnData::NullableInt64(_) => 4,
106        }
107    }
108}
109
110// ──────────────────────────────────────────────────────────────────────────────
111// ORC Integer RLE v2
112// ──────────────────────────────────────────────────────────────────────────────
113
114/// ORC integer RLE v2 encoder/decoder.
115///
116/// This is a simplified subset focusing on DELTA sub-encoding:
117/// 1. Compute deltas from the base value.
118/// 2. Encode with variable-length zigzag encoding (LEB-128 signed).
119pub struct IntRleV2;
120
121impl IntRleV2 {
122    /// Encode a slice of `i64` values using delta + zigzag LEB-128.
123    ///
124    /// Output format:
125    /// ```text
126    /// [count: u32 LE][base: i64 zigzag LEB128][delta_0 ... delta_N-1 zigzag LEB128]
127    /// ```
128    pub fn encode(values: &[i64]) -> Vec<u8> {
129        let mut out = Vec::new();
130
131        // count
132        let count = values.len() as u32;
133        out.extend_from_slice(&count.to_le_bytes());
134
135        if values.is_empty() {
136            return out;
137        }
138
139        // base value
140        encode_zigzag_leb128(values[0], &mut out);
141
142        // deltas
143        for i in 1..values.len() {
144            let delta = values[i].wrapping_sub(values[i - 1]);
145            encode_zigzag_leb128(delta, &mut out);
146        }
147
148        out
149    }
150
151    /// Decode bytes produced by `encode` back into `Vec<i64>`.
152    pub fn decode(data: &[u8]) -> Result<Vec<i64>> {
153        if data.len() < 4 {
154            return Err(IoError::FormatError(
155                "RLE v2: too short to read count".into(),
156            ));
157        }
158        let count = u32::from_le_bytes(
159            data[0..4]
160                .try_into()
161                .map_err(|_| IoError::FormatError("RLE v2: count bytes".into()))?,
162        ) as usize;
163
164        if count == 0 {
165            return Ok(Vec::new());
166        }
167
168        let mut pos = 4;
169        let (base, consumed) = decode_zigzag_leb128(&data[pos..])?;
170        pos += consumed;
171
172        let mut values = Vec::with_capacity(count);
173        values.push(base);
174
175        for _ in 1..count {
176            let (delta, consumed) = decode_zigzag_leb128(&data[pos..])?;
177            pos += consumed;
178            let prev = *values.last().expect("values non-empty");
179            values.push(prev.wrapping_add(delta));
180        }
181
182        Ok(values)
183    }
184}
185
186/// Encode a signed `i64` with zigzag mapping and then LEB-128.
187fn encode_zigzag_leb128(value: i64, out: &mut Vec<u8>) {
188    // Zigzag: 0->0, -1->1, 1->2, -2->3, 2->4 ...
189    let zigzag = ((value << 1) ^ (value >> 63)) as u64;
190    let mut v = zigzag;
191    loop {
192        let low7 = (v & 0x7F) as u8;
193        v >>= 7;
194        if v == 0 {
195            out.push(low7);
196            break;
197        } else {
198            out.push(low7 | 0x80);
199        }
200    }
201}
202
203/// Decode a zigzag LEB-128 value; returns `(value, bytes_consumed)`.
204fn decode_zigzag_leb128(data: &[u8]) -> Result<(i64, usize)> {
205    let mut result: u64 = 0;
206    let mut shift = 0u32;
207    for (i, &byte) in data.iter().enumerate() {
208        result |= ((byte & 0x7F) as u64) << shift;
209        shift += 7;
210        if byte & 0x80 == 0 {
211            // Undo zigzag
212            let signed = ((result >> 1) as i64) ^ -((result & 1) as i64);
213            return Ok((signed, i + 1));
214        }
215        if shift >= 64 {
216            return Err(IoError::FormatError(
217                "zigzag LEB-128: overflow (shift >= 64)".into(),
218            ));
219        }
220    }
221    Err(IoError::FormatError(
222        "zigzag LEB-128: truncated encoding".into(),
223    ))
224}
225
226// ──────────────────────────────────────────────────────────────────────────────
227// Stripe metadata
228// ──────────────────────────────────────────────────────────────────────────────
229
230#[derive(Debug, Clone)]
231struct StripeInfo {
232    /// Byte offset of stripe start in the file.
233    offset: u64,
234    /// Number of rows in the stripe.
235    num_rows: u32,
236    /// Column names in order.
237    column_names: Vec<String>,
238}
239
240// ──────────────────────────────────────────────────────────────────────────────
241// OrcLiteWriter
242// ──────────────────────────────────────────────────────────────────────────────
243
244/// Writer for the ORC-lite columnar format.
245///
246/// Data is written in stripes. Each stripe is self-contained.
247/// After all stripes are written, `finalize()` writes the footer.
248///
249/// # Example
250/// ```no_run
251/// use scirs2_io::formats::orc_lite::{OrcLiteWriter, OrcColumnData};
252/// use std::fs::File;
253///
254/// let file = File::create("/tmp/data.orc").unwrap();
255/// let mut writer = OrcLiteWriter::new(file).unwrap();
256/// let cols = vec![
257///     ("id".to_string(), OrcColumnData::Int64(vec![1, 2, 3])),
258///     ("value".to_string(), OrcColumnData::Float64(vec![1.1, 2.2, 3.3])),
259/// ];
260/// writer.write_stripe(&cols).unwrap();
261/// writer.finalize().unwrap();
262/// ```
263pub struct OrcLiteWriter<W: Write + Seek> {
264    inner: BufWriter<W>,
265    stripes: Vec<StripeInfo>,
266    current_offset: u64,
267}
268
269impl<W: Write + Seek> OrcLiteWriter<W> {
270    /// Create a new writer, writing the file magic header.
271    pub fn new(writer: W) -> Result<Self> {
272        let mut bw = BufWriter::new(writer);
273        bw.write_all(FILE_MAGIC).map_err(IoError::Io)?;
274        Ok(OrcLiteWriter {
275            inner: bw,
276            stripes: Vec::new(),
277            current_offset: FILE_MAGIC_LEN as u64,
278        })
279    }
280
281    /// Write one stripe (a batch of columns of equal length).
282    ///
283    /// All columns must have the same number of rows.
284    pub fn write_stripe(&mut self, columns: &[(String, OrcColumnData)]) -> Result<()> {
285        if columns.is_empty() {
286            return Ok(());
287        }
288
289        // Validate all columns have equal length
290        let num_rows = columns[0].1.len();
291        for (name, col) in columns {
292            if col.len() != num_rows {
293                return Err(IoError::ValidationError(format!(
294                    "column '{}' has {} rows, expected {}",
295                    name,
296                    col.len(),
297                    num_rows
298                )));
299            }
300        }
301
302        let stripe_offset = self.current_offset;
303
304        // STRIPE_MAGIC
305        self.inner.write_all(STRIPE_MAGIC).map_err(IoError::Io)?;
306        self.current_offset += 4;
307
308        // num_rows, num_cols
309        let num_rows_u32 = num_rows as u32;
310        let num_cols_u32 = columns.len() as u32;
311        self.inner
312            .write_all(&num_rows_u32.to_le_bytes())
313            .map_err(IoError::Io)?;
314        self.inner
315            .write_all(&num_cols_u32.to_le_bytes())
316            .map_err(IoError::Io)?;
317        self.current_offset += 8;
318
319        // Encode each column
320        let mut encoded: Vec<(String, ColumnEncoding, u8, Vec<u8>)> = Vec::new();
321        for (name, col) in columns {
322            let (enc, type_tag, bytes) = encode_column(col)?;
323            encoded.push((name.clone(), enc, type_tag, bytes));
324        }
325
326        // Write column headers
327        for (name, enc, type_tag, bytes) in &encoded {
328            let name_bytes = name.as_bytes();
329            let name_len = name_bytes.len() as u16;
330            self.inner
331                .write_all(&name_len.to_le_bytes())
332                .map_err(IoError::Io)?;
333            self.inner.write_all(name_bytes).map_err(IoError::Io)?;
334            self.inner.write_all(&[*enc as u8]).map_err(IoError::Io)?;
335            self.inner.write_all(&[*type_tag]).map_err(IoError::Io)?;
336            let data_len = bytes.len() as u32;
337            self.inner
338                .write_all(&data_len.to_le_bytes())
339                .map_err(IoError::Io)?;
340            self.current_offset += 2 + name_bytes.len() as u64 + 1 + 1 + 4;
341        }
342
343        // Write column data
344        for (_, _, _, bytes) in &encoded {
345            self.inner.write_all(bytes).map_err(IoError::Io)?;
346            self.current_offset += bytes.len() as u64;
347        }
348
349        self.stripes.push(StripeInfo {
350            offset: stripe_offset,
351            num_rows: num_rows_u32,
352            column_names: columns.iter().map(|(n, _)| n.clone()).collect(),
353        });
354
355        Ok(())
356    }
357
358    /// Finalise the file by writing the footer and flushing.
359    pub fn finalize(mut self) -> Result<()> {
360        let footer_offset = self.current_offset;
361
362        // FOOTER_MAGIC
363        self.inner.write_all(FOOTER_MAGIC).map_err(IoError::Io)?;
364
365        // Number of stripes
366        let n_stripes = self.stripes.len() as u32;
367        self.inner
368            .write_all(&n_stripes.to_le_bytes())
369            .map_err(IoError::Io)?;
370
371        // Per-stripe info
372        for stripe in &self.stripes {
373            self.inner
374                .write_all(&stripe.offset.to_le_bytes())
375                .map_err(IoError::Io)?;
376            self.inner
377                .write_all(&stripe.num_rows.to_le_bytes())
378                .map_err(IoError::Io)?;
379            let n_cols = stripe.column_names.len() as u32;
380            self.inner
381                .write_all(&n_cols.to_le_bytes())
382                .map_err(IoError::Io)?;
383            for col_name in &stripe.column_names {
384                let bytes = col_name.as_bytes();
385                let len = bytes.len() as u16;
386                self.inner
387                    .write_all(&len.to_le_bytes())
388                    .map_err(IoError::Io)?;
389                self.inner.write_all(bytes).map_err(IoError::Io)?;
390            }
391        }
392
393        // Footer length (from FOOT_MAGIC to here, not including the length field itself)
394        let footer_length =
395            (self.inner.stream_position().map_err(IoError::Io)? - footer_offset) as u32;
396        self.inner
397            .write_all(&footer_length.to_le_bytes())
398            .map_err(IoError::Io)?;
399
400        self.inner.flush().map_err(IoError::Io)?;
401        Ok(())
402    }
403}
404
405// ──────────────────────────────────────────────────────────────────────────────
406// OrcLiteReader
407// ──────────────────────────────────────────────────────────────────────────────
408
409/// Metadata about a stripe.
410#[derive(Debug, Clone)]
411pub struct StripeMetadata {
412    /// Byte offset of stripe in the file.
413    pub offset: u64,
414    /// Number of rows.
415    pub num_rows: u32,
416    /// Column names.
417    pub column_names: Vec<String>,
418}
419
420/// Reader for the ORC-lite columnar format.
421///
422/// First call `open()` to read and validate the file header and footer.
423/// Then iterate over stripes via `read_stripe()`.
424pub struct OrcLiteReader<R: Read + Seek> {
425    inner: BufReader<R>,
426    pub stripes: Vec<StripeMetadata>,
427}
428
429impl<R: Read + Seek> OrcLiteReader<R> {
430    /// Open an ORC-lite file and read its footer/stripe index.
431    pub fn open(reader: R) -> Result<Self> {
432        let mut br = BufReader::new(reader);
433
434        // Validate magic
435        let mut magic = [0u8; FILE_MAGIC_LEN];
436        br.read_exact(&mut magic).map_err(IoError::Io)?;
437        if &magic != FILE_MAGIC {
438            return Err(IoError::FormatError(
439                "Not an ORC-lite file (bad magic)".into(),
440            ));
441        }
442
443        // Seek to the last 4 bytes to read footer_length
444        let file_size = br.seek(SeekFrom::End(0)).map_err(IoError::Io)?;
445        if file_size < FILE_MAGIC_LEN as u64 + 4 {
446            return Err(IoError::FormatError("ORC-lite file too small".into()));
447        }
448        br.seek(SeekFrom::End(-4)).map_err(IoError::Io)?;
449        let mut fl_bytes = [0u8; 4];
450        br.read_exact(&mut fl_bytes).map_err(IoError::Io)?;
451        let footer_length = u32::from_le_bytes(fl_bytes) as u64;
452
453        // Seek to footer start
454        let footer_start = file_size - 4 - footer_length;
455        br.seek(SeekFrom::Start(footer_start))
456            .map_err(IoError::Io)?;
457
458        // Validate footer magic
459        let mut fmagic = [0u8; 4];
460        br.read_exact(&mut fmagic).map_err(IoError::Io)?;
461        if &fmagic != FOOTER_MAGIC {
462            return Err(IoError::FormatError("ORC-lite: bad footer magic".into()));
463        }
464
465        // Read number of stripes
466        let mut n_stripes_bytes = [0u8; 4];
467        br.read_exact(&mut n_stripes_bytes).map_err(IoError::Io)?;
468        let n_stripes = u32::from_le_bytes(n_stripes_bytes) as usize;
469
470        let mut stripes = Vec::with_capacity(n_stripes);
471        for _ in 0..n_stripes {
472            let mut buf8 = [0u8; 8];
473            br.read_exact(&mut buf8).map_err(IoError::Io)?;
474            let offset = u64::from_le_bytes(buf8);
475
476            let mut buf4 = [0u8; 4];
477            br.read_exact(&mut buf4).map_err(IoError::Io)?;
478            let num_rows = u32::from_le_bytes(buf4);
479
480            br.read_exact(&mut buf4).map_err(IoError::Io)?;
481            let n_cols = u32::from_le_bytes(buf4) as usize;
482
483            let mut column_names = Vec::with_capacity(n_cols);
484            for _ in 0..n_cols {
485                let mut len_bytes = [0u8; 2];
486                br.read_exact(&mut len_bytes).map_err(IoError::Io)?;
487                let name_len = u16::from_le_bytes(len_bytes) as usize;
488                let mut name_bytes = vec![0u8; name_len];
489                br.read_exact(&mut name_bytes).map_err(IoError::Io)?;
490                let name = String::from_utf8(name_bytes).map_err(|e| {
491                    IoError::FormatError(format!("ORC-lite: invalid UTF-8 column name: {}", e))
492                })?;
493                column_names.push(name);
494            }
495
496            stripes.push(StripeMetadata {
497                offset,
498                num_rows,
499                column_names,
500            });
501        }
502
503        Ok(OrcLiteReader { inner: br, stripes })
504    }
505
506    /// Read all columns from stripe `stripe_idx`.
507    pub fn read_stripe(&mut self, stripe_idx: usize) -> Result<HashMap<String, OrcColumnData>> {
508        let stripe = self.stripes.get(stripe_idx).ok_or_else(|| {
509            IoError::NotFound(format!(
510                "ORC-lite: stripe index {} out of range",
511                stripe_idx
512            ))
513        })?;
514        let stripe_offset = stripe.offset;
515        let num_rows = stripe.num_rows as usize;
516        let num_cols = stripe.column_names.len();
517
518        self.inner
519            .seek(SeekFrom::Start(stripe_offset))
520            .map_err(IoError::Io)?;
521
522        // Validate stripe magic
523        let mut smagic = [0u8; 4];
524        self.inner.read_exact(&mut smagic).map_err(IoError::Io)?;
525        if &smagic != STRIPE_MAGIC {
526            return Err(IoError::FormatError("ORC-lite: bad stripe magic".into()));
527        }
528
529        // Skip num_rows and num_cols (already known from footer)
530        let mut skip8 = [0u8; 8];
531        self.inner.read_exact(&mut skip8).map_err(IoError::Io)?;
532
533        // Read column headers
534        let mut col_headers: Vec<(String, ColumnEncoding, u8, u32)> = Vec::with_capacity(num_cols);
535        for _ in 0..num_cols {
536            let mut len_bytes = [0u8; 2];
537            self.inner.read_exact(&mut len_bytes).map_err(IoError::Io)?;
538            let name_len = u16::from_le_bytes(len_bytes) as usize;
539            let mut name_bytes = vec![0u8; name_len];
540            self.inner
541                .read_exact(&mut name_bytes)
542                .map_err(IoError::Io)?;
543            let name = String::from_utf8(name_bytes).map_err(|e| {
544                IoError::FormatError(format!("ORC-lite stripe: invalid UTF-8 column name: {}", e))
545            })?;
546
547            let mut enc_byte = [0u8; 1];
548            self.inner.read_exact(&mut enc_byte).map_err(IoError::Io)?;
549            let enc = ColumnEncoding::from_u8(enc_byte[0])?;
550
551            let mut type_tag_byte = [0u8; 1];
552            self.inner
553                .read_exact(&mut type_tag_byte)
554                .map_err(IoError::Io)?;
555            let type_tag = type_tag_byte[0];
556
557            let mut data_len_bytes = [0u8; 4];
558            self.inner
559                .read_exact(&mut data_len_bytes)
560                .map_err(IoError::Io)?;
561            let data_len = u32::from_le_bytes(data_len_bytes);
562
563            col_headers.push((name, enc, type_tag, data_len));
564        }
565
566        // Read column data
567        let mut result = HashMap::new();
568        for (name, enc, type_tag, data_len) in col_headers {
569            let mut data_bytes = vec![0u8; data_len as usize];
570            self.inner
571                .read_exact(&mut data_bytes)
572                .map_err(IoError::Io)?;
573            let col_data = decode_column(&data_bytes, enc, type_tag, num_rows)?;
574            result.insert(name, col_data);
575        }
576
577        Ok(result)
578    }
579
580    /// Read all stripes and concatenate their columns.
581    pub fn read_all(&mut self) -> Result<HashMap<String, OrcColumnData>> {
582        let n = self.stripes.len();
583        if n == 0 {
584            return Ok(HashMap::new());
585        }
586
587        let mut combined: HashMap<String, OrcColumnData> = HashMap::new();
588
589        for i in 0..n {
590            let stripe_data = self.read_stripe(i)?;
591            for (name, col) in stripe_data {
592                let entry = combined.entry(name).or_insert_with(|| match &col {
593                    OrcColumnData::Int64(_) => OrcColumnData::Int64(Vec::new()),
594                    OrcColumnData::Float64(_) => OrcColumnData::Float64(Vec::new()),
595                    OrcColumnData::Boolean(_) => OrcColumnData::Boolean(Vec::new()),
596                    OrcColumnData::String(_) => OrcColumnData::String(Vec::new()),
597                    OrcColumnData::NullableInt64(_) => OrcColumnData::NullableInt64(Vec::new()),
598                });
599                concat_columns(entry, col)?;
600            }
601        }
602
603        Ok(combined)
604    }
605}
606
607// ──────────────────────────────────────────────────────────────────────────────
608// Encode / decode helpers
609// ──────────────────────────────────────────────────────────────────────────────
610
611fn encode_column(col: &OrcColumnData) -> Result<(ColumnEncoding, u8, Vec<u8>)> {
612    let type_tag = col.type_tag();
613    match col {
614        OrcColumnData::Int64(v) => {
615            // Use RLE v2 for integer columns
616            let encoded = IntRleV2::encode(v);
617            Ok((ColumnEncoding::RleV2, type_tag, encoded))
618        }
619        OrcColumnData::Float64(v) => {
620            // Direct f64 LE encoding
621            let mut bytes = Vec::with_capacity(v.len() * 8);
622            for &x in v {
623                bytes.extend_from_slice(&x.to_le_bytes());
624            }
625            Ok((ColumnEncoding::Direct, type_tag, bytes))
626        }
627        OrcColumnData::Boolean(v) => {
628            // Bit-packed: 8 booleans per byte
629            let mut bytes = Vec::with_capacity((v.len() + 7) / 8 + 4);
630            let count = v.len() as u32;
631            bytes.extend_from_slice(&count.to_le_bytes());
632            let mut byte = 0u8;
633            for (i, &b) in v.iter().enumerate() {
634                if b {
635                    byte |= 1 << (i % 8);
636                }
637                if i % 8 == 7 {
638                    bytes.push(byte);
639                    byte = 0;
640                }
641            }
642            if v.len() % 8 != 0 {
643                bytes.push(byte);
644            }
645            Ok((ColumnEncoding::Direct, type_tag, bytes))
646        }
647        OrcColumnData::String(v) => {
648            // Length-prefixed UTF-8 strings: u32 count + [u32 len + bytes]...
649            let mut bytes = Vec::new();
650            let count = v.len() as u32;
651            bytes.extend_from_slice(&count.to_le_bytes());
652            for s in v {
653                let sb = s.as_bytes();
654                let slen = sb.len() as u32;
655                bytes.extend_from_slice(&slen.to_le_bytes());
656                bytes.extend_from_slice(sb);
657            }
658            Ok((ColumnEncoding::DirectString, type_tag, bytes))
659        }
660        OrcColumnData::NullableInt64(v) => {
661            // Validity bitmap + values (nulls encoded as 0)
662            let count = v.len() as u32;
663            let mut bytes = Vec::new();
664            bytes.extend_from_slice(&count.to_le_bytes());
665            // Validity bitmap
666            let mut bitmap_byte = 0u8;
667            for (i, opt) in v.iter().enumerate() {
668                if opt.is_some() {
669                    bitmap_byte |= 1 << (i % 8);
670                }
671                if i % 8 == 7 {
672                    bytes.push(bitmap_byte);
673                    bitmap_byte = 0;
674                }
675            }
676            if v.len() % 8 != 0 {
677                bytes.push(bitmap_byte);
678            }
679            // Values (0 for null)
680            let vals: Vec<i64> = v.iter().map(|o| o.unwrap_or(0)).collect();
681            let encoded = IntRleV2::encode(&vals);
682            bytes.extend_from_slice(&encoded);
683            Ok((ColumnEncoding::RleV2, type_tag, bytes))
684        }
685    }
686}
687
688fn decode_column(
689    data: &[u8],
690    enc: ColumnEncoding,
691    type_tag: u8,
692    _num_rows: usize,
693) -> Result<OrcColumnData> {
694    match type_tag {
695        0 => {
696            // Int64
697            let values = IntRleV2::decode(data)?;
698            Ok(OrcColumnData::Int64(values))
699        }
700        1 => {
701            // Float64
702            if !data.len().is_multiple_of(8) {
703                return Err(IoError::FormatError(
704                    "ORC-lite: f64 column data not multiple of 8 bytes".into(),
705                ));
706            }
707            let values: Vec<f64> = data
708                .chunks_exact(8)
709                .map(|chunk| {
710                    let arr: [u8; 8] = chunk.try_into().expect("chunk is exactly 8 bytes");
711                    f64::from_le_bytes(arr)
712                })
713                .collect();
714            Ok(OrcColumnData::Float64(values))
715        }
716        2 => {
717            // Boolean
718            if data.len() < 4 {
719                return Err(IoError::FormatError(
720                    "ORC-lite: boolean column too short".into(),
721                ));
722            }
723            let count = u32::from_le_bytes(
724                data[0..4]
725                    .try_into()
726                    .map_err(|_| IoError::FormatError("boolean count bytes".into()))?,
727            ) as usize;
728            let mut values = Vec::with_capacity(count);
729            for i in 0..count {
730                let byte_idx = 4 + i / 8;
731                let bit_idx = i % 8;
732                let byte = data.get(byte_idx).ok_or_else(|| {
733                    IoError::FormatError("ORC-lite: boolean bitmap truncated".into())
734                })?;
735                values.push((byte >> bit_idx) & 1 == 1);
736            }
737            Ok(OrcColumnData::Boolean(values))
738        }
739        3 => {
740            // String
741            if data.len() < 4 {
742                return Err(IoError::FormatError(
743                    "ORC-lite: string column too short".into(),
744                ));
745            }
746            let count = u32::from_le_bytes(
747                data[0..4]
748                    .try_into()
749                    .map_err(|_| IoError::FormatError("string count bytes".into()))?,
750            ) as usize;
751            let mut pos = 4;
752            let mut values = Vec::with_capacity(count);
753            for _ in 0..count {
754                if pos + 4 > data.len() {
755                    return Err(IoError::FormatError(
756                        "ORC-lite: string length truncated".into(),
757                    ));
758                }
759                let slen = u32::from_le_bytes(
760                    data[pos..pos + 4]
761                        .try_into()
762                        .map_err(|_| IoError::FormatError("string len bytes".into()))?,
763                ) as usize;
764                pos += 4;
765                if pos + slen > data.len() {
766                    return Err(IoError::FormatError(
767                        "ORC-lite: string data truncated".into(),
768                    ));
769                }
770                let s = String::from_utf8(data[pos..pos + slen].to_vec()).map_err(|e| {
771                    IoError::FormatError(format!("ORC-lite: invalid UTF-8 string: {}", e))
772                })?;
773                values.push(s);
774                pos += slen;
775            }
776            Ok(OrcColumnData::String(values))
777        }
778        4 => {
779            // NullableInt64
780            if data.len() < 4 {
781                return Err(IoError::FormatError(
782                    "ORC-lite: nullable int64 column too short".into(),
783                ));
784            }
785            let count = u32::from_le_bytes(
786                data[0..4]
787                    .try_into()
788                    .map_err(|_| IoError::FormatError("nullable count bytes".into()))?,
789            ) as usize;
790
791            let bitmap_bytes = (count + 7) / 8;
792            let bitmap_end = 4 + bitmap_bytes;
793            if bitmap_end > data.len() {
794                return Err(IoError::FormatError(
795                    "ORC-lite: nullable bitmap truncated".into(),
796                ));
797            }
798
799            let values_raw = IntRleV2::decode(&data[bitmap_end..])?;
800            let mut values = Vec::with_capacity(count);
801            for i in 0..count {
802                let byte_idx = 4 + i / 8;
803                let bit_idx = i % 8;
804                let byte = data
805                    .get(byte_idx)
806                    .ok_or_else(|| IoError::FormatError("ORC-lite: nullable bitmap read".into()))?;
807                let is_valid = (byte >> bit_idx) & 1 == 1;
808                values.push(if is_valid {
809                    values_raw.get(i).copied().map(Some).unwrap_or(None)
810                } else {
811                    None
812                });
813            }
814            Ok(OrcColumnData::NullableInt64(values))
815        }
816        other => Err(IoError::FormatError(format!(
817            "ORC-lite: unknown type tag {}",
818            other
819        ))),
820    }
821}
822
823/// Append `src` into `dst` (must be same variant).
824fn concat_columns(dst: &mut OrcColumnData, src: OrcColumnData) -> Result<()> {
825    match (dst, src) {
826        (OrcColumnData::Int64(a), OrcColumnData::Int64(b)) => a.extend(b),
827        (OrcColumnData::Float64(a), OrcColumnData::Float64(b)) => a.extend(b),
828        (OrcColumnData::Boolean(a), OrcColumnData::Boolean(b)) => a.extend(b),
829        (OrcColumnData::String(a), OrcColumnData::String(b)) => a.extend(b),
830        (OrcColumnData::NullableInt64(a), OrcColumnData::NullableInt64(b)) => a.extend(b),
831        _ => {
832            return Err(IoError::FormatError(
833                "ORC-lite: column type mismatch across stripes".into(),
834            ))
835        }
836    }
837    Ok(())
838}
839
840// ──────────────────────────────────────────────────────────────────────────────
841// Convenience free functions
842// ──────────────────────────────────────────────────────────────────────────────
843
844/// Write columnar data to an ORC-lite file on disk.
845pub fn write_orc_lite<P: AsRef<Path>>(path: P, columns: &[(String, OrcColumnData)]) -> Result<()> {
846    let file = std::fs::File::create(path.as_ref()).map_err(IoError::Io)?;
847    let mut writer = OrcLiteWriter::new(file)?;
848    writer.write_stripe(columns)?;
849    writer.finalize()
850}
851
852/// Read all data from an ORC-lite file on disk.
853pub fn read_orc_lite<P: AsRef<Path>>(path: P) -> Result<HashMap<String, OrcColumnData>> {
854    let file = std::fs::File::open(path.as_ref()).map_err(IoError::Io)?;
855    let mut reader = OrcLiteReader::open(file)?;
856    reader.read_all()
857}
858
859// ──────────────────────────────────────────────────────────────────────────────
860// Tests
861// ──────────────────────────────────────────────────────────────────────────────
862
863#[cfg(test)]
864mod tests {
865    use super::*;
866    use std::io::Cursor;
867
868    // Helper: write to a Vec<u8> and read back
869    fn roundtrip(cols: &[(String, OrcColumnData)]) -> HashMap<String, OrcColumnData> {
870        let mut buf: Vec<u8> = Vec::new();
871        {
872            let cursor = Cursor::new(&mut buf);
873            let mut writer = OrcLiteWriter::new(cursor).expect("writer create");
874            writer.write_stripe(cols).expect("write stripe");
875            writer.finalize().expect("finalize");
876        }
877        let cursor = Cursor::new(&buf[..]);
878        let mut reader = OrcLiteReader::open(cursor).expect("reader open");
879        reader.read_all().expect("read all")
880    }
881
882    #[test]
883    fn test_int64_roundtrip() {
884        let vals: Vec<i64> = vec![1, 2, 3, 100, -5, 0, 1000];
885        let cols = vec![("ids".to_string(), OrcColumnData::Int64(vals.clone()))];
886        let result = roundtrip(&cols);
887        assert_eq!(result["ids"], OrcColumnData::Int64(vals));
888    }
889
890    #[test]
891    fn test_float64_roundtrip() {
892        let vals: Vec<f64> = vec![1.1, 2.2, -3.3, 0.0, 1e10];
893        let cols = vec![("values".to_string(), OrcColumnData::Float64(vals.clone()))];
894        let result = roundtrip(&cols);
895        assert_eq!(result["values"], OrcColumnData::Float64(vals));
896    }
897
898    #[test]
899    fn test_boolean_roundtrip() {
900        let vals: Vec<bool> = vec![true, false, true, true, false];
901        let cols = vec![("flags".to_string(), OrcColumnData::Boolean(vals.clone()))];
902        let result = roundtrip(&cols);
903        assert_eq!(result["flags"], OrcColumnData::Boolean(vals));
904    }
905
906    #[test]
907    fn test_string_roundtrip() {
908        let vals: Vec<String> = vec!["hello".into(), "world".into(), "ORC-lite".into(), "".into()];
909        let cols = vec![("names".to_string(), OrcColumnData::String(vals.clone()))];
910        let result = roundtrip(&cols);
911        assert_eq!(result["names"], OrcColumnData::String(vals));
912    }
913
914    #[test]
915    fn test_nullable_int64_roundtrip() {
916        let vals: Vec<Option<i64>> = vec![Some(1), None, Some(3), None, Some(99)];
917        let cols = vec![(
918            "opt_ids".to_string(),
919            OrcColumnData::NullableInt64(vals.clone()),
920        )];
921        let result = roundtrip(&cols);
922        assert_eq!(result["opt_ids"], OrcColumnData::NullableInt64(vals));
923    }
924
925    #[test]
926    fn test_multi_column_roundtrip() {
927        let n = 100;
928        let ids: Vec<i64> = (0..n).collect();
929        let vals: Vec<f64> = (0..n).map(|i| i as f64 * 1.5).collect();
930        let labels: Vec<String> = (0..n).map(|i| format!("item_{}", i)).collect();
931
932        let cols = vec![
933            ("id".to_string(), OrcColumnData::Int64(ids.clone())),
934            ("val".to_string(), OrcColumnData::Float64(vals.clone())),
935            ("label".to_string(), OrcColumnData::String(labels.clone())),
936        ];
937        let result = roundtrip(&cols);
938        assert_eq!(result["id"], OrcColumnData::Int64(ids));
939        assert_eq!(result["val"], OrcColumnData::Float64(vals));
940        assert_eq!(result["label"], OrcColumnData::String(labels));
941    }
942
943    #[test]
944    fn test_multi_stripe_roundtrip() {
945        let mut buf: Vec<u8> = Vec::new();
946        {
947            let cursor = Cursor::new(&mut buf);
948            let mut writer = OrcLiteWriter::new(cursor).expect("writer");
949            for stripe_id in 0..3 {
950                let start = stripe_id * 10i64;
951                let ids: Vec<i64> = (start..start + 10).collect();
952                let cols = vec![("id".to_string(), OrcColumnData::Int64(ids))];
953                writer.write_stripe(&cols).expect("write stripe");
954            }
955            writer.finalize().expect("finalize");
956        }
957
958        let cursor = Cursor::new(&buf[..]);
959        let mut reader = OrcLiteReader::open(cursor).expect("reader");
960        assert_eq!(reader.stripes.len(), 3);
961        let all = reader.read_all().expect("read all");
962        if let OrcColumnData::Int64(ids) = &all["id"] {
963            assert_eq!(ids.len(), 30);
964            assert_eq!(ids[0], 0);
965            assert_eq!(ids[29], 29);
966        } else {
967            panic!("expected Int64 column");
968        }
969    }
970
971    #[test]
972    fn test_int_rle_v2_roundtrip() {
973        let values: Vec<i64> = vec![0, 1, 1, 2, 3, 5, 8, 13, -7, 100, -100, i64::MAX / 2];
974        let encoded = IntRleV2::encode(&values);
975        let decoded = IntRleV2::decode(&encoded).expect("decode");
976        assert_eq!(decoded, values);
977    }
978
979    #[test]
980    fn test_int_rle_v2_empty() {
981        let encoded = IntRleV2::encode(&[]);
982        let decoded = IntRleV2::decode(&encoded).expect("decode empty");
983        assert!(decoded.is_empty());
984    }
985
986    #[test]
987    fn test_bad_magic() {
988        let cursor = Cursor::new(b"NOTMAGIC!!");
989        let result = OrcLiteReader::open(cursor);
990        assert!(result.is_err());
991    }
992
993    #[test]
994    fn test_file_roundtrip() {
995        let tmp_dir = std::env::temp_dir();
996        let path = tmp_dir.join("test_orc_lite.orc");
997
998        let ids: Vec<i64> = vec![10, 20, 30];
999        let vals: Vec<f64> = vec![1.0, 2.0, 3.0];
1000        let cols = vec![
1001            ("id".to_string(), OrcColumnData::Int64(ids.clone())),
1002            ("val".to_string(), OrcColumnData::Float64(vals.clone())),
1003        ];
1004
1005        write_orc_lite(&path, &cols).expect("write file");
1006        let result = read_orc_lite(&path).expect("read file");
1007
1008        assert_eq!(result["id"], OrcColumnData::Int64(ids));
1009        assert_eq!(result["val"], OrcColumnData::Float64(vals));
1010
1011        // Cleanup
1012        let _ = std::fs::remove_file(&path);
1013    }
1014}
scirs2_io/formats/orc_lite.rs

scirs2_io/formats/
orc_lite.rs
