densa_arrow/
table.rs

1use std::fmt;
2
3use arrow_array::RecordBatch;
4use arrow_schema::SchemaRef;
5use densa_grid::{Align, CellStyle, CellText, Color, Column, GridSource, Palette, CELL_BYTES};
6
7use crate::reader::{is_numeric, Format, Reader};
8
9/// Rows of the first batch measured to size the columns.
10const SAMPLE_ROWS: usize = 1_000;
11/// Rows of each later batch measured to widen number and time columns (spread over the
12/// batch, its last row included): values that grow, such as counts and row numbers, grow
13/// toward the end.
14const WIDEN_ROWS: usize = 64;
15/// Widest a column is sized from its content; longer text is cut with `…`.
16const MAX_SAMPLED_CHARS: usize = 40;
17/// Width of a column with no rows yet, unless its title is wider.
18const DEFAULT_CHARS: usize = 8;
19
20/// What [`ArrowTable::write_cell`] wrote.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22#[non_exhaustive]
23pub enum CellState {
24    /// The value.
25    Value,
26    /// `null`, for a missing value.
27    Null,
28    /// `?`, for a type this crate does not format (lists, structs, maps, ...).
29    Unsupported,
30    /// `#`, for a value that cannot be shown: longer than a cell holds (as `1e300` without
31    /// fixed decimals), or invalid (a decimal scale beyond 38, a time of day outside a day).
32    Invalid,
33    /// Nothing: the row or column does not exist.
34    OutOfRange,
35}
36
37/// Why a batch could not be added, or a column setting applied.
38#[derive(Debug, Clone, PartialEq, Eq)]
39#[non_exhaustive]
40pub enum TableError {
41    /// No batch was given, so there is no schema.
42    NoBatches,
43    /// The batch has a different number of columns.
44    ColumnCount {
45        /// Columns in the table.
46        expected: usize,
47        /// Columns in the batch.
48        found: usize,
49    },
50    /// The batch's column at this index has another data type.
51    ColumnMismatch {
52        /// Index of the first differing column.
53        column: usize,
54    },
55    /// There is no column at this index.
56    NoSuchColumn(usize),
57}
58
59impl fmt::Display for TableError {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        match self {
62            Self::NoBatches => f.write_str("no record batch was given, so there is no schema"),
63            Self::ColumnCount { expected, found } => {
64                write!(f, "the batch has {found} columns, the table {expected}")
65            }
66            Self::ColumnMismatch { column } => {
67                write!(f, "column {column} of the batch has another data type")
68            }
69            Self::NoSuchColumn(column) => write!(f, "there is no column {column}"),
70        }
71    }
72}
73
74impl std::error::Error for TableError {}
75
76/// Arrow record batches shown as one table: rows run through the batches in order.
77///
78/// Formatting is chosen per column from its data type when a batch is added, so writing a
79/// cell is a lookup and a format into a stack buffer, without allocating. Batches can be
80/// appended while the table is shown (a stream); their arrays are shared, not copied.
81///
82/// Formatted types: booleans, integers, floats, `Decimal128`, strings (`Utf8`, `LargeUtf8`,
83/// `Utf8View`), dictionaries of those, dates, times, timestamps (with a time zone: in UTC,
84/// marked `Z`) and durations. Other types (lists, structs, maps, binary, intervals) show
85/// `?`.
86///
87/// A batch is accepted when its columns have the table's data types in order; field names,
88/// nullability and metadata may differ. A timestamp's time zone is part of its data type.
89#[derive(Debug, Clone)]
90pub struct ArrowTable {
91    schema: SchemaRef,
92    /// Each non-empty batch's columns.
93    batches: Vec<Vec<Reader>>,
94    /// First row of each batch in `batches`.
95    starts: Vec<usize>,
96    rows: usize,
97    formats: Vec<Format>,
98    /// Whether the columns have been sized from a batch's content.
99    sized: bool,
100    /// Bumped whenever content changes a column's width: the first rows, or a later value
101    /// too wide for its number column.
102    widths: u64,
103}
104
105impl ArrowTable {
106    /// An empty table for batches with `schema`; columns are sized by their titles until
107    /// the first rows arrive.
108    pub fn new(schema: SchemaRef) -> Self {
109        let formats = new_formats(&schema);
110        Self {
111            schema,
112            batches: Vec::new(),
113            starts: Vec::new(),
114            rows: 0,
115            formats,
116            sized: false,
117            widths: 0,
118        }
119    }
120
121    /// A table of `batches`, with the first batch's schema.
122    ///
123    /// # Errors
124    /// [`TableError::NoBatches`] for none; a column error if a batch's schema differs.
125    pub fn from_batches(
126        batches: impl IntoIterator<Item = RecordBatch>,
127    ) -> Result<Self, TableError> {
128        let mut batches = batches.into_iter();
129        let first = batches.next().ok_or(TableError::NoBatches)?;
130        let mut table = Self::new(first.schema());
131        for batch in std::iter::once(first).chain(batches) {
132            table.push(batch)?;
133        }
134        Ok(table)
135    }
136
137    /// Appends `batch`'s rows. The first rows to arrive size the columns (their titles and
138    /// up to the first 1,000 values, at most 40 characters); later rows widen a number or time
139    /// column whose values no longer fit (measuring some rows of each batch, the last among
140    /// them), never narrowing it. Columns sized with
141    /// [`set_column_chars`](Self::set_column_chars) keep their width.
142    ///
143    /// # Errors
144    /// A column error if `batch` has other columns or data types; the table is unchanged.
145    pub fn push(&mut self, batch: RecordBatch) -> Result<(), TableError> {
146        self.check(&batch)?;
147        if batch.num_rows() == 0 {
148            return Ok(());
149        }
150        let readers: Vec<Reader> = batch.columns().iter().cloned().map(Reader::new).collect();
151        if self.sized {
152            self.widen_columns(&readers);
153        } else {
154            self.size_columns(&readers);
155        }
156        self.starts.push(self.rows);
157        self.rows += batch.num_rows();
158        self.batches.push(readers);
159        Ok(())
160    }
161
162    /// The table's schema.
163    pub fn schema(&self) -> &SchemaRef {
164        &self.schema
165    }
166
167    /// Rows in all batches.
168    pub fn rows(&self) -> usize {
169        self.rows
170    }
171
172    /// Grid columns: titled by field name, sized by content, numbers and times aligned right.
173    pub fn columns(&self) -> Vec<Column> {
174        columns_of(&self.schema, &self.formats)
175    }
176
177    /// Prints floating-point values of `column` with `decimals` fixed decimal places
178    /// (`None`: the shortest exact form). A column sized by the first rows is sized again
179    /// from them, as now printed.
180    ///
181    /// # Errors
182    /// [`TableError::NoSuchColumn`] if there is no such column.
183    pub fn set_decimals(
184        &mut self,
185        column: usize,
186        decimals: Option<usize>,
187    ) -> Result<(), TableError> {
188        self.format_mut(column)?.decimals = decimals;
189        if let (Some(reader), Some(field)) = (
190            self.batches.first().and_then(|b| b.get(column)),
191            self.schema.fields().get(column),
192        ) {
193            let format = &mut self.formats[column];
194            size_column(format, reader, field.name());
195            // Later batches may have widened it: measure them again as now printed.
196            if is_numeric(field.data_type()) && !format.width_set {
197                let mut text = CellText::new();
198                for readers in &self.batches[1..] {
199                    let rows = spread(readers[column].len(), WIDEN_ROWS);
200                    let widest = widest(&readers[column], format, rows, &mut text);
201                    format.max_chars = format.max_chars.max(widest);
202                }
203            }
204        }
205        Ok(())
206    }
207
208    /// Sets `column` to `chars` characters wide (1 to [`CELL_BYTES`], the most a cell can
209    /// hold): [`columns`](Self::columns) reports it, longer text is cut to fit ending in
210    /// `…`, and the first rows to arrive no longer size it.
211    ///
212    /// # Errors
213    /// [`TableError::NoSuchColumn`] if there is no such column.
214    pub fn set_column_chars(&mut self, column: usize, chars: usize) -> Result<(), TableError> {
215        let format = self.format_mut(column)?;
216        format.max_chars = chars.clamp(1, CELL_BYTES);
217        format.width_set = true;
218        Ok(())
219    }
220
221    /// Appends the value at `row` of `column` to `text`, or a placeholder (see
222    /// [`CellState`]), never a partial value. Never allocates.
223    pub fn write_cell(&self, row: usize, column: usize, text: &mut CellText) -> CellState {
224        if row >= self.rows {
225            return CellState::OutOfRange;
226        }
227        // At least one start is <= row, since row < rows and starts[0] == 0.
228        let batch = self.starts.partition_point(|&start| start <= row) - 1;
229        self.write_at(batch, row - self.starts[batch], column, text)
230    }
231
232    /// [`write_cell`](Self::write_cell) for row `local` of stored batch `batch` (the batches
233    /// with rows, in the order pushed).
234    pub(crate) fn write_at(
235        &self,
236        batch: usize,
237        local: usize,
238        column: usize,
239        text: &mut CellText,
240    ) -> CellState {
241        let reader = self.batches.get(batch).and_then(|b| b.get(column));
242        let (Some(reader), Some(format)) = (reader, self.formats.get(column)) else {
243            return CellState::OutOfRange;
244        };
245        if local >= reader.len() {
246            return CellState::OutOfRange;
247        }
248        write_value(reader, local, format, text)
249    }
250
251    /// A [`GridSource`] showing the table from `first_row` (scrolled), in `palette`'s
252    /// text color, with nulls and unsupported values dimmed.
253    pub fn view<'a>(&'a self, first_row: usize, palette: &Palette) -> ArrowGrid<'a> {
254        ArrowGrid {
255            table: self,
256            first_row,
257            text: palette.text,
258            dim: palette.dim_text,
259        }
260    }
261
262    fn check(&self, batch: &RecordBatch) -> Result<(), TableError> {
263        let ours = self.schema.fields();
264        let theirs = batch.schema_ref().fields();
265        if ours.len() != theirs.len() {
266            return Err(TableError::ColumnCount {
267                expected: ours.len(),
268                found: theirs.len(),
269            });
270        }
271        match ours
272            .iter()
273            .zip(theirs.iter())
274            .position(|(a, b)| a.data_type() != b.data_type())
275        {
276            Some(column) => Err(TableError::ColumnMismatch { column }),
277            None => Ok(()),
278        }
279    }
280
281    /// Widens each column to its longest value among the first rows (and its title).
282    fn size_columns(&mut self, readers: &[Reader]) {
283        for ((reader, format), field) in readers
284            .iter()
285            .zip(&mut self.formats)
286            .zip(self.schema.fields().iter())
287        {
288            size_column(format, reader, field.name());
289        }
290        self.sized = true;
291        self.widths = self.widths.wrapping_add(1);
292    }
293
294    /// Widens each number and time column (not sized by the application) to its widest
295    /// value among some of `readers`' rows, if that is wider than the column.
296    fn widen_columns(&mut self, readers: &[Reader]) {
297        let mut text = CellText::new();
298        let mut widened = false;
299        for ((reader, format), field) in readers
300            .iter()
301            .zip(&mut self.formats)
302            .zip(self.schema.fields().iter())
303        {
304            widened |= widen_column(format, reader, field.data_type(), &mut text);
305        }
306        if widened {
307            self.widths = self.widths.wrapping_add(1);
308        }
309    }
310
311    /// Changes whenever content changed a column's width (a revision, not a count).
312    pub(crate) fn widths(&self) -> u64 {
313        self.widths
314    }
315
316    /// Batches with rows held, in the order pushed.
317    pub(crate) fn stored_batches(&self) -> usize {
318        self.batches.len()
319    }
320
321    /// Drops every batch, keeping the schema and the column formats and widths.
322    pub(crate) fn clear_batches(&mut self) {
323        self.batches.clear();
324        self.starts.clear();
325        self.rows = 0;
326    }
327
328    fn format_mut(&mut self, column: usize) -> Result<&mut Format, TableError> {
329        self.formats
330            .get_mut(column)
331            .ok_or(TableError::NoSuchColumn(column))
332    }
333}
334
335/// Each column's format before any rows: its title's width (at least a few characters).
336pub(crate) fn new_formats(schema: &SchemaRef) -> Vec<Format> {
337    schema
338        .fields()
339        .iter()
340        .map(|field| Format {
341            decimals: None,
342            max_chars: field.name().chars().count().max(DEFAULT_CHARS),
343            width_set: false,
344        })
345        .collect()
346}
347
348/// Grid columns for `schema` in `formats`: titled by field name, as wide as the format,
349/// numbers and times aligned right.
350pub(crate) fn columns_of(schema: &SchemaRef, formats: &[Format]) -> Vec<Column> {
351    schema
352        .fields()
353        .iter()
354        .zip(formats)
355        .map(|(field, format)| {
356            let align = if is_numeric(field.data_type()) {
357                Align::Right
358            } else {
359                Align::Left
360            };
361            let chars = u32::try_from(format.max_chars).unwrap_or(u32::MAX);
362            Column::new(field.name().clone(), chars, align)
363        })
364        .collect()
365}
366
367/// Appends the value at row `local` (within its rows) of `reader` to `text` as `format`
368/// says, or a placeholder: `null`, `?` for a type not formatted, `#` for a value that does
369/// not fit. Never allocates.
370pub(crate) fn write_value(
371    reader: &Reader,
372    local: usize,
373    format: &Format,
374    text: &mut CellText,
375) -> CellState {
376    if reader.is_null(local) {
377        // A cell always has room for `null`.
378        let _ = text.push_str("null");
379        return CellState::Null;
380    }
381    if !reader.is_supported() {
382        reader.write(local, text, format);
383        return CellState::Unsupported;
384    }
385    let mark = text.len();
386    if reader.write(local, text, format) {
387        return CellState::Value;
388    }
389    text.truncate(mark);
390    // A cell always has room for one byte.
391    let _ = text.push_str("#");
392    CellState::Invalid
393}
394
395/// Widens a number or time column (of `data_type`, not sized by the application) to its
396/// widest value among some of `reader`'s rows, if that is wider; whether it did.
397pub(crate) fn widen_column(
398    format: &mut Format,
399    reader: &Reader,
400    data_type: &arrow_schema::DataType,
401    text: &mut CellText,
402) -> bool {
403    if format.width_set || !is_numeric(data_type) {
404        return false;
405    }
406    let widest = widest(reader, format, spread(reader.len(), WIDEN_ROWS), text);
407    let wider = widest > format.max_chars;
408    if wider {
409        format.max_chars = widest;
410    }
411    wider
412}
413
414/// Sizes a column to its longest value among the first rows `reader` holds and its title
415/// `title`, unless its width was set.
416pub(crate) fn size_column(format: &mut Format, reader: &Reader, title: &str) {
417    if format.width_set {
418        return;
419    }
420    let mut text = CellText::new();
421    let rows = 0..reader.len().min(SAMPLE_ROWS);
422    let widest = widest(reader, format, rows, &mut text).max(title.chars().count());
423    format.max_chars = widest.max(1);
424}
425
426/// The widest of `rows`' values in `reader` as `format` writes them, up to 40 characters
427/// (nulls count as nothing).
428fn widest(
429    reader: &Reader,
430    format: &Format,
431    rows: impl Iterator<Item = usize>,
432    text: &mut CellText,
433) -> usize {
434    let sampling = Format {
435        max_chars: MAX_SAMPLED_CHARS,
436        ..*format
437    };
438    let mut widest = 0;
439    for row in rows {
440        text.clear();
441        if !reader.is_null(row) {
442            reader.write(row, text, &sampling);
443        }
444        widest = widest.max(text.as_str().chars().count());
445    }
446    widest.min(MAX_SAMPLED_CHARS)
447}
448
449/// Up to `most` rows of `len`, spread evenly from the first to the last.
450fn spread(len: usize, most: usize) -> impl Iterator<Item = usize> {
451    let count = len.min(most);
452    // The first, then evenly to the last (one row: the first is the last).
453    (0..count).map(move |i| i * (len - 1) / (count - 1).max(1))
454}
455
456/// An [`ArrowTable`] as a grid's data, scrolled to a first row. Cheap to create each frame.
457#[derive(Debug, Clone, Copy)]
458pub struct ArrowGrid<'a> {
459    table: &'a ArrowTable,
460    first_row: usize,
461    text: Color,
462    dim: Color,
463}
464
465impl GridSource for ArrowGrid<'_> {
466    fn rows(&self) -> usize {
467        self.table.rows().saturating_sub(self.first_row)
468    }
469
470    fn cell(&self, row: usize, column: usize, text: &mut CellText) -> CellStyle {
471        let state = self
472            .table
473            .write_cell(self.first_row.saturating_add(row), column, text);
474        let color = if state == CellState::Value {
475            self.text
476        } else {
477            self.dim
478        };
479        CellStyle::new(color)
480    }
481}
densa_arrow/table.rs

densa_arrow/
table.rs
