@@ -27,7 +27,7 @@ use std::fmt::Display;
2727use comfy_table:: { Cell , Table } ;
2828
2929use arrow_array:: { Array , ArrayRef , RecordBatch } ;
30- use arrow_schema:: ArrowError ;
30+ use arrow_schema:: { ArrowError , SchemaRef } ;
3131
3232use 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