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