| 682 | |
| 683 | #[test] |
| 684 | fn test_write_csv() { |
| 685 | let schema = Schema::new(vec![ |
| 686 | Field::new("c1", DataType::Utf8, false), |
| 687 | Field::new("c2", DataType::Float64, true), |
| 688 | Field::new("c3", DataType::UInt32, false), |
| 689 | Field::new("c4", DataType::Boolean, true), |
| 690 | Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true), |
| 691 | Field::new("c6", DataType::Time32(TimeUnit::Second), false), |
| 692 | Field::new_dictionary("c7", DataType::Int32, DataType::Utf8, false), |
| 693 | ]); |
| 694 | |
| 695 | let c1 = StringArray::from(vec![ |
| 696 | "Lorem ipsum dolor sit amet", |
| 697 | "consectetur adipiscing elit", |
| 698 | "sed do eiusmod tempor", |
| 699 | ]); |
| 700 | let c2 = |
| 701 | PrimitiveArray::<Float64Type>::from(vec![Some(123.564532), None, Some(-556132.25)]); |
| 702 | let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]); |
| 703 | let c4 = BooleanArray::from(vec![Some(true), Some(false), None]); |
| 704 | let c5 = |
| 705 | TimestampMillisecondArray::from(vec![None, Some(1555584887378), Some(1555555555555)]); |
| 706 | let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]); |
| 707 | let c7: DictionaryArray<Int32Type> = |
| 708 | vec!["cupcakes", "cupcakes", "foo"].into_iter().collect(); |
| 709 | |
| 710 | let batch = RecordBatch::try_new( |
| 711 | Arc::new(schema), |
| 712 | vec![ |
| 713 | Arc::new(c1), |
| 714 | Arc::new(c2), |
| 715 | Arc::new(c3), |
| 716 | Arc::new(c4), |
| 717 | Arc::new(c5), |
| 718 | Arc::new(c6), |
| 719 | Arc::new(c7), |
| 720 | ], |
| 721 | ) |
| 722 | .unwrap(); |
| 723 | |
| 724 | let mut file = tempfile::tempfile().unwrap(); |
| 725 | |
| 726 | let mut writer = Writer::new(&mut file); |
| 727 | let batches = vec![&batch, &batch]; |
| 728 | for batch in batches { |
| 729 | writer.write(batch).unwrap(); |
| 730 | } |
| 731 | drop(writer); |
| 732 | |
| 733 | // check that file was written successfully |
| 734 | file.rewind().unwrap(); |
| 735 | let mut buffer: Vec<u8> = vec![]; |
| 736 | file.read_to_end(&mut buffer).unwrap(); |
| 737 | |
| 738 | let expected = r#"c1,c2,c3,c4,c5,c6,c7 |
| 739 | Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes |
| 740 | consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378,06:51:20,cupcakes |
| 741 | sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo |