Skip to content

Commit c87a4d9

Browse files
lewiszlwalamb
andauthored
Add pretty_format_batches_with_schema function (#7642)
# Which issue does this PR close? Closes #NNN. # Rationale for this change Improve empty batches format. # What changes are included in this PR? Add new `pretty_format_batches_with_schema` function. # Are there any user-facing changes? Yes. --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent 2f2e705 commit c87a4d9

File tree

1 file changed

+63
-19
lines changed

1 file changed

+63
-19
lines changed

arrow-cast/src/pretty.rs

Lines changed: 63 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use std::fmt::Display;
2727
use comfy_table::{Cell, Table};
2828

2929
use arrow_array::{Array, ArrayRef, RecordBatch};
30-
use arrow_schema::ArrowError;
30+
use arrow_schema::{ArrowError, SchemaRef};
3131

3232
use crate::display::{ArrayFormatter, FormatOptions};
3333

@@ -65,6 +65,38 @@ pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<impl Display, Ar
6565
pretty_format_batches_with_options(results, &options)
6666
}
6767

68+
/// Create a visual representation of [`RecordBatch`]es with a provided schema.
69+
///
70+
/// Useful to display empty batches.
71+
///
72+
/// # Example
73+
/// ```
74+
/// # use std::sync::Arc;
75+
/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
76+
/// # use arrow_cast::pretty::pretty_format_batches_with_schema;
77+
/// # use arrow_schema::{DataType, Field, Schema};
78+
/// let schema = Arc::new(Schema::new(vec![
79+
/// Field::new("a", DataType::Int32, false),
80+
/// Field::new("b", DataType::Utf8, true),
81+
/// ]));
82+
/// // Note, returned object implements `Display`
83+
/// let pretty_table = pretty_format_batches_with_schema(schema, &[]).unwrap();
84+
/// let table_str = format!("Batches:\n{pretty_table}");
85+
/// assert_eq!(table_str,
86+
/// r#"Batches:
87+
/// +---+---+
88+
/// | a | b |
89+
/// +---+---+
90+
/// +---+---+"#);
91+
/// ```
92+
pub fn pretty_format_batches_with_schema(
93+
schema: SchemaRef,
94+
results: &[RecordBatch],
95+
) -> Result<impl Display, ArrowError> {
96+
let options = FormatOptions::default().with_display_error(true);
97+
create_table(Some(schema), results, &options)
98+
}
99+
68100
/// Create a visual representation of [`RecordBatch`]es with formatting options.
69101
///
70102
/// # Arguments
@@ -99,7 +131,7 @@ pub fn pretty_format_batches_with_options(
99131
results: &[RecordBatch],
100132
options: &FormatOptions,
101133
) -> Result<impl Display, ArrowError> {
102-
create_table(results, options)
134+
create_table(None, results, options)
103135
}
104136

105137
/// Create a visual representation of [`ArrayRef`]
@@ -139,29 +171,41 @@ pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<(), ArrowEr
139171
}
140172

141173
/// Convert a series of record batches into a table
142-
fn create_table(results: &[RecordBatch], options: &FormatOptions) -> Result<Table, ArrowError> {
174+
fn create_table(
175+
schema_opt: Option<SchemaRef>,
176+
results: &[RecordBatch],
177+
options: &FormatOptions,
178+
) -> Result<Table, ArrowError> {
143179
let mut table = Table::new();
144180
table.load_preset("||--+-++| ++++++");
145181

146-
if results.is_empty() {
147-
return Ok(table);
148-
}
149-
150-
let schema = results[0].schema();
151-
152-
let mut header = Vec::new();
153-
for field in schema.fields() {
154-
if options.types_info() {
155-
header.push(Cell::new(format!(
156-
"{}\n{}",
157-
field.name(),
158-
field.data_type()
159-
)))
182+
let schema_opt = schema_opt.or_else(|| {
183+
if results.is_empty() {
184+
None
160185
} else {
161-
header.push(Cell::new(field.name()));
186+
Some(results[0].schema())
187+
}
188+
});
189+
190+
if let Some(schema) = schema_opt {
191+
let mut header = Vec::new();
192+
for field in schema.fields() {
193+
if options.types_info() {
194+
header.push(Cell::new(format!(
195+
"{}\n{}",
196+
field.name(),
197+
field.data_type()
198+
)))
199+
} else {
200+
header.push(Cell::new(field.name()));
201+
}
162202
}
203+
table.set_header(header);
204+
}
205+
206+
if results.is_empty() {
207+
return Ok(table);
163208
}
164-
table.set_header(header);
165209

166210
for batch in results {
167211
let formatters = batch

0 commit comments

Comments
 (0)