arrow_cast/
pretty.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Utilities for pretty printing [`RecordBatch`]es and [`Array`]s.
19//!
20//! Note this module is not available unless `feature = "prettyprint"` is enabled.
21//!
22//! [`RecordBatch`]: arrow_array::RecordBatch
23//! [`Array`]: arrow_array::Array
24
25use std::fmt::Display;
26
27use comfy_table::{Cell, Table};
28
29use arrow_array::{Array, ArrayRef, RecordBatch};
30use arrow_schema::{ArrowError, SchemaRef};
31
32use crate::display::{ArrayFormatter, FormatOptions};
33
34/// Create a visual representation of [`RecordBatch`]es
35///
36/// Uses default values for display. See [`pretty_format_batches_with_options`]
37/// for more control.
38///
39/// # Example
40/// ```
41/// # use std::sync::Arc;
42/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
43/// # use arrow_cast::pretty::pretty_format_batches;
44/// # let batch = RecordBatch::try_from_iter(vec![
45/// #       ("a", Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as ArrayRef),
46/// #       ("b", Arc::new(StringArray::from(vec![Some("a"), Some("b"), None, Some("d"), Some("e")]))),
47/// # ]).unwrap();
48/// // Note, returned object implements `Display`
49/// let pretty_table = pretty_format_batches(&[batch]).unwrap();
50/// let table_str = format!("Batches:\n{pretty_table}");
51/// assert_eq!(table_str,
52/// r#"Batches:
53/// +---+---+
54/// | a | b |
55/// +---+---+
56/// | 1 | a |
57/// | 2 | b |
58/// | 3 |   |
59/// | 4 | d |
60/// | 5 | e |
61/// +---+---+"#);
62/// ```
63pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<impl Display, ArrowError> {
64    let options = FormatOptions::default().with_display_error(true);
65    pretty_format_batches_with_options(results, &options)
66}
67
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/// ```
92pub 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
100/// Create a visual representation of [`RecordBatch`]es with formatting options.
101///
102/// # Arguments
103/// * `results` - A slice of record batches to display
104/// * `options` - [`FormatOptions`] that control the resulting display
105///
106/// # Example
107/// ```
108/// # use std::sync::Arc;
109/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
110/// # use arrow_cast::display::FormatOptions;
111/// # use arrow_cast::pretty::{pretty_format_batches, pretty_format_batches_with_options};
112/// # let batch = RecordBatch::try_from_iter(vec![
113/// #       ("a", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef),
114/// #       ("b", Arc::new(StringArray::from(vec![Some("a"), None]))),
115/// # ]).unwrap();
116/// let options = FormatOptions::new()
117///   .with_null("<NULL>");
118/// // Note, returned object implements `Display`
119/// let pretty_table = pretty_format_batches_with_options(&[batch], &options).unwrap();
120/// let table_str = format!("Batches:\n{pretty_table}");
121/// assert_eq!(table_str,
122/// r#"Batches:
123/// +---+--------+
124/// | a | b      |
125/// +---+--------+
126/// | 1 | a      |
127/// | 2 | <NULL> |
128/// +---+--------+"#);
129/// ```
130pub fn pretty_format_batches_with_options(
131    results: &[RecordBatch],
132    options: &FormatOptions,
133) -> Result<impl Display, ArrowError> {
134    create_table(None, results, options)
135}
136
137/// Create a visual representation of [`ArrayRef`]
138///
139/// Uses default values for display. See [`pretty_format_columns_with_options`]
140///
141/// See [`pretty_format_batches`] for an example
142pub fn pretty_format_columns(
143    col_name: &str,
144    results: &[ArrayRef],
145) -> Result<impl Display, ArrowError> {
146    let options = FormatOptions::default().with_display_error(true);
147    pretty_format_columns_with_options(col_name, results, &options)
148}
149
150/// Create a visual representation of [`ArrayRef`] with formatting options.
151///
152/// See [`pretty_format_batches_with_options`] for an example
153pub fn pretty_format_columns_with_options(
154    col_name: &str,
155    results: &[ArrayRef],
156    options: &FormatOptions,
157) -> Result<impl Display, ArrowError> {
158    create_column(col_name, results, options)
159}
160
161/// Prints a visual representation of record batches to stdout
162pub fn print_batches(results: &[RecordBatch]) -> Result<(), ArrowError> {
163    println!("{}", pretty_format_batches(results)?);
164    Ok(())
165}
166
167/// Prints a visual representation of a list of column to stdout
168pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<(), ArrowError> {
169    println!("{}", pretty_format_columns(col_name, results)?);
170    Ok(())
171}
172
173/// Convert a series of record batches into a table
174fn create_table(
175    schema_opt: Option<SchemaRef>,
176    results: &[RecordBatch],
177    options: &FormatOptions,
178) -> Result<Table, ArrowError> {
179    let mut table = Table::new();
180    table.load_preset("||--+-++|    ++++++");
181
182    let schema_opt = schema_opt.or_else(|| {
183        if results.is_empty() {
184            None
185        } else {
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            }
202        }
203        table.set_header(header);
204    }
205
206    if results.is_empty() {
207        return Ok(table);
208    }
209
210    for batch in results {
211        let formatters = batch
212            .columns()
213            .iter()
214            .map(|c| ArrayFormatter::try_new(c.as_ref(), options))
215            .collect::<Result<Vec<_>, ArrowError>>()?;
216
217        for row in 0..batch.num_rows() {
218            let mut cells = Vec::new();
219            for formatter in &formatters {
220                cells.push(Cell::new(formatter.value(row)));
221            }
222            table.add_row(cells);
223        }
224    }
225
226    Ok(table)
227}
228
229fn create_column(
230    field: &str,
231    columns: &[ArrayRef],
232    options: &FormatOptions,
233) -> Result<Table, ArrowError> {
234    let mut table = Table::new();
235    table.load_preset("||--+-++|    ++++++");
236
237    if columns.is_empty() {
238        return Ok(table);
239    }
240
241    let header = vec![Cell::new(field)];
242    table.set_header(header);
243
244    for col in columns {
245        let formatter = ArrayFormatter::try_new(col.as_ref(), options)?;
246        for row in 0..col.len() {
247            let cells = vec![Cell::new(formatter.value(row))];
248            table.add_row(cells);
249        }
250    }
251
252    Ok(table)
253}
254
255#[cfg(test)]
256mod tests {
257    use std::fmt::Write;
258    use std::sync::Arc;
259
260    use half::f16;
261
262    use arrow_array::builder::*;
263    use arrow_array::types::*;
264    use arrow_array::*;
265    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer};
266    use arrow_schema::*;
267
268    use crate::display::{array_value_to_string, DurationFormat};
269
270    use super::*;
271
272    #[test]
273    fn test_pretty_format_batches() {
274        // define a schema.
275        let schema = Arc::new(Schema::new(vec![
276            Field::new("a", DataType::Utf8, true),
277            Field::new("b", DataType::Int32, true),
278        ]));
279
280        // define data.
281        let batch = RecordBatch::try_new(
282            schema,
283            vec![
284                Arc::new(array::StringArray::from(vec![
285                    Some("a"),
286                    Some("b"),
287                    None,
288                    Some("d"),
289                ])),
290                Arc::new(array::Int32Array::from(vec![
291                    Some(1),
292                    None,
293                    Some(10),
294                    Some(100),
295                ])),
296            ],
297        )
298        .unwrap();
299
300        let table = pretty_format_batches(&[batch]).unwrap().to_string();
301
302        let expected = vec![
303            "+---+-----+",
304            "| a | b   |",
305            "+---+-----+",
306            "| a | 1   |",
307            "| b |     |",
308            "|   | 10  |",
309            "| d | 100 |",
310            "+---+-----+",
311        ];
312
313        let actual: Vec<&str> = table.lines().collect();
314
315        assert_eq!(expected, actual, "Actual result:\n{table}");
316    }
317
318    #[test]
319    fn test_pretty_format_columns() {
320        let columns = vec![
321            Arc::new(array::StringArray::from(vec![
322                Some("a"),
323                Some("b"),
324                None,
325                Some("d"),
326            ])) as ArrayRef,
327            Arc::new(array::StringArray::from(vec![Some("e"), None, Some("g")])),
328        ];
329
330        let table = pretty_format_columns("a", &columns).unwrap().to_string();
331
332        let expected = vec![
333            "+---+", "| a |", "+---+", "| a |", "| b |", "|   |", "| d |", "| e |", "|   |",
334            "| g |", "+---+",
335        ];
336
337        let actual: Vec<&str> = table.lines().collect();
338
339        assert_eq!(expected, actual, "Actual result:\n{table}");
340    }
341
342    #[test]
343    fn test_pretty_format_null() {
344        let schema = Arc::new(Schema::new(vec![
345            Field::new("a", DataType::Utf8, true),
346            Field::new("b", DataType::Int32, true),
347            Field::new("c", DataType::Null, true),
348        ]));
349
350        let num_rows = 4;
351        let arrays = schema
352            .fields()
353            .iter()
354            .map(|f| new_null_array(f.data_type(), num_rows))
355            .collect();
356
357        // define data (null)
358        let batch = RecordBatch::try_new(schema, arrays).unwrap();
359
360        let table = pretty_format_batches(&[batch]).unwrap().to_string();
361
362        let expected = vec![
363            "+---+---+---+",
364            "| a | b | c |",
365            "+---+---+---+",
366            "|   |   |   |",
367            "|   |   |   |",
368            "|   |   |   |",
369            "|   |   |   |",
370            "+---+---+---+",
371        ];
372
373        let actual: Vec<&str> = table.lines().collect();
374
375        assert_eq!(expected, actual, "Actual result:\n{table:#?}");
376    }
377
378    #[test]
379    fn test_pretty_format_dictionary() {
380        // define a schema.
381        let field = Field::new_dictionary("d1", DataType::Int32, DataType::Utf8, true);
382        let schema = Arc::new(Schema::new(vec![field]));
383
384        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
385
386        builder.append_value("one");
387        builder.append_null();
388        builder.append_value("three");
389        let array = Arc::new(builder.finish());
390
391        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
392
393        let table = pretty_format_batches(&[batch]).unwrap().to_string();
394
395        let expected = vec![
396            "+-------+",
397            "| d1    |",
398            "+-------+",
399            "| one   |",
400            "|       |",
401            "| three |",
402            "+-------+",
403        ];
404
405        let actual: Vec<&str> = table.lines().collect();
406
407        assert_eq!(expected, actual, "Actual result:\n{table}");
408    }
409
410    #[test]
411    fn test_pretty_format_fixed_size_list() {
412        // define a schema.
413        let field_type =
414            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, true)), 3);
415        let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));
416
417        let keys_builder = Int32Array::builder(3);
418        let mut builder = FixedSizeListBuilder::new(keys_builder, 3);
419
420        builder.values().append_slice(&[1, 2, 3]);
421        builder.append(true);
422        builder.values().append_slice(&[4, 5, 6]);
423        builder.append(false);
424        builder.values().append_slice(&[7, 8, 9]);
425        builder.append(true);
426
427        let array = Arc::new(builder.finish());
428
429        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
430        let table = pretty_format_batches(&[batch]).unwrap().to_string();
431        let expected = vec![
432            "+-----------+",
433            "| d1        |",
434            "+-----------+",
435            "| [1, 2, 3] |",
436            "|           |",
437            "| [7, 8, 9] |",
438            "+-----------+",
439        ];
440
441        let actual: Vec<&str> = table.lines().collect();
442
443        assert_eq!(expected, actual, "Actual result:\n{table}");
444    }
445
446    #[test]
447    fn test_pretty_format_string_view() {
448        let schema = Arc::new(Schema::new(vec![Field::new(
449            "d1",
450            DataType::Utf8View,
451            true,
452        )]));
453
454        // Use a small capacity so we end up with multiple views
455        let mut builder = StringViewBuilder::with_capacity(20);
456        builder.append_value("hello");
457        builder.append_null();
458        builder.append_value("longer than 12 bytes");
459        builder.append_value("another than 12 bytes");
460        builder.append_null();
461        builder.append_value("small");
462
463        let array: ArrayRef = Arc::new(builder.finish());
464        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
465        let table = pretty_format_batches(&[batch]).unwrap().to_string();
466        let expected = vec![
467            "+-----------------------+",
468            "| d1                    |",
469            "+-----------------------+",
470            "| hello                 |",
471            "|                       |",
472            "| longer than 12 bytes  |",
473            "| another than 12 bytes |",
474            "|                       |",
475            "| small                 |",
476            "+-----------------------+",
477        ];
478
479        let actual: Vec<&str> = table.lines().collect();
480
481        assert_eq!(expected, actual, "Actual result:\n{table:#?}");
482    }
483
484    #[test]
485    fn test_pretty_format_binary_view() {
486        let schema = Arc::new(Schema::new(vec![Field::new(
487            "d1",
488            DataType::BinaryView,
489            true,
490        )]));
491
492        // Use a small capacity so we end up with multiple views
493        let mut builder = BinaryViewBuilder::with_capacity(20);
494        builder.append_value(b"hello");
495        builder.append_null();
496        builder.append_value(b"longer than 12 bytes");
497        builder.append_value(b"another than 12 bytes");
498        builder.append_null();
499        builder.append_value(b"small");
500
501        let array: ArrayRef = Arc::new(builder.finish());
502        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
503        let table = pretty_format_batches(&[batch]).unwrap().to_string();
504        let expected = vec![
505            "+--------------------------------------------+",
506            "| d1                                         |",
507            "+--------------------------------------------+",
508            "| 68656c6c6f                                 |",
509            "|                                            |",
510            "| 6c6f6e676572207468616e203132206279746573   |",
511            "| 616e6f74686572207468616e203132206279746573 |",
512            "|                                            |",
513            "| 736d616c6c                                 |",
514            "+--------------------------------------------+",
515        ];
516
517        let actual: Vec<&str> = table.lines().collect();
518
519        assert_eq!(expected, actual, "Actual result:\n\n{table:#?}");
520    }
521
522    #[test]
523    fn test_pretty_format_fixed_size_binary() {
524        // define a schema.
525        let field_type = DataType::FixedSizeBinary(3);
526        let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));
527
528        let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 3);
529
530        builder.append_value([1, 2, 3]).unwrap();
531        builder.append_null();
532        builder.append_value([7, 8, 9]).unwrap();
533
534        let array = Arc::new(builder.finish());
535
536        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
537        let table = pretty_format_batches(&[batch]).unwrap().to_string();
538        let expected = vec![
539            "+--------+",
540            "| d1     |",
541            "+--------+",
542            "| 010203 |",
543            "|        |",
544            "| 070809 |",
545            "+--------+",
546        ];
547
548        let actual: Vec<&str> = table.lines().collect();
549
550        assert_eq!(expected, actual, "Actual result:\n{table}");
551    }
552
553    /// Generate an array with type $ARRAYTYPE with a numeric value of
554    /// $VALUE, and compare $EXPECTED_RESULT to the output of
555    /// formatting that array with `pretty_format_batches`
556    macro_rules! check_datetime {
557        ($ARRAYTYPE:ident, $VALUE:expr, $EXPECTED_RESULT:expr) => {
558            let mut builder = $ARRAYTYPE::builder(10);
559            builder.append_value($VALUE);
560            builder.append_null();
561            let array = builder.finish();
562
563            let schema = Arc::new(Schema::new(vec![Field::new(
564                "f",
565                array.data_type().clone(),
566                true,
567            )]));
568            let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
569
570            let table = pretty_format_batches(&[batch])
571                .expect("formatting batches")
572                .to_string();
573
574            let expected = $EXPECTED_RESULT;
575            let actual: Vec<&str> = table.lines().collect();
576
577            assert_eq!(expected, actual, "Actual result:\n\n{actual:#?}\n\n");
578        };
579    }
580
581    fn timestamp_batch<T: ArrowTimestampType>(timezone: &str, value: T::Native) -> RecordBatch {
582        let mut builder = PrimitiveBuilder::<T>::with_capacity(10);
583        builder.append_value(value);
584        builder.append_null();
585        let array = builder.finish();
586        let array = array.with_timezone(timezone);
587
588        let schema = Arc::new(Schema::new(vec![Field::new(
589            "f",
590            array.data_type().clone(),
591            true,
592        )]));
593        RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
594    }
595
596    #[test]
597    fn test_pretty_format_timestamp_second_with_fixed_offset_timezone() {
598        let batch = timestamp_batch::<TimestampSecondType>("+08:00", 11111111);
599        let table = pretty_format_batches(&[batch]).unwrap().to_string();
600
601        let expected = vec![
602            "+---------------------------+",
603            "| f                         |",
604            "+---------------------------+",
605            "| 1970-05-09T22:25:11+08:00 |",
606            "|                           |",
607            "+---------------------------+",
608        ];
609        let actual: Vec<&str> = table.lines().collect();
610        assert_eq!(expected, actual, "Actual result:\n\n{actual:#?}\n\n");
611    }
612
613    #[test]
614    fn test_pretty_format_timestamp_second() {
615        let expected = vec![
616            "+---------------------+",
617            "| f                   |",
618            "+---------------------+",
619            "| 1970-05-09T14:25:11 |",
620            "|                     |",
621            "+---------------------+",
622        ];
623        check_datetime!(TimestampSecondArray, 11111111, expected);
624    }
625
626    #[test]
627    fn test_pretty_format_timestamp_millisecond() {
628        let expected = vec![
629            "+-------------------------+",
630            "| f                       |",
631            "+-------------------------+",
632            "| 1970-01-01T03:05:11.111 |",
633            "|                         |",
634            "+-------------------------+",
635        ];
636        check_datetime!(TimestampMillisecondArray, 11111111, expected);
637    }
638
639    #[test]
640    fn test_pretty_format_timestamp_microsecond() {
641        let expected = vec![
642            "+----------------------------+",
643            "| f                          |",
644            "+----------------------------+",
645            "| 1970-01-01T00:00:11.111111 |",
646            "|                            |",
647            "+----------------------------+",
648        ];
649        check_datetime!(TimestampMicrosecondArray, 11111111, expected);
650    }
651
652    #[test]
653    fn test_pretty_format_timestamp_nanosecond() {
654        let expected = vec![
655            "+-------------------------------+",
656            "| f                             |",
657            "+-------------------------------+",
658            "| 1970-01-01T00:00:00.011111111 |",
659            "|                               |",
660            "+-------------------------------+",
661        ];
662        check_datetime!(TimestampNanosecondArray, 11111111, expected);
663    }
664
665    #[test]
666    fn test_pretty_format_date_32() {
667        let expected = vec![
668            "+------------+",
669            "| f          |",
670            "+------------+",
671            "| 1973-05-19 |",
672            "|            |",
673            "+------------+",
674        ];
675        check_datetime!(Date32Array, 1234, expected);
676    }
677
678    #[test]
679    fn test_pretty_format_date_64() {
680        let expected = vec![
681            "+---------------------+",
682            "| f                   |",
683            "+---------------------+",
684            "| 2005-03-18T01:58:20 |",
685            "|                     |",
686            "+---------------------+",
687        ];
688        check_datetime!(Date64Array, 1111111100000, expected);
689    }
690
691    #[test]
692    fn test_pretty_format_time_32_second() {
693        let expected = vec![
694            "+----------+",
695            "| f        |",
696            "+----------+",
697            "| 00:18:31 |",
698            "|          |",
699            "+----------+",
700        ];
701        check_datetime!(Time32SecondArray, 1111, expected);
702    }
703
704    #[test]
705    fn test_pretty_format_time_32_millisecond() {
706        let expected = vec![
707            "+--------------+",
708            "| f            |",
709            "+--------------+",
710            "| 03:05:11.111 |",
711            "|              |",
712            "+--------------+",
713        ];
714        check_datetime!(Time32MillisecondArray, 11111111, expected);
715    }
716
717    #[test]
718    fn test_pretty_format_time_64_microsecond() {
719        let expected = vec![
720            "+-----------------+",
721            "| f               |",
722            "+-----------------+",
723            "| 00:00:11.111111 |",
724            "|                 |",
725            "+-----------------+",
726        ];
727        check_datetime!(Time64MicrosecondArray, 11111111, expected);
728    }
729
730    #[test]
731    fn test_pretty_format_time_64_nanosecond() {
732        let expected = vec![
733            "+--------------------+",
734            "| f                  |",
735            "+--------------------+",
736            "| 00:00:00.011111111 |",
737            "|                    |",
738            "+--------------------+",
739        ];
740        check_datetime!(Time64NanosecondArray, 11111111, expected);
741    }
742
743    #[test]
744    fn test_int_display() {
745        let array = Arc::new(Int32Array::from(vec![6, 3])) as ArrayRef;
746        let actual_one = array_value_to_string(&array, 0).unwrap();
747        let expected_one = "6";
748
749        let actual_two = array_value_to_string(&array, 1).unwrap();
750        let expected_two = "3";
751        assert_eq!(actual_one, expected_one);
752        assert_eq!(actual_two, expected_two);
753    }
754
755    #[test]
756    fn test_decimal_display() {
757        let precision = 10;
758        let scale = 2;
759
760        let array = [Some(101), None, Some(200), Some(3040)]
761            .into_iter()
762            .collect::<Decimal128Array>()
763            .with_precision_and_scale(precision, scale)
764            .unwrap();
765
766        let dm = Arc::new(array) as ArrayRef;
767
768        let schema = Arc::new(Schema::new(vec![Field::new(
769            "f",
770            dm.data_type().clone(),
771            true,
772        )]));
773
774        let batch = RecordBatch::try_new(schema, vec![dm]).unwrap();
775
776        let table = pretty_format_batches(&[batch]).unwrap().to_string();
777
778        let expected = vec![
779            "+-------+",
780            "| f     |",
781            "+-------+",
782            "| 1.01  |",
783            "|       |",
784            "| 2.00  |",
785            "| 30.40 |",
786            "+-------+",
787        ];
788
789        let actual: Vec<&str> = table.lines().collect();
790        assert_eq!(expected, actual, "Actual result:\n{table}");
791    }
792
793    #[test]
794    fn test_decimal_display_zero_scale() {
795        let precision = 5;
796        let scale = 0;
797
798        let array = [Some(101), None, Some(200), Some(3040)]
799            .into_iter()
800            .collect::<Decimal128Array>()
801            .with_precision_and_scale(precision, scale)
802            .unwrap();
803
804        let dm = Arc::new(array) as ArrayRef;
805
806        let schema = Arc::new(Schema::new(vec![Field::new(
807            "f",
808            dm.data_type().clone(),
809            true,
810        )]));
811
812        let batch = RecordBatch::try_new(schema, vec![dm]).unwrap();
813
814        let table = pretty_format_batches(&[batch]).unwrap().to_string();
815        let expected = vec![
816            "+------+", "| f    |", "+------+", "| 101  |", "|      |", "| 200  |", "| 3040 |",
817            "+------+",
818        ];
819
820        let actual: Vec<&str> = table.lines().collect();
821        assert_eq!(expected, actual, "Actual result:\n{table}");
822    }
823
824    #[test]
825    fn test_pretty_format_struct() {
826        let schema = Schema::new(vec![
827            Field::new_struct(
828                "c1",
829                vec![
830                    Field::new("c11", DataType::Int32, true),
831                    Field::new_struct(
832                        "c12",
833                        vec![Field::new("c121", DataType::Utf8, false)],
834                        false,
835                    ),
836                ],
837                false,
838            ),
839            Field::new("c2", DataType::Utf8, false),
840        ]);
841
842        let c1 = StructArray::from(vec![
843            (
844                Arc::new(Field::new("c11", DataType::Int32, true)),
845                Arc::new(Int32Array::from(vec![Some(1), None, Some(5)])) as ArrayRef,
846            ),
847            (
848                Arc::new(Field::new_struct(
849                    "c12",
850                    vec![Field::new("c121", DataType::Utf8, false)],
851                    false,
852                )),
853                Arc::new(StructArray::from(vec![(
854                    Arc::new(Field::new("c121", DataType::Utf8, false)),
855                    Arc::new(StringArray::from(vec![Some("e"), Some("f"), Some("g")])) as ArrayRef,
856                )])) as ArrayRef,
857            ),
858        ]);
859        let c2 = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
860
861        let batch =
862            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
863
864        let table = pretty_format_batches(&[batch]).unwrap().to_string();
865        let expected = vec![
866            "+--------------------------+----+",
867            "| c1                       | c2 |",
868            "+--------------------------+----+",
869            "| {c11: 1, c12: {c121: e}} | a  |",
870            "| {c11: , c12: {c121: f}}  | b  |",
871            "| {c11: 5, c12: {c121: g}} | c  |",
872            "+--------------------------+----+",
873        ];
874
875        let actual: Vec<&str> = table.lines().collect();
876        assert_eq!(expected, actual, "Actual result:\n{table}");
877    }
878
879    #[test]
880    fn test_pretty_format_dense_union() {
881        let mut builder = UnionBuilder::new_dense();
882        builder.append::<Int32Type>("a", 1).unwrap();
883        builder.append::<Float64Type>("b", 3.2234).unwrap();
884        builder.append_null::<Float64Type>("b").unwrap();
885        builder.append_null::<Int32Type>("a").unwrap();
886        let union = builder.build().unwrap();
887
888        let schema = Schema::new(vec![Field::new_union(
889            "Teamsters",
890            vec![0, 1],
891            vec![
892                Field::new("a", DataType::Int32, false),
893                Field::new("b", DataType::Float64, false),
894            ],
895            UnionMode::Dense,
896        )]);
897
898        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(union)]).unwrap();
899        let table = pretty_format_batches(&[batch]).unwrap().to_string();
900        let actual: Vec<&str> = table.lines().collect();
901        let expected = vec![
902            "+------------+",
903            "| Teamsters  |",
904            "+------------+",
905            "| {a=1}      |",
906            "| {b=3.2234} |",
907            "| {b=}       |",
908            "| {a=}       |",
909            "+------------+",
910        ];
911
912        assert_eq!(expected, actual);
913    }
914
915    #[test]
916    fn test_pretty_format_sparse_union() {
917        let mut builder = UnionBuilder::new_sparse();
918        builder.append::<Int32Type>("a", 1).unwrap();
919        builder.append::<Float64Type>("b", 3.2234).unwrap();
920        builder.append_null::<Float64Type>("b").unwrap();
921        builder.append_null::<Int32Type>("a").unwrap();
922        let union = builder.build().unwrap();
923
924        let schema = Schema::new(vec![Field::new_union(
925            "Teamsters",
926            vec![0, 1],
927            vec![
928                Field::new("a", DataType::Int32, false),
929                Field::new("b", DataType::Float64, false),
930            ],
931            UnionMode::Sparse,
932        )]);
933
934        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(union)]).unwrap();
935        let table = pretty_format_batches(&[batch]).unwrap().to_string();
936        let actual: Vec<&str> = table.lines().collect();
937        let expected = vec![
938            "+------------+",
939            "| Teamsters  |",
940            "+------------+",
941            "| {a=1}      |",
942            "| {b=3.2234} |",
943            "| {b=}       |",
944            "| {a=}       |",
945            "+------------+",
946        ];
947
948        assert_eq!(expected, actual);
949    }
950
951    #[test]
952    fn test_pretty_format_nested_union() {
953        //Inner UnionArray
954        let mut builder = UnionBuilder::new_dense();
955        builder.append::<Int32Type>("b", 1).unwrap();
956        builder.append::<Float64Type>("c", 3.2234).unwrap();
957        builder.append_null::<Float64Type>("c").unwrap();
958        builder.append_null::<Int32Type>("b").unwrap();
959        builder.append_null::<Float64Type>("c").unwrap();
960        let inner = builder.build().unwrap();
961
962        let inner_field = Field::new_union(
963            "European Union",
964            vec![0, 1],
965            vec![
966                Field::new("b", DataType::Int32, false),
967                Field::new("c", DataType::Float64, false),
968            ],
969            UnionMode::Dense,
970        );
971
972        // Can't use UnionBuilder with non-primitive types, so manually build outer UnionArray
973        let a_array = Int32Array::from(vec![None, None, None, Some(1234), Some(23)]);
974        let type_ids = [1, 1, 0, 0, 1].into_iter().collect::<ScalarBuffer<i8>>();
975
976        let children = vec![Arc::new(a_array) as Arc<dyn Array>, Arc::new(inner)];
977
978        let union_fields = [
979            (0, Arc::new(Field::new("a", DataType::Int32, true))),
980            (1, Arc::new(inner_field.clone())),
981        ]
982        .into_iter()
983        .collect();
984
985        let outer = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
986
987        let schema = Schema::new(vec![Field::new_union(
988            "Teamsters",
989            vec![0, 1],
990            vec![Field::new("a", DataType::Int32, true), inner_field],
991            UnionMode::Sparse,
992        )]);
993
994        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(outer)]).unwrap();
995        let table = pretty_format_batches(&[batch]).unwrap().to_string();
996        let actual: Vec<&str> = table.lines().collect();
997        let expected = vec![
998            "+-----------------------------+",
999            "| Teamsters                   |",
1000            "+-----------------------------+",
1001            "| {European Union={b=1}}      |",
1002            "| {European Union={c=3.2234}} |",
1003            "| {a=}                        |",
1004            "| {a=1234}                    |",
1005            "| {European Union={c=}}       |",
1006            "+-----------------------------+",
1007        ];
1008        assert_eq!(expected, actual);
1009    }
1010
1011    #[test]
1012    fn test_writing_formatted_batches() {
1013        // define a schema.
1014        let schema = Arc::new(Schema::new(vec![
1015            Field::new("a", DataType::Utf8, true),
1016            Field::new("b", DataType::Int32, true),
1017        ]));
1018
1019        // define data.
1020        let batch = RecordBatch::try_new(
1021            schema,
1022            vec![
1023                Arc::new(array::StringArray::from(vec![
1024                    Some("a"),
1025                    Some("b"),
1026                    None,
1027                    Some("d"),
1028                ])),
1029                Arc::new(array::Int32Array::from(vec![
1030                    Some(1),
1031                    None,
1032                    Some(10),
1033                    Some(100),
1034                ])),
1035            ],
1036        )
1037        .unwrap();
1038
1039        let mut buf = String::new();
1040        write!(&mut buf, "{}", pretty_format_batches(&[batch]).unwrap()).unwrap();
1041
1042        let s = [
1043            "+---+-----+",
1044            "| a | b   |",
1045            "+---+-----+",
1046            "| a | 1   |",
1047            "| b |     |",
1048            "|   | 10  |",
1049            "| d | 100 |",
1050            "+---+-----+",
1051        ];
1052        let expected = s.join("\n");
1053        assert_eq!(expected, buf);
1054    }
1055
1056    #[test]
1057    fn test_float16_display() {
1058        let values = vec![
1059            Some(f16::from_f32(f32::NAN)),
1060            Some(f16::from_f32(4.0)),
1061            Some(f16::from_f32(f32::NEG_INFINITY)),
1062        ];
1063        let array = Arc::new(values.into_iter().collect::<Float16Array>()) as ArrayRef;
1064
1065        let schema = Arc::new(Schema::new(vec![Field::new(
1066            "f16",
1067            array.data_type().clone(),
1068            true,
1069        )]));
1070
1071        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
1072
1073        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1074
1075        let expected = vec![
1076            "+------+", "| f16  |", "+------+", "| NaN  |", "| 4    |", "| -inf |", "+------+",
1077        ];
1078
1079        let actual: Vec<&str> = table.lines().collect();
1080        assert_eq!(expected, actual, "Actual result:\n{table}");
1081    }
1082
1083    #[test]
1084    fn test_pretty_format_interval_day_time() {
1085        let arr = Arc::new(arrow_array::IntervalDayTimeArray::from(vec![
1086            Some(IntervalDayTime::new(-1, -600_000)),
1087            Some(IntervalDayTime::new(0, -1001)),
1088            Some(IntervalDayTime::new(0, -1)),
1089            Some(IntervalDayTime::new(0, 1)),
1090            Some(IntervalDayTime::new(0, 10)),
1091            Some(IntervalDayTime::new(0, 100)),
1092        ]));
1093
1094        let schema = Arc::new(Schema::new(vec![Field::new(
1095            "IntervalDayTime",
1096            arr.data_type().clone(),
1097            true,
1098        )]));
1099
1100        let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
1101
1102        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1103
1104        let expected = vec![
1105            "+------------------+",
1106            "| IntervalDayTime  |",
1107            "+------------------+",
1108            "| -1 days -10 mins |",
1109            "| -1.001 secs      |",
1110            "| -0.001 secs      |",
1111            "| 0.001 secs       |",
1112            "| 0.010 secs       |",
1113            "| 0.100 secs       |",
1114            "+------------------+",
1115        ];
1116
1117        let actual: Vec<&str> = table.lines().collect();
1118
1119        assert_eq!(expected, actual, "Actual result:\n{table}");
1120    }
1121
1122    #[test]
1123    fn test_pretty_format_interval_month_day_nano_array() {
1124        let arr = Arc::new(arrow_array::IntervalMonthDayNanoArray::from(vec![
1125            Some(IntervalMonthDayNano::new(-1, -1, -600_000_000_000)),
1126            Some(IntervalMonthDayNano::new(0, 0, -1_000_000_001)),
1127            Some(IntervalMonthDayNano::new(0, 0, -1)),
1128            Some(IntervalMonthDayNano::new(0, 0, 1)),
1129            Some(IntervalMonthDayNano::new(0, 0, 10)),
1130            Some(IntervalMonthDayNano::new(0, 0, 100)),
1131            Some(IntervalMonthDayNano::new(0, 0, 1_000)),
1132            Some(IntervalMonthDayNano::new(0, 0, 10_000)),
1133            Some(IntervalMonthDayNano::new(0, 0, 100_000)),
1134            Some(IntervalMonthDayNano::new(0, 0, 1_000_000)),
1135            Some(IntervalMonthDayNano::new(0, 0, 10_000_000)),
1136            Some(IntervalMonthDayNano::new(0, 0, 100_000_000)),
1137            Some(IntervalMonthDayNano::new(0, 0, 1_000_000_000)),
1138        ]));
1139
1140        let schema = Arc::new(Schema::new(vec![Field::new(
1141            "IntervalMonthDayNano",
1142            arr.data_type().clone(),
1143            true,
1144        )]));
1145
1146        let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
1147
1148        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1149
1150        let expected = vec![
1151            "+--------------------------+",
1152            "| IntervalMonthDayNano     |",
1153            "+--------------------------+",
1154            "| -1 mons -1 days -10 mins |",
1155            "| -1.000000001 secs        |",
1156            "| -0.000000001 secs        |",
1157            "| 0.000000001 secs         |",
1158            "| 0.000000010 secs         |",
1159            "| 0.000000100 secs         |",
1160            "| 0.000001000 secs         |",
1161            "| 0.000010000 secs         |",
1162            "| 0.000100000 secs         |",
1163            "| 0.001000000 secs         |",
1164            "| 0.010000000 secs         |",
1165            "| 0.100000000 secs         |",
1166            "| 1.000000000 secs         |",
1167            "+--------------------------+",
1168        ];
1169
1170        let actual: Vec<&str> = table.lines().collect();
1171
1172        assert_eq!(expected, actual, "Actual result:\n{table}");
1173    }
1174
1175    #[test]
1176    fn test_format_options() {
1177        let options = FormatOptions::default()
1178            .with_null("null")
1179            .with_types_info(true);
1180        let int32_array = Int32Array::from(vec![Some(1), Some(2), None, Some(3), Some(4)]);
1181        let string_array =
1182            StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz"), None]);
1183
1184        let batch = RecordBatch::try_from_iter([
1185            ("my_int32_name", Arc::new(int32_array) as _),
1186            ("my_string_name", Arc::new(string_array) as _),
1187        ])
1188        .unwrap();
1189
1190        let column = pretty_format_columns_with_options(
1191            "my_column_name",
1192            &[batch.column(0).clone()],
1193            &options,
1194        )
1195        .unwrap()
1196        .to_string();
1197
1198        let expected_column = vec![
1199            "+----------------+",
1200            "| my_column_name |",
1201            "+----------------+",
1202            "| 1              |",
1203            "| 2              |",
1204            "| null           |",
1205            "| 3              |",
1206            "| 4              |",
1207            "+----------------+",
1208        ];
1209
1210        let actual: Vec<&str> = column.lines().collect();
1211        assert_eq!(expected_column, actual, "Actual result:\n{column}");
1212
1213        let batch = pretty_format_batches_with_options(&[batch], &options)
1214            .unwrap()
1215            .to_string();
1216
1217        let expected_table = vec![
1218            "+---------------+----------------+",
1219            "| my_int32_name | my_string_name |",
1220            "| Int32         | Utf8           |",
1221            "+---------------+----------------+",
1222            "| 1             | foo            |",
1223            "| 2             | bar            |",
1224            "| null          | null           |",
1225            "| 3             | baz            |",
1226            "| 4             | null           |",
1227            "+---------------+----------------+",
1228        ];
1229
1230        let actual: Vec<&str> = batch.lines().collect();
1231        assert_eq!(expected_table, actual, "Actual result:\n{batch}");
1232    }
1233
1234    #[test]
1235    fn duration_pretty_and_iso_extremes() {
1236        // Build [MIN, MAX, 3661, NULL]
1237        let arr = DurationSecondArray::from(vec![Some(i64::MIN), Some(i64::MAX), Some(3661), None]);
1238        let array: ArrayRef = Arc::new(arr);
1239
1240        // Pretty formatting
1241        let opts = FormatOptions::default().with_null("null");
1242        let opts = opts.with_duration_format(DurationFormat::Pretty);
1243        let pretty = pretty_format_columns_with_options("pretty", &[array.clone()], &opts)
1244            .unwrap()
1245            .to_string();
1246
1247        // Expected output
1248        let expected_pretty = vec![
1249            "+------------------------------+",
1250            "| pretty                       |",
1251            "+------------------------------+",
1252            "| <invalid>                    |",
1253            "| <invalid>                    |",
1254            "| 0 days 1 hours 1 mins 1 secs |",
1255            "| null                         |",
1256            "+------------------------------+",
1257        ];
1258
1259        let actual: Vec<&str> = pretty.lines().collect();
1260        assert_eq!(expected_pretty, actual, "Actual result:\n{pretty}");
1261
1262        // ISO8601 formatting
1263        let opts_iso = FormatOptions::default()
1264            .with_null("null")
1265            .with_duration_format(DurationFormat::ISO8601);
1266        let iso = pretty_format_columns_with_options("iso", &[array], &opts_iso)
1267            .unwrap()
1268            .to_string();
1269
1270        // Expected output
1271        let expected_iso = vec![
1272            "+-----------+",
1273            "| iso       |",
1274            "+-----------+",
1275            "| <invalid> |",
1276            "| <invalid> |",
1277            "| PT3661S   |",
1278            "| null      |",
1279            "+-----------+",
1280        ];
1281
1282        let actual: Vec<&str> = iso.lines().collect();
1283        assert_eq!(expected_iso, actual, "Actual result:\n{iso}");
1284    }
1285}