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