//! Tiny Arrow column/batch builders shared by the collectors — keep the parsers
//! readable (build `Vec<Option<…>>` columns, hand them here).
use std::sync::Arc;
use anyhow::Result;
use arrow_array::builder::{
BinaryBuilder, BooleanBuilder, Date32Builder, Decimal128Builder, Float64Builder, Int64Builder,
StringBuilder, TimestampMicrosecondBuilder,
};
use arrow_array::{ArrayRef, RecordBatch};
use arrow_schema::{DataType, Field, Schema, TimeUnit};
/// A nullable Utf8 column from owned optional strings.
pub fn str_col(values: &[Option<String>]) -> ArrayRef {
let mut b = StringBuilder::with_capacity(values.len(), values.len() * 16);
for v in values {
match v {
Some(s) => b.append_value(s),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable Int64 column.
pub fn i64_col(values: &[Option<i64>]) -> ArrayRef {
let mut b = Int64Builder::with_capacity(values.len());
for v in values {
match v {
Some(n) => b.append_value(*n),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable Binary column (WKB geometry blobs for the geo connectors).
pub fn bin_col(values: &[Option<Vec<u8>>]) -> ArrayRef {
let mut b = BinaryBuilder::with_capacity(values.len(), values.len() * 64);
for v in values {
match v {
Some(bytes) => b.append_value(bytes),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable Float64 column.
pub fn f64_col(values: &[Option<f64>]) -> ArrayRef {
let mut b = Float64Builder::with_capacity(values.len());
for v in values {
match v {
Some(x) => b.append_value(*x),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable Boolean column (CSV bool-sniff, mainframe level-88 flags).
pub fn bool_col(values: &[Option<bool>]) -> ArrayRef {
let mut b = BooleanBuilder::with_capacity(values.len());
for v in values {
match v {
Some(x) => b.append_value(*x),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable `Date32` column — values are **days since the Unix epoch**
/// (1970-01-01), the Arrow `Date32` convention. Used by the CSV date sniff and
/// COBOL `PIC 9(8)` date fields.
pub fn date32_col(values: &[Option<i32>]) -> ArrayRef {
let mut b = Date32Builder::with_capacity(values.len());
for v in values {
match v {
Some(d) => b.append_value(*d),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable `Timestamp(Microsecond, None)` column — values are **microseconds
/// since the Unix epoch** (timezone-naive). Used by the CSV datetime sniff.
pub fn timestamp_col(values: &[Option<i64>]) -> ArrayRef {
let mut b = TimestampMicrosecondBuilder::with_capacity(values.len());
for v in values {
match v {
Some(us) => b.append_value(*us),
None => b.append_null(),
}
}
Arc::new(b.finish())
}
/// A nullable `Decimal128(precision, scale)` column — values are the **unscaled**
/// `i128` mantissa (i.e. `value * 10^scale`), the Arrow `Decimal128` convention.
/// This is the shared money/packed-decimal builder: mainframe COMP-3/zoned and
/// CSV money columns all land here (e.g. `-42.50` at scale 2 → unscaled `-4250`).
pub fn decimal128_col(values: &[Option<i128>], precision: u8, scale: i8) -> ArrayRef {
let mut b = Decimal128Builder::with_capacity(values.len());
for v in values {
match v {
Some(n) => b.append_value(*n),
None => b.append_null(),
}
}
let arr = b
.finish()
.with_precision_and_scale(precision, scale)
.expect("valid Decimal128 precision/scale");
Arc::new(arr)
}
/// The Arrow `DataType` for a naive microsecond timestamp column (matches
/// [`timestamp_col`]). Handy when declaring a schema field for it.
pub fn timestamp_type() -> DataType {
DataType::Timestamp(TimeUnit::Microsecond, None)
}
/// Build a schema of `(name, DataType, nullable)` fields.
pub fn schema(fields: &[(&str, DataType, bool)]) -> Schema {
Schema::new(
fields
.iter()
.map(|(n, dt, null)| Field::new(*n, dt.clone(), *null))
.collect::<Vec<_>>(),
)
}
/// Zip a schema with already-built columns into one `RecordBatch`.
pub fn batch(schema: &Schema, columns: Vec<ArrayRef>) -> Result<RecordBatch> {
Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?)
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::{
Array, BooleanArray, Date32Array, Decimal128Array, TimestampMicrosecondArray,
};
#[test]
fn bool_col_carries_values_and_nulls() {
let a = bool_col(&[Some(true), None, Some(false)]);
let a = a.as_any().downcast_ref::<BooleanArray>().unwrap();
assert_eq!(a.len(), 3);
assert!(a.value(0));
assert!(a.is_null(1));
assert!(!a.value(2));
}
#[test]
fn date32_col_is_days_since_epoch() {
// 2020-01-01 = 18262 days since 1970-01-01.
let a = date32_col(&[Some(18262), None]);
let a = a.as_any().downcast_ref::<Date32Array>().unwrap();
assert_eq!(a.value(0), 18262);
assert!(a.is_null(1));
assert_eq!(a.data_type(), &DataType::Date32);
}
#[test]
fn timestamp_col_is_micros_and_naive() {
let a = timestamp_col(&[Some(1_000_000), None]);
assert_eq!(a.data_type(), ×tamp_type());
let a = a
.as_any()
.downcast_ref::<TimestampMicrosecondArray>()
.unwrap();
assert_eq!(a.value(0), 1_000_000);
assert!(a.is_null(1));
}
#[test]
fn decimal128_col_carries_unscaled_precision_scale() {
// -42.50 at scale 2 → unscaled -4250; precision 7.
let a = decimal128_col(&[Some(-4250), None, Some(4250)], 7, 2);
assert_eq!(a.data_type(), &DataType::Decimal128(7, 2));
let a = a.as_any().downcast_ref::<Decimal128Array>().unwrap();
assert_eq!(a.value(0), -4250);
assert!(a.is_null(1));
assert_eq!(a.value(2), 4250);
// human-readable render honours the scale
assert_eq!(a.value_as_string(0), "-42.50");
}
}