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/// Widest a column is sized from its content; longer text is cut with `…`.
12const MAX_SAMPLED_CHARS: usize = 40;
13/// Width of a column with no rows yet, unless its title is wider.
14const DEFAULT_CHARS: usize = 8;
15
16/// What [`ArrowTable::write_cell`] wrote.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18#[non_exhaustive]
19pub enum CellState {
20    /// The value.
21    Value,
22    /// `null`, for a missing value.
23    Null,
24    /// `?`, for a type this crate does not format (lists, structs, maps, ...).
25    Unsupported,
26    /// `#`, for a value that cannot be shown: longer than a cell holds (as `1e300` without
27    /// fixed decimals), or invalid (a decimal scale beyond 38, a time of day outside a day).
28    Invalid,
29    /// Nothing: the row or column does not exist.
30    OutOfRange,
31}
32
33/// Why a batch could not be added, or a column setting applied.
34#[derive(Debug, Clone, PartialEq, Eq)]
35#[non_exhaustive]
36pub enum TableError {
37    /// No batch was given, so there is no schema.
38    NoBatches,
39    /// The batch has a different number of columns.
40    ColumnCount {
41        /// Columns in the table.
42        expected: usize,
43        /// Columns in the batch.
44        found: usize,
45    },
46    /// The batch's column at this index has another data type.
47    ColumnMismatch {
48        /// Index of the first differing column.
49        column: usize,
50    },
51    /// There is no column at this index.
52    NoSuchColumn(usize),
53}
54
55impl fmt::Display for TableError {
56    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
57        match self {
58            Self::NoBatches => f.write_str("no record batch was given, so there is no schema"),
59            Self::ColumnCount { expected, found } => {
60                write!(f, "the batch has {found} columns, the table {expected}")
61            }
62            Self::ColumnMismatch { column } => {
63                write!(f, "column {column} of the batch has another data type")
64            }
65            Self::NoSuchColumn(column) => write!(f, "there is no column {column}"),
66        }
67    }
68}
69
70impl std::error::Error for TableError {}
71
72/// Arrow record batches shown as one table: rows run through the batches in order.
73///
74/// Formatting is chosen per column from its data type when a batch is added, so writing a
75/// cell is a lookup and a format into a stack buffer, without allocating. Batches can be
76/// appended while the table is shown (a stream); their arrays are shared, not copied.
77///
78/// Formatted types: booleans, integers, floats, `Decimal128`, strings (`Utf8`, `LargeUtf8`,
79/// `Utf8View`), dictionaries of those, dates, times, timestamps (with a time zone: in UTC,
80/// marked `Z`) and durations. Other types (lists, structs, maps, binary, intervals) show
81/// `?`.
82///
83/// A batch is accepted when its columns have the table's data types in order; field names,
84/// nullability and metadata may differ. A timestamp's time zone is part of its data type.
85#[derive(Debug, Clone)]
86pub struct ArrowTable {
87    schema: SchemaRef,
88    /// Each non-empty batch's columns.
89    batches: Vec<Vec<Reader>>,
90    /// First row of each batch in `batches`.
91    starts: Vec<usize>,
92    rows: usize,
93    formats: Vec<Format>,
94    /// Whether the columns have been sized from a batch's content.
95    sized: bool,
96}
97
98impl ArrowTable {
99    /// An empty table for batches with `schema`; columns are sized by their titles until
100    /// the first rows arrive.
101    pub fn new(schema: SchemaRef) -> Self {
102        let formats = schema
103            .fields()
104            .iter()
105            .map(|field| Format {
106                decimals: None,
107                max_chars: field.name().chars().count().max(DEFAULT_CHARS),
108                width_set: false,
109            })
110            .collect();
111        Self {
112            schema,
113            batches: Vec::new(),
114            starts: Vec::new(),
115            rows: 0,
116            formats,
117            sized: false,
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), except columns sized with
139    /// [`set_column_chars`](Self::set_column_chars).
140    ///
141    /// # Errors
142    /// A column error if `batch` has other columns or data types; the table is unchanged.
143    pub fn push(&mut self, batch: RecordBatch) -> Result<(), TableError> {
144        self.check(&batch)?;
145        if batch.num_rows() == 0 {
146            return Ok(());
147        }
148        let readers: Vec<Reader> = batch.columns().iter().cloned().map(Reader::new).collect();
149        if !self.sized {
150            self.size_columns(&readers);
151        }
152        self.starts.push(self.rows);
153        self.rows += batch.num_rows();
154        self.batches.push(readers);
155        Ok(())
156    }
157
158    /// The table's schema.
159    pub fn schema(&self) -> &SchemaRef {
160        &self.schema
161    }
162
163    /// Rows in all batches.
164    pub fn rows(&self) -> usize {
165        self.rows
166    }
167
168    /// Grid columns: titled by field name, sized by content, numbers and times aligned right.
169    pub fn columns(&self) -> Vec<Column> {
170        self.schema
171            .fields()
172            .iter()
173            .zip(&self.formats)
174            .map(|(field, format)| {
175                let align = if is_numeric(field.data_type()) {
176                    Align::Right
177                } else {
178                    Align::Left
179                };
180                let chars = u32::try_from(format.max_chars).unwrap_or(u32::MAX);
181                Column::new(field.name().clone(), chars, align)
182            })
183            .collect()
184    }
185
186    /// Prints floating-point values of `column` with `decimals` fixed decimal places
187    /// (`None`: the shortest exact form). A column sized by the first rows is sized again
188    /// from them, as now printed.
189    ///
190    /// # Errors
191    /// [`TableError::NoSuchColumn`] if there is no such column.
192    pub fn set_decimals(
193        &mut self,
194        column: usize,
195        decimals: Option<usize>,
196    ) -> Result<(), TableError> {
197        self.format_mut(column)?.decimals = decimals;
198        if let (Some(reader), Some(field)) = (
199            self.batches.first().and_then(|b| b.get(column)),
200            self.schema.fields().get(column),
201        ) {
202            let format = &mut self.formats[column];
203            size_column(format, reader, field.name());
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        if reader.is_null(local) {
249            // A cell always has room for `null`.
250            let _ = text.push_str("null");
251            return CellState::Null;
252        }
253        if !reader.is_supported() {
254            reader.write(local, text, format);
255            return CellState::Unsupported;
256        }
257        let mark = text.len();
258        if reader.write(local, text, format) {
259            return CellState::Value;
260        }
261        text.truncate(mark);
262        // A cell always has room for one byte.
263        let _ = text.push_str("#");
264        CellState::Invalid
265    }
266
267    /// A [`GridSource`] showing the table from `first_row` (scrolled), in `palette`'s
268    /// text color, with nulls and unsupported values dimmed.
269    pub fn view<'a>(&'a self, first_row: usize, palette: &Palette) -> ArrowGrid<'a> {
270        ArrowGrid {
271            table: self,
272            first_row,
273            text: palette.text,
274            dim: palette.dim_text,
275        }
276    }
277
278    fn check(&self, batch: &RecordBatch) -> Result<(), TableError> {
279        let ours = self.schema.fields();
280        let theirs = batch.schema_ref().fields();
281        if ours.len() != theirs.len() {
282            return Err(TableError::ColumnCount {
283                expected: ours.len(),
284                found: theirs.len(),
285            });
286        }
287        match ours
288            .iter()
289            .zip(theirs.iter())
290            .position(|(a, b)| a.data_type() != b.data_type())
291        {
292            Some(column) => Err(TableError::ColumnMismatch { column }),
293            None => Ok(()),
294        }
295    }
296
297    /// Widens each column to its longest value among the first rows (and its title).
298    fn size_columns(&mut self, readers: &[Reader]) {
299        for ((reader, format), field) in readers
300            .iter()
301            .zip(&mut self.formats)
302            .zip(self.schema.fields().iter())
303        {
304            size_column(format, reader, field.name());
305        }
306        self.sized = true;
307    }
308
309    /// Whether rows have sized the columns.
310    pub(crate) fn is_sized(&self) -> bool {
311        self.sized
312    }
313
314    /// Batches with rows held, in the order pushed.
315    pub(crate) fn stored_batches(&self) -> usize {
316        self.batches.len()
317    }
318
319    /// Drops every batch, keeping the schema and the column formats and widths.
320    pub(crate) fn clear_batches(&mut self) {
321        self.batches.clear();
322        self.starts.clear();
323        self.rows = 0;
324    }
325
326    fn format_mut(&mut self, column: usize) -> Result<&mut Format, TableError> {
327        self.formats
328            .get_mut(column)
329            .ok_or(TableError::NoSuchColumn(column))
330    }
331}
332
333/// Sizes a column to its longest value among the first rows `reader` holds and its title
334/// `title`, unless its width was set.
335fn size_column(format: &mut Format, reader: &Reader, title: &str) {
336    if format.width_set {
337        return;
338    }
339    let sampling = Format {
340        max_chars: MAX_SAMPLED_CHARS,
341        ..*format
342    };
343    let mut text = CellText::new();
344    let mut widest = title.chars().count();
345    for row in 0..reader.len().min(SAMPLE_ROWS) {
346        text.clear();
347        if !reader.is_null(row) {
348            reader.write(row, &mut text, &sampling);
349        }
350        widest = widest.max(text.as_str().chars().count());
351    }
352    format.max_chars = widest.max(1);
353}
354
355/// An [`ArrowTable`] as a grid's data, scrolled to a first row. Cheap to create each frame.
356#[derive(Debug, Clone, Copy)]
357pub struct ArrowGrid<'a> {
358    table: &'a ArrowTable,
359    first_row: usize,
360    text: Color,
361    dim: Color,
362}
363
364impl GridSource for ArrowGrid<'_> {
365    fn rows(&self) -> usize {
366        self.table.rows().saturating_sub(self.first_row)
367    }
368
369    fn cell(&self, row: usize, column: usize, text: &mut CellText) -> CellStyle {
370        let state = self
371            .table
372            .write_cell(self.first_row.saturating_add(row), column, text);
373        let color = if state == CellState::Value {
374            self.text
375        } else {
376            self.dim
377        };
378        CellStyle::new(color)
379    }
380}
densa_arrow/table.rs

densa_arrow/
table.rs
