arrs/output/
table.rs

1use std::io::Write;
2
3use arrow_array::RecordBatch;
4use arrow_schema::SchemaRef;
5use comfy_table::{Table, presets};
6
7use crate::Result;
8use crate::cli::BinaryFormat;
9use crate::output::{RowWriter, value};
10
11#[derive(Debug, Copy, Clone, PartialEq, Eq)]
12pub enum TableStyle {
13    Pretty,
14    Plain,
15}
16
17/// Writes a comfy-table-rendered grid to `W` once `finish()` is called.
18///
19/// Buffers every row from `write_batch` because the renderer needs the full
20/// data set to compute column widths. This is fine for the metadata commands
21/// (small row counts by construction), but is the reason `Format::Table` is
22/// not the default for `cat`/`head`/etc., which want streaming behaviour.
23pub struct TableRowWriter<W: Write> {
24    writer: W,
25    binary_format: BinaryFormat,
26    table_style: TableStyle,
27    schema: Option<SchemaRef>,
28    rows: Vec<Vec<String>>,
29}
30
31impl<W: Write> TableRowWriter<W> {
32    pub fn new(writer: W, binary_format: BinaryFormat, table_style: TableStyle) -> Self {
33        Self {
34            writer,
35            binary_format,
36            table_style,
37            schema: None,
38            rows: Vec::new(),
39        }
40    }
41}
42
43impl<W: Write> RowWriter for TableRowWriter<W> {
44    fn start(&mut self, schema: &SchemaRef) -> Result<()> {
45        self.schema = Some(schema.clone());
46        Ok(())
47    }
48
49    fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
50        let rows = (0..batch.num_rows())
51            .map(|row| {
52                batch
53                    .columns()
54                    .iter()
55                    .map(|column| value::table_cell(column.as_ref(), row, self.binary_format))
56                    .collect::<Result<Vec<_>>>()
57            })
58            .collect::<Result<Vec<_>>>()?;
59        self.rows.extend(rows);
60        Ok(())
61    }
62
63    fn finish(&mut self) -> Result<()> {
64        let schema = self
65            .schema
66            .as_ref()
67            .expect("start() must be called before finish()");
68
69        let mut table = Table::new();
70        // Pretty borders only when stdout is a real terminal — pipelines and
71        // captured-output test runs get an ASCII grid that's grep-friendly.
72        let preset = match self.table_style {
73            TableStyle::Plain => presets::ASCII_FULL,
74            TableStyle::Pretty => presets::UTF8_FULL,
75        };
76        table.load_preset(preset);
77
78        let headers: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
79        table.set_header(headers);
80
81        table.add_rows(self.rows.drain(..));
82
83        writeln!(self.writer, "{table}")?;
84        self.writer.flush()?;
85        Ok(())
86    }
87}
arrs/output/table.rs

arrs/output/
table.rs
