faucet_core/
columnar.rs

1//! Opt-in Apache Arrow columnar record path (feature `arrow`, RFC 0002 / #375).
2//!
3//! This is the **escape hatch** described in #324: a columnar representation a
4//! connector may produce or consume at the page boundary, *additive* to and
5//! coexisting with the default `serde_json::Value` row model. It touches neither
6//! [`StreamPage`](crate::StreamPage) nor any existing connector — an Arrow-native
7//! source overrides [`Source::stream_batches`](crate::Source::stream_batches) and
8//! an Arrow-native sink overrides
9//! [`Sink::write_batch_columnar`](crate::Sink::write_batch_columnar); the pipeline
10//! uses the columnar path only when **both** sides advertise support (and no
11//! `Value`-shaped stage needs to observe the records), so a
12//! `parquet → parquet` chain never materializes `Value`.
13//!
14//! The `RecordBatch ↔ Value` conversions here are the single source of truth for
15//! the shim (they match `faucet-transform-sql`'s `shovel` byte-for-byte, incl.
16//! `with_explicit_nulls(true)` so an explicit-null field round-trips as
17//! `"key": null` rather than being silently dropped — audit #321 H6).
18
19use crate::FaucetError;
20use arrow::array::RecordBatch;
21use arrow::datatypes::{Schema, SchemaRef};
22use serde_json::Value;
23use std::sync::Arc;
24
25/// A page of records in **columnar** (Arrow) form, the columnar analogue of
26/// [`StreamPage`](crate::StreamPage).
27///
28/// `bookmark` carries the exact same checkpoint semantics as `StreamPage`:
29/// `Some` triggers flush + bookmark-persist after the batch is durably written;
30/// most sources emit `Some` only on the final batch, CDC-style sources per
31/// committed transaction.
32#[derive(Debug, Clone)]
33pub struct ColumnarPage {
34    /// The record batch to write to the sink for this page.
35    pub batch: RecordBatch,
36    /// Optional bookmark to checkpoint after this batch is durably written.
37    pub bookmark: Option<Value>,
38}
39
40impl ColumnarPage {
41    /// Construct a columnar page from a batch and optional bookmark.
42    pub fn new(batch: RecordBatch, bookmark: Option<Value>) -> Self {
43        Self { batch, bookmark }
44    }
45
46    /// Number of rows in the batch.
47    pub fn num_rows(&self) -> usize {
48        self.batch.num_rows()
49    }
50}
51
52/// Map any display-able error into a [`FaucetError::Transform`] with context.
53fn te<E: std::fmt::Display>(ctx: &str, e: E) -> FaucetError {
54    FaucetError::Transform(format!("columnar shim: {ctx}: {e}"))
55}
56
57/// Infer an Arrow [`Schema`] from a slice of JSON records (each a JSON object).
58///
59/// Returns the inferred schema in an [`Arc`]. An empty slice yields a schema
60/// with no fields.
61pub fn infer_arrow_schema(records: &[Value]) -> Result<SchemaRef, FaucetError> {
62    let iter = records
63        .iter()
64        .map(|v| Ok::<_, arrow::error::ArrowError>(v.clone()));
65    let schema = arrow_json::reader::infer_json_schema_from_iterator(iter)
66        .map_err(|e| te("schema inference", e))?;
67    Ok(Arc::new(schema))
68}
69
70/// Encode a slice of JSON records into a single [`RecordBatch`] against `schema`.
71///
72/// Returns an empty batch if `records` is empty.
73pub fn values_to_record_batch(
74    records: &[Value],
75    schema: SchemaRef,
76) -> Result<RecordBatch, FaucetError> {
77    let mut decoder = arrow_json::ReaderBuilder::new(schema.clone())
78        .build_decoder()
79        .map_err(|e| te("decoder build", e))?;
80    decoder.serialize(records).map_err(|e| te("encode", e))?;
81    let mut batches = Vec::new();
82    while let Some(b) = decoder.flush().map_err(|e| te("flush", e))? {
83        batches.push(b);
84    }
85    if batches.is_empty() {
86        return Ok(RecordBatch::new_empty(schema));
87    }
88    if batches.len() == 1 {
89        return Ok(batches.pop().unwrap());
90    }
91    arrow::compute::concat_batches(&schema, &batches).map_err(|e| te("concat", e))
92}
93
94/// Convenience: infer the schema from `records` and encode them into a batch.
95pub fn values_to_record_batch_inferred(records: &[Value]) -> Result<RecordBatch, FaucetError> {
96    let schema = infer_arrow_schema(records)?;
97    values_to_record_batch(records, schema)
98}
99
100/// Decode a [`RecordBatch`] into JSON objects (one per row).
101///
102/// Uses `arrow-json`'s array writer with **explicit nulls enabled**, so a
103/// null-valued column is emitted as `"key": null` rather than omitted — without
104/// this a `SELECT *`-style identity would silently delete every explicit-null
105/// field (audit #321 H6). An empty batch returns an empty `Vec`.
106pub fn record_batch_to_values(batch: &RecordBatch) -> Result<Vec<Value>, FaucetError> {
107    let mut buf = Vec::new();
108    {
109        let mut writer = arrow_json::writer::WriterBuilder::new()
110            .with_explicit_nulls(true)
111            .build::<_, arrow_json::writer::JsonArray>(&mut buf);
112        writer.write(batch).map_err(|e| te("json write", e))?;
113        writer.finish().map_err(|e| te("json finish", e))?;
114    }
115    serde_json::from_slice(&buf).map_err(|e| te("json parse", e))
116}
117
118/// Compare two schemas for field-level equality (name + data-type + nullability).
119pub fn schema_eq(a: &Schema, b: &Schema) -> bool {
120    a.fields() == b.fields()
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use serde_json::json;
127
128    #[test]
129    fn round_trip_scalars_nulls_nested() {
130        let recs = vec![
131            json!({"id": 1, "name": "a", "score": 1.5, "ok": true, "tags": ["x", "y"]}),
132            json!({"id": 2, "name": null, "score": null, "ok": false, "tags": []}),
133        ];
134        let batch = values_to_record_batch_inferred(&recs).unwrap();
135        assert_eq!(batch.num_rows(), 2);
136        let back = record_batch_to_values(&batch).unwrap();
137        assert_eq!(back[0]["id"], json!(1));
138        assert_eq!(back[0]["tags"], json!(["x", "y"]));
139        // #321 H6: an explicit-null field survives the round-trip.
140        assert!(back[1].as_object().unwrap().contains_key("name"));
141        assert_eq!(back[1]["name"], json!(null));
142    }
143
144    #[test]
145    fn empty_records_yield_empty_batch_and_back() {
146        let schema = infer_arrow_schema(&[json!({"a": 1})]).unwrap();
147        let batch = values_to_record_batch(&[], schema).unwrap();
148        assert_eq!(batch.num_rows(), 0);
149        assert!(record_batch_to_values(&batch).unwrap().is_empty());
150    }
151
152    #[test]
153    fn columnar_page_reports_rows() {
154        let batch = values_to_record_batch_inferred(&[json!({"a": 1}), json!({"a": 2})]).unwrap();
155        let page = ColumnarPage::new(batch, Some(json!({"lsn": 42})));
156        assert_eq!(page.num_rows(), 2);
157        assert_eq!(page.bookmark, Some(json!({"lsn": 42})));
158    }
159}
faucet_core/columnar.rs

faucet_core/
columnar.rs
