1use arrow_array::*;
177use arrow_cast::display::*;
178use arrow_schema::*;
179use csv::ByteRecord;
180use std::io::Write;
181
182use crate::map_csv_error;
183const DEFAULT_NULL_VALUE: &str = "";
184
185pub use csv::QuoteStyle;
200
201#[derive(Debug)]
205pub struct Writer<W: Write> {
206 writer: csv::Writer<W>,
208 has_headers: bool,
210 date_format: Option<String>,
212 datetime_format: Option<String>,
214 timestamp_format: Option<String>,
216 timestamp_tz_format: Option<String>,
218 time_format: Option<String>,
220 beginning: bool,
222 null_value: Option<String>,
224 ignore_leading_whitespace: bool,
226 ignore_trailing_whitespace: bool,
228}
229
230impl<W: Write> Writer<W> {
231 pub fn new(writer: W) -> Self {
236 let delimiter = b',';
237 WriterBuilder::new().with_delimiter(delimiter).build(writer)
238 }
239
240 pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
242 let num_columns = batch.num_columns();
243 if self.beginning {
244 if self.has_headers {
245 let mut headers: Vec<String> = Vec::with_capacity(num_columns);
246 batch
247 .schema()
248 .fields()
249 .iter()
250 .for_each(|field| headers.push(field.name().to_string()));
251 self.writer
252 .write_record(&headers[..])
253 .map_err(map_csv_error)?;
254 }
255 self.beginning = false;
256 }
257
258 let options = FormatOptions::default()
259 .with_null(self.null_value.as_deref().unwrap_or(DEFAULT_NULL_VALUE))
260 .with_date_format(self.date_format.as_deref())
261 .with_datetime_format(self.datetime_format.as_deref())
262 .with_timestamp_format(self.timestamp_format.as_deref())
263 .with_timestamp_tz_format(self.timestamp_tz_format.as_deref())
264 .with_time_format(self.time_format.as_deref());
265
266 let converters = batch
267 .columns()
268 .iter()
269 .map(|a| {
270 if a.data_type().is_nested() {
271 Err(ArrowError::CsvError(format!(
272 "Nested type {} is not supported in CSV",
273 a.data_type()
274 )))
275 } else {
276 ArrayFormatter::try_new(a.as_ref(), &options)
277 }
278 })
279 .collect::<Result<Vec<_>, ArrowError>>()?;
280
281 let mut buffer = String::with_capacity(1024);
282 let mut byte_record = ByteRecord::with_capacity(1024, converters.len());
283
284 for row_idx in 0..batch.num_rows() {
285 byte_record.clear();
286 for (col_idx, converter) in converters.iter().enumerate() {
287 buffer.clear();
288 converter.value(row_idx).write(&mut buffer).map_err(|e| {
289 ArrowError::CsvError(format!(
290 "Error processing row {}, col {}: {e}",
291 row_idx + 1,
292 col_idx + 1
293 ))
294 })?;
295
296 let field_bytes =
297 self.get_trimmed_field_bytes(&buffer, batch.column(col_idx).data_type());
298 byte_record.push_field(field_bytes);
299 }
300
301 self.writer
302 .write_byte_record(&byte_record)
303 .map_err(map_csv_error)?;
304 }
305 self.writer.flush()?;
306
307 Ok(())
308 }
309
310 fn get_trimmed_field_bytes<'a>(&self, buffer: &'a str, data_type: &DataType) -> &'a [u8] {
312 let should_trim = (self.ignore_leading_whitespace || self.ignore_trailing_whitespace)
314 && matches!(
315 data_type,
316 DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
317 );
318
319 if !should_trim {
320 return buffer.as_bytes();
321 }
322
323 let mut trimmed = buffer;
324 if self.ignore_leading_whitespace {
325 trimmed = trimmed.trim_start();
326 }
327 if self.ignore_trailing_whitespace {
328 trimmed = trimmed.trim_end();
329 }
330 trimmed.as_bytes()
331 }
332
333 pub fn into_inner(self) -> W {
335 self.writer.into_inner().unwrap()
337 }
338}
339
340impl<W: Write> RecordBatchWriter for Writer<W> {
341 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
342 self.write(batch)
343 }
344
345 fn close(self) -> Result<(), ArrowError> {
346 Ok(())
347 }
348}
349
350#[derive(Clone, Debug)]
352pub struct WriterBuilder {
353 delimiter: u8,
355 has_header: bool,
357 quote: u8,
359 escape: u8,
361 terminator: Terminator,
363 double_quote: bool,
365 date_format: Option<String>,
367 datetime_format: Option<String>,
369 timestamp_format: Option<String>,
371 timestamp_tz_format: Option<String>,
373 time_format: Option<String>,
375 null_value: Option<String>,
377 ignore_leading_whitespace: bool,
379 ignore_trailing_whitespace: bool,
381 quote_style: QuoteStyle,
383}
384
385#[derive(Clone, Debug)]
387pub enum Terminator {
388 CRLF,
390 Any(u8),
392}
393
394impl Default for WriterBuilder {
395 fn default() -> Self {
396 WriterBuilder {
397 delimiter: b',',
398 has_header: true,
399 quote: b'"',
400 escape: b'\\',
401 terminator: Terminator::Any(b'\n'),
402 double_quote: true,
403 date_format: None,
404 datetime_format: None,
405 timestamp_format: None,
406 timestamp_tz_format: None,
407 time_format: None,
408 null_value: None,
409 ignore_leading_whitespace: false,
410 ignore_trailing_whitespace: false,
411 quote_style: QuoteStyle::default(),
412 }
413 }
414}
415
416impl WriterBuilder {
417 pub fn new() -> Self {
439 Self::default()
440 }
441
442 pub fn with_header(mut self, header: bool) -> Self {
444 self.has_header = header;
445 self
446 }
447
448 pub fn header(&self) -> bool {
450 self.has_header
451 }
452
453 pub fn with_delimiter(mut self, delimiter: u8) -> Self {
455 self.delimiter = delimiter;
456 self
457 }
458
459 pub fn delimiter(&self) -> u8 {
461 self.delimiter
462 }
463
464 pub fn with_quote(mut self, quote: u8) -> Self {
466 self.quote = quote;
467 self
468 }
469
470 pub fn quote(&self) -> u8 {
472 self.quote
473 }
474
475 pub fn with_escape(mut self, escape: u8) -> Self {
483 self.escape = escape;
484 self
485 }
486
487 pub fn escape(&self) -> u8 {
489 self.escape
490 }
491
492 pub fn with_double_quote(mut self, double_quote: bool) -> Self {
500 self.double_quote = double_quote;
501 self
502 }
503
504 pub fn double_quote(&self) -> bool {
506 self.double_quote
507 }
508
509 pub fn with_date_format(mut self, format: String) -> Self {
511 self.date_format = Some(format);
512 self
513 }
514
515 pub fn date_format(&self) -> Option<&str> {
517 self.date_format.as_deref()
518 }
519
520 pub fn with_datetime_format(mut self, format: String) -> Self {
522 self.datetime_format = Some(format);
523 self
524 }
525
526 pub fn datetime_format(&self) -> Option<&str> {
528 self.datetime_format.as_deref()
529 }
530
531 pub fn with_time_format(mut self, format: String) -> Self {
533 self.time_format = Some(format);
534 self
535 }
536
537 pub fn time_format(&self) -> Option<&str> {
539 self.time_format.as_deref()
540 }
541
542 pub fn with_timestamp_format(mut self, format: String) -> Self {
544 self.timestamp_format = Some(format);
545 self
546 }
547
548 pub fn timestamp_format(&self) -> Option<&str> {
550 self.timestamp_format.as_deref()
551 }
552
553 pub fn with_timestamp_tz_format(mut self, tz_format: String) -> Self {
555 self.timestamp_tz_format = Some(tz_format);
556 self
557 }
558
559 pub fn timestamp_tz_format(&self) -> Option<&str> {
561 self.timestamp_tz_format.as_deref()
562 }
563
564 pub fn with_null(mut self, null_value: String) -> Self {
566 self.null_value = Some(null_value);
567 self
568 }
569
570 pub fn null(&self) -> &str {
572 self.null_value.as_deref().unwrap_or(DEFAULT_NULL_VALUE)
573 }
574
575 pub fn with_ignore_leading_whitespace(mut self, ignore: bool) -> Self {
578 self.ignore_leading_whitespace = ignore;
579 self
580 }
581
582 pub fn ignore_leading_whitespace(&self) -> bool {
584 self.ignore_leading_whitespace
585 }
586
587 pub fn with_ignore_trailing_whitespace(mut self, ignore: bool) -> Self {
590 self.ignore_trailing_whitespace = ignore;
591 self
592 }
593
594 pub fn ignore_trailing_whitespace(&self) -> bool {
596 self.ignore_trailing_whitespace
597 }
598
599 pub fn with_quote_style(mut self, quote_style: QuoteStyle) -> Self {
615 self.quote_style = quote_style;
616 self
617 }
618
619 pub fn quote_style(&self) -> QuoteStyle {
621 self.quote_style
622 }
623
624 pub fn with_line_terminator(mut self, terminator: Terminator) -> Self {
626 self.terminator = terminator;
627 self
628 }
629
630 pub fn line_terminator(&self) -> &Terminator {
632 &self.terminator
633 }
634
635 pub fn build<W: Write>(self, writer: W) -> Writer<W> {
637 let mut builder = csv::WriterBuilder::new();
638
639 let terminator = match self.terminator {
640 Terminator::CRLF => csv::Terminator::CRLF,
641 Terminator::Any(byte) => csv::Terminator::Any(byte),
642 };
643
644 let writer = builder
645 .delimiter(self.delimiter)
646 .quote(self.quote)
647 .quote_style(self.quote_style)
648 .double_quote(self.double_quote)
649 .escape(self.escape)
650 .terminator(terminator)
651 .from_writer(writer);
652 Writer {
653 writer,
654 beginning: true,
655 has_headers: self.has_header,
656 date_format: self.date_format,
657 datetime_format: self.datetime_format,
658 time_format: self.time_format,
659 timestamp_format: self.timestamp_format,
660 timestamp_tz_format: self.timestamp_tz_format,
661 null_value: self.null_value,
662 ignore_leading_whitespace: self.ignore_leading_whitespace,
663 ignore_trailing_whitespace: self.ignore_trailing_whitespace,
664 }
665 }
666}
667
668#[cfg(test)]
669#[expect(
670 clippy::verbose_file_reads,
671 reason = "`tempfile::tempfile()` has no path, so the contents can only be read back through the handle"
672)]
673mod tests {
674 use super::*;
675
676 use crate::ReaderBuilder;
677 use arrow_array::builder::{
678 BinaryBuilder, Decimal32Builder, Decimal64Builder, Decimal128Builder, Decimal256Builder,
679 FixedSizeBinaryBuilder, LargeBinaryBuilder,
680 };
681 use arrow_array::types::*;
682 use arrow_buffer::i256;
683 use core::str;
684 use std::io::{Cursor, Read, Seek};
685 use std::sync::Arc;
686
687 #[test]
688 fn test_write_csv() {
689 let schema = Schema::new(vec![
690 Field::new("c1", DataType::Utf8, false),
691 Field::new("c2", DataType::Float64, true),
692 Field::new("c3", DataType::UInt32, false),
693 Field::new("c4", DataType::Boolean, true),
694 Field::new("c5", DataType::Timestamp(TimeUnit::Millisecond, None), true),
695 Field::new("c6", DataType::Time32(TimeUnit::Second), false),
696 Field::new_dictionary("c7", DataType::Int32, DataType::Utf8, false),
697 ]);
698
699 let c1 = StringArray::from(vec![
700 "Lorem ipsum dolor sit amet",
701 "consectetur adipiscing elit",
702 "sed do eiusmod tempor",
703 ]);
704 let c2 =
705 PrimitiveArray::<Float64Type>::from(vec![Some(123.564532), None, Some(-556132.25)]);
706 let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
707 let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
708 let c5 =
709 TimestampMillisecondArray::from(vec![None, Some(1555584887378), Some(1555555555555)]);
710 let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
711 let c7: DictionaryArray<Int32Type> =
712 vec!["cupcakes", "cupcakes", "foo"].into_iter().collect();
713
714 let batch = RecordBatch::try_new(
715 Arc::new(schema),
716 vec![
717 Arc::new(c1),
718 Arc::new(c2),
719 Arc::new(c3),
720 Arc::new(c4),
721 Arc::new(c5),
722 Arc::new(c6),
723 Arc::new(c7),
724 ],
725 )
726 .unwrap();
727
728 let mut file = tempfile::tempfile().unwrap();
729
730 let mut writer = Writer::new(&mut file);
731 let batches = vec![&batch, &batch];
732 for batch in batches {
733 writer.write(batch).unwrap();
734 }
735 drop(writer);
736
737 file.rewind().unwrap();
739 let mut buffer: Vec<u8> = vec![];
740 file.read_to_end(&mut buffer).unwrap();
741
742 let expected = r"c1,c2,c3,c4,c5,c6,c7
743Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes
744consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378,06:51:20,cupcakes
745sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo
746Lorem ipsum dolor sit amet,123.564532,3,true,,00:20:34,cupcakes
747consectetur adipiscing elit,,2,false,2019-04-18T10:54:47.378,06:51:20,cupcakes
748sed do eiusmod tempor,-556132.25,1,,2019-04-18T02:45:55.555,23:46:03,foo
749";
750 assert_eq!(expected, str::from_utf8(&buffer).unwrap());
751 }
752
753 #[test]
754 fn test_write_csv_decimal() {
755 let schema = Schema::new(vec![
756 Field::new("c1", DataType::Decimal32(9, 6), true),
757 Field::new("c2", DataType::Decimal64(17, 6), true),
758 Field::new("c3", DataType::Decimal128(38, 6), true),
759 Field::new("c4", DataType::Decimal256(76, 6), true),
760 ]);
761
762 let mut c1_builder = Decimal32Builder::new().with_data_type(DataType::Decimal32(9, 6));
763 c1_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]);
764 let c1 = c1_builder.finish();
765
766 let mut c2_builder = Decimal64Builder::new().with_data_type(DataType::Decimal64(17, 6));
767 c2_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]);
768 let c2 = c2_builder.finish();
769
770 let mut c3_builder = Decimal128Builder::new().with_data_type(DataType::Decimal128(38, 6));
771 c3_builder.extend(vec![Some(-3335724), Some(2179404), None, Some(290472)]);
772 let c3 = c3_builder.finish();
773
774 let mut c4_builder = Decimal256Builder::new().with_data_type(DataType::Decimal256(76, 6));
775 c4_builder.extend(vec![
776 Some(i256::from_i128(-3335724)),
777 Some(i256::from_i128(2179404)),
778 None,
779 Some(i256::from_i128(290472)),
780 ]);
781 let c4 = c4_builder.finish();
782
783 let batch = RecordBatch::try_new(
784 Arc::new(schema),
785 vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
786 )
787 .unwrap();
788
789 let mut file = tempfile::tempfile().unwrap();
790
791 let mut writer = Writer::new(&mut file);
792 let batches = vec![&batch, &batch];
793 for batch in batches {
794 writer.write(batch).unwrap();
795 }
796 drop(writer);
797
798 file.rewind().unwrap();
800 let mut buffer: Vec<u8> = vec![];
801 file.read_to_end(&mut buffer).unwrap();
802
803 let expected = r"c1,c2,c3,c4
804-3.335724,-3.335724,-3.335724,-3.335724
8052.179404,2.179404,2.179404,2.179404
806,,,
8070.290472,0.290472,0.290472,0.290472
808-3.335724,-3.335724,-3.335724,-3.335724
8092.179404,2.179404,2.179404,2.179404
810,,,
8110.290472,0.290472,0.290472,0.290472
812";
813 assert_eq!(expected, str::from_utf8(&buffer).unwrap());
814 }
815
816 #[test]
817 fn test_write_csv_custom_options() {
818 let schema = Schema::new(vec![
819 Field::new("c1", DataType::Utf8, false),
820 Field::new("c2", DataType::Float64, true),
821 Field::new("c3", DataType::UInt32, false),
822 Field::new("c4", DataType::Boolean, true),
823 Field::new("c6", DataType::Time32(TimeUnit::Second), false),
824 ]);
825
826 let c1 = StringArray::from(vec![
827 "Lorem ipsum \ndolor sit amet",
828 "consectetur \"adipiscing\" elit",
829 "sed do eiusmod tempor",
830 ]);
831 let c2 =
832 PrimitiveArray::<Float64Type>::from(vec![Some(123.564532), None, Some(-556132.25)]);
833 let c3 = PrimitiveArray::<UInt32Type>::from(vec![3, 2, 1]);
834 let c4 = BooleanArray::from(vec![Some(true), Some(false), None]);
835 let c6 = Time32SecondArray::from(vec![1234, 24680, 85563]);
836
837 let batch = RecordBatch::try_new(
838 Arc::new(schema),
839 vec![
840 Arc::new(c1),
841 Arc::new(c2),
842 Arc::new(c3),
843 Arc::new(c4),
844 Arc::new(c6),
845 ],
846 )
847 .unwrap();
848
849 let mut file = tempfile::tempfile().unwrap();
850
851 let builder = WriterBuilder::new()
852 .with_header(false)
853 .with_delimiter(b'|')
854 .with_quote(b'\'')
855 .with_null("NULL".to_string())
856 .with_time_format("%r".to_string());
857 let mut writer = builder.build(&mut file);
858 let batches = vec![&batch];
859 for batch in batches {
860 writer.write(batch).unwrap();
861 }
862 drop(writer);
863
864 file.rewind().unwrap();
866 let mut buffer: Vec<u8> = vec![];
867 file.read_to_end(&mut buffer).unwrap();
868
869 assert_eq!(
870 "'Lorem ipsum \ndolor sit amet'|123.564532|3|true|12:20:34 AM\nconsectetur \"adipiscing\" elit|NULL|2|false|06:51:20 AM\nsed do eiusmod tempor|-556132.25|1|NULL|11:46:03 PM\n"
871 .to_string(),
872 String::from_utf8(buffer).unwrap()
873 );
874
875 let mut file = tempfile::tempfile().unwrap();
876
877 let builder = WriterBuilder::new()
878 .with_header(true)
879 .with_double_quote(false)
880 .with_escape(b'$');
881 let mut writer = builder.build(&mut file);
882 let batches = vec![&batch];
883 for batch in batches {
884 writer.write(batch).unwrap();
885 }
886 drop(writer);
887
888 file.rewind().unwrap();
889 let mut buffer: Vec<u8> = vec![];
890 file.read_to_end(&mut buffer).unwrap();
891
892 assert_eq!(
893 "c1,c2,c3,c4,c6\n\"Lorem ipsum \ndolor sit amet\",123.564532,3,true,00:20:34\n\"consectetur $\"adipiscing$\" elit\",,2,false,06:51:20\nsed do eiusmod tempor,-556132.25,1,,23:46:03\n"
894 .to_string(),
895 String::from_utf8(buffer).unwrap()
896 );
897 }
898
899 #[test]
900 fn test_conversion_consistency() {
901 let schema = Schema::new(vec![
904 Field::new("c1", DataType::Date32, false),
905 Field::new("c2", DataType::Date64, false),
906 Field::new("c3", DataType::Timestamp(TimeUnit::Nanosecond, None), false),
907 ]);
908
909 let nanoseconds = vec![
910 1599566300000000000,
911 1599566200000000000,
912 1599566100000000000,
913 ];
914 let c1 = Date32Array::from(vec![3, 2, 1]);
915 let c2 = Date64Array::from(vec![3, 2, 1]);
916 let c3 = TimestampNanosecondArray::from(nanoseconds.clone());
917
918 let batch = RecordBatch::try_new(
919 Arc::new(schema.clone()),
920 vec![Arc::new(c1), Arc::new(c2), Arc::new(c3)],
921 )
922 .unwrap();
923
924 let builder = WriterBuilder::new().with_header(false);
925
926 let mut buf: Cursor<Vec<u8>> = Default::default();
927 {
929 let mut writer = builder.build(&mut buf);
930 writer.write(&batch).unwrap();
931 }
932 buf.set_position(0);
933
934 let mut reader = ReaderBuilder::new(Arc::new(schema))
935 .with_batch_size(3)
936 .build_buffered(buf)
937 .unwrap();
938
939 let rb = reader.next().unwrap().unwrap();
940 let c1 = rb.column(0).as_any().downcast_ref::<Date32Array>().unwrap();
941 let c2 = rb.column(1).as_any().downcast_ref::<Date64Array>().unwrap();
942 let c3 = rb
943 .column(2)
944 .as_any()
945 .downcast_ref::<TimestampNanosecondArray>()
946 .unwrap();
947
948 let actual = c1.into_iter().collect::<Vec<_>>();
949 let expected = vec![Some(3), Some(2), Some(1)];
950 assert_eq!(actual, expected);
951 let actual = c2.into_iter().collect::<Vec<_>>();
952 let expected = vec![Some(3), Some(2), Some(1)];
953 assert_eq!(actual, expected);
954 let actual = c3.into_iter().collect::<Vec<_>>();
955 let expected = nanoseconds.into_iter().map(Some).collect::<Vec<_>>();
956 assert_eq!(actual, expected);
957 }
958
959 #[test]
960 fn test_write_csv_invalid_cast() {
961 let schema = Schema::new(vec![
962 Field::new("c0", DataType::UInt32, false),
963 Field::new("c1", DataType::Date64, false),
964 ]);
965
966 let c0 = UInt32Array::from(vec![Some(123), Some(234)]);
967 let c1 = Date64Array::from(vec![Some(1926632005177), Some(1926632005177685347)]);
968 let batch =
969 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c0), Arc::new(c1)]).unwrap();
970
971 let mut file = tempfile::tempfile().unwrap();
972 let mut writer = Writer::new(&mut file);
973 let batches = vec![&batch, &batch];
974
975 for batch in batches {
976 let err = writer.write(batch).unwrap_err().to_string();
977 assert_eq!(
978 err,
979 "Csv error: Error processing row 2, col 2: Cast error: Failed to convert 1926632005177685347 to temporal for Date64"
980 )
981 }
982 drop(writer);
983 }
984
985 #[test]
986 fn test_write_csv_using_rfc3339() {
987 let schema = Schema::new(vec![
988 Field::new(
989 "c1",
990 DataType::Timestamp(TimeUnit::Millisecond, Some("+00:00".into())),
991 true,
992 ),
993 Field::new("c2", DataType::Timestamp(TimeUnit::Millisecond, None), true),
994 Field::new("c3", DataType::Date32, false),
995 Field::new("c4", DataType::Time32(TimeUnit::Second), false),
996 ]);
997
998 let c1 = TimestampMillisecondArray::from(vec![Some(1555584887378), Some(1635577147000)])
999 .with_timezone("+00:00".to_string());
1000 let c2 = TimestampMillisecondArray::from(vec![Some(1555584887378), Some(1635577147000)]);
1001 let c3 = Date32Array::from(vec![3, 2]);
1002 let c4 = Time32SecondArray::from(vec![1234, 24680]);
1003
1004 let batch = RecordBatch::try_new(
1005 Arc::new(schema),
1006 vec![Arc::new(c1), Arc::new(c2), Arc::new(c3), Arc::new(c4)],
1007 )
1008 .unwrap();
1009
1010 let mut file = tempfile::tempfile().unwrap();
1011
1012 let builder = WriterBuilder::new();
1013 let mut writer = builder.build(&mut file);
1014 let batches = vec![&batch];
1015 for batch in batches {
1016 writer.write(batch).unwrap();
1017 }
1018 drop(writer);
1019
1020 file.rewind().unwrap();
1021 let mut buffer: Vec<u8> = vec![];
1022 file.read_to_end(&mut buffer).unwrap();
1023
1024 assert_eq!(
1025 "c1,c2,c3,c4
10262019-04-18T10:54:47.378Z,2019-04-18T10:54:47.378,1970-01-04,00:20:34
10272021-10-30T06:59:07Z,2021-10-30T06:59:07,1970-01-03,06:51:20\n",
1028 String::from_utf8(buffer).unwrap()
1029 );
1030 }
1031
1032 #[test]
1033 fn test_write_csv_tz_format() {
1034 let schema = Schema::new(vec![
1035 Field::new(
1036 "c1",
1037 DataType::Timestamp(TimeUnit::Millisecond, Some("+02:00".into())),
1038 true,
1039 ),
1040 Field::new(
1041 "c2",
1042 DataType::Timestamp(TimeUnit::Second, Some("+04:00".into())),
1043 true,
1044 ),
1045 ]);
1046 let c1 = TimestampMillisecondArray::from(vec![Some(1_000), Some(2_000)])
1047 .with_timezone("+02:00".to_string());
1048 let c2 = TimestampSecondArray::from(vec![Some(1_000_000), None])
1049 .with_timezone("+04:00".to_string());
1050 let batch =
1051 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
1052
1053 let mut file = tempfile::tempfile().unwrap();
1054 let mut writer = WriterBuilder::new()
1055 .with_timestamp_tz_format("%M:%H".to_string())
1056 .build(&mut file);
1057 writer.write(&batch).unwrap();
1058
1059 drop(writer);
1060 file.rewind().unwrap();
1061 let mut buffer: Vec<u8> = vec![];
1062 file.read_to_end(&mut buffer).unwrap();
1063
1064 let output = String::from_utf8(buffer).unwrap();
1065 assert_eq!(output, "c1,c2\n00:02,46:17\n00:02,\n");
1066 }
1067
1068 #[test]
1069 fn test_write_csv_with_lf_terminator() {
1070 let output = write_batch_with_terminator(Terminator::Any(b'\n'));
1071 assert_eq!(output, "c1,c2\nhello,1\nworld,2\n");
1072 }
1073
1074 #[test]
1075 fn test_write_csv_with_crlf_terminator() {
1076 let output = write_batch_with_terminator(Terminator::CRLF);
1077 assert_eq!(output, "c1,c2\r\nhello,1\r\nworld,2\r\n");
1078 }
1079
1080 #[test]
1081 fn test_write_csv_with_any_terminator() {
1082 let output = write_batch_with_terminator(Terminator::Any(b'|'));
1083 assert_eq!(output, "c1,c2|hello,1|world,2|");
1084 }
1085
1086 fn write_batch_with_terminator(terminator: Terminator) -> String {
1087 let schema = Schema::new(vec![
1088 Field::new("c1", DataType::Utf8, false),
1089 Field::new("c2", DataType::UInt32, false),
1090 ]);
1091
1092 let c1 = StringArray::from(vec!["hello", "world"]);
1093 let c2 = PrimitiveArray::<UInt32Type>::from(vec![1, 2]);
1094
1095 let batch =
1096 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]).unwrap();
1097
1098 let mut buf = Vec::new();
1099 let mut writer = WriterBuilder::new()
1100 .with_line_terminator(terminator)
1101 .build(&mut buf);
1102 writer.write(&batch).unwrap();
1103 drop(writer);
1104
1105 String::from_utf8(buf).unwrap()
1106 }
1107
1108 #[test]
1109 fn test_write_csv_binary() {
1110 let fixed_size = 8;
1111 let schema = SchemaRef::new(Schema::new(vec![
1112 Field::new("c1", DataType::Binary, true),
1113 Field::new("c2", DataType::FixedSizeBinary(fixed_size), true),
1114 Field::new("c3", DataType::LargeBinary, true),
1115 ]));
1116 let mut c1_builder = BinaryBuilder::new();
1117 c1_builder.append_value(b"Homer");
1118 c1_builder.append_value(b"Bart");
1119 c1_builder.append_null();
1120 c1_builder.append_value(b"Ned");
1121 let mut c2_builder = FixedSizeBinaryBuilder::new(fixed_size);
1122 c2_builder.append_value(b"Simpson ").unwrap();
1123 c2_builder.append_value(b"Simpson ").unwrap();
1124 c2_builder.append_null();
1125 c2_builder.append_value(b"Flanders").unwrap();
1126 let mut c3_builder = LargeBinaryBuilder::new();
1127 c3_builder.append_null();
1128 c3_builder.append_null();
1129 c3_builder.append_value(b"Comic Book Guy");
1130 c3_builder.append_null();
1131
1132 let batch = RecordBatch::try_new(
1133 schema,
1134 vec![
1135 Arc::new(c1_builder.finish()) as ArrayRef,
1136 Arc::new(c2_builder.finish()) as ArrayRef,
1137 Arc::new(c3_builder.finish()) as ArrayRef,
1138 ],
1139 )
1140 .unwrap();
1141
1142 let mut buf = Vec::new();
1143 let builder = WriterBuilder::new();
1144 let mut writer = builder.build(&mut buf);
1145 writer.write(&batch).unwrap();
1146 drop(writer);
1147 assert_eq!(
1148 "\
1149 c1,c2,c3\n\
1150 486f6d6572,53696d70736f6e20,\n\
1151 42617274,53696d70736f6e20,\n\
1152 ,,436f6d696320426f6f6b20477579\n\
1153 4e6564,466c616e64657273,\n\
1154 ",
1155 String::from_utf8(buf).unwrap()
1156 );
1157 }
1158
1159 #[test]
1160 fn test_write_csv_whitespace_handling() {
1161 let schema = Schema::new(vec![
1162 Field::new("c1", DataType::Utf8, false),
1163 Field::new("c2", DataType::Float64, true),
1164 Field::new("c3", DataType::Utf8, true),
1165 ]);
1166
1167 let c1 = StringArray::from(vec![
1168 " leading space",
1169 "trailing space ",
1170 " both spaces ",
1171 "no spaces",
1172 ]);
1173 let c2 = PrimitiveArray::<Float64Type>::from(vec![
1174 Some(123.45),
1175 Some(678.90),
1176 None,
1177 Some(111.22),
1178 ]);
1179 let c3 = StringArray::from(vec![
1180 Some(" test "),
1181 Some("value "),
1182 None,
1183 Some(" another"),
1184 ]);
1185
1186 let batch = RecordBatch::try_new(
1187 Arc::new(schema),
1188 vec![Arc::new(c1), Arc::new(c2), Arc::new(c3)],
1189 )
1190 .unwrap();
1191
1192 let mut buf = Vec::new();
1194 let builder = WriterBuilder::new();
1195 let mut writer = builder.build(&mut buf);
1196 writer.write(&batch).unwrap();
1197 drop(writer);
1198 assert_eq!(
1199 "c1,c2,c3\n leading space,123.45, test \ntrailing space ,678.9,value \n both spaces ,,\nno spaces,111.22, another\n",
1200 String::from_utf8(buf).unwrap()
1201 );
1202
1203 let mut buf = Vec::new();
1205 let builder = WriterBuilder::new().with_ignore_leading_whitespace(true);
1206 let mut writer = builder.build(&mut buf);
1207 writer.write(&batch).unwrap();
1208 drop(writer);
1209 assert_eq!(
1210 "c1,c2,c3\nleading space,123.45,test \ntrailing space ,678.9,value \nboth spaces ,,\nno spaces,111.22,another\n",
1211 String::from_utf8(buf).unwrap()
1212 );
1213
1214 let mut buf = Vec::new();
1216 let builder = WriterBuilder::new().with_ignore_trailing_whitespace(true);
1217 let mut writer = builder.build(&mut buf);
1218 writer.write(&batch).unwrap();
1219 drop(writer);
1220 assert_eq!(
1221 "c1,c2,c3\n leading space,123.45, test\ntrailing space,678.9,value\n both spaces,,\nno spaces,111.22, another\n",
1222 String::from_utf8(buf).unwrap()
1223 );
1224
1225 let mut buf = Vec::new();
1227 let builder = WriterBuilder::new()
1228 .with_ignore_leading_whitespace(true)
1229 .with_ignore_trailing_whitespace(true);
1230 let mut writer = builder.build(&mut buf);
1231 writer.write(&batch).unwrap();
1232 drop(writer);
1233 assert_eq!(
1234 "c1,c2,c3\nleading space,123.45,test\ntrailing space,678.9,value\nboth spaces,,\nno spaces,111.22,another\n",
1235 String::from_utf8(buf).unwrap()
1236 );
1237 }
1238
1239 #[test]
1240 fn test_write_csv_whitespace_with_special_chars() {
1241 let schema = Schema::new(vec![Field::new("c1", DataType::Utf8, false)]);
1242
1243 let c1 = StringArray::from(vec![
1244 " quoted \"value\" ",
1245 " new\nline ",
1246 " comma,value ",
1247 "\ttab\tvalue\t",
1248 ]);
1249
1250 let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1)]).unwrap();
1251
1252 let mut buf = Vec::new();
1254 let builder = WriterBuilder::new()
1255 .with_ignore_leading_whitespace(true)
1256 .with_ignore_trailing_whitespace(true);
1257 let mut writer = builder.build(&mut buf);
1258 writer.write(&batch).unwrap();
1259 drop(writer);
1260
1261 assert_eq!(
1263 "c1\n\"quoted \"\"value\"\"\"\n\"new\nline\"\n\"comma,value\"\ntab\tvalue\n",
1264 String::from_utf8(buf).unwrap()
1265 );
1266 }
1267
1268 #[test]
1269 fn test_write_csv_whitespace_all_string_types() {
1270 use arrow_array::{LargeStringArray, StringViewArray};
1271
1272 let schema = Schema::new(vec![
1273 Field::new("utf8", DataType::Utf8, false),
1274 Field::new("large_utf8", DataType::LargeUtf8, false),
1275 Field::new("utf8_view", DataType::Utf8View, false),
1276 ]);
1277
1278 let utf8 = StringArray::from(vec![" leading", "trailing ", " both ", "no_spaces"]);
1279
1280 let large_utf8 =
1281 LargeStringArray::from(vec![" leading", "trailing ", " both ", "no_spaces"]);
1282
1283 let utf8_view =
1284 StringViewArray::from(vec![" leading", "trailing ", " both ", "no_spaces"]);
1285
1286 let batch = RecordBatch::try_new(
1287 Arc::new(schema),
1288 vec![Arc::new(utf8), Arc::new(large_utf8), Arc::new(utf8_view)],
1289 )
1290 .unwrap();
1291
1292 let mut buf = Vec::new();
1294 let builder = WriterBuilder::new();
1295 let mut writer = builder.build(&mut buf);
1296 writer.write(&batch).unwrap();
1297 drop(writer);
1298 assert_eq!(
1299 "utf8,large_utf8,utf8_view\n leading, leading, leading\ntrailing ,trailing ,trailing \n both , both , both \nno_spaces,no_spaces,no_spaces\n",
1300 String::from_utf8(buf).unwrap()
1301 );
1302
1303 let mut buf = Vec::new();
1305 let builder = WriterBuilder::new()
1306 .with_ignore_leading_whitespace(true)
1307 .with_ignore_trailing_whitespace(true);
1308 let mut writer = builder.build(&mut buf);
1309 writer.write(&batch).unwrap();
1310 drop(writer);
1311 assert_eq!(
1312 "utf8,large_utf8,utf8_view\nleading,leading,leading\ntrailing,trailing,trailing\nboth,both,both\nno_spaces,no_spaces,no_spaces\n",
1313 String::from_utf8(buf).unwrap()
1314 );
1315
1316 let mut buf = Vec::new();
1318 let builder = WriterBuilder::new().with_ignore_leading_whitespace(true);
1319 let mut writer = builder.build(&mut buf);
1320 writer.write(&batch).unwrap();
1321 drop(writer);
1322 assert_eq!(
1323 "utf8,large_utf8,utf8_view\nleading,leading,leading\ntrailing ,trailing ,trailing \nboth ,both ,both \nno_spaces,no_spaces,no_spaces\n",
1324 String::from_utf8(buf).unwrap()
1325 );
1326
1327 let mut buf = Vec::new();
1329 let builder = WriterBuilder::new().with_ignore_trailing_whitespace(true);
1330 let mut writer = builder.build(&mut buf);
1331 writer.write(&batch).unwrap();
1332 drop(writer);
1333 assert_eq!(
1334 "utf8,large_utf8,utf8_view\n leading, leading, leading\ntrailing,trailing,trailing\n both, both, both\nno_spaces,no_spaces,no_spaces\n",
1335 String::from_utf8(buf).unwrap()
1336 );
1337 }
1338
1339 fn write_quote_style(batch: &RecordBatch, quote_style: QuoteStyle) -> String {
1340 let mut buf = Vec::new();
1341 let mut writer = WriterBuilder::new()
1342 .with_quote_style(quote_style)
1343 .build(&mut buf);
1344 writer.write(batch).unwrap();
1345 drop(writer);
1346 String::from_utf8(buf).unwrap()
1347 }
1348
1349 fn write_quote_style_with_null(
1350 batch: &RecordBatch,
1351 quote_style: QuoteStyle,
1352 null_value: &str,
1353 ) -> String {
1354 let mut buf = Vec::new();
1355 let mut writer = WriterBuilder::new()
1356 .with_quote_style(quote_style)
1357 .with_null(null_value.to_string())
1358 .build(&mut buf);
1359 writer.write(batch).unwrap();
1360 drop(writer);
1361 String::from_utf8(buf).unwrap()
1362 }
1363
1364 #[test]
1365 fn test_write_csv_quote_style() {
1366 let schema = Schema::new(vec![
1367 Field::new("text", DataType::Utf8, false),
1368 Field::new("number", DataType::Int32, false),
1369 Field::new("float", DataType::Float64, false),
1370 ]);
1371
1372 let text = StringArray::from(vec!["hello", "world", "comma,value", "quote\"test"]);
1373 let number = Int32Array::from(vec![1, 2, 3, 4]);
1374 let float = Float64Array::from(vec![1.1, 2.2, 3.3, 4.4]);
1375
1376 let batch = RecordBatch::try_new(
1377 Arc::new(schema),
1378 vec![Arc::new(text), Arc::new(number), Arc::new(float)],
1379 )
1380 .unwrap();
1381
1382 assert_eq!(
1384 "text,number,float\nhello,1,1.1\nworld,2,2.2\n\"comma,value\",3,3.3\n\"quote\"\"test\",4,4.4\n",
1385 write_quote_style(&batch, QuoteStyle::Necessary)
1386 );
1387
1388 assert_eq!(
1390 "\"text\",\"number\",\"float\"\n\"hello\",\"1\",\"1.1\"\n\"world\",\"2\",\"2.2\"\n\"comma,value\",\"3\",\"3.3\"\n\"quote\"\"test\",\"4\",\"4.4\"\n",
1391 write_quote_style(&batch, QuoteStyle::Always)
1392 );
1393
1394 assert_eq!(
1396 "\"text\",\"number\",\"float\"\n\"hello\",1,1.1\n\"world\",2,2.2\n\"comma,value\",3,3.3\n\"quote\"\"test\",4,4.4\n",
1397 write_quote_style(&batch, QuoteStyle::NonNumeric)
1398 );
1399
1400 assert_eq!(
1403 "text,number,float\nhello,1,1.1\nworld,2,2.2\ncomma,value,3,3.3\nquote\"test,4,4.4\n",
1404 write_quote_style(&batch, QuoteStyle::Never)
1405 );
1406 }
1407
1408 #[test]
1409 fn test_write_csv_quote_style_with_nulls() {
1410 let schema = Schema::new(vec![
1411 Field::new("text", DataType::Utf8, true),
1412 Field::new("number", DataType::Int32, true),
1413 ]);
1414
1415 let text = StringArray::from(vec![Some("hello"), None, Some("world")]);
1416 let number = Int32Array::from(vec![Some(1), Some(2), None]);
1417
1418 let batch =
1419 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(text), Arc::new(number)]).unwrap();
1420
1421 assert_eq!(
1423 "\"text\",\"number\"\n\"hello\",\"1\"\n\"\",\"2\"\n\"world\",\"\"\n",
1424 write_quote_style(&batch, QuoteStyle::Always)
1425 );
1426
1427 assert_eq!(
1429 "\"text\",\"number\"\n\"hello\",\"1\"\n\"NULL\",\"2\"\n\"world\",\"NULL\"\n",
1430 write_quote_style_with_null(&batch, QuoteStyle::Always, "NULL")
1431 );
1432 }
1433}