Skip to content

Commit b7b90fc

Browse files
nevi-mealamb
authored andcommitted
ARROW-11898: [Rust] Pretty print columns
Adds a way to also pretty-print a slice of columns, as a convenience to aviod creating a single-column record batch. Closes #9650 from nevi-me/pretty-print-col Authored-by: Neville Dipale <nevilledips@gmail.com> Signed-off-by: Andrew Lamb <andrew@nerdnetworks.org>
1 parent 46161d2 commit b7b90fc

1 file changed

Lines changed: 60 additions & 1 deletion

File tree

rust/arrow/src/util/pretty.rs

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
//! Utilities for printing record batches. Note this module is not
1919
//! available unless `feature = "prettyprint"` is enabled.
2020
21-
use crate::record_batch::RecordBatch;
21+
use crate::{array::ArrayRef, record_batch::RecordBatch};
2222

2323
use prettytable::format;
2424
use prettytable::{Cell, Row, Table};
@@ -32,12 +32,23 @@ pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<String> {
3232
Ok(create_table(results)?.to_string())
3333
}
3434

35+
///! Create a visual representation of columns
36+
pub fn pretty_format_columns(col_name: &str, results: &[ArrayRef]) -> Result<String> {
37+
Ok(create_column(col_name, results)?.to_string())
38+
}
39+
3540
///! Prints a visual representation of record batches to stdout
3641
pub fn print_batches(results: &[RecordBatch]) -> Result<()> {
3742
create_table(results)?.printstd();
3843
Ok(())
3944
}
4045

46+
///! Prints a visual representation of a list of column to stdout
47+
pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<()> {
48+
create_column(col_name, results)?.printstd();
49+
Ok(())
50+
}
51+
4152
///! Convert a series of record batches into a table
4253
fn create_table(results: &[RecordBatch]) -> Result<Table> {
4354
let mut table = Table::new();
@@ -69,6 +80,28 @@ fn create_table(results: &[RecordBatch]) -> Result<Table> {
6980
Ok(table)
7081
}
7182

83+
fn create_column(field: &str, columns: &[ArrayRef]) -> Result<Table> {
84+
let mut table = Table::new();
85+
table.set_format(*format::consts::FORMAT_NO_LINESEP_WITH_TITLE);
86+
87+
if columns.is_empty() {
88+
return Ok(table);
89+
}
90+
91+
let header = vec![Cell::new(field)];
92+
table.set_titles(Row::new(header));
93+
94+
for col in columns {
95+
for row in 0..col.len() {
96+
let mut cells = Vec::new();
97+
cells.push(Cell::new(&array_value_to_string(&col, row)?));
98+
table.add_row(Row::new(cells));
99+
}
100+
}
101+
102+
Ok(table)
103+
}
104+
72105
#[cfg(test)]
73106
mod tests {
74107
use crate::{
@@ -132,6 +165,32 @@ mod tests {
132165
Ok(())
133166
}
134167

168+
#[test]
169+
fn test_pretty_format_columns() -> Result<()> {
170+
let columns = vec![
171+
Arc::new(array::StringArray::from(vec![
172+
Some("a"),
173+
Some("b"),
174+
None,
175+
Some("d"),
176+
])) as ArrayRef,
177+
Arc::new(array::StringArray::from(vec![Some("e"), None, Some("g")])),
178+
];
179+
180+
let table = pretty_format_columns("a", &columns)?;
181+
182+
let expected = vec![
183+
"+---+", "| a |", "+---+", "| a |", "| b |", "| |", "| d |", "| e |",
184+
"| |", "| g |", "+---+",
185+
];
186+
187+
let actual: Vec<&str> = table.lines().collect();
188+
189+
assert_eq!(expected, actual, "Actual result:\n{}", table);
190+
191+
Ok(())
192+
}
193+
135194
#[test]
136195
fn test_pretty_format_null() {
137196
let schema = Arc::new(Schema::new(vec![

0 commit comments

Comments
 (0)