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 = schema
110            .fields()
111            .iter()
112            .map(|field| Format {
113                decimals: None,
114                max_chars: field.name().chars().count().max(DEFAULT_CHARS),
115                width_set: false,
116            })
117            .collect();
118        Self {
119            schema,
120            batches: Vec::new(),
121            starts: Vec::new(),
122            rows: 0,
123            formats,
124            sized: false,
125            widths: 0,
126        }
127    }
128
129    /// A table of `batches`, with the first batch's schema.
130    ///
131    /// # Errors
132    /// [`TableError::NoBatches`] for none; a column error if a batch's schema differs.
133    pub fn from_batches(
134        batches: impl IntoIterator<Item = RecordBatch>,
135    ) -> Result<Self, TableError> {
136        let mut batches = batches.into_iter();
137        let first = batches.next().ok_or(TableError::NoBatches)?;
138        let mut table = Self::new(first.schema());
139        for batch in std::iter::once(first).chain(batches) {
140            table.push(batch)?;
141        }
142        Ok(table)
143    }
144
145    /// Appends `batch`'s rows. The first rows to arrive size the columns (their titles and
146    /// up to the first 1,000 values, at most 40 characters); later rows widen a number or time
147    /// column whose values no longer fit (measuring some rows of each batch, the last among
148    /// them), never narrowing it. Columns sized with
149    /// [`set_column_chars`](Self::set_column_chars) keep their width.
150    ///
151    /// # Errors
152    /// A column error if `batch` has other columns or data types; the table is unchanged.
153    pub fn push(&mut self, batch: RecordBatch) -> Result<(), TableError> {
154        self.check(&batch)?;
155        if batch.num_rows() == 0 {
156            return Ok(());
157        }
158        let readers: Vec<Reader> = batch.columns().iter().cloned().map(Reader::new).collect();
159        if self.sized {
160            self.widen_columns(&readers);
161        } else {
162            self.size_columns(&readers);
163        }
164        self.starts.push(self.rows);
165        self.rows += batch.num_rows();
166        self.batches.push(readers);
167        Ok(())
168    }
169
170    /// The table's schema.
171    pub fn schema(&self) -> &SchemaRef {
172        &self.schema
173    }
174
175    /// Rows in all batches.
176    pub fn rows(&self) -> usize {
177        self.rows
178    }
179
180    /// Grid columns: titled by field name, sized by content, numbers and times aligned right.
181    pub fn columns(&self) -> Vec<Column> {
182        self.schema
183            .fields()
184            .iter()
185            .zip(&self.formats)
186            .map(|(field, format)| {
187                let align = if is_numeric(field.data_type()) {
188                    Align::Right
189                } else {
190                    Align::Left
191                };
192                let chars = u32::try_from(format.max_chars).unwrap_or(u32::MAX);
193                Column::new(field.name().clone(), chars, align)
194            })
195            .collect()
196    }
197
198    /// Prints floating-point values of `column` with `decimals` fixed decimal places
199    /// (`None`: the shortest exact form). A column sized by the first rows is sized again
200    /// from them, as now printed.
201    ///
202    /// # Errors
203    /// [`TableError::NoSuchColumn`] if there is no such column.
204    pub fn set_decimals(
205        &mut self,
206        column: usize,
207        decimals: Option<usize>,
208    ) -> Result<(), TableError> {
209        self.format_mut(column)?.decimals = decimals;
210        if let (Some(reader), Some(field)) = (
211            self.batches.first().and_then(|b| b.get(column)),
212            self.schema.fields().get(column),
213        ) {
214            let format = &mut self.formats[column];
215            size_column(format, reader, field.name());
216            // Later batches may have widened it: measure them again as now printed.
217            if is_numeric(field.data_type()) && !format.width_set {
218                let mut text = CellText::new();
219                for readers in &self.batches[1..] {
220                    let rows = spread(readers[column].len(), WIDEN_ROWS);
221                    let widest = widest(&readers[column], format, rows, &mut text);
222                    format.max_chars = format.max_chars.max(widest);
223                }
224            }
225        }
226        Ok(())
227    }
228
229    /// Sets `column` to `chars` characters wide (1 to [`CELL_BYTES`], the most a cell can
230    /// hold): [`columns`](Self::columns) reports it, longer text is cut to fit ending in
231    /// `…`, and the first rows to arrive no longer size it.
232    ///
233    /// # Errors
234    /// [`TableError::NoSuchColumn`] if there is no such column.
235    pub fn set_column_chars(&mut self, column: usize, chars: usize) -> Result<(), TableError> {
236        let format = self.format_mut(column)?;
237        format.max_chars = chars.clamp(1, CELL_BYTES);
238        format.width_set = true;
239        Ok(())
240    }
241
242    /// Appends the value at `row` of `column` to `text`, or a placeholder (see
243    /// [`CellState`]), never a partial value. Never allocates.
244    pub fn write_cell(&self, row: usize, column: usize, text: &mut CellText) -> CellState {
245        if row >= self.rows {
246            return CellState::OutOfRange;
247        }
248        // At least one start is <= row, since row < rows and starts[0] == 0.
249        let batch = self.starts.partition_point(|&start| start <= row) - 1;
250        self.write_at(batch, row - self.starts[batch], column, text)
251    }
252
253    /// [`write_cell`](Self::write_cell) for row `local` of stored batch `batch` (the batches
254    /// with rows, in the order pushed).
255    pub(crate) fn write_at(
256        &self,
257        batch: usize,
258        local: usize,
259        column: usize,
260        text: &mut CellText,
261    ) -> CellState {
262        let reader = self.batches.get(batch).and_then(|b| b.get(column));
263        let (Some(reader), Some(format)) = (reader, self.formats.get(column)) else {
264            return CellState::OutOfRange;
265        };
266        if local >= reader.len() {
267            return CellState::OutOfRange;
268        }
269        if reader.is_null(local) {
270            // A cell always has room for `null`.
271            let _ = text.push_str("null");
272            return CellState::Null;
273        }
274        if !reader.is_supported() {
275            reader.write(local, text, format);
276            return CellState::Unsupported;
277        }
278        let mark = text.len();
279        if reader.write(local, text, format) {
280            return CellState::Value;
281        }
282        text.truncate(mark);
283        // A cell always has room for one byte.
284        let _ = text.push_str("#");
285        CellState::Invalid
286    }
287
288    /// A [`GridSource`] showing the table from `first_row` (scrolled), in `palette`'s
289    /// text color, with nulls and unsupported values dimmed.
290    pub fn view<'a>(&'a self, first_row: usize, palette: &Palette) -> ArrowGrid<'a> {
291        ArrowGrid {
292            table: self,
293            first_row,
294            text: palette.text,
295            dim: palette.dim_text,
296        }
297    }
298
299    fn check(&self, batch: &RecordBatch) -> Result<(), TableError> {
300        let ours = self.schema.fields();
301        let theirs = batch.schema_ref().fields();
302        if ours.len() != theirs.len() {
303            return Err(TableError::ColumnCount {
304                expected: ours.len(),
305                found: theirs.len(),
306            });
307        }
308        match ours
309            .iter()
310            .zip(theirs.iter())
311            .position(|(a, b)| a.data_type() != b.data_type())
312        {
313            Some(column) => Err(TableError::ColumnMismatch { column }),
314            None => Ok(()),
315        }
316    }
317
318    /// Widens each column to its longest value among the first rows (and its title).
319    fn size_columns(&mut self, readers: &[Reader]) {
320        for ((reader, format), field) in readers
321            .iter()
322            .zip(&mut self.formats)
323            .zip(self.schema.fields().iter())
324        {
325            size_column(format, reader, field.name());
326        }
327        self.sized = true;
328        self.widths = self.widths.wrapping_add(1);
329    }
330
331    /// Widens each number and time column (not sized by the application) to its widest
332    /// value among some of `readers`' rows, if that is wider than the column.
333    fn widen_columns(&mut self, readers: &[Reader]) {
334        let mut text = CellText::new();
335        let mut widened = false;
336        for ((reader, format), field) in readers
337            .iter()
338            .zip(&mut self.formats)
339            .zip(self.schema.fields().iter())
340        {
341            if format.width_set || !is_numeric(field.data_type()) {
342                continue;
343            }
344            let widest = widest(reader, format, spread(reader.len(), WIDEN_ROWS), &mut text);
345            if widest > format.max_chars {
346                format.max_chars = widest;
347                widened = true;
348            }
349        }
350        if widened {
351            self.widths = self.widths.wrapping_add(1);
352        }
353    }
354
355    /// Changes whenever content changed a column's width (a revision, not a count).
356    pub(crate) fn widths(&self) -> u64 {
357        self.widths
358    }
359
360    /// Batches with rows held, in the order pushed.
361    pub(crate) fn stored_batches(&self) -> usize {
362        self.batches.len()
363    }
364
365    /// Drops every batch, keeping the schema and the column formats and widths.
366    pub(crate) fn clear_batches(&mut self) {
367        self.batches.clear();
368        self.starts.clear();
369        self.rows = 0;
370    }
371
372    fn format_mut(&mut self, column: usize) -> Result<&mut Format, TableError> {
373        self.formats
374            .get_mut(column)
375            .ok_or(TableError::NoSuchColumn(column))
376    }
377}
378
379/// Sizes a column to its longest value among the first rows `reader` holds and its title
380/// `title`, unless its width was set.
381fn size_column(format: &mut Format, reader: &Reader, title: &str) {
382    if format.width_set {
383        return;
384    }
385    let mut text = CellText::new();
386    let rows = 0..reader.len().min(SAMPLE_ROWS);
387    let widest = widest(reader, format, rows, &mut text).max(title.chars().count());
388    format.max_chars = widest.max(1);
389}
390
391/// The widest of `rows`' values in `reader` as `format` writes them, up to 40 characters
392/// (nulls count as nothing).
393fn widest(
394    reader: &Reader,
395    format: &Format,
396    rows: impl Iterator<Item = usize>,
397    text: &mut CellText,
398) -> usize {
399    let sampling = Format {
400        max_chars: MAX_SAMPLED_CHARS,
401        ..*format
402    };
403    let mut widest = 0;
404    for row in rows {
405        text.clear();
406        if !reader.is_null(row) {
407            reader.write(row, text, &sampling);
408        }
409        widest = widest.max(text.as_str().chars().count());
410    }
411    widest.min(MAX_SAMPLED_CHARS)
412}
413
414/// Up to `most` rows of `len`, spread evenly from the first to the last.
415fn spread(len: usize, most: usize) -> impl Iterator<Item = usize> {
416    let count = len.min(most);
417    // The first, then evenly to the last (one row: the first is the last).
418    (0..count).map(move |i| i * (len - 1) / (count - 1).max(1))
419}
420
421/// An [`ArrowTable`] as a grid's data, scrolled to a first row. Cheap to create each frame.
422#[derive(Debug, Clone, Copy)]
423pub struct ArrowGrid<'a> {
424    table: &'a ArrowTable,
425    first_row: usize,
426    text: Color,
427    dim: Color,
428}
429
430impl GridSource for ArrowGrid<'_> {
431    fn rows(&self) -> usize {
432        self.table.rows().saturating_sub(self.first_row)
433    }
434
435    fn cell(&self, row: usize, column: usize, text: &mut CellText) -> CellStyle {
436        let state = self
437            .table
438            .write_cell(self.first_row.saturating_add(row), column, text);
439        let color = if state == CellState::Value {
440            self.text
441        } else {
442            self.dim
443        };
444        CellStyle::new(color)
445    }
446}
densa_arrow/table.rs

densa_arrow/
table.rs
