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 ]));
1112
1113 let schema = Arc::new(Schema::new(vec![Field::new(
1114 "IntervalDayTime",
1115 arr.data_type().clone(),
1116 true,
1117 )]));
1118
1119 let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
1120
1121 let table = pretty_format_batches(&[batch]).unwrap().to_string();
1122
1123 let expected = vec![
1124 "+------------------+",
1125 "| IntervalDayTime |",
1126 "+------------------+",
1127 "| -1 days -10 mins |",
1128 "| -1.001 secs |",
1129 "| -0.001 secs |",
1130 "| 0.001 secs |",
1131 "| 0.010 secs |",
1132 "| 0.100 secs |",
1133 "+------------------+",
1134 ];
1135
1136 let actual: Vec<&str> = table.lines().collect();
1137
1138 assert_eq!(expected, actual, "Actual result:\n{table}");
1139 }
1140
1141 #[test]
1142 fn test_pretty_format_interval_month_day_nano_array() {
1143 let arr = Arc::new(arrow_array::IntervalMonthDayNanoArray::from(vec![
1144 Some(IntervalMonthDayNano::new(-1, -1, -600_000_000_000)),
1145 Some(IntervalMonthDayNano::new(0, 0, -1_000_000_001)),
1146 Some(IntervalMonthDayNano::new(0, 0, -1)),
1147 Some(IntervalMonthDayNano::new(0, 0, 1)),
1148 Some(IntervalMonthDayNano::new(0, 0, 10)),
1149 Some(IntervalMonthDayNano::new(0, 0, 100)),
1150 Some(IntervalMonthDayNano::new(0, 0, 1_000)),
1151 Some(IntervalMonthDayNano::new(0, 0, 10_000)),
1152 Some(IntervalMonthDayNano::new(0, 0, 100_000)),
1153 Some(IntervalMonthDayNano::new(0, 0, 1_000_000)),
1154 Some(IntervalMonthDayNano::new(0, 0, 10_000_000)),
1155 Some(IntervalMonthDayNano::new(0, 0, 100_000_000)),
1156 Some(IntervalMonthDayNano::new(0, 0, 1_000_000_000)),
1157 ]));
1158
1159 let schema = Arc::new(Schema::new(vec![Field::new(
1160 "IntervalMonthDayNano",
1161 arr.data_type().clone(),
1162 true,
1163 )]));
1164
1165 let batch = RecordBatch::try_new(schema, vec![arr]).unwrap();
1166
1167 let table = pretty_format_batches(&[batch]).unwrap().to_string();
1168
1169 let expected = vec![
1170 "+--------------------------+",
1171 "| IntervalMonthDayNano |",
1172 "+--------------------------+",
1173 "| -1 mons -1 days -10 mins |",
1174 "| -1.000000001 secs |",
1175 "| -0.000000001 secs |",
1176 "| 0.000000001 secs |",
1177 "| 0.000000010 secs |",
1178 "| 0.000000100 secs |",
1179 "| 0.000001000 secs |",
1180 "| 0.000010000 secs |",
1181 "| 0.000100000 secs |",
1182 "| 0.001000000 secs |",
1183 "| 0.010000000 secs |",
1184 "| 0.100000000 secs |",
1185 "| 1.000000000 secs |",
1186 "+--------------------------+",
1187 ];
1188
1189 let actual: Vec<&str> = table.lines().collect();
1190
1191 assert_eq!(expected, actual, "Actual result:\n{table}");
1192 }
1193
1194 #[test]
1195 fn test_format_options() {
1196 let options = FormatOptions::default()
1197 .with_null("null")
1198 .with_types_info(true);
1199 let int32_array = Int32Array::from(vec![Some(1), Some(2), None, Some(3), Some(4)]);
1200 let string_array =
1201 StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz"), None]);
1202
1203 let batch = RecordBatch::try_from_iter([
1204 ("my_int32_name", Arc::new(int32_array) as _),
1205 ("my_string_name", Arc::new(string_array) as _),
1206 ])
1207 .unwrap();
1208
1209 let column = pretty_format_columns_with_options(
1210 "my_column_name",
1211 &[batch.column(0).clone()],
1212 &options,
1213 )
1214 .unwrap()
1215 .to_string();
1216
1217 let expected_column = vec![
1218 "+----------------+",
1219 "| my_column_name |",
1220 "+----------------+",
1221 "| 1 |",
1222 "| 2 |",
1223 "| null |",
1224 "| 3 |",
1225 "| 4 |",
1226 "+----------------+",
1227 ];
1228
1229 let actual: Vec<&str> = column.lines().collect();
1230 assert_eq!(expected_column, actual, "Actual result:\n{column}");
1231
1232 let batch = pretty_format_batches_with_options(&[batch], &options)
1233 .unwrap()
1234 .to_string();
1235
1236 let expected_table = vec![
1237 "+---------------+----------------+",
1238 "| my_int32_name | my_string_name |",
1239 "| Int32 | Utf8 |",
1240 "+---------------+----------------+",
1241 "| 1 | foo |",
1242 "| 2 | bar |",
1243 "| null | null |",
1244 "| 3 | baz |",
1245 "| 4 | null |",
1246 "+---------------+----------------+",
1247 ];
1248
1249 let actual: Vec<&str> = batch.lines().collect();
1250 assert_eq!(expected_table, actual, "Actual result:\n{batch}");
1251 }
1252
1253 #[test]
1254 fn duration_pretty_and_iso_extremes() {
1255 let arr = DurationSecondArray::from(vec![Some(i64::MIN), Some(i64::MAX), Some(3661), None]);
1257 let array: ArrayRef = Arc::new(arr);
1258
1259 let opts = FormatOptions::default().with_null("null");
1261 let opts = opts.with_duration_format(DurationFormat::Pretty);
1262 let pretty =
1263 pretty_format_columns_with_options("pretty", std::slice::from_ref(&array), &opts)
1264 .unwrap()
1265 .to_string();
1266
1267 let expected_pretty = vec![
1269 "+------------------------------+",
1270 "| pretty |",
1271 "+------------------------------+",
1272 "| <invalid> |",
1273 "| <invalid> |",
1274 "| 0 days 1 hours 1 mins 1 secs |",
1275 "| null |",
1276 "+------------------------------+",
1277 ];
1278
1279 let actual: Vec<&str> = pretty.lines().collect();
1280 assert_eq!(expected_pretty, actual, "Actual result:\n{pretty}");
1281
1282 let opts_iso = FormatOptions::default()
1284 .with_null("null")
1285 .with_duration_format(DurationFormat::ISO8601);
1286 let iso = pretty_format_columns_with_options("iso", &[array], &opts_iso)
1287 .unwrap()
1288 .to_string();
1289
1290 let expected_iso = vec![
1292 "+-----------+",
1293 "| iso |",
1294 "+-----------+",
1295 "| <invalid> |",
1296 "| <invalid> |",
1297 "| PT3661S |",
1298 "| null |",
1299 "+-----------+",
1300 ];
1301
1302 let actual: Vec<&str> = iso.lines().collect();
1303 assert_eq!(expected_iso, actual, "Actual result:\n{iso}");
1304 }
1305
1306 #[derive(Debug)]
1312 struct TestFormatters {}
1313
1314 impl ArrayFormatterFactory for TestFormatters {
1315 fn create_array_formatter<'formatter>(
1316 &self,
1317 array: &'formatter dyn Array,
1318 options: &FormatOptions<'formatter>,
1319 field: Option<&'formatter Field>,
1320 ) -> Result<Option<ArrayFormatter<'formatter>>, ArrowError> {
1321 if field
1322 .map(|f| f.extension_type_name() == Some("my_money"))
1323 .unwrap_or(false)
1324 {
1325 let array = array.as_primitive();
1327 let display_index = Box::new(MyMoneyFormatter {
1328 array,
1329 options: options.clone(),
1330 });
1331 return Ok(Some(ArrayFormatter::new(display_index, options.safe())));
1332 }
1333
1334 if array.data_type() == &DataType::Int32 {
1335 let array = array.as_primitive();
1336 let display_index = Box::new(MyInt32Formatter {
1337 array,
1338 options: options.clone(),
1339 });
1340 return Ok(Some(ArrayFormatter::new(display_index, options.safe())));
1341 }
1342
1343 Ok(None)
1344 }
1345 }
1346
1347 struct MyMoneyFormatter<'a> {
1349 array: &'a Int32Array,
1350 options: FormatOptions<'a>,
1351 }
1352
1353 impl<'a> DisplayIndex for MyMoneyFormatter<'a> {
1354 fn write(&self, idx: usize, f: &mut dyn Write) -> crate::display::FormatResult {
1355 match self.array.is_valid(idx) {
1356 true => write!(f, "{} €", self.array.value(idx))?,
1357 false => write!(f, "{}", self.options.null())?,
1358 }
1359
1360 Ok(())
1361 }
1362 }
1363
1364 struct MyInt32Formatter<'a> {
1366 array: &'a Int32Array,
1367 options: FormatOptions<'a>,
1368 }
1369
1370 impl<'a> DisplayIndex for MyInt32Formatter<'a> {
1371 fn write(&self, idx: usize, f: &mut dyn Write) -> crate::display::FormatResult {
1372 match self.array.is_valid(idx) {
1373 true => write!(f, "{} (32-Bit)", self.array.value(idx))?,
1374 false => write!(f, "{}", self.options.null())?,
1375 }
1376
1377 Ok(())
1378 }
1379 }
1380
1381 #[test]
1382 fn test_format_batches_with_custom_formatters() {
1383 let options = FormatOptions::new()
1385 .with_null("<NULL>")
1386 .with_formatter_factory(Some(&TestFormatters {}));
1387 let money_metadata = HashMap::from([(
1388 extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1389 "my_money".to_owned(),
1390 )]);
1391 let schema = Arc::new(Schema::new(vec![
1392 Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone()),
1393 ]));
1394
1395 let batch = RecordBatch::try_new(
1397 schema,
1398 vec![Arc::new(array::Int32Array::from(vec![
1399 Some(1),
1400 None,
1401 Some(10),
1402 Some(100),
1403 ]))],
1404 )
1405 .unwrap();
1406
1407 let mut buf = String::new();
1408 write!(
1409 &mut buf,
1410 "{}",
1411 pretty_format_batches_with_options(&[batch], &options).unwrap()
1412 )
1413 .unwrap();
1414
1415 let s = [
1416 "+--------+",
1417 "| income |",
1418 "+--------+",
1419 "| 1 € |",
1420 "| <NULL> |",
1421 "| 10 € |",
1422 "| 100 € |",
1423 "+--------+",
1424 ];
1425 let expected = s.join("\n");
1426 assert_eq!(expected, buf);
1427 }
1428
1429 #[test]
1430 fn test_format_batches_with_custom_formatters_multi_nested_list() {
1431 let options = FormatOptions::new()
1433 .with_null("<NULL>")
1434 .with_formatter_factory(Some(&TestFormatters {}));
1435 let money_metadata = HashMap::from([(
1436 extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1437 "my_money".to_owned(),
1438 )]);
1439 let nested_field = Arc::new(
1440 Field::new_list_field(DataType::Int32, true).with_metadata(money_metadata.clone()),
1441 );
1442
1443 let inner_list = ListBuilder::new(Int32Builder::new()).with_field(nested_field);
1445 let mut outer_list = FixedSizeListBuilder::new(inner_list, 2);
1446 outer_list.values().append_value([Some(1)]);
1447 outer_list.values().append_null();
1448 outer_list.append(true);
1449 outer_list.values().append_value([Some(2), Some(8)]);
1450 outer_list
1451 .values()
1452 .append_value([Some(50), Some(25), Some(25)]);
1453 outer_list.append(true);
1454 let outer_list = outer_list.finish();
1455
1456 let schema = Arc::new(Schema::new(vec![Field::new(
1457 "income",
1458 outer_list.data_type().clone(),
1459 true,
1460 )]));
1461
1462 let batch = RecordBatch::try_new(schema, vec![Arc::new(outer_list)]).unwrap();
1464
1465 let mut buf = String::new();
1466 write!(
1467 &mut buf,
1468 "{}",
1469 pretty_format_batches_with_options(&[batch], &options).unwrap()
1470 )
1471 .unwrap();
1472
1473 let s = [
1474 "+----------------------------------+",
1475 "| income |",
1476 "+----------------------------------+",
1477 "| [[1 €], <NULL>] |",
1478 "| [[2 €, 8 €], [50 €, 25 €, 25 €]] |",
1479 "+----------------------------------+",
1480 ];
1481 let expected = s.join("\n");
1482 assert_eq!(expected, buf);
1483 }
1484
1485 #[test]
1486 fn test_format_batches_with_custom_formatters_nested_struct() {
1487 let options = FormatOptions::new()
1489 .with_null("<NULL>")
1490 .with_formatter_factory(Some(&TestFormatters {}));
1491 let money_metadata = HashMap::from([(
1492 extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1493 "my_money".to_owned(),
1494 )]);
1495 let fields = Fields::from(vec![
1496 Field::new("name", DataType::Utf8, true),
1497 Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone()),
1498 ]);
1499
1500 let schema = Arc::new(Schema::new(vec![Field::new(
1501 "income",
1502 DataType::Struct(fields.clone()),
1503 true,
1504 )]));
1505
1506 let mut nested_data = StructBuilder::new(
1508 fields,
1509 vec![
1510 Box::new(StringBuilder::new()),
1511 Box::new(Int32Builder::new()),
1512 ],
1513 );
1514 nested_data
1515 .field_builder::<StringBuilder>(0)
1516 .unwrap()
1517 .extend([Some("Gimli"), Some("Legolas"), Some("Aragorn")]);
1518 nested_data
1519 .field_builder::<Int32Builder>(1)
1520 .unwrap()
1521 .extend([Some(10), None, Some(30)]);
1522 nested_data.append(true);
1523 nested_data.append(true);
1524 nested_data.append(true);
1525
1526 let batch = RecordBatch::try_new(schema, vec![Arc::new(nested_data.finish())]).unwrap();
1528
1529 let mut buf = String::new();
1530 write!(
1531 &mut buf,
1532 "{}",
1533 pretty_format_batches_with_options(&[batch], &options).unwrap()
1534 )
1535 .unwrap();
1536
1537 let s = [
1538 "+---------------------------------+",
1539 "| income |",
1540 "+---------------------------------+",
1541 "| {name: Gimli, income: 10 €} |",
1542 "| {name: Legolas, income: <NULL>} |",
1543 "| {name: Aragorn, income: 30 €} |",
1544 "+---------------------------------+",
1545 ];
1546 let expected = s.join("\n");
1547 assert_eq!(expected, buf);
1548 }
1549
1550 #[test]
1551 fn test_format_batches_with_custom_formatters_nested_map() {
1552 let options = FormatOptions::new()
1554 .with_null("<NULL>")
1555 .with_formatter_factory(Some(&TestFormatters {}));
1556 let money_metadata = HashMap::from([(
1557 extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1558 "my_money".to_owned(),
1559 )]);
1560
1561 let mut array = MapBuilder::<StringBuilder, Int32Builder>::new(
1562 None,
1563 StringBuilder::new(),
1564 Int32Builder::new(),
1565 )
1566 .with_values_field(
1567 Field::new("values", DataType::Int32, true).with_metadata(money_metadata.clone()),
1568 );
1569 array
1570 .keys()
1571 .extend([Some("Gimli"), Some("Legolas"), Some("Aragorn")]);
1572 array.values().extend([Some(10), None, Some(30)]);
1573 array.append(true).unwrap();
1574 let array = array.finish();
1575
1576 let schema = Arc::new(Schema::new(vec![Field::new(
1578 "income",
1579 array.data_type().clone(),
1580 true,
1581 )]));
1582 let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
1583
1584 let mut buf = String::new();
1585 write!(
1586 &mut buf,
1587 "{}",
1588 pretty_format_batches_with_options(&[batch], &options).unwrap()
1589 )
1590 .unwrap();
1591
1592 let s = [
1593 "+-----------------------------------------------+",
1594 "| income |",
1595 "+-----------------------------------------------+",
1596 "| {Gimli: 10 €, Legolas: <NULL>, Aragorn: 30 €} |",
1597 "+-----------------------------------------------+",
1598 ];
1599 let expected = s.join("\n");
1600 assert_eq!(expected, buf);
1601 }
1602
1603 #[test]
1604 fn test_format_batches_with_custom_formatters_nested_union() {
1605 let options = FormatOptions::new()
1607 .with_null("<NULL>")
1608 .with_formatter_factory(Some(&TestFormatters {}));
1609 let money_metadata = HashMap::from([(
1610 extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1611 "my_money".to_owned(),
1612 )]);
1613 let fields = UnionFields::new(
1614 vec![0],
1615 vec![Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone())],
1616 );
1617
1618 let mut array_builder = UnionBuilder::new_dense();
1620 array_builder.append::<Int32Type>("income", 1).unwrap();
1621 let (_, type_ids, offsets, children) = array_builder.build().unwrap().into_parts();
1622 let array = UnionArray::try_new(fields, type_ids, offsets, children).unwrap();
1623
1624 let schema = Arc::new(Schema::new(vec![Field::new(
1625 "income",
1626 array.data_type().clone(),
1627 true,
1628 )]));
1629
1630 let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
1632
1633 let mut buf = String::new();
1634 write!(
1635 &mut buf,
1636 "{}",
1637 pretty_format_batches_with_options(&[batch], &options).unwrap()
1638 )
1639 .unwrap();
1640
1641 let s = [
1642 "+--------------+",
1643 "| income |",
1644 "+--------------+",
1645 "| {income=1 €} |",
1646 "+--------------+",
1647 ];
1648 let expected = s.join("\n");
1649 assert_eq!(expected, buf);
1650 }
1651
1652 #[test]
1653 fn test_format_batches_with_custom_formatters_custom_schema_overrules_batch_schema() {
1654 let options = FormatOptions::new().with_formatter_factory(Some(&TestFormatters {}));
1656 let money_metadata = HashMap::from([(
1657 extension::EXTENSION_TYPE_NAME_KEY.to_owned(),
1658 "my_money".to_owned(),
1659 )]);
1660 let schema = Arc::new(Schema::new(vec![
1661 Field::new("income", DataType::Int32, true).with_metadata(money_metadata.clone()),
1662 ]));
1663
1664 let batch = RecordBatch::try_new(
1666 schema,
1667 vec![Arc::new(array::Int32Array::from(vec![
1668 Some(1),
1669 None,
1670 Some(10),
1671 Some(100),
1672 ]))],
1673 )
1674 .unwrap();
1675
1676 let mut buf = String::new();
1677 write!(
1678 &mut buf,
1679 "{}",
1680 create_table(
1681 Some(Arc::new(Schema::new(vec![Field::new(
1683 "income",
1684 DataType::Int32,
1685 true
1686 ),]))),
1687 &[batch],
1688 &options,
1689 )
1690 .unwrap()
1691 )
1692 .unwrap();
1693
1694 let s = [
1696 "+--------------+",
1697 "| income |",
1698 "+--------------+",
1699 "| 1 (32-Bit) |",
1700 "| |",
1701 "| 10 (32-Bit) |",
1702 "| 100 (32-Bit) |",
1703 "+--------------+",
1704 ];
1705 let expected = s.join("\n");
1706 assert_eq!(expected, buf);
1707 }
1708
1709 #[test]
1710 fn test_format_column_with_custom_formatters() {
1711 let array = Arc::new(array::Int32Array::from(vec![
1713 Some(1),
1714 None,
1715 Some(10),
1716 Some(100),
1717 ]));
1718
1719 let mut buf = String::new();
1720 write!(
1721 &mut buf,
1722 "{}",
1723 pretty_format_columns_with_options(
1724 "income",
1725 &[array],
1726 &FormatOptions::default().with_formatter_factory(Some(&TestFormatters {}))
1727 )
1728 .unwrap()
1729 )
1730 .unwrap();
1731
1732 let s = [
1733 "+--------------+",
1734 "| income |",
1735 "+--------------+",
1736 "| 1 (32-Bit) |",
1737 "| |",
1738 "| 10 (32-Bit) |",
1739 "| 100 (32-Bit) |",
1740 "+--------------+",
1741 ];
1742 let expected = s.join("\n");
1743 assert_eq!(expected, buf);
1744 }
1745
1746 #[test]
1747 fn test_pretty_format_batches_with_schema_with_wrong_number_of_fields() {
1748 let schema_a = Arc::new(Schema::new(vec![
1749 Field::new("a", DataType::Int32, true),
1750 Field::new("b", DataType::Utf8, true),
1751 ]));
1752 let schema_b = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
1753
1754 let batch = RecordBatch::try_new(
1756 schema_b,
1757 vec![Arc::new(array::Int32Array::from(vec![
1758 Some(1),
1759 None,
1760 Some(10),
1761 Some(100),
1762 ]))],
1763 )
1764 .unwrap();
1765
1766 let error = pretty_format_batches_with_schema(schema_a, &[batch])
1767 .err()
1768 .unwrap();
1769 assert_eq!(
1770 &error.to_string(),
1771 "Invalid argument error: Expected the same number of columns in a record batch (1) as the number of fields (2) in the schema"
1772 );
1773 }
1774}