1use std::io::Write;
2
3use arrow::array::Time64MicrosecondArray;
4use arrow::array::types::Decimal128Type;
5use arrow::array::*;
6use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
7use arrow::record_batch::RecordBatch;
8
9use crate::error::Result;
10use crate::types::decimal::scaled_i128_to_decimal_str;
11
12pub struct CsvFormat;
13
14pub struct CsvFormatWriter {
15 writer: Box<dyn Write + Send>,
16 bytes_written: u64,
17}
18
19impl super::Format for CsvFormat {
20 fn create_writer(
21 &self,
22 schema: &SchemaRef,
23 mut writer: Box<dyn Write + Send>,
24 ) -> Result<Box<dyn super::FormatWriter + Send>> {
25 if let Some(field) = schema
30 .fields()
31 .iter()
32 .find(|f| !csv_serializable(f.data_type()))
33 {
34 anyhow::bail!(
35 "CSV cannot serialize column '{}' (Arrow type {:?}); use `format: parquet` \
36 or drop the column from the query",
37 field.name(),
38 field.data_type()
39 );
40 }
41 let header = schema
42 .fields()
43 .iter()
44 .map(|f| f.name().as_str())
45 .collect::<Vec<_>>()
46 .join(",");
47 let header_bytes = header.len() as u64 + 1; writeln!(writer, "{}", header)?;
49 Ok(Box::new(CsvFormatWriter {
50 writer,
51 bytes_written: header_bytes,
52 }))
53 }
54
55 fn file_extension(&self) -> &str {
56 "csv"
57 }
58}
59
60impl super::FormatWriter for CsvFormatWriter {
61 fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
62 let mut buf = Vec::with_capacity(batch.num_rows() * batch.num_columns() * 8);
63 for row_idx in 0..batch.num_rows() {
64 for col_idx in 0..batch.num_columns() {
65 if col_idx > 0 {
66 buf.push(b',');
67 }
68 write_csv_value(&mut buf, batch.column(col_idx), row_idx)?;
69 }
70 buf.push(b'\n');
71 }
72 self.bytes_written += buf.len() as u64;
73 self.writer.write_all(&buf)?;
74 Ok(())
75 }
76
77 fn finish(self: Box<Self>) -> Result<()> {
78 Ok(())
79 }
80
81 fn bytes_written(&self) -> u64 {
82 self.bytes_written
83 }
84}
85
86pub(crate) fn csv_serializable(dt: &DataType) -> bool {
91 matches!(
92 dt,
93 DataType::Boolean
94 | DataType::Int16
95 | DataType::Int32
96 | DataType::Int64
97 | DataType::UInt64
98 | DataType::Decimal128(_, _)
99 | DataType::Float32
100 | DataType::Float64
101 | DataType::Utf8
102 | DataType::Binary
103 | DataType::FixedSizeBinary(16)
104 | DataType::Date32
105 | DataType::Time64(TimeUnit::Microsecond)
106 | DataType::Timestamp(TimeUnit::Microsecond, _)
107 )
108}
109
110fn write_lower_hex(writer: &mut dyn Write, bytes: &[u8]) -> Result<()> {
116 const HEX: &[u8; 16] = b"0123456789abcdef";
117 let mut chunk = [0u8; 1024];
118 for slab in bytes.chunks(chunk.len() / 2) {
119 let mut n = 0;
120 for &b in slab {
121 chunk[n] = HEX[(b >> 4) as usize];
122 chunk[n + 1] = HEX[(b & 0x0f) as usize];
123 n += 2;
124 }
125 writer.write_all(&chunk[..n])?;
126 }
127 Ok(())
128}
129
130fn write_csv_value(writer: &mut dyn Write, array: &dyn Array, idx: usize) -> Result<()> {
131 if array.is_null(idx) {
132 return Ok(());
133 }
134
135 match array.data_type() {
136 DataType::Boolean => {
137 let arr = array
138 .as_any()
139 .downcast_ref::<BooleanArray>()
140 .expect("DataType/Array mismatch");
141 write!(writer, "{}", arr.value(idx))?;
142 }
143 DataType::Int16 => {
144 let arr = array
145 .as_any()
146 .downcast_ref::<Int16Array>()
147 .expect("DataType/Array mismatch");
148 write!(writer, "{}", arr.value(idx))?;
149 }
150 DataType::Int32 => {
151 let arr = array
152 .as_any()
153 .downcast_ref::<Int32Array>()
154 .expect("DataType/Array mismatch");
155 write!(writer, "{}", arr.value(idx))?;
156 }
157 DataType::Int64 => {
158 let arr = array
159 .as_any()
160 .downcast_ref::<Int64Array>()
161 .expect("DataType/Array mismatch");
162 write!(writer, "{}", arr.value(idx))?;
163 }
164 DataType::UInt64 => {
165 let arr = array
166 .as_any()
167 .downcast_ref::<UInt64Array>()
168 .expect("DataType/Array mismatch");
169 write!(writer, "{}", arr.value(idx))?;
170 }
171 DataType::Decimal128(_, scale) => {
172 let arr = array.as_primitive::<Decimal128Type>();
173 let text = scaled_i128_to_decimal_str(arr.value(idx), *scale);
174 writer.write_all(text.as_bytes())?;
175 }
176 DataType::Float32 => {
177 let arr = array
178 .as_any()
179 .downcast_ref::<Float32Array>()
180 .expect("DataType/Array mismatch");
181 write!(writer, "{}", arr.value(idx))?;
182 }
183 DataType::Float64 => {
184 let arr = array
185 .as_any()
186 .downcast_ref::<Float64Array>()
187 .expect("DataType/Array mismatch");
188 write!(writer, "{}", arr.value(idx))?;
189 }
190 DataType::Utf8 => {
191 let arr = array
192 .as_any()
193 .downcast_ref::<StringArray>()
194 .expect("DataType/Array mismatch");
195 let val = arr.value(idx);
196 if val
200 .bytes()
201 .any(|b| matches!(b, b',' | b'"' | b'\n' | b'\r'))
202 {
203 writer.write_all(b"\"")?;
204 let mut rest = val;
205 while let Some(pos) = rest.find('"') {
206 writer.write_all(&rest.as_bytes()[..pos])?;
207 writer.write_all(b"\"\"")?;
208 rest = &rest[pos + 1..];
209 }
210 writer.write_all(rest.as_bytes())?;
211 writer.write_all(b"\"")?;
212 } else {
213 writer.write_all(val.as_bytes())?;
214 }
215 }
216 DataType::Binary => {
217 let arr = array
218 .as_any()
219 .downcast_ref::<BinaryArray>()
220 .expect("DataType/Array mismatch");
221 write_lower_hex(writer, arr.value(idx))?;
222 }
223 DataType::FixedSizeBinary(16) => {
231 let arr = array
232 .as_any()
233 .downcast_ref::<FixedSizeBinaryArray>()
234 .expect("DataType/Array mismatch");
235 let val = arr.value(idx);
236 let mut bytes = [0u8; 16];
237 bytes.copy_from_slice(val);
238 write!(writer, "{}", uuid::Uuid::from_bytes(bytes).to_hyphenated())?;
239 }
240 DataType::Date32 => {
241 let arr = array
242 .as_any()
243 .downcast_ref::<Date32Array>()
244 .expect("DataType/Array mismatch");
245 let days = arr.value(idx);
246 let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).expect("epoch is valid");
252 let date =
253 chrono::Duration::try_days(days as i64).and_then(|d| epoch.checked_add_signed(d));
254 if let Some(date) = date {
255 write!(writer, "{}", date)?;
256 }
257 }
258 DataType::Time64(TimeUnit::Microsecond) => {
259 let arr = array
260 .as_any()
261 .downcast_ref::<Time64MicrosecondArray>()
262 .expect("DataType/Array mismatch");
263 let micros = arr.value(idx);
264 let secs = micros / 1_000_000;
265 let frac_us = micros % 1_000_000;
266 write!(
267 writer,
268 "{:02}:{:02}:{:02}.{:06}",
269 secs / 3600,
270 (secs % 3600) / 60,
271 secs % 60,
272 frac_us
273 )?;
274 }
275 DataType::Timestamp(TimeUnit::Microsecond, _) => {
276 let arr = array
277 .as_any()
278 .downcast_ref::<TimestampMicrosecondArray>()
279 .expect("DataType/Array mismatch");
280 let micros = arr.value(idx);
281 let secs = micros / 1_000_000;
282 let nsecs = ((micros % 1_000_000) * 1_000) as u32;
283 if let Some(dt) = chrono::DateTime::from_timestamp(secs, nsecs) {
284 use chrono::{Datelike as _, Timelike as _};
285 let y = dt.year();
286 if (0..=9999).contains(&y) {
292 write!(
293 writer,
294 "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:06}",
295 y,
296 dt.month(),
297 dt.day(),
298 dt.hour(),
299 dt.minute(),
300 dt.second(),
301 dt.nanosecond() / 1_000
302 )?;
303 } else {
304 write!(writer, "{}", dt.format("%Y-%m-%dT%H:%M:%S%.6f"))?;
305 }
306 }
307 }
308 other => {
309 anyhow::bail!(
313 "CSV: no serializer for Arrow type {other:?} (column should have been rejected at writer creation)"
314 );
315 }
316 }
317
318 Ok(())
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324 use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
325 use std::sync::Arc;
326
327 fn cell<A: Array + 'static>(array: A, idx: usize) -> String {
329 let mut buf = Vec::new();
330 write_csv_value(&mut buf, &array, idx).unwrap();
331 String::from_utf8(buf).unwrap()
332 }
333
334 fn null_cell(dt: DataType) -> String {
336 use arrow::array::new_null_array;
337 let arr = new_null_array(&dt, 1);
338 let mut buf = Vec::new();
339 write_csv_value(&mut buf, arr.as_ref(), 0).unwrap();
340 String::from_utf8(buf).unwrap()
341 }
342
343 #[test]
346 fn time64_renders_exact_wall_clock() {
347 use arrow::array::Time64MicrosecondArray;
350 assert_eq!(
351 cell(Time64MicrosecondArray::from(vec![86_399_999_999_i64]), 0),
352 "23:59:59.999999"
353 );
354 assert_eq!(
355 cell(Time64MicrosecondArray::from(vec![3_661_000_001_i64]), 0),
356 "01:01:01.000001"
357 );
358 assert_eq!(
359 cell(Time64MicrosecondArray::from(vec![0_i64]), 0),
360 "00:00:00.000000"
361 );
362 }
363
364 #[test]
365 fn write_batch_layout_and_byte_accounting_are_exact() {
366 use arrow::array::Int64Array;
370 use std::sync::{Arc as SArc, Mutex};
371
372 #[derive(Clone)]
373 struct SharedBuf(SArc<Mutex<Vec<u8>>>);
374 impl std::io::Write for SharedBuf {
375 fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
376 self.0.lock().unwrap().extend_from_slice(b);
377 Ok(b.len())
378 }
379 fn flush(&mut self) -> std::io::Result<()> {
380 Ok(())
381 }
382 }
383
384 let schema: SchemaRef = Arc::new(Schema::new(vec![
385 Field::new("a", DataType::Int64, false),
386 Field::new("b", DataType::Int64, false),
387 ]));
388 let sink = SharedBuf(SArc::new(Mutex::new(Vec::new())));
389 use crate::format::Format as _;
390 let mut w = CsvFormat
391 .create_writer(&schema, Box::new(sink.clone()))
392 .unwrap();
393 let batch = RecordBatch::try_new(
394 schema,
395 vec![
396 Arc::new(Int64Array::from(vec![1, 2])),
397 Arc::new(Int64Array::from(vec![10, 20])),
398 ],
399 )
400 .unwrap();
401 w.write_batch(&batch).unwrap();
402
403 let text = String::from_utf8(sink.0.lock().unwrap().clone()).unwrap();
404 assert_eq!(
405 text, "a,b\n1,10\n2,20\n",
406 "exact CSV layout (no leading commas)"
407 );
408 assert_eq!(
409 w.bytes_written(),
410 text.len() as u64,
411 "byte accounting must equal the physical output"
412 );
413 }
414
415 #[test]
418 fn null_value_writes_empty_string() {
419 assert_eq!(null_cell(DataType::Int64), "");
420 assert_eq!(null_cell(DataType::Utf8), "");
421 assert_eq!(null_cell(DataType::Boolean), "");
422 }
423
424 #[test]
427 fn bool_true_writes_true() {
428 assert_eq!(cell(BooleanArray::from(vec![true]), 0), "true");
429 }
430
431 #[test]
432 fn bool_false_writes_false() {
433 assert_eq!(cell(BooleanArray::from(vec![false]), 0), "false");
434 }
435
436 #[test]
437 fn int16_value() {
438 assert_eq!(cell(Int16Array::from(vec![42i16]), 0), "42");
439 }
440
441 #[test]
442 fn int32_negative() {
443 assert_eq!(cell(Int32Array::from(vec![-7i32]), 0), "-7");
444 }
445
446 #[test]
447 fn decimal128_writes_exact_text() {
448 let arr = Decimal128Array::from(vec![10i128])
449 .with_precision_and_scale(18, 2)
450 .unwrap();
451 assert_eq!(cell(arr, 0), "0.10");
452 let scaled =
453 crate::types::decimal::decimal_str_to_scaled_i128("999999999999.99", 2).unwrap();
454 let arr = Decimal128Array::from(vec![scaled])
455 .with_precision_and_scale(18, 2)
456 .unwrap();
457 assert_eq!(cell(arr, 0), "999999999999.99");
458 }
459
460 #[test]
461 fn int64_large() {
462 assert_eq!(
463 cell(Int64Array::from(vec![9_999_999_999i64]), 0),
464 "9999999999"
465 );
466 }
467
468 #[test]
469 fn float32_value() {
470 let result = cell(Float32Array::from(vec![1.5f32]), 0);
471 assert!(result.starts_with("1.5"), "got: {result}");
472 }
473
474 #[test]
475 fn float64_value() {
476 let result = cell(Float64Array::from(vec![std::f64::consts::PI]), 0);
477 assert!(result.starts_with("3.14"), "got: {result}");
478 }
479
480 #[test]
491 fn float_special_values_emit_literals_not_empty() {
492 assert_eq!(cell(Float64Array::from(vec![f64::NAN]), 0), "NaN");
493 assert_eq!(cell(Float64Array::from(vec![f64::INFINITY]), 0), "inf");
494 assert_eq!(cell(Float64Array::from(vec![f64::NEG_INFINITY]), 0), "-inf");
495 assert_eq!(cell(Float32Array::from(vec![f32::NAN]), 0), "NaN");
496 assert_eq!(cell(Float32Array::from(vec![f32::INFINITY]), 0), "inf");
497 assert_eq!(cell(Float64Array::from(vec![-0.0f64]), 0), "-0");
499 }
500
501 #[test]
504 fn plain_string_no_quoting() {
505 assert_eq!(cell(StringArray::from(vec!["hello"]), 0), "hello");
506 }
507
508 #[test]
509 fn string_with_comma_is_quoted() {
510 assert_eq!(cell(StringArray::from(vec!["a,b"]), 0), "\"a,b\"");
511 }
512
513 #[test]
514 fn string_with_double_quote_is_escaped() {
515 let result = cell(StringArray::from(vec![r#"say "hi""#]), 0);
517 assert_eq!(result, r#""say ""hi""""#);
518 }
519
520 #[test]
521 fn string_with_newline_is_quoted() {
522 let result = cell(StringArray::from(vec!["line1\nline2"]), 0);
523 assert!(
524 result.starts_with('"') && result.ends_with('"'),
525 "got: {result}"
526 );
527 assert!(result.contains("line1\nline2"), "got: {result}");
528 }
529
530 #[test]
536 fn roast_string_with_carriage_return_is_quoted() {
537 let result = cell(StringArray::from(vec!["a\rb"]), 0);
538 assert_eq!(
539 result, "\"a\rb\"",
540 "lone CR must force quoting per RFC 4180, but got unquoted cell {result:?}"
541 );
542 }
543
544 #[test]
547 fn binary_is_written_as_hex() {
548 let arr = BinaryArray::from_vec(vec![&[0xDE, 0xAD, 0xBE, 0xEF][..]]);
549 assert_eq!(cell(arr, 0), "deadbeef");
550 }
551
552 #[test]
553 fn binary_empty_writes_empty() {
554 let arr = BinaryArray::from_vec(vec![&[][..]]);
555 assert_eq!(cell(arr, 0), "");
556 }
557
558 #[test]
561 fn date32_epoch_is_1970_01_01() {
562 assert_eq!(cell(Date32Array::from(vec![0i32]), 0), "1970-01-01");
563 }
564
565 #[test]
566 fn date32_positive_offset() {
567 assert_eq!(cell(Date32Array::from(vec![365i32]), 0), "1971-01-01");
569 }
570
571 #[test]
574 fn timestamp_micros_formats_as_iso() {
575 let micros: i64 = 1_672_531_200 * 1_000_000;
577 let _schema = Arc::new(Schema::new(vec![Field::new(
578 "ts",
579 DataType::Timestamp(TimeUnit::Microsecond, None),
580 true,
581 )]));
582 let arr = TimestampMicrosecondArray::from(vec![micros]);
583 let result = cell(arr, 0);
584 assert!(result.starts_with("2023-01-01T"), "got: {result}");
585 assert!(result.contains("00:00:00"), "got: {result}");
586 }
587
588 #[test]
591 fn csv_format_write_batch_tracks_bytes_and_succeeds() {
592 use crate::format::Format;
593
594 let schema = Arc::new(Schema::new(vec![
595 Field::new("id", DataType::Int64, false),
596 Field::new("name", DataType::Utf8, true),
597 ]));
598 let batch = arrow::record_batch::RecordBatch::try_new(
599 schema.clone(),
600 vec![
601 Arc::new(Int64Array::from(vec![1i64, 2])),
602 Arc::new(StringArray::from(vec![Some("alice"), None])),
603 ],
604 )
605 .unwrap();
606
607 let fmt = CsvFormat;
609 let mut writer = fmt
610 .create_writer(&schema, Box::new(Vec::<u8>::new()))
611 .unwrap();
612 writer.write_batch(&batch).unwrap();
613 assert!(
615 writer.bytes_written() > 10,
616 "expected >10 bytes, got {}",
617 writer.bytes_written()
618 );
619 writer.finish().unwrap();
620 }
621
622 #[test]
625 fn csv_rejects_array_columns_loudly() {
626 use crate::format::Format;
627 let schema = Arc::new(Schema::new(vec![
628 Field::new("id", DataType::Int64, false),
629 Field::new(
630 "tags",
631 DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
632 true,
633 ),
634 ]));
635 let Err(err) = CsvFormat.create_writer(&schema, Box::new(Vec::<u8>::new())) else {
636 panic!("CSV must reject array columns, not silently drop them");
637 };
638 let msg = format!("{err:#}");
639 assert!(msg.contains("tags"), "error must name the column: {msg}");
640 assert!(msg.to_lowercase().contains("csv"), "{msg}");
641 }
642
643 #[test]
647 fn every_serializable_type_is_actually_written() {
648 use crate::format::Format;
649 let cols: Vec<(&str, ArrayRef)> = vec![
650 ("b", Arc::new(BooleanArray::from(vec![true]))),
651 ("i16", Arc::new(Int16Array::from(vec![1i16]))),
652 ("i32", Arc::new(Int32Array::from(vec![1i32]))),
653 ("i64", Arc::new(Int64Array::from(vec![1i64]))),
654 ("u64", Arc::new(UInt64Array::from(vec![1u64]))),
655 (
656 "dec",
657 Arc::new(
658 Decimal128Array::from(vec![100i128])
659 .with_precision_and_scale(18, 2)
660 .unwrap(),
661 ),
662 ),
663 ("f32", Arc::new(Float32Array::from(vec![1.0f32]))),
664 ("f64", Arc::new(Float64Array::from(vec![1.0f64]))),
665 ("s", Arc::new(StringArray::from(vec!["x"]))),
666 ("bin", Arc::new(BinaryArray::from_vec(vec![&[1u8][..]]))),
667 (
668 "uuid",
669 Arc::new(
670 FixedSizeBinaryArray::try_from_iter(std::iter::once(vec![0u8; 16])).unwrap(),
671 ),
672 ),
673 ("d", Arc::new(Date32Array::from(vec![0i32]))),
674 ("t", Arc::new(Time64MicrosecondArray::from(vec![0i64]))),
675 ("ts", Arc::new(TimestampMicrosecondArray::from(vec![0i64]))),
676 ];
677 let fields: Vec<Field> = cols
678 .iter()
679 .map(|(n, a)| Field::new(*n, a.data_type().clone(), true))
680 .collect();
681 for f in &fields {
683 assert!(
684 csv_serializable(f.data_type()),
685 "test type {:?} not in csv_serializable",
686 f.data_type()
687 );
688 }
689 let schema = Arc::new(Schema::new(fields));
690 let arrays: Vec<ArrayRef> = cols.into_iter().map(|(_, a)| a).collect();
691 let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap();
692 let mut w = CsvFormat
693 .create_writer(&schema, Box::new(Vec::<u8>::new()))
694 .unwrap();
695 w.write_batch(&batch)
696 .expect("every serializable type must write without hitting the fallthrough");
697 }
698
699 #[test]
704 fn binary_hex_matches_per_byte_format_for_all_byte_values() {
705 let all: Vec<u8> = (0..=255u8).collect();
708 for case in [&all[..], &[][..], &[0x00, 0xff, 0x10, 0x0a]] {
709 let expected: String = case.iter().map(|b| format!("{b:02x}")).collect();
710 let got = cell(BinaryArray::from_vec(vec![case]), 0);
711 assert_eq!(got, expected, "hex mismatch for {case:?}");
712 }
713 }
714
715 #[test]
716 fn binary_hex_spans_chunk_boundary() {
717 let big: Vec<u8> = (0..2000u32).map(|i| (i % 256) as u8).collect();
719 let expected: String = big.iter().map(|b| format!("{b:02x}")).collect();
720 let got = cell(BinaryArray::from_vec(vec![&big[..]]), 0);
721 assert_eq!(got, expected);
722 }
723
724 #[test]
725 fn timestamp_fast_path_matches_chrono_format() {
726 let cases: [i64; 6] = [
730 0,
731 1_700_000_000_123_456, 1_000_000 * 86_399 + 999_999, -1, 253_402_300_799_000_000, 300_000_000_000_000_000, ];
737 for micros in cases {
738 let got = cell(TimestampMicrosecondArray::from(vec![micros]), 0);
739 let secs = micros / 1_000_000;
740 let nsecs = ((micros % 1_000_000) * 1_000) as u32;
741 let expected = match chrono::DateTime::from_timestamp(secs, nsecs) {
742 Some(dt) => format!("{}", dt.format("%Y-%m-%dT%H:%M:%S%.6f")),
743 None => String::new(),
744 };
745 assert_eq!(got, expected, "timestamp mismatch for micros={micros}");
746 }
747 }
748}