Skip to main content

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