Line data Source code
1 : //! Renders Arrow IPC results as box-drawn tables.
2 : //!
3 : //! Cell values are formatted by arrow-cast (`ArrayFormatter`): floats with their shortest round-trip digits,
4 : //! temporal values in ISO 8601, binaries in hex, nested values like `[1, 2]` or `{a: 1}`.
5 :
6 : use super::unicode::{display_width, graphemes};
7 : use arrow_array::timezone::Tz;
8 : use arrow_array::{make_array, Array, ArrayRef, RecordBatch};
9 : use arrow_cast::display::{ArrayFormatter, FormatOptions};
10 : use arrow_schema::{DataType, Schema};
11 : use std::fmt::Write;
12 : use std::io::Cursor;
13 :
14 : #[derive(Clone, Copy, Debug, PartialEq, Eq)]
15 : enum Alignment {
16 : Left,
17 : Right,
18 : }
19 :
20 : const CELL_PADDING: usize = 2;
21 : const MIN_CONTENT_WIDTH: usize = 1;
22 : const MAX_ROW_HEIGHT: usize = 20;
23 : const TRUNCATION_MARKER: &[u8] = b"...";
24 : const NO_RESULTS_MESSAGE: &[u8] = b"No results";
25 : /// Null cells stay empty, nulls inside nested values print as `null`
26 : const FORMAT_OPTIONS: FormatOptions<'static> = FormatOptions::new().with_null("null");
27 :
28 : /// A cell of a rendered table
29 : struct Cell {
30 : /// The escaped cell text
31 : text: Vec<u8>,
32 : /// The horizontal alignment of the text
33 : alignment: Alignment,
34 : }
35 :
36 : /// A result table with formatted cells
37 : struct Table {
38 : /// The escaped column names
39 : headers: Vec<Vec<u8>>,
40 : /// The alignment of every column, right for numeric and boolean columns
41 : alignments: Vec<Alignment>,
42 : /// The rows of cells
43 : rows: Vec<Vec<Cell>>,
44 : }
45 :
46 : /// Make cell text safe for a terminal
47 116 : fn escape_cell_text(text: &[u8]) -> Vec<u8> {
48 : const HEX: &[u8; 16] = b"0123456789abcdef";
49 116 : let mut output = Vec::with_capacity(text.len());
50 1011 : for (i, &value) in text.iter().enumerate() {
51 1011 : if value == b'\n' {
52 25 : output.push(b'\n');
53 986 : } else if value == b'\r' {
54 3 : if i + 1 >= text.len() || text[i + 1] != b'\n' {
55 2 : output.push(b'\n');
56 2 : }
57 983 : } else if value == b'\t' {
58 1 : output.extend_from_slice(b" ");
59 982 : } else if value < 0x20 || value == 0x7f {
60 3 : output.extend_from_slice(b"\\x");
61 3 : output.push(HEX[(value >> 4) as usize]);
62 3 : output.push(HEX[(value & 0x0f) as usize]);
63 979 : } else {
64 979 : output.push(value);
65 979 : }
66 : }
67 116 : output
68 116 : }
69 :
70 : /// Numbers and booleans are right-aligned
71 42 : fn is_right_aligned(data_type: &DataType) -> bool {
72 42 : data_type.is_numeric() || *data_type == DataType::Boolean
73 42 : }
74 :
75 20 : fn create_table(schema: &Schema) -> Table {
76 : Table {
77 42 : headers: schema.fields().iter().map(|f| escape_cell_text(f.name().as_bytes())).collect(),
78 20 : alignments: schema
79 20 : .fields()
80 20 : .iter()
81 42 : .map(|f| if is_right_aligned(f.data_type()) { Alignment::Right } else { Alignment::Left })
82 20 : .collect(),
83 20 : rows: Vec::new(),
84 : }
85 20 : }
86 :
87 : /// arrow-cast resolves named time zones only with chrono-tz, which the shell does not link. Timestamp columns
88 : /// with such a time zone are shown in UTC.
89 42 : fn displayable(column: &ArrayRef) -> ArrayRef {
90 42 : match column.data_type() {
91 2 : DataType::Timestamp(unit, Some(tz)) if tz.parse::<Tz>().is_err() => {
92 1 : let utc = DataType::Timestamp(*unit, Some("+00:00".into()));
93 1 : make_array(column.to_data().into_builder().data_type(utc).build().expect("timestamps keep their layout"))
94 : }
95 41 : _ => column.clone(),
96 : }
97 42 : }
98 :
99 20 : fn append_batch(table: &mut Table, batch: &RecordBatch) {
100 20 : let columns: Vec<ArrayRef> = batch.columns().iter().map(displayable).collect();
101 42 : let nulls: Vec<_> = columns.iter().map(|column| column.logical_nulls()).collect();
102 : // A column that arrow-cast cannot format shows its type
103 20 : let formatters: Vec<_> =
104 42 : columns.iter().map(|column| ArrayFormatter::try_new(column.as_ref(), &FORMAT_OPTIONS).ok()).collect();
105 20 : let mut text = String::new();
106 35 : for row in 0..batch.num_rows() {
107 35 : let mut cells = Vec::with_capacity(columns.len());
108 74 : for column in 0..columns.len() {
109 74 : text.clear();
110 74 : if !nulls[column].as_ref().is_some_and(|nulls| nulls.is_null(row)) {
111 61 : let _ = match &formatters[column] {
112 61 : Some(formatter) => write!(text, "{}", formatter.value(row)),
113 0 : None => write!(text, "{}", columns[column].data_type()),
114 : };
115 13 : }
116 74 : cells.push(Cell { text: escape_cell_text(text.as_bytes()), alignment: table.alignments[column] });
117 : }
118 35 : table.rows.push(cells);
119 : }
120 20 : }
121 :
122 : /// Read an Arrow IPC file, or an IPC stream if the data is no file. Errors are the arrow-rs error after
123 : /// `Invalid: `, empty buffers have a message of their own.
124 44 : fn read_ipc(data: &[u8]) -> Result<Table, String> {
125 44 : if data.is_empty() {
126 15 : return Err("Invalid: Arrow IPC buffer is empty".into());
127 29 : }
128 29 : if let Ok(reader) = arrow_ipc::reader::FileReader::try_new(Cursor::new(data), None) {
129 19 : let mut table = create_table(&reader.schema());
130 19 : for batch in reader {
131 19 : let batch = batch.map_err(|e| format!("Invalid: {}", e))?;
132 19 : append_batch(&mut table, &batch);
133 : }
134 19 : return Ok(table);
135 10 : }
136 10 : check_stream_framing(data)?;
137 1 : let reader =
138 2 : arrow_ipc::reader::StreamReader::try_new(Cursor::new(data), None).map_err(|e| format!("Invalid: {}", e))?;
139 1 : let mut table = create_table(&reader.schema());
140 1 : for batch in reader {
141 1 : let batch = batch.map_err(|e| format!("Invalid: {}", e))?;
142 1 : append_batch(&mut table, &batch);
143 : }
144 1 : Ok(table)
145 44 : }
146 :
147 : /// Check the message framing of an IPC stream: arrow-rs ends a stream at a truncated message prefix, this
148 : /// reports it as an error. Messages are read up to the end-of-stream marker.
149 10 : fn check_stream_framing(data: &[u8]) -> Result<(), String> {
150 20 : let read_i32 = |pos: usize| i32::from_le_bytes(data[pos..pos + 4].try_into().unwrap());
151 10 : let mut pos = 0;
152 : loop {
153 13 : let remaining = data.len() - pos;
154 13 : if remaining == 0 {
155 0 : return Ok(());
156 13 : }
157 13 : if remaining < 4 {
158 1 : return Err(format!("Invalid: Corrupted message, only {} bytes available", remaining));
159 12 : }
160 12 : let continuation = read_i32(pos);
161 12 : pos += 4;
162 7 : let metadata_length = match continuation {
163 : -1 => {
164 9 : let remaining = data.len() - pos;
165 9 : if remaining < 4 {
166 1 : return Err(format!("Invalid: Corrupted metadata length, only {} bytes available", remaining));
167 8 : }
168 8 : let length = read_i32(pos);
169 8 : pos += 4;
170 6 : match length {
171 2 : 0 => return Ok(()),
172 6 : l if l < 0 => return Err("IOError: Invalid IPC message: negative metadata length".into()),
173 5 : l => l as usize,
174 : }
175 : }
176 0 : 0 => return Ok(()),
177 3 : l if l > 0 => l as usize,
178 1 : _ => return Err("IOError: Invalid IPC stream: negative continuation token".into()),
179 : };
180 7 : let remaining = data.len() - pos;
181 7 : if remaining < metadata_length {
182 3 : return Err(format!(
183 3 : "Invalid: Expected to read {} metadata bytes, but only read {}",
184 3 : metadata_length, remaining
185 3 : ));
186 4 : }
187 4 : let message =
188 4 : arrow_ipc::root_as_message(&data[pos..pos + metadata_length]).map_err(|e| format!("Invalid: {}", e))?;
189 4 : pos += metadata_length;
190 4 : let body_length = message.bodyLength().max(0) as usize;
191 4 : let remaining = data.len() - pos;
192 4 : if remaining < body_length {
193 1 : return Err(format!(
194 1 : "IOError: Expected to be able to read {} bytes for message body, got {}",
195 1 : body_length, remaining
196 1 : ));
197 3 : }
198 3 : pos += body_length;
199 : }
200 10 : }
201 :
202 : /// A grapheme cluster of a text
203 : struct Grapheme {
204 : /// The byte offset of the cluster
205 : begin: usize,
206 : /// The byte offset after the cluster
207 : end: usize,
208 : /// The display width in terminal columns
209 : width: usize,
210 : }
211 :
212 200 : fn text_graphemes(text: &[u8]) -> Vec<Grapheme> {
213 200 : graphemes(text)
214 200 : .into_iter()
215 1501 : .map(|(begin, end)| Grapheme { begin, end, width: display_width(&text[begin..end]) })
216 200 : .collect()
217 200 : }
218 :
219 : /// Cut a text to a width and append "..."
220 5 : fn mark_truncated(text: &[u8], width: usize) -> Vec<u8> {
221 5 : let marker_width = width.min(TRUNCATION_MARKER.len());
222 5 : let content_width = width - marker_width;
223 5 : let mut end = 0;
224 5 : let mut rendered = 0;
225 7 : for grapheme in text_graphemes(text) {
226 7 : if rendered + grapheme.width > content_width {
227 5 : break;
228 2 : }
229 2 : rendered += grapheme.width;
230 2 : end = grapheme.end;
231 : }
232 5 : let mut output = text[..end].to_vec();
233 5 : output.extend_from_slice(&TRUNCATION_MARKER[..marker_width]);
234 5 : output
235 5 : }
236 :
237 : /// ASCII whitespace, including vertical tab and form feed
238 1707 : fn is_space(c: u8) -> bool {
239 1707 : matches!(c, b' ' | b'\t' | b'\n' | 0x0b | 0x0c | b'\r')
240 1707 : }
241 :
242 223 : fn wrap_line(text: &[u8], width: usize) -> Vec<Vec<u8>> {
243 223 : let width = width.max(1);
244 223 : if text.is_empty() {
245 28 : return vec![Vec::new()];
246 195 : }
247 195 : let graphemes = text_graphemes(text);
248 1723 : let is_space_grapheme = |g: &Grapheme| g.end - g.begin == 1 && is_space(text[g.begin]);
249 195 : let mut lines = Vec::new();
250 195 : let mut begin = 0;
251 447 : while begin < graphemes.len() {
252 252 : let mut end = begin;
253 252 : let mut rendered: usize = 0;
254 252 : let mut whitespace = None;
255 1718 : while end < graphemes.len() && (rendered == 0 || rendered.saturating_add(graphemes[end].width) <= width) {
256 1466 : rendered += graphemes[end].width;
257 1466 : if is_space_grapheme(&graphemes[end]) {
258 78 : whitespace = Some(end);
259 1388 : }
260 1466 : end += 1;
261 : }
262 252 : if let Some(whitespace) = whitespace {
263 20 : if end < graphemes.len() && whitespace >= begin {
264 5 : end = whitespace + 1;
265 15 : }
266 232 : }
267 252 : let mut content_end = end;
268 257 : while content_end > begin && is_space_grapheme(&graphemes[content_end - 1]) {
269 5 : content_end -= 1;
270 5 : }
271 252 : let byte_begin = graphemes[begin].begin;
272 252 : let byte_end = if content_end > begin { graphemes[content_end - 1].end } else { byte_begin };
273 252 : lines.push(text[byte_begin..byte_end].to_vec());
274 252 : begin = end;
275 : }
276 195 : lines
277 223 : }
278 :
279 169 : fn wrap_text(text: &[u8], width: usize) -> Vec<Vec<u8>> {
280 169 : let mut lines = Vec::new();
281 169 : let mut begin = 0;
282 : loop {
283 1525 : let end = text[begin..].iter().position(|&c| c == b'\n').map(|i| i + begin);
284 223 : lines.extend(wrap_line(&text[begin..end.unwrap_or(text.len())], width));
285 223 : match end {
286 54 : Some(end) => begin = end + 1,
287 169 : None => break,
288 : }
289 : }
290 169 : lines
291 169 : }
292 :
293 34 : fn frame_width(column_count: usize) -> usize {
294 34 : if column_count == 0 {
295 0 : 0
296 : } else {
297 34 : (CELL_PADDING + 1) * column_count + 1
298 : }
299 34 : }
300 :
301 18 : fn visible_column_count(mut column_count: usize, terminal_columns: usize) -> usize {
302 25 : while column_count > 1 && frame_width(column_count) + column_count * MIN_CONTENT_WIDTH > terminal_columns {
303 7 : column_count -= 1;
304 7 : }
305 18 : column_count
306 18 : }
307 :
308 18 : fn resolve_column_widths(table: &Table, terminal_columns: usize) -> Vec<usize> {
309 18 : let column_count = visible_column_count(table.headers.len(), terminal_columns);
310 35 : let mut widths: Vec<usize> = (0..column_count).map(|c| display_width(&table.headers[c]).max(1)).collect();
311 33 : for row in &table.rows {
312 67 : for column in 0..row.len().min(column_count) {
313 94 : for line in wrap_text(&row[column].text, usize::MAX) {
314 94 : widths[column] = widths[column].max(display_width(&line));
315 94 : }
316 : }
317 : }
318 18 : if column_count == 0 {
319 0 : return widths;
320 18 : }
321 18 : let frame = frame_width(column_count);
322 18 : let available = if terminal_columns > frame { terminal_columns - frame } else { column_count };
323 18 : if available < column_count {
324 0 : return vec![MIN_CONTENT_WIDTH; column_count];
325 18 : }
326 : // Only the initial total is compared with the available width, the loop does not check it again
327 18 : let total: usize = widths.iter().sum();
328 18 : if total <= available {
329 9 : return widths;
330 9 : }
331 9 : let mut fixed = vec![false; column_count];
332 : loop {
333 10 : let mut remaining_width = available;
334 10 : let mut remaining_columns = 0;
335 16 : for column in 0..column_count {
336 16 : if fixed[column] {
337 1 : remaining_width = remaining_width.saturating_sub(widths[column]);
338 15 : } else {
339 15 : remaining_columns += 1;
340 15 : }
341 : }
342 10 : if remaining_columns == 0 {
343 0 : break;
344 10 : }
345 10 : let share = MIN_CONTENT_WIDTH.max(remaining_width / remaining_columns);
346 10 : let mut found_small = false;
347 16 : for column in 0..column_count {
348 16 : if !fixed[column] && widths[column] <= share {
349 1 : fixed[column] = true;
350 1 : found_small = true;
351 15 : }
352 : }
353 10 : if found_small {
354 1 : continue;
355 9 : }
356 9 : let base = MIN_CONTENT_WIDTH.max(remaining_width / remaining_columns);
357 9 : let mut excess = remaining_width.saturating_sub(base * remaining_columns);
358 14 : for column in 0..column_count {
359 14 : if !fixed[column] {
360 13 : widths[column] = base + if excess > 0 { 1 } else { 0 };
361 13 : excess = excess.saturating_sub(1);
362 1 : }
363 : }
364 9 : break;
365 : }
366 9 : widths
367 18 : }
368 :
369 54 : fn append_rule(output: &mut Vec<u8>, left: &str, middle: &str, right: &str, fill: &str, widths: &[usize]) {
370 54 : output.extend_from_slice(left.as_bytes());
371 105 : for (column, width) in widths.iter().enumerate() {
372 105 : if column != 0 {
373 51 : output.extend_from_slice(middle.as_bytes());
374 54 : }
375 1281 : for _ in 0..width + 2 {
376 1281 : output.extend_from_slice(fill.as_bytes());
377 1281 : }
378 : }
379 54 : output.extend_from_slice(right.as_bytes());
380 54 : output.push(b'\n');
381 54 : }
382 :
383 193 : fn append_cell(output: &mut Vec<u8>, text: &[u8], width: usize, alignment: Alignment) {
384 193 : let padding = width.saturating_sub(display_width(text));
385 193 : if alignment == Alignment::Right {
386 20 : output.extend(std::iter::repeat_n(b' ', padding));
387 173 : }
388 193 : output.extend_from_slice(text);
389 193 : if alignment == Alignment::Left {
390 173 : output.extend(std::iter::repeat_n(b' ', padding));
391 173 : }
392 193 : }
393 :
394 51 : fn append_row(output: &mut Vec<u8>, cells: &[Cell], widths: &[usize], separator: &str) {
395 51 : let mut wrapped = Vec::with_capacity(widths.len());
396 51 : let mut row_height = 1;
397 102 : for (column, &width) in widths.iter().enumerate() {
398 102 : let text: &[u8] = cells.get(column).map_or(&[], |c| &c.text);
399 102 : let mut lines = wrap_text(text, width);
400 102 : if lines.len() > MAX_ROW_HEIGHT {
401 1 : lines.truncate(MAX_ROW_HEIGHT);
402 1 : let last = lines.last_mut().unwrap();
403 1 : *last = mark_truncated(last, width);
404 101 : }
405 102 : row_height = row_height.max(lines.len());
406 102 : wrapped.push(lines);
407 : }
408 117 : for line in 0..row_height {
409 117 : output.extend_from_slice("│ ".as_bytes());
410 193 : for (column, &width) in widths.iter().enumerate() {
411 193 : if column != 0 {
412 76 : output.push(b' ');
413 76 : output.extend_from_slice(separator.as_bytes());
414 76 : output.push(b' ');
415 117 : }
416 193 : let text: &[u8] = wrapped[column].get(line).map_or(&[], |l| l);
417 193 : let alignment = cells.get(column).map_or(Alignment::Left, |c| c.alignment);
418 193 : append_cell(output, text, width, alignment);
419 : }
420 117 : output.extend_from_slice(" │\n".as_bytes());
421 : }
422 51 : }
423 :
424 20 : fn render_table(table: &Table, terminal_columns: usize) -> Vec<u8> {
425 20 : if table.headers.is_empty() {
426 2 : return NO_RESULTS_MESSAGE.to_vec();
427 18 : }
428 18 : let widths = resolve_column_widths(table, terminal_columns);
429 18 : let mut output = Vec::new();
430 18 : append_rule(&mut output, "╭", "┬", "╮", "─", &widths);
431 18 : let headers: Vec<Cell> = (0..widths.len())
432 35 : .map(|column| {
433 35 : let mut header = table.headers[column].clone();
434 35 : if column + 1 == widths.len() && widths.len() < table.headers.len() {
435 4 : header = mark_truncated(&header, widths[column]);
436 31 : }
437 35 : Cell { text: header, alignment: Alignment::Left }
438 35 : })
439 18 : .collect();
440 18 : append_row(&mut output, &headers, &widths, "│");
441 18 : append_rule(&mut output, "╞", "╪", "╡", "═", &widths);
442 33 : for row in &table.rows {
443 33 : append_row(&mut output, row, &widths, "│");
444 33 : }
445 18 : append_rule(&mut output, "╰", "┴", "╯", "─", &widths);
446 18 : output.pop();
447 18 : output
448 20 : }
449 :
450 : /// Renders Arrow IPC results for a terminal width
451 : #[derive(Clone, Debug)]
452 : pub struct ArrowRenderer {
453 : /// The terminal width in columns, at least 1
454 : terminal_columns: u32,
455 : }
456 :
457 : impl ArrowRenderer {
458 140 : pub fn new(terminal_columns: u32) -> Self {
459 140 : Self { terminal_columns: terminal_columns.max(1) }
460 140 : }
461 :
462 3 : pub fn resize(&mut self, terminal_columns: u32) {
463 3 : self.terminal_columns = terminal_columns.max(1);
464 3 : }
465 :
466 547 : pub fn terminal_columns(&self) -> u32 {
467 547 : self.terminal_columns
468 547 : }
469 :
470 : /// Render an Arrow IPC file or stream. Errors are Arrow error texts.
471 44 : pub fn render_ipc(&self, data: &[u8]) -> Result<Vec<u8>, String> {
472 44 : let table = read_ipc(data)?;
473 20 : Ok(render_table(&table, self.terminal_columns as usize))
474 44 : }
475 : }
476 :
477 : impl Default for ArrowRenderer {
478 1 : fn default() -> Self {
479 1 : Self::new(100)
480 1 : }
481 : }
|
