rich_data/
arrow.rs

1//! Arrow `RecordBatch`es as rows, and Arrow schemas in the
2//! format-neutral model (the `arrow` feature).
3//!
4//! [`schema`] maps an Arrow schema into [`Schema`]: integers of every width
5//! are `integer`, `Decimal*(p, s)` is `decimal(p,s)`, timestamps keep their
6//! time zone, dictionaries (and run-end encoded arrays) take their value
7//! type, and lists, structs and maps nest. Each field keeps Arrow's own
8//! type name as its [native type](Field::native_type) (`Int32`,
9//! `Timestamp(ms, "UTC")`, `Dictionary(Int32, Utf8)`; `List`, `Struct` and
10//! `Map` for nested types, whose children say the rest), is
11//! [required](Field::is_required) when it is not nullable, and carries its
12//! metadata, sorted by key. Types the model has no name for (times,
13//! durations, intervals, unions) are `unknown`, with the native type saying
14//! what they are.
15//!
16//! The schema explorer (#342): [`tree`] draws an Arrow schema as a
17//! [`SchemaTree`], and [`diff`] compares two as a [`SchemaDiff`], fields
18//! added, removed and changed with breaking changes marked.
19//!
20//! [`read_schema`] reads the schema of an Arrow IPC file (Feather v2,
21//! `.arrow`) or stream (`.arrows`) without reading its record batches.
22//!
23//! [`rows`] reads one batch, and [`BatchSource`] streams many. Cells are
24//! [`Value::Int`] and [`Value::Float`] for numbers (a `UInt64` past
25//! `i64::MAX` becomes a float), [`Value::Null`] for nulls, and text for the
26//! rest: booleans as `true`/`false`, decimals exactly, dates as
27//! `YYYY-MM-DD`, timestamps as `YYYY-MM-DDTHH:MM:SS[.fff]` in UTC (with a
28//! `Z` when the type has a time zone), binary as `0x…` hex, and nested
29//! values as compact JSON.
30//!
31//! ```
32//! use std::sync::Arc;
33//! use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
34//! use rich_data::{arrow, Value};
35//!
36//! let batch = RecordBatch::try_from_iter([
37//!     ("id", Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef),
38//!     ("name", Arc::new(StringArray::from(vec!["ada", "alan"])) as ArrayRef),
39//! ])
40//! .unwrap();
41//! let rows = arrow::rows(&batch);
42//! assert_eq!(rows.columns(), ["id", "name"]);
43//! assert_eq!(rows.rows()[1], [Value::Null, Value::from("alan")]);
44//! let id = &rows.schema().unwrap().fields()[0];
45//! assert_eq!((id.data_type().to_string(), id.native_type()), ("integer".into(), Some("Int32")));
46//! ```
47
48use arrow_array::cast::AsArray;
49use arrow_array::types::{
50    Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type,
51    DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, DurationSecondType,
52    Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
53    Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
54    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
55    TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
56};
57use arrow_array::{Array, RecordBatch};
58use arrow_schema::{DataType as Arrow, Field as ArrowField, TimeUnit};
59use serde_json::{Map, Value as Json};
60
61use rich_ext::schema::{SchemaDiff, SchemaTree};
62
63use crate::record::MAX_DEPTH;
64use crate::{DataError, DataType, Field, Row, RowSource, Rows, Schema, Value};
65
66/// An Arrow schema in the format-neutral model, with its metadata.
67pub fn schema(schema: &arrow_schema::Schema) -> Schema {
68    let mut model = Schema::new(schema.fields().iter().map(|f| field(f, 0)));
69    for (key, value) in sorted(schema.metadata()) {
70        model = model.with_metadata(key, value);
71    }
72    model
73}
74
75/// One Arrow field in the model.
76pub fn field(field: &ArrowField, depth: usize) -> Field {
77    let mut model = Field::new(field.name().clone(), data_type(field.data_type(), depth))
78        .nullable(field.is_nullable())
79        .required(!field.is_nullable())
80        .with_native_type(native_type(field.data_type()));
81    for (key, value) in sorted(field.metadata()) {
82        model = model.with_metadata(key, value);
83    }
84    model
85}
86
87/// Metadata in key order, so the model (and what draws it) is the same
88/// every time.
89fn sorted<'a>(
90    metadata: impl IntoIterator<Item = (&'a String, &'a String)>,
91) -> Vec<(&'a String, &'a String)> {
92    let mut entries: Vec<(&String, &String)> = metadata.into_iter().collect();
93    entries.sort();
94    entries
95}
96
97/// Arrow's name for a type: its own for scalars, and only the kind for
98/// nested types (`List`, `FixedSizeList(3)`, `Struct`, `Map`), whose
99/// children are fields of their own.
100pub fn native_type(arrow: &Arrow) -> String {
101    match arrow {
102        Arrow::List(_) => "List".into(),
103        Arrow::LargeList(_) => "LargeList".into(),
104        Arrow::ListView(_) => "ListView".into(),
105        Arrow::LargeListView(_) => "LargeListView".into(),
106        Arrow::FixedSizeList(_, size) => format!("FixedSizeList({size})"),
107        Arrow::Struct(_) => "Struct".into(),
108        Arrow::Map(_, sorted) if *sorted => "Map(sorted)".into(),
109        Arrow::Map(..) => "Map".into(),
110        Arrow::Dictionary(key, values) if values.is_nested() => {
111            format!("Dictionary({key}, {})", native_type(values))
112        }
113        Arrow::RunEndEncoded(ends, values) if values.data_type().is_nested() => format!(
114            "RunEndEncoded({}, {})",
115            ends.data_type(),
116            native_type(values.data_type())
117        ),
118        other => other.to_string(),
119    }
120}
121
122/// An Arrow schema drawn as a tree: each field with its Arrow type,
123/// `(required)` when it is not nullable, and its metadata; nested fields
124/// under their parents.
125///
126/// ```
127/// use arrow_schema::{DataType, Field, Schema, TimeUnit};
128/// use rich::Console;
129///
130/// let schema = Schema::new(vec![
131///     Field::new("id", DataType::Int64, false),
132///     Field::new("at", DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), true),
133///     Field::new("tags", DataType::new_list(DataType::Utf8, true), true),
134/// ]);
135/// let console = Console::builder().width(60).color_system(None).build();
136/// assert_eq!(
137///     console.render_to_string(&rich_data::arrow::tree(&schema)),
138///     "schema  3 fields\n\
139///      ├── id (required)  Int64\n\
140///      ├── at  Timestamp(ms, \"UTC\")\n\
141///      └── tags  List\n    \
142///          └── item  Utf8"
143/// );
144/// ```
145pub fn tree(arrow: &arrow_schema::Schema) -> SchemaTree {
146    SchemaTree::from_model(schema(arrow))
147}
148
149/// What changed between two Arrow schemas: fields added and removed, type
150/// changes, fields that became (or stopped being) nullable, and metadata,
151/// with breaking changes marked.
152///
153/// ```
154/// use arrow_schema::{DataType, Field, Schema};
155///
156/// let old = Schema::new(vec![Field::new("id", DataType::Int32, true)]);
157/// let new = Schema::new(vec![
158///     Field::new("id", DataType::Int64, false),
159///     Field::new("name", DataType::Utf8, true),
160/// ]);
161/// let diff = rich_data::arrow::diff(&old, &new);
162/// let lines: Vec<String> = diff
163///     .changes()
164///     .iter()
165///     .map(|c| format!("{} {} {}{}", c.kind.marker(), c.path, c.detail, if c.breaking { " !" } else { "" }))
166///     .collect();
167/// assert_eq!(
168///     lines,
169///     ["~ id became required !", "~ id type Int32 → Int64 !", "+ name field added (Utf8)"]
170/// );
171/// ```
172pub fn diff(old: &arrow_schema::Schema, new: &arrow_schema::Schema) -> SchemaDiff {
173    SchemaDiff::models(&schema(old), &schema(new))
174}
175
176/// The magic bytes an Arrow IPC file starts (and ends) with.
177const FILE_MAGIC: &[u8; 6] = b"ARROW1";
178
179/// The schema of an Arrow IPC file (`ARROW1`…, Feather v2) or stream,
180/// whichever `reader` holds: the file's footer, or the stream's first
181/// message. Neither record batches nor dictionary batches are read: only
182/// the footer's schema, its length checked against the input's, so a
183/// hostile footer is an error rather than a panic or a large allocation.
184///
185/// ```
186/// use std::io::Cursor;
187/// use arrow_schema::{DataType, Field, Schema};
188///
189/// let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]);
190/// let mut bytes = Vec::new();
191/// arrow_ipc::writer::FileWriter::try_new(&mut bytes, &schema)
192///     .unwrap()
193///     .finish()
194///     .unwrap();
195/// let read = rich_data::arrow::read_schema(Cursor::new(bytes)).unwrap();
196/// assert_eq!(read.field(0).name(), "id");
197/// ```
198pub fn read_schema<R: std::io::Read + std::io::Seek>(
199    mut reader: R,
200) -> Result<arrow_schema::Schema, DataError> {
201    use std::io::SeekFrom;
202    let io = |err: std::io::Error| DataError::new(err.to_string());
203    let start = reader.stream_position().map_err(io)?;
204    let mut magic = [0u8; 6];
205    let read = read_up_to(&mut reader, &mut magic).map_err(io)?;
206    reader.seek(SeekFrom::Start(start)).map_err(io)?;
207    if read == 0 {
208        return Err(DataError::new(
209            "empty input: not an Arrow IPC file or stream",
210        ));
211    }
212    let not_arrow = |err: arrow_schema::ArrowError| {
213        DataError::new(format!("not an Arrow IPC file or stream: {err}"))
214    };
215    if &magic == FILE_MAGIC {
216        file_schema(reader, start).map_err(not_arrow)
217    } else {
218        let stream = arrow_ipc::reader::StreamReader::try_new(reader, None).map_err(not_arrow)?;
219        Ok(stream.schema().as_ref().clone())
220    }
221}
222
223/// Fill as much of `buf` as `reader` has, returning how much that was.
224/// The schema in an IPC file's footer: `FileReader::try_new` without the
225/// dictionary batches it also reads, whose block lengths come from the file
226/// and which it unwraps (a negative one panics).
227fn file_schema<R: std::io::Read + std::io::Seek>(
228    mut reader: R,
229    start: u64,
230) -> Result<arrow_schema::Schema, arrow_schema::ArrowError> {
231    use arrow_schema::ArrowError;
232    use std::io::SeekFrom;
233    // The leading magic and its padding, then the trailing length and magic.
234    const FRAME: u64 = 8 + 10;
235    let len = reader.seek(SeekFrom::End(0))?.saturating_sub(start);
236    if len < FRAME {
237        return Err(ArrowError::ParseError(
238            "Arrow file does not contain correct footer".into(),
239        ));
240    }
241    let mut trailer = [0u8; 10];
242    reader.seek(SeekFrom::End(-10))?;
243    reader.read_exact(&mut trailer)?;
244    let footer_len = arrow_ipc::reader::read_footer_length(trailer)?;
245    if footer_len as u64 > len - FRAME {
246        return Err(ArrowError::ParseError(format!(
247            "Invalid footer length: {footer_len}"
248        )));
249    }
250    let mut footer = vec![0; footer_len];
251    reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
252    reader.read_exact(&mut footer)?;
253    let footer = arrow_ipc::root_as_footer(&footer)
254        .map_err(|err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")))?;
255    let schema = footer
256        .schema()
257        .ok_or_else(|| ArrowError::ParseError("Unable to get schema from IPC Footer".into()))?;
258    arrow_ipc::convert::try_fb_to_schema(schema)
259}
260
261fn read_up_to(reader: &mut impl std::io::Read, buf: &mut [u8]) -> std::io::Result<usize> {
262    let mut filled = 0;
263    while filled < buf.len() {
264        match reader.read(&mut buf[filled..]) {
265            Ok(0) => break,
266            Ok(n) => filled += n,
267            Err(err) if err.kind() == std::io::ErrorKind::Interrupted => {}
268            Err(err) => return Err(err),
269        }
270    }
271    Ok(filled)
272}
273
274/// One Arrow type in the model.
275pub fn data_type(arrow: &Arrow, depth: usize) -> DataType {
276    if depth >= MAX_DEPTH {
277        return DataType::Any;
278    }
279    match arrow {
280        Arrow::Null => DataType::Unknown,
281        Arrow::Boolean => DataType::Boolean,
282        Arrow::Int8
283        | Arrow::Int16
284        | Arrow::Int32
285        | Arrow::Int64
286        | Arrow::UInt8
287        | Arrow::UInt16
288        | Arrow::UInt32
289        | Arrow::UInt64 => DataType::Integer,
290        Arrow::Float16 | Arrow::Float32 | Arrow::Float64 => DataType::Float,
291        Arrow::Decimal32(p, s)
292        | Arrow::Decimal64(p, s)
293        | Arrow::Decimal128(p, s)
294        | Arrow::Decimal256(p, s) => DataType::Decimal {
295            precision: Some(u32::from(*p)),
296            scale: Some(i32::from(*s)),
297        },
298        Arrow::Utf8 | Arrow::LargeUtf8 | Arrow::Utf8View => DataType::String,
299        Arrow::Binary | Arrow::LargeBinary | Arrow::BinaryView | Arrow::FixedSizeBinary(_) => {
300            DataType::Binary
301        }
302        Arrow::Date32 | Arrow::Date64 => DataType::Date,
303        Arrow::Timestamp(_, timezone) => DataType::Timestamp {
304            timezone: timezone.as_deref().map(str::to_string),
305        },
306        Arrow::List(item)
307        | Arrow::LargeList(item)
308        | Arrow::ListView(item)
309        | Arrow::LargeListView(item)
310        | Arrow::FixedSizeList(item, _) => DataType::List(Box::new(field(item, depth + 1))),
311        Arrow::Struct(fields) => {
312            DataType::Struct(fields.iter().map(|f| field(f, depth + 1)).collect())
313        }
314        Arrow::Map(entries, _) => match entries.data_type() {
315            Arrow::Struct(kv) if kv.len() == 2 => DataType::Map {
316                key: Box::new(field(&kv[0], depth + 1)),
317                value: Box::new(field(&kv[1], depth + 1)),
318            },
319            _ => DataType::Unknown,
320        },
321        Arrow::Dictionary(_, values) => data_type(values, depth),
322        Arrow::RunEndEncoded(_, values) => data_type(values.data_type(), depth),
323        _ => DataType::Unknown,
324    }
325}
326
327/// One batch as rows, with its schema.
328pub fn rows(batch: &RecordBatch) -> Rows {
329    let arrow_schema = batch.schema();
330    let mut rows = Rows::new(arrow_schema.fields().iter().map(|f| f.name().clone()))
331        .with_schema(schema(&arrow_schema));
332    for index in 0..batch.num_rows() {
333        rows.push(batch.columns().iter().map(|c| cell(c.as_ref(), index)));
334    }
335    rows
336}
337
338/// A [`RowSource`] over record batches that share a schema: the first
339/// batch's schema is the source's, and a later batch with different
340/// columns is an error.
341#[derive(Debug)]
342pub struct BatchSource<I> {
343    batches: I,
344    columns: Vec<String>,
345    schema: Schema,
346    arrow_schema: arrow_schema::SchemaRef,
347    current: Option<RecordBatch>,
348    row: usize,
349    number: usize,
350}
351
352impl<I: Iterator<Item = RecordBatch>> BatchSource<I> {
353    /// Read `batches`; `None` when there are none (no batch, no schema).
354    pub fn new(batches: impl IntoIterator<IntoIter = I>) -> Option<Self> {
355        let mut batches = batches.into_iter();
356        let first = batches.next()?;
357        let arrow_schema = first.schema();
358        Some(BatchSource {
359            batches,
360            columns: arrow_schema
361                .fields()
362                .iter()
363                .map(|f| f.name().clone())
364                .collect(),
365            schema: schema(&arrow_schema),
366            arrow_schema,
367            current: Some(first),
368            row: 0,
369            number: 1,
370        })
371    }
372}
373
374impl<I: Iterator<Item = RecordBatch>> RowSource for BatchSource<I> {
375    fn columns(&self) -> &[String] {
376        &self.columns
377    }
378
379    fn schema(&self) -> Option<&Schema> {
380        Some(&self.schema)
381    }
382
383    fn next_row(&mut self) -> Option<Result<Row, DataError>> {
384        loop {
385            let batch = self.current.as_ref()?;
386            if self.row < batch.num_rows() {
387                let row = batch
388                    .columns()
389                    .iter()
390                    .map(|c| cell(c.as_ref(), self.row))
391                    .collect();
392                self.row += 1;
393                return Some(Ok(row));
394            }
395            let next = self.batches.next()?;
396            self.number += 1;
397            if next.schema().fields() != self.arrow_schema.fields() {
398                self.current = None;
399                return Some(Err(DataError::at(
400                    self.number,
401                    "batch schema differs from the first batch's",
402                )));
403            }
404            self.current = Some(next);
405            self.row = 0;
406        }
407    }
408}
409
410macro_rules! int {
411    ($array:expr, $index:expr, $($variant:ident => $ty:ty),*) => {
412        match $array.data_type() {
413            $(Arrow::$variant => Some(i128::from($array.as_primitive::<$ty>().value($index))),)*
414            _ => None,
415        }
416    };
417}
418
419/// A run-end encoded array's values and the physical index that holds
420/// logical row `index`; `None` for any other array.
421fn run(array: &dyn Array, index: usize) -> Option<(&dyn Array, usize)> {
422    let Arrow::RunEndEncoded(ends, _) = array.data_type() else {
423        return None;
424    };
425    Some(match ends.data_type() {
426        Arrow::Int16 => {
427            let run = array.as_run::<Int16Type>();
428            (run.values().as_ref(), run.get_physical_index(index))
429        }
430        Arrow::Int32 => {
431            let run = array.as_run::<Int32Type>();
432            (run.values().as_ref(), run.get_physical_index(index))
433        }
434        Arrow::Int64 => {
435            let run = array.as_run::<Int64Type>();
436            (run.values().as_ref(), run.get_physical_index(index))
437        }
438        _ => return None,
439    })
440}
441
442/// The cell at `index` of `array`.
443pub fn cell(array: &dyn Array, index: usize) -> Value {
444    // A run-end encoded array keeps its nulls in its values, so it resolves
445    // to the value that holds the row before the null check.
446    if let Some((values, physical)) = run(array, index) {
447        return cell(values, physical);
448    }
449    if array.is_null(index) {
450        return Value::Null;
451    }
452    let integer = int!(array, index,
453        Int8 => Int8Type, Int16 => Int16Type, Int32 => Int32Type, Int64 => Int64Type,
454        UInt8 => UInt8Type, UInt16 => UInt16Type, UInt32 => UInt32Type, UInt64 => UInt64Type);
455    if let Some(n) = integer {
456        return i64::try_from(n).map_or(Value::Float(n as f64), Value::Int);
457    }
458    match array.data_type() {
459        Arrow::Float16 => Value::Float(array.as_primitive::<Float16Type>().value(index).to_f64()),
460        Arrow::Float32 => Value::Float(f64::from(array.as_primitive::<Float32Type>().value(index))),
461        Arrow::Float64 => Value::Float(array.as_primitive::<Float64Type>().value(index)),
462        Arrow::Duration(unit) => Value::Int(match unit {
463            TimeUnit::Second => array.as_primitive::<DurationSecondType>().value(index),
464            TimeUnit::Millisecond => array.as_primitive::<DurationMillisecondType>().value(index),
465            TimeUnit::Microsecond => array.as_primitive::<DurationMicrosecondType>().value(index),
466            TimeUnit::Nanosecond => array.as_primitive::<DurationNanosecondType>().value(index),
467        }),
468        Arrow::List(_)
469        | Arrow::LargeList(_)
470        | Arrow::ListView(_)
471        | Arrow::LargeListView(_)
472        | Arrow::FixedSizeList(..)
473        | Arrow::Struct(_)
474        | Arrow::Map(..) => Value::Str(json(array, index, 0).to_string()),
475        Arrow::Dictionary(..) => {
476            let dictionary = array.as_any_dictionary();
477            let key = dictionary.normalized_keys()[index];
478            cell(dictionary.values().as_ref(), key)
479        }
480        _ => Value::Str(text(array, index)),
481    }
482}
483
484/// A scalar cell's text: booleans, decimals, strings, binary, dates, times
485/// and timestamps.
486fn text(array: &dyn Array, index: usize) -> String {
487    match array.data_type() {
488        Arrow::Boolean => array.as_boolean().value(index).to_string(),
489        Arrow::Decimal32(_, s) => decimal(
490            array
491                .as_primitive::<Decimal32Type>()
492                .value(index)
493                .to_string(),
494            *s,
495        ),
496        Arrow::Decimal64(_, s) => decimal(
497            array
498                .as_primitive::<Decimal64Type>()
499                .value(index)
500                .to_string(),
501            *s,
502        ),
503        Arrow::Decimal128(_, s) => decimal(
504            array
505                .as_primitive::<Decimal128Type>()
506                .value(index)
507                .to_string(),
508            *s,
509        ),
510        Arrow::Decimal256(_, s) => decimal(
511            array
512                .as_primitive::<Decimal256Type>()
513                .value(index)
514                .to_string(),
515            *s,
516        ),
517        Arrow::Utf8 => array.as_string::<i32>().value(index).to_string(),
518        Arrow::LargeUtf8 => array.as_string::<i64>().value(index).to_string(),
519        Arrow::Utf8View => array.as_string_view().value(index).to_string(),
520        Arrow::Binary => hex(array.as_binary::<i32>().value(index)),
521        Arrow::LargeBinary => hex(array.as_binary::<i64>().value(index)),
522        Arrow::BinaryView => hex(array.as_binary_view().value(index)),
523        Arrow::FixedSizeBinary(_) => hex(array.as_fixed_size_binary().value(index)),
524        Arrow::Date32 => date(i64::from(array.as_primitive::<Date32Type>().value(index))),
525        Arrow::Date64 => date(
526            array
527                .as_primitive::<Date64Type>()
528                .value(index)
529                .div_euclid(86_400_000),
530        ),
531        Arrow::Timestamp(unit, timezone) => {
532            let raw = match unit {
533                TimeUnit::Second => array.as_primitive::<TimestampSecondType>().value(index),
534                TimeUnit::Millisecond => array
535                    .as_primitive::<TimestampMillisecondType>()
536                    .value(index),
537                TimeUnit::Microsecond => array
538                    .as_primitive::<TimestampMicrosecondType>()
539                    .value(index),
540                TimeUnit::Nanosecond => {
541                    array.as_primitive::<TimestampNanosecondType>().value(index)
542                }
543            };
544            let (seconds, fraction) = split(i128::from(raw), *unit);
545            let days = seconds.div_euclid(86_400) as i64;
546            let mut out = format!(
547                "{}T{}",
548                date(days),
549                time(seconds.rem_euclid(86_400), fraction)
550            );
551            if timezone.is_some() {
552                out.push('Z');
553            }
554            out
555        }
556        Arrow::Time32(unit) | Arrow::Time64(unit) => {
557            let raw = match unit {
558                TimeUnit::Second => {
559                    i64::from(array.as_primitive::<Time32SecondType>().value(index))
560                }
561                TimeUnit::Millisecond => {
562                    i64::from(array.as_primitive::<Time32MillisecondType>().value(index))
563                }
564                TimeUnit::Microsecond => array.as_primitive::<Time64MicrosecondType>().value(index),
565                TimeUnit::Nanosecond => array.as_primitive::<Time64NanosecondType>().value(index),
566            };
567            let (seconds, fraction) = split(i128::from(raw), *unit);
568            time(seconds, fraction)
569        }
570        other => format!("<{other}>"),
571    }
572}
573
574/// Seconds and the sub-second digits (as written, without trailing zeros).
575fn split(raw: i128, unit: TimeUnit) -> (i128, String) {
576    let (per_second, width) = match unit {
577        TimeUnit::Second => return (raw, String::new()),
578        TimeUnit::Millisecond => (1_000, 3),
579        TimeUnit::Microsecond => (1_000_000, 6),
580        TimeUnit::Nanosecond => (1_000_000_000, 9),
581    };
582    let fraction = format!("{:0width$}", raw.rem_euclid(per_second));
583    (
584        raw.div_euclid(per_second),
585        fraction.trim_end_matches('0').to_string(),
586    )
587}
588
589/// `HH:MM:SS[.fff]` for seconds into a day.
590fn time(seconds: i128, fraction: String) -> String {
591    let mut out = format!(
592        "{:02}:{:02}:{:02}",
593        seconds / 3600,
594        seconds / 60 % 60,
595        seconds % 60
596    );
597    if !fraction.is_empty() {
598        out.push('.');
599        out.push_str(&fraction);
600    }
601    out
602}
603
604/// `YYYY-MM-DD` for days since 1970-01-01 (Howard Hinnant's
605/// `civil_from_days`).
606fn date(days: i64) -> String {
607    let z = days + 719_468;
608    let era = z.div_euclid(146_097);
609    let doe = z.rem_euclid(146_097);
610    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
611    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
612    let mp = (5 * doy + 2) / 153;
613    let day = doy - (153 * mp + 2) / 5 + 1;
614    let month = if mp < 10 { mp + 3 } else { mp - 9 };
615    let year = yoe + era * 400 + i64::from(month <= 2);
616    format!("{year:04}-{month:02}-{day:02}")
617}
618
619/// An unscaled decimal's digits with the point `scale` places from the end.
620fn decimal(digits: String, scale: i8) -> String {
621    let (sign, digits) = match digits.strip_prefix('-') {
622        Some(rest) => ("-", rest.to_string()),
623        None => ("", digits),
624    };
625    if scale <= 0 {
626        let zeros = "0".repeat(usize::from(scale.unsigned_abs()));
627        return format!("{sign}{digits}{zeros}");
628    }
629    let scale = usize::from(scale.unsigned_abs());
630    let digits = format!("{digits:0>width$}", width = scale + 1);
631    let (whole, fraction) = digits.split_at(digits.len() - scale);
632    format!("{sign}{whole}.{fraction}")
633}
634
635fn hex(bytes: &[u8]) -> String {
636    let mut out = String::with_capacity(2 + bytes.len() * 2);
637    out.push_str("0x");
638    for byte in bytes {
639        out.push_str(&format!("{byte:02x}"));
640    }
641    out
642}
643
644/// A value as JSON, for nested cells.
645fn json(array: &dyn Array, index: usize, depth: usize) -> Json {
646    if let Some((values, physical)) = run(array, index) {
647        return json(values, physical, depth);
648    }
649    if array.is_null(index) {
650        return Json::Null;
651    }
652    if depth >= MAX_DEPTH {
653        return Json::String("…".into());
654    }
655    let items = |values: &dyn Array, depth: usize| {
656        Json::Array(
657            (0..values.len())
658                .map(|i| json(values, i, depth + 1))
659                .collect(),
660        )
661    };
662    match array.data_type() {
663        Arrow::List(_) => items(array.as_list::<i32>().value(index).as_ref(), depth),
664        Arrow::LargeList(_) => items(array.as_list::<i64>().value(index).as_ref(), depth),
665        Arrow::ListView(_) => items(array.as_list_view::<i32>().value(index).as_ref(), depth),
666        Arrow::LargeListView(_) => items(array.as_list_view::<i64>().value(index).as_ref(), depth),
667        Arrow::FixedSizeList(..) => items(array.as_fixed_size_list().value(index).as_ref(), depth),
668        Arrow::Struct(_) => {
669            let array = array.as_struct();
670            let map: Map<String, Json> = array
671                .column_names()
672                .into_iter()
673                .zip(array.columns())
674                .map(|(name, column)| (name.to_string(), json(column.as_ref(), index, depth + 1)))
675                .collect();
676            Json::Object(map)
677        }
678        Arrow::Map(..) => {
679            let entries = array.as_map().value(index);
680            let (keys, values) = (entries.column(0), entries.column(1));
681            let map: Map<String, Json> = (0..entries.len())
682                .map(|i| {
683                    (
684                        cell(keys.as_ref(), i).plain(),
685                        json(values.as_ref(), i, depth + 1),
686                    )
687                })
688                .collect();
689            Json::Object(map)
690        }
691        _ => match cell(array, index) {
692            Value::Null => Json::Null,
693            Value::Int(n) => Json::from(n),
694            Value::Float(f) => serde_json::Number::from_f64(f).map_or(Json::Null, Json::Number),
695            Value::Str(s) if array.data_type() == &Arrow::Boolean => Json::Bool(s == "true"),
696            other => Json::String(other.plain()),
697        },
698    }
699}
700
701#[cfg(test)]
702mod tests {
703    use std::sync::Arc;
704
705    use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StringBuilder};
706    use arrow_array::{
707        ArrayRef, BooleanArray, Date32Array, Decimal128Array, DictionaryArray, Float64Array,
708        StructArray, TimestampMillisecondArray, UInt64Array,
709    };
710    use arrow_schema::Field as ArrowField;
711
712    use super::*;
713
714    #[test]
715    fn scalars_become_cells() {
716        let batch = RecordBatch::try_from_iter([
717            (
718                "big",
719                Arc::new(UInt64Array::from(vec![u64::MAX])) as ArrayRef,
720            ),
721            ("f", Arc::new(Float64Array::from(vec![1.5])) as ArrayRef),
722            ("ok", Arc::new(BooleanArray::from(vec![true])) as ArrayRef),
723            (
724                "price",
725                Arc::new(
726                    Decimal128Array::from(vec![-1205])
727                        .with_precision_and_scale(10, 2)
728                        .unwrap(),
729                ) as ArrayRef,
730            ),
731            ("day", Arc::new(Date32Array::from(vec![19_782])) as ArrayRef),
732            (
733                "at",
734                Arc::new(
735                    TimestampMillisecondArray::from(vec![1_709_164_800_250]).with_timezone("UTC"),
736                ) as ArrayRef,
737            ),
738            (
739                "kind",
740                Arc::new(DictionaryArray::<arrow_array::types::Int8Type>::from_iter(
741                    ["x"],
742                )) as ArrayRef,
743            ),
744        ])
745        .unwrap();
746        let rows = rows(&batch);
747        assert_eq!(
748            rows.rows()[0],
749            [
750                Value::Float(u64::MAX as f64),
751                Value::Float(1.5),
752                Value::from("true"),
753                Value::from("-12.05"),
754                Value::from("2024-02-29"),
755                Value::from("2024-02-29T00:00:00.25Z"),
756                Value::from("x"),
757            ]
758        );
759        let types: Vec<String> = rows
760            .schema()
761            .unwrap()
762            .fields()
763            .iter()
764            .map(|f| f.data_type().to_string())
765            .collect();
766        assert_eq!(
767            types,
768            [
769                "integer",
770                "float",
771                "boolean",
772                "decimal(10,2)",
773                "date",
774                "timestamp[UTC]",
775                "string"
776            ]
777        );
778    }
779
780    #[test]
781    fn nested_values_are_json_and_nest_in_the_schema() {
782        let mut list = ListBuilder::new(Int32Builder::new());
783        list.values().append_value(1);
784        list.values().append_null();
785        list.append(true);
786        let list = Arc::new(list.finish()) as ArrayRef;
787        let mut map = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
788        map.keys().append_value("a");
789        map.values().append_value(2);
790        map.append(true).unwrap();
791        let map = Arc::new(map.finish()) as ArrayRef;
792        let inner = Arc::new(StructArray::from(vec![(
793            Arc::new(ArrowField::new("x", Arrow::Int32, true)),
794            Arc::new(arrow_array::Int32Array::from(vec![7])) as ArrayRef,
795        )])) as ArrayRef;
796        let batch = RecordBatch::try_from_iter([("l", list), ("m", map), ("s", inner)]).unwrap();
797        let rows = rows(&batch);
798        assert_eq!(
799            rows.rows()[0],
800            [
801                Value::from("[1,null]"),
802                Value::from(r#"{"a":2}"#),
803                Value::from(r#"{"x":7}"#)
804            ]
805        );
806        let schema = rows.schema().unwrap();
807        assert_eq!(schema.fields()[0].data_type().to_string(), "list<integer>");
808        assert_eq!(
809            schema.fields()[1].data_type().to_string(),
810            "map<string, integer>"
811        );
812        assert_eq!(schema.fields()[2].children().next().unwrap().name(), "x");
813    }
814
815    #[test]
816    fn a_batch_source_streams_and_checks_schemas() {
817        let batch = |name: &str| {
818            RecordBatch::try_from_iter([(
819                name,
820                Arc::new(arrow_array::Int32Array::from(vec![1, 2])) as ArrayRef,
821            )])
822            .unwrap()
823        };
824        let mut source = BatchSource::new([batch("a"), batch("a"), batch("b")]).unwrap();
825        for _ in 0..4 {
826            assert!(source.next_row().unwrap().is_ok());
827        }
828        let error = source.next_row().unwrap().unwrap_err();
829        assert_eq!(error.line(), Some(3));
830        assert!(source.next_row().is_none());
831        assert!(BatchSource::new(Vec::<RecordBatch>::new()).is_none());
832    }
833
834    #[test]
835    fn list_views_and_run_end_arrays_decode_to_their_values() {
836        use arrow_array::{Int32Array, ListViewArray, RunArray, StringArray};
837
838        let item = Arc::new(ArrowField::new("item", Arrow::Int32, true));
839        let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef;
840        let views = ListViewArray::new(item, vec![2, 0].into(), vec![2, 1].into(), values, None);
841        assert_eq!(cell(&views, 0), Value::Str("[3,4]".into()));
842        assert_eq!(cell(&views, 1), Value::Str("[1]".into()));
843
844        let ends = Int32Array::from(vec![2, 3, 5]);
845        let runs = StringArray::from(vec![Some("a"), None, Some("b")]);
846        let encoded = RunArray::<Int32Type>::try_new(&ends, &runs).unwrap();
847        let cells: Vec<Value> = (0..5).map(|i| cell(&encoded, i)).collect();
848        assert_eq!(
849            cells,
850            [
851                Value::Str("a".into()),
852                Value::Str("a".into()),
853                Value::Null,
854                Value::Str("b".into()),
855                Value::Str("b".into()),
856            ]
857        );
858    }
859
860    #[test]
861    fn dates_and_decimals_format_exactly() {
862        assert_eq!(date(0), "1970-01-01");
863        assert_eq!(date(-1), "1969-12-31");
864        assert_eq!(decimal("5".into(), 3), "0.005");
865        assert_eq!(decimal("-123".into(), 0), "-123");
866        assert_eq!(decimal("12".into(), -2), "1200");
867        assert_eq!(split(-1, TimeUnit::Millisecond), (-1, "999".to_string()));
868        assert_eq!(time(3661, "5".into()), "01:01:01.5");
869    }
870}
rich_data/arrow.rs

rich_data/
arrow.rs
