Skip to main content

arrow_cast/
pretty.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Utilities for pretty printing [`RecordBatch`]es and [`Array`]s.
6//!
7//! Note this module is not available unless `feature = "prettyprint"` is enabled.
8//!
9//! [`RecordBatch`]: arrow_array::RecordBatch
10//! [`Array`]: arrow_array::Array
11
12use arrow_array::{Array, ArrayRef, RecordBatch};
13use arrow_schema::{ArrowError, SchemaRef};
14use comfy_table::{Cell, LineStyle, Table, TableStyle, presets::ASCII_FULL_CONDENSED};
15use std::fmt::Display;
16
17use crate::display::{ArrayFormatter, FormatOptions, make_array_formatter};
18
19const TABLE_STYLE: TableStyle =
20    ASCII_FULL_CONDENSED.header_separator(LineStyle::new('+', '-', '+', '+'));
21
22/// Create a visual representation of [`RecordBatch`]es
23///
24/// Uses default values for display. See [`pretty_format_batches_with_options`]
25/// for more control.
26///
27/// # Example
28/// ```
29/// # use std::sync::Arc;
30/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
31/// # use arrow_cast::pretty::pretty_format_batches;
32/// # let batch = RecordBatch::try_from_iter(vec![
33/// #       ("a", Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as ArrayRef),
34/// #       ("b", Arc::new(StringArray::from(vec![Some("a"), Some("b"), None, Some("d"), Some("e")]))),
35/// # ]).unwrap();
36/// // Note, returned object implements `Display`
37/// let pretty_table = pretty_format_batches(&[batch]).unwrap();
38/// let table_str = format!("Batches:\n{pretty_table}");
39/// assert_eq!(table_str,
40/// r#"Batches:
41/// +---+---+
42/// | a | b |
43/// +---+---+
44/// | 1 | a |
45/// | 2 | b |
46/// | 3 |   |
47/// | 4 | d |
48/// | 5 | e |
49/// +---+---+"#);
50/// ```
51pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<impl Display + use<>, ArrowError> {
52    let options = FormatOptions::default().with_display_error(true);
53    pretty_format_batches_with_options(results, &options)
54}
55
56/// Create a visual representation of [`RecordBatch`]es with a provided schema.
57///
58/// Useful to display empty batches.
59///
60/// # Example
61/// ```
62/// # use std::sync::Arc;
63/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
64/// # use arrow_cast::pretty::pretty_format_batches_with_schema;
65/// # use arrow_schema::{DataType, Field, Schema};
66/// let schema = Arc::new(Schema::new(vec![
67///     Field::new("a", DataType::Int32, false),
68///     Field::new("b", DataType::Utf8, true),
69/// ]));
70/// // Note, returned object implements `Display`
71/// let pretty_table = pretty_format_batches_with_schema(schema, &[]).unwrap();
72/// let table_str = format!("Batches:\n{pretty_table}");
73/// assert_eq!(table_str,
74/// r#"Batches:
75/// +---+---+
76/// | a | b |
77/// +---+---+
78/// +---+---+"#);
79/// ```
80pub fn pretty_format_batches_with_schema(
81    schema: SchemaRef,
82    results: &[RecordBatch],
83) -> Result<impl Display + use<>, ArrowError> {
84    let options = FormatOptions::default().with_display_error(true);
85    create_table(Some(schema), results, &options)
86}
87
88/// Create a visual representation of [`RecordBatch`]es with formatting options.
89///
90/// # Arguments
91/// * `results` - A slice of record batches to display
92/// * `options` - [`FormatOptions`] that control the resulting display
93///
94/// # Example
95/// ```
96/// # use std::sync::Arc;
97/// # use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray};
98/// # use arrow_cast::display::FormatOptions;
99/// # use arrow_cast::pretty::{pretty_format_batches, pretty_format_batches_with_options};
100/// # let batch = RecordBatch::try_from_iter(vec![
101/// #       ("a", Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef),
102/// #       ("b", Arc::new(StringArray::from(vec![Some("a"), None]))),
103/// # ]).unwrap();
104/// let options = FormatOptions::new()
105///   .with_null("<NULL>");
106/// // Note, returned object implements `Display`
107/// let pretty_table = pretty_format_batches_with_options(&[batch], &options).unwrap();
108/// let table_str = format!("Batches:\n{pretty_table}");
109/// assert_eq!(table_str,
110/// r#"Batches:
111/// +---+--------+
112/// | a | b      |
113/// +---+--------+
114/// | 1 | a      |
115/// | 2 | <NULL> |
116/// +---+--------+"#);
117/// ```
118pub fn pretty_format_batches_with_options(
119    results: &[RecordBatch],
120    options: &FormatOptions,
121) -> Result<impl Display + use<>, ArrowError> {
122    create_table(None, results, options)
123}
124
125/// Create a visual representation of [`ArrayRef`]
126///
127/// Uses default values for display. See [`pretty_format_columns_with_options`]
128///
129/// See [`pretty_format_batches`] for an example
130pub fn pretty_format_columns(
131    col_name: &str,
132    results: &[ArrayRef],
133) -> Result<impl Display + use<>, ArrowError> {
134    let options = FormatOptions::default().with_display_error(true);
135    pretty_format_columns_with_options(col_name, results, &options)
136}
137
138/// Create a visual representation of [`ArrayRef`] with formatting options.
139///
140/// See [`pretty_format_batches_with_options`] for an example
141pub fn pretty_format_columns_with_options(
142    col_name: &str,
143    results: &[ArrayRef],
144    options: &FormatOptions,
145) -> Result<impl Display + use<>, ArrowError> {
146    create_column(col_name, results, options)
147}
148
149/// Prints a visual representation of record batches to stdout
150pub fn print_batches(results: &[RecordBatch]) -> Result<(), ArrowError> {
151    println!("{}", pretty_format_batches(results)?);
152    Ok(())
153}
154
155/// Prints a visual representation of a list of column to stdout
156pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<(), ArrowError> {
157    println!("{}", pretty_format_columns(col_name, results)?);
158    Ok(())
159}
160
161/// Convert a series of record batches into a table
162fn create_table(
163    schema_opt: Option<SchemaRef>,
164    results: &[RecordBatch],
165    options: &FormatOptions,
166) -> Result<Table, ArrowError> {
167    let mut table = Table::new();
168    table.load_style(TABLE_STYLE);
169
170    let schema_opt = schema_opt.or_else(|| {
171        if results.is_empty() {
172            None
173        } else {
174            Some(results[0].schema())
175        }
176    });
177
178    if let Some(schema) = &schema_opt {
179        let mut header = Vec::new();
180        for field in schema.fields() {
181            if options.types_info() {
182                header.push(Cell::new(format!(
183                    "{}\n{}",
184                    field.name(),
185                    field.data_type()
186                )))
187            } else {
188                header.push(Cell::new(field.name()));
189            }
190        }
191        table.set_header(header);
192    }
193
194    if results.is_empty() {
195        return Ok(table);
196    }
197
198    for batch in results {
199        let schema = schema_opt.as_ref().unwrap_or_else(|| batch.schema_ref());
200
201        // Could be a custom schema that was provided.
202        if batch.columns().len() != schema.fields().len() {
203            return Err(ArrowError::InvalidArgumentError(format!(
204                "Expected the same number of columns in a record batch ({}) as the number of fields ({}) in the schema",
205                batch.columns().len(),
206                schema.fields.len()
207            )));
208        }
209
210        let formatters = batch
211            .columns()
212            .iter()
213            .zip(schema.fields().iter())
214            .map(|(c, field)| make_array_formatter(c, options, Some(field)))
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_style(TABLE_STYLE);
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 = match options.formatter_factory() {
246            None => ArrayFormatter::try_new(col.as_ref(), options)?,
247            Some(formatters) => formatters
248                .create_array_formatter(col.as_ref(), options, None)
249                .transpose()
250                .unwrap_or_else(|| ArrayFormatter::try_new(col.as_ref(), options))?,
251        };
252        for row in 0..col.len() {
253            let cells = vec![Cell::new(formatter.value(row))];
254            table.add_row(cells);
255        }
256    }
257
258    Ok(table)
259}
260
261#[cfg(test)]
262mod tests {
263    use std::collections::HashMap;
264    use std::fmt::Write;
265    use std::sync::Arc;
266
267    use arrow_array::builder::*;
268    use arrow_array::cast::AsArray;
269    use arrow_array::types::*;
270    use arrow_array::*;
271    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, ScalarBuffer};
272    use arrow_schema::*;
273    use half::f16;
274
275    use crate::display::{
276        ArrayFormatterFactory, DisplayIndex, DurationFormat, array_value_to_string,
277    };
278
279    use super::*;
280
281    #[test]
282    fn test_pretty_format_batches() {
283        // define a schema.
284        let schema = Arc::new(Schema::new(vec![
285            Field::new("a", DataType::Utf8, true),
286            Field::new("b", DataType::Int32, true),
287        ]));
288
289        // define data.
290        let batch = RecordBatch::try_new(
291            schema,
292            vec![
293                Arc::new(array::StringArray::from(vec![
294                    Some("a"),
295                    Some("b"),
296                    None,
297                    Some("d"),
298                ])),
299                Arc::new(array::Int32Array::from(vec![
300                    Some(1),
301                    None,
302                    Some(10),
303                    Some(100),
304                ])),
305            ],
306        )
307        .unwrap();
308
309        let table = pretty_format_batches(&[batch]).unwrap().to_string();
310
311        insta::assert_snapshot!(table, @"
312        +---+-----+
313        | a | b   |
314        +---+-----+
315        | a | 1   |
316        | b |     |
317        |   | 10  |
318        | d | 100 |
319        +---+-----+
320        ");
321    }
322
323    #[test]
324    fn test_pretty_format_columns() {
325        let columns = vec![
326            Arc::new(array::StringArray::from(vec![
327                Some("a"),
328                Some("b"),
329                None,
330                Some("d"),
331            ])) as ArrayRef,
332            Arc::new(array::StringArray::from(vec![Some("e"), None, Some("g")])),
333        ];
334
335        let table = pretty_format_columns("a", &columns).unwrap().to_string();
336
337        insta::assert_snapshot!(table, @"
338        +---+
339        | a |
340        +---+
341        | a |
342        | b |
343        |   |
344        | d |
345        | e |
346        |   |
347        | g |
348        +---+
349        ");
350    }
351
352    #[test]
353    fn test_pretty_format_null() {
354        let schema = Arc::new(Schema::new(vec![
355            Field::new("a", DataType::Utf8, true),
356            Field::new("b", DataType::Int32, true),
357            Field::new("c", DataType::Null, true),
358        ]));
359
360        let num_rows = 4;
361        let arrays = schema
362            .fields()
363            .iter()
364            .map(|f| new_null_array(f.data_type(), num_rows))
365            .collect();
366
367        // define data (null)
368        let batch = RecordBatch::try_new(schema, arrays).unwrap();
369
370        let table = pretty_format_batches(&[batch]).unwrap().to_string();
371
372        insta::assert_snapshot!(table, @"
373        +---+---+---+
374        | a | b | c |
375        +---+---+---+
376        |   |   |   |
377        |   |   |   |
378        |   |   |   |
379        |   |   |   |
380        +---+---+---+
381        ");
382    }
383
384    #[test]
385    fn test_pretty_format_dictionary() {
386        // define a schema.
387        let field = Field::new_dictionary("d1", DataType::Int32, DataType::Utf8, true);
388        let schema = Arc::new(Schema::new(vec![field]));
389
390        let mut builder = StringDictionaryBuilder::<Int32Type>::new();
391
392        builder.append_value("one");
393        builder.append_null();
394        builder.append_value("three");
395        let array = Arc::new(builder.finish());
396
397        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
398
399        let table = pretty_format_batches(&[batch]).unwrap().to_string();
400
401        insta::assert_snapshot!(table, @"
402        +-------+
403        | d1    |
404        +-------+
405        | one   |
406        |       |
407        | three |
408        +-------+
409        ");
410    }
411
412    #[test]
413    fn test_pretty_format_fixed_size_list() {
414        // define a schema.
415        let field_type =
416            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, true)), 3);
417        let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));
418
419        let keys_builder = Int32Array::builder(3);
420        let mut builder = FixedSizeListBuilder::new(keys_builder, 3);
421
422        builder.values().append_slice(&[1, 2, 3]);
423        builder.append(true);
424        builder.values().append_slice(&[4, 5, 6]);
425        builder.append(false);
426        builder.values().append_slice(&[7, 8, 9]);
427        builder.append(true);
428
429        let array = Arc::new(builder.finish());
430
431        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
432        let table = pretty_format_batches(&[batch]).unwrap().to_string();
433
434        insta::assert_snapshot!(table, @"
435        +-----------+
436        | d1        |
437        +-----------+
438        | [1, 2, 3] |
439        |           |
440        | [7, 8, 9] |
441        +-----------+
442        ");
443    }
444
445    #[test]
446    fn test_pretty_format_string_view() {
447        let schema = Arc::new(Schema::new(vec![Field::new(
448            "d1",
449            DataType::Utf8View,
450            true,
451        )]));
452
453        // Use a small capacity so we end up with multiple views
454        let mut builder = StringViewBuilder::with_capacity(20);
455        builder.append_value("hello");
456        builder.append_null();
457        builder.append_value("longer than 12 bytes");
458        builder.append_value("another than 12 bytes");
459        builder.append_null();
460        builder.append_value("small");
461
462        let array: ArrayRef = Arc::new(builder.finish());
463        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
464        let table = pretty_format_batches(&[batch]).unwrap().to_string();
465
466        insta::assert_snapshot!(table, @"
467        +-----------------------+
468        | d1                    |
469        +-----------------------+
470        | hello                 |
471        |                       |
472        | longer than 12 bytes  |
473        | another than 12 bytes |
474        |                       |
475        | small                 |
476        +-----------------------+
477        ");
478    }
479
480    #[test]
481    fn test_pretty_format_binary_view() {
482        let schema = Arc::new(Schema::new(vec![Field::new(
483            "d1",
484            DataType::BinaryView,
485            true,
486        )]));
487
488        // Use a small capacity so we end up with multiple views
489        let mut builder = BinaryViewBuilder::with_capacity(20);
490        builder.append_value(b"hello");
491        builder.append_null();
492        builder.append_value(b"longer than 12 bytes");
493        builder.append_value(b"another than 12 bytes");
494        builder.append_null();
495        builder.append_value(b"small");
496
497        let array: ArrayRef = Arc::new(builder.finish());
498        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
499        let table = pretty_format_batches(&[batch]).unwrap().to_string();
500
501        insta::assert_snapshot!(table, @"
502        +--------------------------------------------+
503        | d1                                         |
504        +--------------------------------------------+
505        | 68656c6c6f                                 |
506        |                                            |
507        | 6c6f6e676572207468616e203132206279746573   |
508        | 616e6f74686572207468616e203132206279746573 |
509        |                                            |
510        | 736d616c6c                                 |
511        +--------------------------------------------+
512        ");
513    }
514
515    #[test]
516    fn test_pretty_format_fixed_size_binary() {
517        // define a schema.
518        let field_type = DataType::FixedSizeBinary(3);
519        let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)]));
520
521        let mut builder = FixedSizeBinaryBuilder::with_capacity(3, 3);
522
523        builder.append_value([1, 2, 3]).unwrap();
524        builder.append_null();
525        builder.append_value([7, 8, 9]).unwrap();
526
527        let array = Arc::new(builder.finish());
528
529        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
530        let table = pretty_format_batches(&[batch]).unwrap().to_string();
531
532        insta::assert_snapshot!(table, @"
533        +--------+
534        | d1     |
535        +--------+
536        | 010203 |
537        |        |
538        | 070809 |
539        +--------+
540        ");
541    }
542
543    /// Generate an array of [`ArrowPrimitiveType`] with a numeric `value`,
544    /// then format it with `pretty_format_batches`.
545    fn format_primitive_batch<T: ArrowPrimitiveType>(value: T::Native) -> String {
546        let mut builder = PrimitiveBuilder::<T>::with_capacity(10);
547        builder.append_value(value);
548        builder.append_null();
549        let array = builder.finish();
550        let schema = Arc::new(Schema::new(vec![Field::new(
551            "f",
552            array.data_type().clone(),
553            true,
554        )]));
555        let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
556        pretty_format_batches(&[batch])
557            .expect("formatting batches")
558            .to_string()
559    }
560
561    fn timestamp_batch<T: ArrowTimestampType>(timezone: &str, value: T::Native) -> RecordBatch {
562        let mut builder = PrimitiveBuilder::<T>::with_capacity(10);
563        builder.append_value(value);
564        builder.append_null();
565        let array = builder.finish();
566        let array = array.with_timezone(timezone);
567
568        let schema = Arc::new(Schema::new(vec![Field::new(
569            "f",
570            array.data_type().clone(),
571            true,
572        )]));
573        RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap()
574    }
575
576    #[test]
577    fn test_pretty_format_timestamp_second_with_fixed_offset_timezone() {
578        let batch = timestamp_batch::<TimestampSecondType>("+08:00", 11111111);
579        let table = pretty_format_batches(&[batch]).unwrap().to_string();
580
581        insta::assert_snapshot!(table, @"
582        +---------------------------+
583        | f                         |
584        +---------------------------+
585        | 1970-05-09T22:25:11+08:00 |
586        |                           |
587        +---------------------------+
588        ");
589    }
590
591    #[test]
592    fn test_pretty_format_timestamp_second() {
593        let table = format_primitive_batch::<TimestampSecondType>(11111111);
594        insta::assert_snapshot!(table, @"
595        +---------------------+
596        | f                   |
597        +---------------------+
598        | 1970-05-09T14:25:11 |
599        |                     |
600        +---------------------+
601        ");
602    }
603
604    #[test]
605    fn test_pretty_format_timestamp_millisecond() {
606        let table = format_primitive_batch::<TimestampMillisecondType>(11111111);
607        insta::assert_snapshot!(table, @"
608        +-------------------------+
609        | f                       |
610        +-------------------------+
611        | 1970-01-01T03:05:11.111 |
612        |                         |
613        +-------------------------+
614        ");
615    }
616
617    #[test]
618    fn test_pretty_format_timestamp_microsecond() {
619        let table = format_primitive_batch::<TimestampMicrosecondType>(11111111);
620        insta::assert_snapshot!(table, @"
621        +----------------------------+
622        | f                          |
623        +----------------------------+
624        | 1970-01-01T00:00:11.111111 |
625        |                            |
626        +----------------------------+
627        ");
628    }
629
630    #[test]
631    fn test_pretty_format_timestamp_nanosecond() {
632        let table = format_primitive_batch::<TimestampNanosecondType>(11111111);
633        insta::assert_snapshot!(table, @"
634        +-------------------------------+
635        | f                             |
636        +-------------------------------+
637        | 1970-01-01T00:00:00.011111111 |
638        |                               |
639        +-------------------------------+
640        ");
641    }
642
643    #[test]
644    fn test_pretty_format_date_32() {
645        let table = format_primitive_batch::<Date32Type>(1234);
646        insta::assert_snapshot!(table, @"
647        +------------+
648        | f          |
649        +------------+
650        | 1973-05-19 |
651        |            |
652        +------------+
653        ");
654    }
655
656    #[test]
657    fn test_pretty_format_date_64() {
658        let table = format_primitive_batch::<Date64Type>(1111111100000);
659        insta::assert_snapshot!(table, @"
660        +---------------------+
661        | f                   |
662        +---------------------+
663        | 2005-03-18T01:58:20 |
664        |                     |
665        +---------------------+
666        ");
667    }
668
669    #[test]
670    fn test_pretty_format_time_32_second() {
671        let table = format_primitive_batch::<Time32SecondType>(1111);
672        insta::assert_snapshot!(table, @"
673        +----------+
674        | f        |
675        +----------+
676        | 00:18:31 |
677        |          |
678        +----------+
679        ");
680    }
681
682    #[test]
683    fn test_pretty_format_time_32_millisecond() {
684        let table = format_primitive_batch::<Time32MillisecondType>(11111111);
685        insta::assert_snapshot!(table, @"
686        +--------------+
687        | f            |
688        +--------------+
689        | 03:05:11.111 |
690        |              |
691        +--------------+
692        ");
693    }
694
695    #[test]
696    fn test_pretty_format_time_64_microsecond() {
697        let table = format_primitive_batch::<Time64MicrosecondType>(11111111);
698        insta::assert_snapshot!(table, @"
699        +-----------------+
700        | f               |
701        +-----------------+
702        | 00:00:11.111111 |
703        |                 |
704        +-----------------+
705        ");
706    }
707
708    #[test]
709    fn test_pretty_format_time_64_nanosecond() {
710        let table = format_primitive_batch::<Time64NanosecondType>(11111111);
711        insta::assert_snapshot!(table, @"
712        +--------------------+
713        | f                  |
714        +--------------------+
715        | 00:00:00.011111111 |
716        |                    |
717        +--------------------+
718        ");
719    }
720
721    #[test]
722    fn test_int_display() {
723        let array = Arc::new(Int32Array::from(vec![6, 3])) as ArrayRef;
724        insta::assert_snapshot!(array_value_to_string(&array, 0).unwrap(), @"6");
725        insta::assert_snapshot!(array_value_to_string(&array, 1).unwrap(), @"3");
726    }
727
728    #[test]
729    fn test_decimal_display() {
730        let precision = 10;
731        let scale = 2;
732
733        let array = [Some(101), None, Some(200), Some(3040)]
734            .into_iter()
735            .collect::<Decimal128Array>()
736            .with_precision_and_scale(precision, scale)
737            .unwrap();
738
739        let dm = Arc::new(array) as ArrayRef;
740
741        let schema = Arc::new(Schema::new(vec![Field::new(
742            "f",
743            dm.data_type().clone(),
744            true,
745        )]));
746
747        let batch = RecordBatch::try_new(schema, vec![dm]).unwrap();
748
749        let table = pretty_format_batches(&[batch]).unwrap().to_string();
750
751        insta::assert_snapshot!(table, @"
752        +-------+
753        | f     |
754        +-------+
755        | 1.01  |
756        |       |
757        | 2.00  |
758        | 30.40 |
759        +-------+
760        ");
761    }
762
763    #[test]
764    fn test_decimal_display_zero_scale() {
765        let precision = 5;
766        let scale = 0;
767
768        let array = [Some(101), None, Some(200), Some(3040)]
769            .into_iter()
770            .collect::<Decimal128Array>()
771            .with_precision_and_scale(precision, scale)
772            .unwrap();
773
774        let dm = Arc::new(array) as ArrayRef;
775
776        let schema = Arc::new(Schema::new(vec![Field::new(
777            "f",
778            dm.data_type().clone(),
779            true,
780        )]));
781
782        let batch = RecordBatch::try_new(schema, vec![dm]).unwrap();
783
784        let table = pretty_format_batches(&[batch]).unwrap().to_string();
785
786        insta::assert_snapshot!(table, @"
787        +------+
788        | f    |
789        +------+
790        | 101  |
791        |      |
792        | 200  |
793        | 3040 |
794        +------+
795        ");
796    }
797
798    #[test]
799    fn test_pretty_format_struct() {
800        let schema = Schema::new(vec![
801            Field::new_struct(
802                "c1",
803                vec![
804                    Field::new("c11", DataType::Int32, true),
805                    Field::new_struct(
806                        "c12",
807                        vec![Field::new("c121", DataType::Utf8, false)],
808                        false,
809                    ),
810                ],
811                false,
812            ),
813            Field::new("c2", DataType::Utf8, false),
814        ]);
815
816        let c1 = StructArray::from(vec![
817            (
818                Arc::new(Field::new("c11", DataType::Int32, true)),
819                Arc::new(Int32Array::from(vec![Some(1), None, Some(5)])) as ArrayRef,
820            ),
821            (
822                Arc::new(Field::new_struct(
823                    "c12",
824                    vec![Field::new("c121", DataType::Utf8, false)],
825                    false,
826                )),
827                Arc::new(StructArray::from(vec![(
828                    Arc::new(Field::new("c121", DataType::Utf8, false)),
829                    Arc::new(StringArray::from(vec![Some("e"), Some("f"), Some("g")])) as ArrayRef,
830                )])) as ArrayRef,
831            ),
832        ]);
833        let c2 = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
834
835        let batch =
836            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
837
838        let table = pretty_format_batches(&[batch]).unwrap().to_string();
839
840        insta::assert_snapshot!(table, @"
841        +--------------------------+----+
842        | c1                       | c2 |
843        +--------------------------+----+
844        | {c11: 1, c12: {c121: e}} | a  |
845        | {c11: , c12: {c121: f}}  | b  |
846        | {c11: 5, c12: {c121: g}} | c  |
847        +--------------------------+----+
848        ");
849    }
850
851    #[test]
852    fn test_pretty_format_dense_union() {
853        let mut builder = UnionBuilder::new_dense();
854        builder.append::<Int32Type>("a", 1).unwrap();
855        builder.append::<Float64Type>("b", 3.2234).unwrap();
856        builder.append_null::<Float64Type>("b").unwrap();
857        builder.append_null::<Int32Type>("a").unwrap();
858        let union = builder.build().unwrap();
859
860        let schema = Schema::new(vec![Field::new_union(
861            "Teamsters",
862            vec![0, 1],
863            vec![
864                Field::new("a", DataType::Int32, false),
865                Field::new("b", DataType::Float64, false),
866            ],
867            UnionMode::Dense,
868        )]);
869
870        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(union)]).unwrap();
871        let table = pretty_format_batches(&[batch]).unwrap().to_string();
872
873        insta::assert_snapshot!(table, @"
874        +------------+
875        | Teamsters  |
876        +------------+
877        | {a=1}      |
878        | {b=3.2234} |
879        | {b=}       |
880        | {a=}       |
881        +------------+
882        ");
883    }
884
885    #[test]
886    fn test_pretty_format_sparse_union() {
887        let mut builder = UnionBuilder::new_sparse();
888        builder.append::<Int32Type>("a", 1).unwrap();
889        builder.append::<Float64Type>("b", 3.2234).unwrap();
890        builder.append_null::<Float64Type>("b").unwrap();
891        builder.append_null::<Int32Type>("a").unwrap();
892        let union = builder.build().unwrap();
893
894        let schema = Schema::new(vec![Field::new_union(
895            "Teamsters",
896            vec![0, 1],
897            vec![
898                Field::new("a", DataType::Int32, false),
899                Field::new("b", DataType::Float64, false),
900            ],
901            UnionMode::Sparse,
902        )]);
903
904        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(union)]).unwrap();
905        let table = pretty_format_batches(&[batch]).unwrap().to_string();
906
907        insta::assert_snapshot!(table, @"
908        +------------+
909        | Teamsters  |
910        +------------+
911        | {a=1}      |
912        | {b=3.2234} |
913        | {b=}       |
914        | {a=}       |
915        +------------+
916        ");
917    }
918
919    #[test]
920    fn test_pretty_format_nested_union() {
921        //Inner UnionArray
922        let mut builder = UnionBuilder::new_dense();
923        builder.append::<Int32Type>("b", 1).unwrap();
924        builder.append::<Float64Type>("c", 3.2234).unwrap();
925        builder.append_null::<Float64Type>("c").unwrap();
926        builder.append_null::<Int32Type>("b").unwrap();
927        builder.append_null::<Float64Type>("c").unwrap();
928        let inner = builder.build().unwrap();
929
930        let inner_field = Field::new_union(
931            "European Union",
932            vec![0, 1],
933            vec![
934                Field::new("b", DataType::Int32, false),
935                Field::new("c", DataType::Float64, false),
936            ],
937            UnionMode::Dense,
938        );
939
940        // Can't use UnionBuilder with non-primitive types, so manually build outer UnionArray
941        let a_array = Int32Array::from(vec![None, None, None, Some(1234), Some(23)]);
942        let type_ids = [1, 1, 0, 0, 1].into_iter().collect::<ScalarBuffer<i8>>();
943
944        let children = vec![Arc::new(a_array) as Arc<dyn Array>, Arc::new(inner)];
945
946        let union_fields = [
947            (0, Arc::new(Field::new("a", DataType::Int32, true))),
948            (1, Arc::new(inner_field.clone())),
949        ]
950        .into_iter()
951        .collect();
952
953        let outer = UnionArray::try_new(union_fields, type_ids, None, children).unwrap();
954
955        let schema = Schema::new(vec![Field::new_union(
956            "Teamsters",
957            vec![0, 1],
958            vec![Field::new("a", DataType::Int32, true), inner_field],
959            UnionMode::Sparse,
960        )]);
961
962        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(outer)]).unwrap();
963        let table = pretty_format_batches(&[batch]).unwrap().to_string();
964
965        insta::assert_snapshot!(table, @"
966        +-----------------------------+
967        | Teamsters                   |
968        +-----------------------------+
969        | {European Union={b=1}}      |
970        | {European Union={c=3.2234}} |
971        | {a=}                        |
972        | {a=1234}                    |
973        | {European Union={c=}}       |
974        +-----------------------------+
975        ");
976    }
977
978    #[test]
979    fn test_writing_formatted_batches() {
980        // define a schema.
981        let schema = Arc::new(Schema::new(vec![
982            Field::new("a", DataType::Utf8, true),
983            Field::new("b", DataType::Int32, true),
984        ]));
985
986        // define data.
987        let batch = RecordBatch::try_new(
988            schema,
989            vec![
990                Arc::new(array::StringArray::from(vec![
991                    Some("a"),
992                    Some("b"),
993                    None,
994                    Some("d"),
995                ])),
996                Arc::new(array::Int32Array::from(vec![
997                    Some(1),
998                    None,
999                    Some(10),
1000                    Some(100),
1001                ])),
1002            ],
1003        )
1004        .unwrap();
1005
1006        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1007
1008        insta::assert_snapshot!(table, @"
1009        +---+-----+
1010        | a | b   |
1011        +---+-----+
1012        | a | 1   |
1013        | b |     |
1014        |   | 10  |
1015        | d | 100 |
1016        +---+-----+
1017        ");
1018    }
1019
1020    #[test]
1021    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
1022    fn test_float16_display() {
1023        let values = vec![
1024            Some(f16::from_f32(f32::NAN)),
1025            Some(f16::from_f32(4.0)),
1026            Some(f16::from_f32(f32::NEG_INFINITY)),
1027        ];
1028        let array = Arc::new(values.into_iter().collect::<Float16Array>()) as ArrayRef;
1029
1030        let schema = Arc::new(Schema::new(vec![Field::new(
1031            "f16",
1032            array.data_type().clone(),
1033            true,
1034        )]));
1035
1036        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
1037
1038        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1039
1040        insta::assert_snapshot!(table, @"
1041        +------+
1042        | f16  |
1043        +------+
1044        | NaN  |
1045        | 4    |
1046        | -inf |
1047        +------+
1048        ");
1049    }
1050
1051    #[test]
1052    fn test_pretty_format_interval_day_time() {
1053        let arr = Arc::new(arrow_array::IntervalDayTimeArray::from(vec![
1054            Some(IntervalDayTime::new(-1, -600_000)),
1055            Some(IntervalDayTime::new(0, -1001)),
1056            Some(IntervalDayTime::new(0, -1)),
1057            Some(IntervalDayTime::new(0, 1)),
1058            Some(IntervalDayTime::new(0, 10)),
1059            Some(IntervalDayTime::new(0, 100)),
1060            Some(IntervalDayTime::new(0, 0)),
1061        ]));
1062
1063        let schema = Arc::new(Schema::new(vec![Field::new(
1064            "IntervalDayTime",
1065            arr.data_type().clone(),
1066            true,
1067        )]));
1068
1069        let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
1070
1071        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1072
1073        insta::assert_snapshot!(table, @"
1074        +------------------+
1075        | IntervalDayTime  |
1076        +------------------+
1077        | -1 days -10 mins |
1078        | -1.001 secs      |
1079        | -0.001 secs      |
1080        | 0.001 secs       |
1081        | 0.010 secs       |
1082        | 0.100 secs       |
1083        | 0 secs           |
1084        +------------------+
1085        ");
1086    }
1087
1088    #[test]
1089    fn test_pretty_format_interval_month_day_nano_array() {
1090        let arr = Arc::new(arrow_array::IntervalMonthDayNanoArray::from(vec![
1091            Some(IntervalMonthDayNano::new(-1, -1, -600_000_000_000)),
1092            Some(IntervalMonthDayNano::new(0, 0, -1_000_000_001)),
1093            Some(IntervalMonthDayNano::new(0, 0, -1)),
1094            Some(IntervalMonthDayNano::new(0, 0, 1)),
1095            Some(IntervalMonthDayNano::new(0, 0, 10)),
1096            Some(IntervalMonthDayNano::new(0, 0, 100)),
1097            Some(IntervalMonthDayNano::new(0, 0, 1_000)),
1098            Some(IntervalMonthDayNano::new(0, 0, 10_000)),
1099            Some(IntervalMonthDayNano::new(0, 0, 100_000)),
1100            Some(IntervalMonthDayNano::new(0, 0, 1_000_000)),
1101            Some(IntervalMonthDayNano::new(0, 0, 10_000_000)),
1102            Some(IntervalMonthDayNano::new(0, 0, 100_000_000)),
1103            Some(IntervalMonthDayNano::new(0, 0, 1_000_000_000)),
1104            Some(IntervalMonthDayNano::new(0, 0, 0)),
1105        ]));
1106
1107        let schema = Arc::new(Schema::new(vec![Field::new(
1108            "IntervalMonthDayNano",
1109            arr.data_type().clone(),
1110            true,
1111        )]));
1112
1113        let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
1114
1115        let table = pretty_format_batches(&[batch]).unwrap().to_string();
1116
1117        insta::assert_snapshot!(table, @"
1118        +--------------------------+
1119        | IntervalMonthDayNano     |
1120        +--------------------------+
1121        | -1 mons -1 days -10 mins |
1122        | -1.000000001 secs        |
1123        | -0.000000001 secs        |
1124        | 0.000000001 secs         |
1125        | 0.000000010 secs         |
1126        | 0.000000100 secs         |
1127        | 0.000001000 secs         |
1128        | 0.000010000 secs         |
1129        | 0.000100000 secs         |
1130        | 0.001000000 secs         |
1131        | 0.010000000 secs         |
1132        | 0.100000000 secs         |
1133        | 1.000000000 secs         |
1134        | 0 secs                   |
1135        +--------------------------+
1136        ");
1137    }
1138
1139    #[test]
1140    fn test_format_options() {
1141        let options = FormatOptions::default()
1142            .with_null("null")
1143            .with_types_info(true);
1144        let int32_array = Int32Array::from(vec![Some(1), Some(2), None, Some(3), Some(4)]);
1145        let string_array =
1146            StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz"), None]);
1147
1148        let batch = RecordBatch::try_from_iter([
1149            ("my_int32_name", Arc::new(int32_array) as _),
1150            ("my_string_name", Arc::new(string_array) as _),
1151        ])
1152        .unwrap();
1153
1154        let column = pretty_format_columns_with_options(
1155            "my_column_name",
1156            &[batch.column(0).clone()],
1157            &options,
1158        )
1159        .unwrap()
1160        .to_string();
1161
1162        insta::assert_snapshot!(column, @"
1163        +----------------+
1164        | my_column_name |
1165        +----------------+
1166        | 1              |
1167        | 2              |
1168        | null           |
1169        | 3              |
1170        | 4              |
1171        +----------------+
1172        ");
1173
1174        let table = pretty_format_batches_with_options(&[batch], &options)
1175            .unwrap()
1176            .to_string();
1177
1178        insta::assert_snapshot!(table, @"
1179        +---------------+----------------+
1180        | my_int32_name | my_string_name |
1181        | Int32         | Utf8           |
1182        +---------------+----------------+
1183        | 1             | foo            |
1184        | 2             | bar            |
1185        | null          | null           |
1186        | 3             | baz            |
1187        | 4             | null           |
1188        +---------------+----------------+
1189        ");
1190    }
1191
1192    #[test]
1193    fn duration_pretty_and_iso_extremes() {
1194        // Build [MIN, MAX, 3661, NULL]
1195        let arr = DurationSecondArray::from(vec![Some(i64::MIN), Some(i64::MAX), Some(3661), None]);
1196        let array: ArrayRef = Arc::new(arr);
1197
1198        // Pretty formatting
1199        let opts = FormatOptions::default().with_null("null");
1200        let opts = opts.with_duration_format(DurationFormat::Pretty);
1201        let pretty =
1202            pretty_format_columns_with_options("pretty", std::slice::from_ref(&array), &opts)
1203                .unwrap()
1204                .to_string();
1205
1206        insta::assert_snapshot!(pretty, @"
1207        +------------------------------+
1208        | pretty                       |
1209        +------------------------------+
1210        | <invalid>                    |
1211        | <invalid>                    |
1212        | 0 days 1 hours 1 mins 1 secs |
1213        | null                         |
1214        +------------------------------+
1215        ");
1216
1217        // ISO8601 formatting
1218        let opts_iso = FormatOptions::default()
1219            .with_null("null")
1220            .with_duration_format(DurationFormat::ISO8601);
1221        let iso = pretty_format_columns_with_options("iso", &[array], &opts_iso)
1222            .unwrap()
1223            .to_string();
1224
1225        insta::assert_snapshot!(iso, @"
1226        +-----------+
1227        | iso       |
1228        +-----------+
1229        | <invalid> |
1230        | <invalid> |
1231        | PT3661S   |
1232        | null      |
1233        +-----------+
1234        ");
1235    }
1236
1237    //
1238    // Custom Formatting
1239    //
1240
1241    /// The factory that will create the [`ArrayFormatter`]s.
1242    #[derive(Debug)]
1243    struct TestFormatters {}
1244
1245    impl ArrayFormatterFactory for TestFormatters {
1246        fn create_array_formatter<'formatter>(
1247            &self,
1248            array: &'formatter dyn Array,
1249            options: &FormatOptions<'formatter>,
1250            field: Option<&'formatter Field>,
1251        ) -> Result<Option<ArrayFormatter<'formatter>>, ArrowError> {
1252            if field
1253                .map(|f| f.extension_type_name() == Some("my_money"))
1254                .unwrap_or(false)
1255            {
1256                // We assume that my_money always is an Int32.
1257                let array = array.as_primitive();
1258                let display_index = Box::new(MyMoneyFormatter {
1259                    array,
1260                    options: options.clone(),
1261                });
1262                return Ok(Some(ArrayFormatter::new(display_index, options.safe())));
1263            }
1264
1265            if array.data_type() == &DataType::Int32 {
1266                let array = array.as_primitive();
1267                let display_index = Box::new(MyInt32Formatter {
1268                    array,
1269                    options: options.clone(),
1270                });
1271                return Ok(Some(ArrayFormatter::new(display_index, options.safe())));
1272            }
1273
1274            Ok(None)
1275        }
1276    }
1277
1278    /// A format that will append a "€" sign to the end of the Int32 values.
1279    struct MyMoneyFormatter<'a> {
1280        array: &'a Int32Array,
1281        options: FormatOptions<'a>,
1282    }
1283
1284    impl DisplayIndex for MyMoneyFormatter<'_> {
1285        fn write(&self, idx: usize, f: &mut dyn Write) -> crate::display::FormatResult {
1286            match self.array.is_valid(idx) {
1287                true => write!(f, "{} €", self.array.value(idx))?,
1288                false => write!(f, "{}", self.options.null())?,
1289            }
1290
1291            Ok(())
1292        }
1293    }
1294
1295    /// The actual formatter
1296    struct MyInt32Formatter<'a> {
1297        array: &'a Int32Array,
1298        options: FormatOptions<'a>,
1299    }
1300
1301    impl DisplayIndex for MyInt32Formatter<'_> {
1302        fn write(&self, idx: usize, f: &mut dyn Write) -> crate::display::FormatResult {
1303            match self.array.is_valid(idx) {
1304                true => write!(f, "{} (32-Bit)", self.array.value(idx))?,
1305                false => write!(f, "{}", self.options.null())?,
1306            }
1307
1308            Ok(())
1309        }
1310    }
1311
1312    #[test]
1313    fn test_format_batches_with_custom_formatters() {
1314        // define a schema.
1315        let options = FormatOptions::new()
1316            .with_null("<NULL>")
1317            .with_formatter_factory(Some(&TestFormatters {}));
1318        let money_metadata = HashMap::from([(
1319            extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1320            "my_money".to_owned(),
1321        )]);
1322        let schema = Arc::new(Schema::new(vec![
1323            Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone()),
1324        ]));
1325
1326        // define data.
1327        let batch = RecordBatch::try_new(
1328            schema,
1329            vec![Arc::new(array::Int32Array::from(vec![
1330                Some(1),
1331                None,
1332                Some(10),
1333                Some(100),
1334            ]))],
1335        )
1336        .unwrap();
1337
1338        let table = pretty_format_batches_with_options(&[batch], &options)
1339            .unwrap()
1340            .to_string();
1341
1342        insta::assert_snapshot!(table, @"
1343        +--------+
1344        | income |
1345        +--------+
1346        | 1 €    |
1347        | <NULL> |
1348        | 10 €   |
1349        | 100 €  |
1350        +--------+
1351        ");
1352    }
1353
1354    #[test]
1355    fn test_format_batches_with_custom_formatters_multi_nested_list() {
1356        // define a schema.
1357        let options = FormatOptions::new()
1358            .with_null("<NULL>")
1359            .with_formatter_factory(Some(&TestFormatters {}));
1360        let money_metadata = HashMap::from([(
1361            extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1362            "my_money".to_owned(),
1363        )]);
1364        let nested_field = Arc::new(
1365            Field::new_list_field(DataType::Int32, true).with_metadata(money_metadata.clone()),
1366        );
1367
1368        // Create nested data
1369        let inner_list = ListBuilder::new(Int32Builder::new()).with_field(nested_field);
1370        let mut outer_list = FixedSizeListBuilder::new(inner_list, 2);
1371        outer_list.values().append_value([Some(1)]);
1372        outer_list.values().append_null();
1373        outer_list.append(true);
1374        outer_list.values().append_value([Some(2), Some(8)]);
1375        outer_list
1376            .values()
1377            .append_value([Some(50), Some(25), Some(25)]);
1378        outer_list.append(true);
1379        let outer_list = outer_list.finish();
1380
1381        let schema = Arc::new(Schema::new(vec![Field::new(
1382            "income",
1383            outer_list.data_type().clone(),
1384            true,
1385        )]));
1386
1387        // define data.
1388        let batch = RecordBatch::try_new(schema, vec![Arc::new(outer_list)]).unwrap();
1389
1390        let table = pretty_format_batches_with_options(&[batch], &options)
1391            .unwrap()
1392            .to_string();
1393
1394        insta::assert_snapshot!(table, @"
1395        +----------------------------------+
1396        | income                           |
1397        +----------------------------------+
1398        | [[1 €], <NULL>]                  |
1399        | [[2 €, 8 €], [50 €, 25 €, 25 €]] |
1400        +----------------------------------+
1401        ");
1402    }
1403
1404    #[test]
1405    fn test_format_batches_with_custom_formatters_nested_struct() {
1406        // define a schema.
1407        let options = FormatOptions::new()
1408            .with_null("<NULL>")
1409            .with_formatter_factory(Some(&TestFormatters {}));
1410        let money_metadata = HashMap::from([(
1411            extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1412            "my_money".to_owned(),
1413        )]);
1414        let fields = Fields::from(vec![
1415            Field::new("name", DataType::Utf8, true),
1416            Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone()),
1417        ]);
1418
1419        let schema = Arc::new(Schema::new(vec![Field::new(
1420            "income",
1421            DataType::Struct(fields.clone()),
1422            true,
1423        )]));
1424
1425        // Create nested data
1426        let mut nested_data = StructBuilder::new(
1427            fields,
1428            vec![
1429                Box::new(StringBuilder::new()),
1430                Box::new(Int32Builder::new()),
1431            ],
1432        );
1433        nested_data
1434            .field_builder::<StringBuilder>(0)
1435            .unwrap()
1436            .extend([Some("Gimli"), Some("Legolas"), Some("Aragorn")]);
1437        nested_data
1438            .field_builder::<Int32Builder>(1)
1439            .unwrap()
1440            .extend([Some(10), None, Some(30)]);
1441        nested_data.append(true);
1442        nested_data.append(true);
1443        nested_data.append(true);
1444
1445        // define data.
1446        let batch = RecordBatch::try_new(schema, vec![Arc::new(nested_data.finish())]).unwrap();
1447
1448        let table = pretty_format_batches_with_options(&[batch], &options)
1449            .unwrap()
1450            .to_string();
1451
1452        insta::assert_snapshot!(table, @"
1453        +---------------------------------+
1454        | income                          |
1455        +---------------------------------+
1456        | {name: Gimli, income: 10 €}     |
1457        | {name: Legolas, income: <NULL>} |
1458        | {name: Aragorn, income: 30 €}   |
1459        +---------------------------------+
1460        ");
1461    }
1462
1463    #[test]
1464    fn test_format_batches_with_custom_formatters_nested_map() {
1465        // define a schema.
1466        let options = FormatOptions::new()
1467            .with_null("<NULL>")
1468            .with_formatter_factory(Some(&TestFormatters {}));
1469        let money_metadata = HashMap::from([(
1470            extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1471            "my_money".to_owned(),
1472        )]);
1473
1474        let mut array = MapBuilder::<StringBuilder, Int32Builder>::new(
1475            None,
1476            StringBuilder::new(),
1477            Int32Builder::new(),
1478        )
1479        .with_values_field(
1480            Field::new("my_values", DataType::Int32, true).with_metadata(money_metadata.clone()),
1481        );
1482        array
1483            .keys()
1484            .extend([Some("Gimli"), Some("Legolas"), Some("Aragorn")]);
1485        array.values().extend([Some(10), None, Some(30)]);
1486        array.append(true).unwrap();
1487        let array = array.finish();
1488
1489        // define data.
1490        let schema = Arc::new(Schema::new(vec![Field::new(
1491            "income",
1492            array.data_type().clone(),
1493            true,
1494        )]));
1495        let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
1496
1497        let table = pretty_format_batches_with_options(&[batch], &options)
1498            .unwrap()
1499            .to_string();
1500
1501        insta::assert_snapshot!(table, @"
1502        +-----------------------------------------------+
1503        | income                                        |
1504        +-----------------------------------------------+
1505        | {Gimli: 10 €, Legolas: <NULL>, Aragorn: 30 €} |
1506        +-----------------------------------------------+
1507        ");
1508    }
1509
1510    #[test]
1511    fn test_format_batches_with_custom_formatters_nested_union() {
1512        // define a schema.
1513        let options = FormatOptions::new()
1514            .with_null("<NULL>")
1515            .with_formatter_factory(Some(&TestFormatters {}));
1516        let money_metadata = HashMap::from([(
1517            extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1518            "my_money".to_owned(),
1519        )]);
1520        let fields = UnionFields::try_new(
1521            vec![0],
1522            vec![Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone())],
1523        )
1524        .unwrap();
1525
1526        // Create nested data and construct it with the correct metadata
1527        let mut array_builder = UnionBuilder::new_dense();
1528        array_builder.append::<Int32Type>("income", 1).unwrap();
1529        let (_, type_ids, offsets, children) = array_builder.build().unwrap().into_parts();
1530        let array = UnionArray::try_new(fields, type_ids, offsets, children).unwrap();
1531
1532        let schema = Arc::new(Schema::new(vec![Field::new(
1533            "income",
1534            array.data_type().clone(),
1535            true,
1536        )]));
1537
1538        // define data.
1539        let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
1540
1541        let table = pretty_format_batches_with_options(&[batch], &options)
1542            .unwrap()
1543            .to_string();
1544
1545        insta::assert_snapshot!(table, @"
1546        +--------------+
1547        | income       |
1548        +--------------+
1549        | {income=1 €} |
1550        +--------------+
1551        ");
1552    }
1553
1554    #[test]
1555    fn test_format_batches_with_custom_formatters_custom_schema_overrules_batch_schema() {
1556        // define a schema.
1557        let options = FormatOptions::new().with_formatter_factory(Some(&TestFormatters {}));
1558        let money_metadata = HashMap::from([(
1559            extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1560            "my_money".to_owned(),
1561        )]);
1562        let schema = Arc::new(Schema::new(vec![
1563            Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone()),
1564        ]));
1565
1566        // define data.
1567        let batch = RecordBatch::try_new(
1568            schema,
1569            vec![Arc::new(array::Int32Array::from(vec![
1570                Some(1),
1571                None,
1572                Some(10),
1573                Some(100),
1574            ]))],
1575        )
1576        .unwrap();
1577
1578        let table = create_table(
1579            // No metadata compared to test_format_batches_with_custom_formatters
1580            Some(Arc::new(Schema::new(vec![Field::new(
1581                "income",
1582                DataType::Int32,
1583                true,
1584            )]))),
1585            &[batch],
1586            &options,
1587        )
1588        .unwrap()
1589        .to_string();
1590
1591        // No € formatting as in test_format_batches_with_custom_formatters
1592        insta::assert_snapshot!(table, @"
1593        +--------------+
1594        | income       |
1595        +--------------+
1596        | 1 (32-Bit)   |
1597        |              |
1598        | 10 (32-Bit)  |
1599        | 100 (32-Bit) |
1600        +--------------+
1601        ");
1602    }
1603
1604    #[test]
1605    fn test_format_column_with_custom_formatters() {
1606        // define data.
1607        let array = Arc::new(array::Int32Array::from(vec![
1608            Some(1),
1609            None,
1610            Some(10),
1611            Some(100),
1612        ]));
1613
1614        let table = pretty_format_columns_with_options(
1615            "income",
1616            &[array],
1617            &FormatOptions::default().with_formatter_factory(Some(&TestFormatters {})),
1618        )
1619        .unwrap()
1620        .to_string();
1621
1622        insta::assert_snapshot!(table, @"
1623        +--------------+
1624        | income       |
1625        +--------------+
1626        | 1 (32-Bit)   |
1627        |              |
1628        | 10 (32-Bit)  |
1629        | 100 (32-Bit) |
1630        +--------------+
1631        ");
1632    }
1633
1634    #[test]
1635    fn test_pretty_format_batches_with_schema_with_wrong_number_of_fields() {
1636        let schema_a = Arc::new(Schema::new(vec![
1637            Field::new("a", DataType::Int32, true),
1638            Field::new("b", DataType::Utf8, true),
1639        ]));
1640        let schema_b = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
1641
1642        // define data.
1643        let batch = RecordBatch::try_new(
1644            schema_b,
1645            vec![Arc::new(array::Int32Array::from(vec![
1646                Some(1),
1647                None,
1648                Some(10),
1649                Some(100),
1650            ]))],
1651        )
1652        .unwrap();
1653
1654        let error = pretty_format_batches_with_schema(schema_a, &[batch])
1655            .err()
1656            .unwrap();
1657        insta::assert_snapshot!(error, @"Invalid argument error: Expected the same number of columns in a record batch (1) as the number of fields (2) in the schema");
1658    }
1659
1660    #[test]
1661    fn test_quoted_strings() {
1662        let schema = Arc::new(Schema::new(vec![Field::new(
1663            "strings",
1664            DataType::Utf8,
1665            true,
1666        )]));
1667
1668        let string_array = StringArray::from(vec![
1669            Some("hello"),
1670            Some("world"),
1671            Some(""),
1672            Some("tab\there"),
1673            Some("newline\ntest"),
1674            Some("quote\"test"),
1675            Some("backslash\\test"),
1676            None,
1677        ]);
1678
1679        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(string_array)]).unwrap();
1680
1681        let options_none = FormatOptions::new().with_null("NULL");
1682        let table = pretty_format_batches_with_options(std::slice::from_ref(&batch), &options_none)
1683            .unwrap()
1684            .to_string();
1685
1686        insta::assert_snapshot!(table, @"
1687        +----------------+
1688        | strings        |
1689        +----------------+
1690        | hello          |
1691        | world          |
1692        |                |
1693        | tab	here       |
1694        | newline        |
1695        | test           |
1696        | quote\"test     |
1697        | backslash\\test |
1698        | NULL           |
1699        +----------------+
1700        ");
1701
1702        let options_quoted = FormatOptions::new()
1703            .with_null("NULL")
1704            .with_quoted_strings(true);
1705
1706        let table = pretty_format_batches_with_options(&[batch], &options_quoted)
1707            .unwrap()
1708            .to_string();
1709
1710        insta::assert_snapshot!(table, @r#"
1711        +-------------------+
1712        | strings           |
1713        +-------------------+
1714        | "hello"           |
1715        | "world"           |
1716        | ""                |
1717        | "tab\there"       |
1718        | "newline\ntest"   |
1719        | "quote\"test"     |
1720        | "backslash\\test" |
1721        | NULL              |
1722        +-------------------+
1723        "#);
1724    }
1725
1726    #[test]
1727    fn test_string_view_quoted() {
1728        let schema = Arc::new(Schema::new(vec![Field::new(
1729            "view_strings",
1730            DataType::Utf8View,
1731            true,
1732        )]));
1733
1734        let mut builder = StringViewBuilder::new();
1735        builder.append_value("hello");
1736        builder.append_null();
1737        builder.append_value("quote\"test");
1738
1739        let array: ArrayRef = Arc::new(builder.finish());
1740        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
1741
1742        let options = FormatOptions::new().with_quoted_strings(true);
1743
1744        let table = pretty_format_batches_with_options(&[batch], &options)
1745            .unwrap()
1746            .to_string();
1747
1748        insta::assert_snapshot!(table, @"
1749        +---------------+
1750        | view_strings  |
1751        +---------------+
1752        | \"hello\"       |
1753        |               |
1754        | \"quote\\\"test\" |
1755        +---------------+
1756        ");
1757    }
1758
1759    #[test]
1760    fn test_quoted_strings_in_struct() {
1761        let string_builder = StringBuilder::new();
1762        let mut name_builder = string_builder;
1763        name_builder.append_value("Alice");
1764        name_builder.append_value("");
1765        name_builder.append_value("Bob");
1766
1767        let fields = vec![Field::new("name", DataType::Utf8, false)];
1768        let mut struct_builder = StructBuilder::new(fields, vec![Box::new(name_builder)]);
1769        struct_builder.append(true);
1770        struct_builder.append(true);
1771        struct_builder.append(true);
1772
1773        let struct_array = struct_builder.finish();
1774
1775        let schema = Arc::new(Schema::new(vec![Field::new(
1776            "person",
1777            struct_array.data_type().clone(),
1778            false,
1779        )]));
1780
1781        let batch = RecordBatch::try_new(schema, vec![Arc::new(struct_array)]).unwrap();
1782
1783        let options_none = FormatOptions::new();
1784        let table = pretty_format_batches_with_options(std::slice::from_ref(&batch), &options_none)
1785            .unwrap()
1786            .to_string();
1787
1788        insta::assert_snapshot!(table, @"
1789        +---------------+
1790        | person        |
1791        +---------------+
1792        | {name: Alice} |
1793        | {name: }      |
1794        | {name: Bob}   |
1795        +---------------+
1796        ");
1797
1798        let options_quoted = FormatOptions::new().with_quoted_strings(true);
1799        let table = pretty_format_batches_with_options(&[batch], &options_quoted)
1800            .unwrap()
1801            .to_string();
1802
1803        insta::assert_snapshot!(table, @"
1804        +-----------------+
1805        | person          |
1806        +-----------------+
1807        | {name: \"Alice\"} |
1808        | {name: \"\"}      |
1809        | {name: \"Bob\"}   |
1810        +-----------------+
1811        ");
1812    }
1813}