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