facett-grid/src/arrow.rs
//! **Arrow ingest** — share primitive and string buffers with the source batch.
//! Arrays are downcast once; sparse edits leave the source buffers unchanged.
//! Dates retain the grid's formatted-text representation.
//!
//! Supported source types: signed/unsigned ints (→ `Int`), `Float32/64` (→
//! `Float`), `Utf8/LargeUtf8` (→ `Text`, promoted to `Enum` if low-cardinality),
//! `Boolean` (→ `Bool`), `Date32/Date64` (→ `Date` text),
//! `Timestamp(Second/Millisecond/Microsecond/Nanosecond, _)` (→ `Date` text,
//! rendered as ISO-8601 datetime). Anything else falls back to blank `Text`.
use arrow_array::{
Array, BooleanArray, Date32Array, Date64Array, Float32Array, Float64Array, Int16Array, Int32Array,
Int64Array, Int8Array, LargeStringArray, RecordBatch, StringArray,
TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray,
UInt16Array, UInt32Array, UInt64Array, UInt8Array,
};
use arrow_schema::{DataType, TimeUnit};
use crate::data::{CellKind, Column, GridData, Keys, Value};
use std::collections::BTreeMap;
/// Build a [`GridData`] from an Arrow [`RecordBatch`]. Each Arrow column becomes a
/// typed [`Column`]; lengths come from the batch. Never fails (unknown types fall
/// back to text).
pub fn grid_data_from_batch(batch: &RecordBatch) -> GridData {
let nrows = batch.num_rows();
let schema = batch.schema();
let mut named: Vec<(String, Column)> = Vec::with_capacity(batch.num_columns());
for (ci, field) in schema.fields().iter().enumerate() {
let arr = batch.column(ci).as_ref();
named.push((field.name().clone(), column_from_array(arr, nrows)));
}
GridData::new(named)
}
fn nulls_of(arr: &dyn Array, n: usize) -> Vec<bool> {
(0..n).map(|i| arr.is_null(i)).collect()
}
#[cfg(test)]
thread_local! {
/// Test-only: cells read one at a time through [`ArrowColumn::value`] on this
/// thread. Whole-column work (sort, Σ) must not read cells this way.
pub(crate) static VALUE_READS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
}
/// A typed Arrow array with sparse, validated cell edits. Cloning shares buffers.
#[derive(Clone, Debug)]
pub struct ArrowColumn {
values: Values,
edits: BTreeMap<usize, String>,
}
// Downcast once on ingest, not on each cell access.
macro_rules! arrays {
($($variant:ident: $ty:ty => $kind:ident, $value:ident, $convert:expr;)*) => {
#[derive(Clone, Debug)]
enum Values { $($variant($ty),)* }
impl ArrowColumn {
fn from_array(array: &dyn Array) -> Option<Self> {
$(if let Some(a) = array.as_any().downcast_ref::<$ty>() {
return Some(Self { values: Values::$variant(a.clone()), edits: BTreeMap::new() });
})*
None
}
pub(crate) fn len(&self) -> usize {
match &self.values { $(Values::$variant(a) => a.len(),)* }
}
pub(crate) fn kind(&self) -> CellKind {
match &self.values { $(Values::$variant(_) => CellKind::$kind,)* }
}
pub(crate) fn value(&self, row: usize) -> Option<Value<'_>> {
#[cfg(test)]
VALUE_READS.with(|n| n.set(n.get() + 1));
if row >= self.len() { return None; }
if let Some(s) = self.edits.get(&row) {
return Some(match self.kind() {
CellKind::Int if self.is_unsigned() => Value::UInt(s.parse().ok()?),
CellKind::Int => Value::Int(s.parse().ok()?),
CellKind::Float => Value::Float(s.parse().ok()?),
CellKind::Bool => Value::Bool(matches!(s.as_str(), "true" | "1")),
_ => Value::Text(s),
});
}
match &self.values {
$(Values::$variant(a) => (!a.is_null(row)).then(|| Value::$value(($convert)(a.value(row)))),)*
}
}
/// Fold every row's value, in row order with edits applied, in ONE
/// sequential pass over the typed buffer. Whole-column readers use this
/// rather than [`value`](Self::value) per row, which probes the edit map
/// and re-dispatches on the array type on every call. The accumulator
/// goes by value, so a sum stays in registers.
pub(crate) fn fold_values<'a, A>(&'a self, init: A, mut f: impl FnMut(A, usize, Option<Value<'a>>) -> A) -> A {
// No edits (the common case) gets a loop of its own: with the
// overlay's `value` call beside it the loop runs ~4x slower.
if self.edits.is_empty() {
return match &self.values {
$(Values::$variant(a) => a.iter().enumerate().fold(init, |acc, (row, v)| {
f(acc, row, v.map(|v| Value::$value(($convert)(v))))
}),)*
};
}
// The next edited row; `usize::MAX` once there are no more.
let mut edits = self.edits.keys().copied();
let mut next_edit = edits.next().unwrap_or(usize::MAX);
match &self.values {
$(Values::$variant(a) => a.iter().enumerate().fold(init, |acc, (row, v)| {
if row == next_edit {
next_edit = edits.next().unwrap_or(usize::MAX);
f(acc, row, self.value(row))
} else {
f(acc, row, v.map(|v| Value::$value(($convert)(v))))
}
}),)*
}
}
/// The column's sort [`Keys`]: the typed buffer copied out once, then
/// each edited cell parsed once, exactly as [`value`](Self::value) reads it.
pub(crate) fn keys(&self) -> Keys<'_> {
match &self.values {
$(Values::$variant(a) => {
let mut keys: Vec<_> = a.iter().map(|v| v.map($convert)).collect();
for &row in self.edits.keys() {
keys[row] = match self.value(row) {
Some(Value::$value(v)) => Some(v),
_ => None,
};
}
Keys::$value(keys)
})*
}
}
}
}
}
arrays! {
I8: Int8Array => Int, Int, i64::from;
I16: Int16Array => Int, Int, i64::from;
I32: Int32Array => Int, Int, i64::from;
I64: Int64Array => Int, Int, std::convert::identity;
U8: UInt8Array => Int, Int, i64::from;
U16: UInt16Array => Int, Int, i64::from;
U32: UInt32Array => Int, Int, i64::from;
U64: UInt64Array => Int, UInt, std::convert::identity;
F32: Float32Array => Float, Float, f64::from;
F64: Float64Array => Float, Float, std::convert::identity;
Bool: BooleanArray => Bool, Bool, std::convert::identity;
Text: StringArray => Text, Text, std::convert::identity;
LargeText: LargeStringArray => Text, Text, std::convert::identity;
}
impl ArrowColumn {
pub(crate) fn is_unsigned(&self) -> bool {
matches!(self.values, Values::U64(_))
}
pub(crate) fn edit(&mut self, row: usize, value: &str) {
if row < self.len() {
let value = if self.kind() == CellKind::Bool {
value.to_ascii_lowercase()
} else {
value.to_owned()
};
self.edits.insert(row, value);
}
}
}
fn column_from_array(arr: &dyn Array, n: usize) -> Column {
if let Some(column) = ArrowColumn::from_array(arr) {
return Column::Arrow(column);
}
match arr.data_type() {
DataType::Date32 => {
let a = arr.as_any().downcast_ref::<Date32Array>().unwrap();
let v = (0..n)
.map(|i| if a.is_null(i) { String::new() } else { date_from_days(a.value(i)) })
.collect();
Column::Date { v, nulls: nulls_of(arr, n) }
}
DataType::Date64 => {
let a = arr.as_any().downcast_ref::<Date64Array>().unwrap();
let v = (0..n)
.map(|i| if a.is_null(i) { String::new() } else { date_from_days((a.value(i) / 86_400_000) as i32) })
.collect();
Column::Date { v, nulls: nulls_of(arr, n) }
}
// Timestamp variants: render as ISO-8601 datetime strings (YYYY-MM-DDTHH:MM:SS).
// All four TimeUnit sub-types map to the same Column::Date kind so they align
// and sort consistently with Date32/Date64 in mixed schemas.
DataType::Timestamp(unit, _tz) => {
// Convert raw timestamp value to microseconds-since-epoch, then format.
let v: Vec<String> = match unit {
TimeUnit::Second => {
let a = arr.as_any().downcast_ref::<TimestampSecondArray>().unwrap();
(0..n).map(|i| if a.is_null(i) { String::new() } else { datetime_from_us(a.value(i) * 1_000_000) }).collect()
}
TimeUnit::Millisecond => {
let a = arr.as_any().downcast_ref::<TimestampMillisecondArray>().unwrap();
(0..n).map(|i| if a.is_null(i) { String::new() } else { datetime_from_us(a.value(i) * 1_000) }).collect()
}
TimeUnit::Microsecond => {
let a = arr.as_any().downcast_ref::<TimestampMicrosecondArray>().unwrap();
(0..n).map(|i| if a.is_null(i) { String::new() } else { datetime_from_us(a.value(i)) }).collect()
}
TimeUnit::Nanosecond => {
let a = arr.as_any().downcast_ref::<TimestampNanosecondArray>().unwrap();
(0..n).map(|i| if a.is_null(i) { String::new() } else { datetime_from_us(a.value(i) / 1_000) }).collect()
}
};
Column::Date { v, nulls: nulls_of(arr, n) }
}
// Unknown / nested → best-effort text. We can't generically stringify an
// arbitrary arrow scalar without `arrow-cast`, so emit a typed placeholder
// (the type name) rather than pull another crate. Common types are all
// handled above; this is the rare tail.
_ => {
let v = (0..n).map(|_| String::new()).collect();
Column::Text { v, nulls: nulls_of(arr, n) }
}
}
}
/// Microseconds-since-epoch (Unix epoch 1970-01-01T00:00:00Z) → `YYYY-MM-DDTHH:MM:SS`.
/// Pure integer arithmetic — no chrono dependency. Sub-second precision is dropped
/// (the common analytics display form; the raw value is still available via state_json).
fn datetime_from_us(us: i64) -> String {
// Split into whole seconds and remaining microseconds.
let secs = if us >= 0 { us / 1_000_000 } else { (us - 999_999) / 1_000_000 };
let day_secs = secs.rem_euclid(86_400);
let days = ((secs - day_secs) / 86_400) as i32;
let h = day_secs / 3600;
let m = (day_secs % 3600) / 60;
let s = day_secs % 60;
format!("{}T{h:02}:{m:02}:{s:02}", date_from_days(days))
}
/// Days-since-epoch (1970-01-01) → `YYYY-MM-DD`. Pure integer civil-date math
/// (Howard Hinnant's algorithm) — no chrono dependency.
fn date_from_days(days: i32) -> String {
let z = days as i64 + 719_468;
let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
let doe = z - era * 146_097; // [0, 146096]
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146_096) / 365; // [0, 399]
let y = yoe + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); // [0, 365]
let mp = (5 * doy + 2) / 153; // [0, 11]
let d = doy - (153 * mp + 2) / 5 + 1; // [1, 31]
let m = if mp < 10 { mp + 3 } else { mp - 9 }; // [1, 12]
let y = if m <= 2 { y + 1 } else { y };
format!("{y:04}-{m:02}-{d:02}")
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_schema::{Field, Schema};
use std::sync::Arc;
#[test]
fn maps_arrow_types_to_kinds() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("i", DataType::Int64, false),
Field::new("f", DataType::Float64, false),
Field::new("s", DataType::Utf8, false),
Field::new("b", DataType::Boolean, false),
Field::new("status", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 2, 3, 4])),
Arc::new(Float64Array::from(vec![1.5, 2.5, 3.5, 4.5])),
Arc::new(StringArray::from(vec!["a-1", "a-2", "a-3", "a-4"])),
Arc::new(BooleanArray::from(vec![true, false, true, false])),
Arc::new(StringArray::from(vec!["ok", "warn", "ok", "warn"])),
],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.nrows(), 4);
assert_eq!(g.ncols(), 5);
assert_eq!(g.kinds[0], CellKind::Int);
assert_eq!(g.kinds[1], CellKind::Float);
assert_eq!(g.kinds[2], CellKind::Text, "unique strings stay text");
assert_eq!(g.kinds[3], CellKind::Bool);
assert_eq!(g.kinds[4], CellKind::Enum, "low-cardinality → enum chip");
assert_eq!(g.cell_string(0, 0), "1");
assert_eq!(g.cell_string(0, 1), "1.5");
}
#[test]
fn date32_renders_iso() {
// 2021-01-01 is 18628 days after epoch.
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("d", DataType::Date32, false)])),
vec![Arc::new(Date32Array::from(vec![18628]))],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.kinds[0], CellKind::Date);
assert_eq!(g.cell_string(0, 0), "2021-01-01");
}
#[test]
fn nulls_round_trip_as_empty() {
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("i", DataType::Int64, true)])),
vec![Arc::new(Int64Array::from(vec![Some(7), None, Some(9)]))],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.cell_string(0, 0), "7");
assert_eq!(g.cell_string(1, 0), "", "null → empty string");
assert!(g.columns[0].is_null(1));
}
#[test]
fn timestamp_microsecond_renders_iso_datetime() {
use arrow_array::TimestampMicrosecondArray;
use arrow_schema::TimeUnit;
// 2021-01-01T00:00:00Z = 18628 days × 86400 s/day × 1_000_000 µs/s
let epoch_us: i64 = 18628_i64 * 86_400 * 1_000_000;
// 2021-06-15T13:45:30Z = epoch + (165 days + 13h45m30s) in µs
let dt_us: i64 = (18628 + 165) * 86_400 * 1_000_000 + (13 * 3600 + 45 * 60 + 30) * 1_000_000;
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("ts", DataType::Timestamp(TimeUnit::Microsecond, None), true),
])),
vec![Arc::new(
TimestampMicrosecondArray::from(vec![Some(epoch_us), None, Some(dt_us)]),
)],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.kinds[0], CellKind::Date, "Timestamp maps to Date kind");
assert_eq!(g.cell_string(0, 0), "2021-01-01T00:00:00", "epoch renders as ISO datetime");
assert_eq!(g.cell_string(1, 0), "", "null Timestamp → empty string");
assert_eq!(g.cell_string(2, 0), "2021-06-15T13:45:30", "non-epoch timestamp renders correctly");
}
#[test]
fn timestamp_second_renders_iso_datetime() {
use arrow_array::TimestampSecondArray;
use arrow_schema::TimeUnit;
// 2020-03-01T12:00:00Z
let secs: i64 = 18322_i64 * 86_400 + 12 * 3600; // 18322 days from epoch = 2020-03-01
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("ts", DataType::Timestamp(TimeUnit::Second, None), false),
])),
vec![Arc::new(TimestampSecondArray::from(vec![secs]))],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.kinds[0], CellKind::Date);
assert_eq!(g.cell_string(0, 0), "2020-03-01T12:00:00");
}
#[test]
fn timestamp_millisecond_renders_iso_datetime() {
use arrow_array::TimestampMillisecondArray;
use arrow_schema::TimeUnit;
// epoch: 2021-01-01T00:00:00Z = 18628 days × 86400 s/day × 1000 ms/s.
let epoch_ms: i64 = 18628_i64 * 86_400 * 1_000;
// non-epoch: 2021-06-15T13:45:30Z = epoch + 165 days + 13h45m30s, in ms.
// The 500 ms remainder must be DROPPED (sub-second precision is not shown).
let dt_ms: i64 = (18628 + 165) * 86_400 * 1_000 + (13 * 3600 + 45 * 60 + 30) * 1_000 + 500;
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"ts",
DataType::Timestamp(TimeUnit::Millisecond, None),
true,
)])),
vec![Arc::new(TimestampMillisecondArray::from(vec![Some(epoch_ms), None, Some(dt_ms)]))],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.kinds[0], CellKind::Date, "Timestamp(ms) maps to Date kind");
assert_eq!(g.cell_string(0, 0), "2021-01-01T00:00:00", "epoch ms renders as ISO datetime");
assert_eq!(g.cell_string(1, 0), "", "null Timestamp(ms) → empty string");
assert_eq!(g.cell_string(2, 0), "2021-06-15T13:45:30", "non-epoch ms renders, sub-second dropped");
}
#[test]
fn timestamp_nanosecond_renders_iso_datetime() {
use arrow_array::TimestampNanosecondArray;
use arrow_schema::TimeUnit;
// epoch: 2021-01-01T00:00:00Z in ns.
let epoch_ns: i64 = 18628_i64 * 86_400 * 1_000_000_000;
// non-epoch: 2021-06-15T13:45:30Z + 123_456_789 ns; the sub-second part must
// be DROPPED (datetime_from_us divides ns→µs then formats whole seconds).
let dt_ns: i64 = (18628 + 165) * 86_400 * 1_000_000_000
+ (13 * 3600 + 45 * 60 + 30) * 1_000_000_000
+ 123_456_789;
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new(
"ts",
DataType::Timestamp(TimeUnit::Nanosecond, None),
true,
)])),
vec![Arc::new(TimestampNanosecondArray::from(vec![Some(epoch_ns), None, Some(dt_ns)]))],
)
.unwrap();
let g = grid_data_from_batch(&batch);
assert_eq!(g.kinds[0], CellKind::Date, "Timestamp(ns) maps to Date kind");
assert_eq!(g.cell_string(0, 0), "2021-01-01T00:00:00", "epoch ns renders as ISO datetime");
assert_eq!(g.cell_string(1, 0), "", "null Timestamp(ns) → empty string");
assert_eq!(g.cell_string(2, 0), "2021-06-15T13:45:30", "non-epoch ns renders, sub-second dropped");
}
#[test]
fn primitive_and_text_import_share_sliced_buffers_and_edits_stay_sparse() {
let ints = Int64Array::from(vec![99, 1, 2, 3]).slice(1, 3);
let text =
StringArray::from(vec![Some("skip"), Some("first"), None, Some("last")]).slice(1, 3);
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("n", DataType::Int64, false),
Field::new("s", DataType::Utf8, true),
])),
vec![Arc::new(ints.clone()), Arc::new(text.clone())],
)
.unwrap();
let mut data = grid_data_from_batch(&batch);
let Column::Arrow(c) = &data.columns[0] else {
panic!("shared column")
};
let Values::I64(a) = &c.values else {
panic!("typed column")
};
assert_eq!(a.values().as_ptr(), ints.values().as_ptr());
let Some(Value::Text(s)) = data.columns[1].value(0) else {
panic!("text")
};
assert_eq!(s.as_ptr(), text.value(0).as_ptr());
assert!(data.columns[1].is_null(1));
let Column::Arrow(c) = &mut data.columns[1] else {
panic!("shared column")
};
c.edit(1, "edited");
assert_eq!(c.edits.len(), 1);
assert_eq!(data.cell_raw(1, 1), "edited");
assert!(text.is_null(1), "the original source is immutable");
assert!(!data.columns[1].is_null(1));
assert_eq!(data.cell_raw(2, 1), "last");
}
}