faucet_core/file_format/
orc.rs

1//! Apache ORC (#719) — **read-only**.
2//!
3//! Decoded with `orc-rust`, which reads straight into Arrow `RecordBatch`es,
4//! stripe by stripe: memory is bounded by one stripe's decoded batch (capped at
5//! the configured batch size), not by the file. `orc.columns` projects
6//! top-level columns before any stripe is decoded.
7//!
8//! There is no ORC writer here by design. `orc-rust`'s writer is experimental:
9//! it covers primitive columns only, panics on nested and temporal types, and
10//! writes no compression — not something to put behind a sink a pipeline
11//! trusts. A sink asked for `format: orc` refuses at encode time; Parquet is
12//! the columnar format to write.
13//!
14//! On the record path the batches go through
15//! [`record_batch_to_values`](crate::columnar::record_batch_to_values), so an
16//! ORC file and a Parquet file with the same content produce the same records.
17
18use super::OrcOptions;
19use crate::error::FaucetError;
20use arrow::array::RecordBatch;
21use arrow::datatypes::SchemaRef;
22use orc_rust::projection::ProjectionMask;
23use orc_rust::reader::ChunkReader;
24use serde_json::Value;
25
26/// Where an ORC file's bytes come from.
27///
28/// A closed set rather than a generic reader so `orc-rust`'s trait never
29/// appears in this crate's public API.
30#[derive(Debug)]
31pub enum OrcInput {
32    /// A local file, read by stripe with positioned reads.
33    File(std::fs::File),
34    /// A whole object already in memory.
35    Bytes(bytes::Bytes),
36}
37
38/// Decode a whole ORC object into records.
39pub fn decode(bytes: &[u8], opts: &OrcOptions) -> Result<Vec<Value>, FaucetError> {
40    decode_input(OrcInput::Bytes(bytes::Bytes::copy_from_slice(bytes)), opts)
41}
42
43/// [`decode`] taking ownership of the body, so it is read in place rather
44/// than copied (CORE-73).
45pub fn decode_owned(bytes: Vec<u8>, opts: &OrcOptions) -> Result<Vec<Value>, FaucetError> {
46    decode_input(OrcInput::Bytes(bytes::Bytes::from(bytes)), opts)
47}
48
49fn decode_input(input: OrcInput, opts: &OrcOptions) -> Result<Vec<Value>, FaucetError> {
50    let mut out = Vec::new();
51    read_batches(input, opts, 0, &mut |b| {
52        out.extend(crate::columnar::record_batch_to_values(&b)?);
53        Ok(())
54    })?;
55    Ok(out)
56}
57
58/// Stream an ORC file's batches (at most `batch_size` rows each; `0` means the
59/// reader's own stripe-sized batches). Returns the (projected) Arrow schema,
60/// known even for a file with no rows.
61pub fn read_batches(
62    input: OrcInput,
63    opts: &OrcOptions,
64    batch_size: usize,
65    f: &mut dyn FnMut(RecordBatch) -> Result<(), FaucetError>,
66) -> Result<SchemaRef, FaucetError> {
67    read_batches_checked(input, opts, batch_size, &mut |_| Ok(()), f)
68}
69
70/// [`read_batches`], calling `check` with the file's (projected) schema
71/// before any stripe is decoded — so a caller can refuse a file whose shape
72/// it cannot accept without having emitted any of its rows.
73pub fn read_batches_checked(
74    input: OrcInput,
75    opts: &OrcOptions,
76    batch_size: usize,
77    check: &mut dyn FnMut(&SchemaRef) -> Result<(), FaucetError>,
78    f: &mut dyn FnMut(RecordBatch) -> Result<(), FaucetError>,
79) -> Result<SchemaRef, FaucetError> {
80    match input {
81        OrcInput::File(file) => read_from(file, opts, batch_size, check, f),
82        OrcInput::Bytes(b) => read_from(b, opts, batch_size, check, f),
83    }
84}
85
86fn read_from<R: ChunkReader>(
87    reader: R,
88    opts: &OrcOptions,
89    batch_size: usize,
90    check: &mut dyn FnMut(&SchemaRef) -> Result<(), FaucetError>,
91    f: &mut dyn FnMut(RecordBatch) -> Result<(), FaucetError>,
92) -> Result<SchemaRef, FaucetError> {
93    let mut builder = orc_rust::ArrowReaderBuilder::try_new(reader)
94        .map_err(|e| FaucetError::Source(format!("orc: not a readable ORC file: {e}")))?;
95    if let Some(cols) = &opts.columns {
96        let root = builder.file_metadata().root_data_type();
97        let known: Vec<&str> = root.children().iter().map(|c| c.name()).collect();
98        if let Some(missing) = cols.iter().find(|c| !known.contains(&c.as_str())) {
99            return Err(FaucetError::Source(format!(
100                "orc: column {missing:?} is not in the file (columns: {})",
101                known.join(", ")
102            )));
103        }
104        let mask = ProjectionMask::named_roots(root, cols);
105        builder = builder.with_projection(mask);
106    }
107    if batch_size > 0 {
108        builder = builder.with_batch_size(batch_size);
109    }
110    let schema = builder.schema();
111    check(&schema)?;
112    for batch in builder.build() {
113        let batch = batch.map_err(|e| FaucetError::Source(format!("orc: decode error: {e}")))?;
114        if batch.num_rows() > 0 {
115            f(batch)?;
116        }
117    }
118    Ok(schema)
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    const FIXTURE: &[u8] = include_bytes!("../../tests/fixtures/orc/people.orc");
126
127    #[test]
128    fn reads_every_column_and_row() {
129        let rows = decode(FIXTURE, &OrcOptions::default()).expect("decode");
130        assert_eq!(rows.len(), 3);
131        assert_eq!(rows[0]["id"], serde_json::json!(1));
132        assert_eq!(rows[0]["name"], serde_json::json!("ada"));
133        assert!(rows[0].get("score").is_some());
134    }
135
136    #[test]
137    fn projection_keeps_only_the_named_columns() {
138        let opts = OrcOptions {
139            columns: Some(vec!["name".into()]),
140        };
141        let mut schemas = Vec::new();
142        let mut rows = Vec::new();
143        let schema = read_batches(
144            OrcInput::Bytes(bytes::Bytes::from_static(FIXTURE)),
145            &opts,
146            2,
147            &mut |b| {
148                schemas.push(b.schema());
149                rows.extend(crate::columnar::record_batch_to_values(&b)?);
150                Ok(())
151            },
152        )
153        .expect("read");
154        assert_eq!(schema.fields().len(), 1);
155        assert_eq!(schemas.len(), 2, "batch size 2 over 3 rows");
156        assert_eq!(rows[2], serde_json::json!({"name": "grace"}));
157    }
158
159    #[test]
160    fn unknown_columns_and_garbage_are_errors() {
161        let opts = OrcOptions {
162            columns: Some(vec!["nope".into()]),
163        };
164        let err = decode(FIXTURE, &opts).expect_err("unknown column");
165        assert!(err.to_string().contains("nope"), "{err}");
166        let err = decode(b"definitely not orc", &OrcOptions::default()).expect_err("garbage");
167        assert!(err.to_string().contains("orc"), "{err}");
168    }
169
170    #[test]
171    fn reads_from_a_local_file() {
172        let dir = tempfile::tempdir().expect("tmp");
173        let path = dir.path().join("p.orc");
174        std::fs::write(&path, FIXTURE).expect("write");
175        let mut n = 0;
176        read_batches(
177            OrcInput::File(std::fs::File::open(&path).expect("open")),
178            &OrcOptions::default(),
179            0,
180            &mut |b| {
181                n += b.num_rows();
182                Ok(())
183            },
184        )
185        .expect("read");
186        assert_eq!(n, 3);
187    }
188}
faucet_core/file_format/orc.rs

faucet_core/file_format/
orc.rs
