Skip to main content

rivet/format/
csv.rs

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        // Fail loud: arrays and other nested/wide Arrow types have no CSV cell
26        // representation. Reject them up front, naming the column, instead of
27        // silently writing empty values for every row — `format: parquet` or
28        // excluding the column from the query is the fix.
29        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        // RFC-4180 quote each column NAME with the SAME rule the Utf8 data arm
42        // applies to cells (quote on , " CR LF, double embedded quotes). Data
43        // cells were quoted but the header was a raw `join(",")`, so a curated-
44        // query alias with a comma (`... AS "Amount, USD"`) or quote split the
45        // header across columns and silently mis-aligned EVERY downstream reader
46        // from the data. Built once (not the per-cell hot path), so plain String.
47        let header = schema
48            .fields()
49            .iter()
50            .map(|f| {
51                let n = f.name();
52                if n.bytes().any(|b| matches!(b, b',' | b'"' | b'\n' | b'\r')) {
53                    format!("\"{}\"", n.replace('"', "\"\""))
54                } else {
55                    n.clone()
56                }
57            })
58            .collect::<Vec<String>>()
59            .join(",");
60        let header_bytes = header.len() as u64 + 1; // +1 for newline
61        writeln!(writer, "{}", header)?;
62        Ok(Box::new(CsvFormatWriter {
63            writer,
64            bytes_written: header_bytes,
65        }))
66    }
67
68    fn file_extension(&self) -> &str {
69        "csv"
70    }
71}
72
73impl super::FormatWriter for CsvFormatWriter {
74    fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
75        let mut buf = Vec::with_capacity(batch.num_rows() * batch.num_columns() * 8);
76        for row_idx in 0..batch.num_rows() {
77            for col_idx in 0..batch.num_columns() {
78                if col_idx > 0 {
79                    buf.push(b',');
80                }
81                write_csv_value(&mut buf, batch.column(col_idx), row_idx)?;
82            }
83            buf.push(b'\n');
84        }
85        self.bytes_written += buf.len() as u64;
86        self.writer.write_all(&buf)?;
87        Ok(())
88    }
89
90    fn finish(self: Box<Self>) -> Result<()> {
91        Ok(())
92    }
93
94    fn bytes_written(&self) -> u64 {
95        self.bytes_written
96    }
97}
98
99/// Arrow types `write_csv_value` can serialize. Everything else — lists,
100/// structs, maps, `Decimal256`, non-UUID fixed binary, … — has no CSV cell
101/// representation and is rejected at writer creation rather than silently
102/// emitted as an empty value.
103pub(crate) fn csv_serializable(dt: &DataType) -> bool {
104    matches!(
105        dt,
106        DataType::Boolean
107            | DataType::Int16
108            | DataType::Int32
109            | DataType::Int64
110            | DataType::UInt64
111            | DataType::Decimal128(_, _)
112            | DataType::Float32
113            | DataType::Float64
114            | DataType::Utf8
115            | DataType::Binary
116            | DataType::FixedSizeBinary(16)
117            | DataType::Date32
118            | DataType::Time64(TimeUnit::Microsecond)
119            | DataType::Timestamp(TimeUnit::Microsecond, _)
120    )
121}
122
123/// Lowercase hex-encode `bytes` into `writer`. Replaces a per-byte
124/// `write!("{:02x}")` (which drove `core::fmt` through the `dyn Write` vtable
125/// once per byte) with a table lookup batched through a stack buffer — zero
126/// allocation, byte-identical output. On a binary/blob column this is the
127/// difference between N format-machinery calls and a handful of `write_all`s.
128fn write_lower_hex(writer: &mut dyn Write, bytes: &[u8]) -> Result<()> {
129    const HEX: &[u8; 16] = b"0123456789abcdef";
130    let mut chunk = [0u8; 1024];
131    for slab in bytes.chunks(chunk.len() / 2) {
132        let mut n = 0;
133        for &b in slab {
134            chunk[n] = HEX[(b >> 4) as usize];
135            chunk[n + 1] = HEX[(b & 0x0f) as usize];
136            n += 2;
137        }
138        writer.write_all(&chunk[..n])?;
139    }
140    Ok(())
141}
142
143fn write_csv_value(writer: &mut dyn Write, array: &dyn Array, idx: usize) -> Result<()> {
144    if array.is_null(idx) {
145        return Ok(());
146    }
147
148    match array.data_type() {
149        DataType::Boolean => {
150            let arr = array
151                .as_any()
152                .downcast_ref::<BooleanArray>()
153                .expect("DataType/Array mismatch");
154            write!(writer, "{}", arr.value(idx))?;
155        }
156        DataType::Int16 => {
157            let arr = array
158                .as_any()
159                .downcast_ref::<Int16Array>()
160                .expect("DataType/Array mismatch");
161            write!(writer, "{}", arr.value(idx))?;
162        }
163        DataType::Int32 => {
164            let arr = array
165                .as_any()
166                .downcast_ref::<Int32Array>()
167                .expect("DataType/Array mismatch");
168            write!(writer, "{}", arr.value(idx))?;
169        }
170        DataType::Int64 => {
171            let arr = array
172                .as_any()
173                .downcast_ref::<Int64Array>()
174                .expect("DataType/Array mismatch");
175            write!(writer, "{}", arr.value(idx))?;
176        }
177        DataType::UInt64 => {
178            let arr = array
179                .as_any()
180                .downcast_ref::<UInt64Array>()
181                .expect("DataType/Array mismatch");
182            write!(writer, "{}", arr.value(idx))?;
183        }
184        DataType::Decimal128(_, scale) => {
185            let arr = array.as_primitive::<Decimal128Type>();
186            let text = scaled_i128_to_decimal_str(arr.value(idx), *scale);
187            writer.write_all(text.as_bytes())?;
188        }
189        DataType::Float32 => {
190            let arr = array
191                .as_any()
192                .downcast_ref::<Float32Array>()
193                .expect("DataType/Array mismatch");
194            write!(writer, "{}", arr.value(idx))?;
195        }
196        DataType::Float64 => {
197            let arr = array
198                .as_any()
199                .downcast_ref::<Float64Array>()
200                .expect("DataType/Array mismatch");
201            write!(writer, "{}", arr.value(idx))?;
202        }
203        DataType::Utf8 => {
204            let arr = array
205                .as_any()
206                .downcast_ref::<StringArray>()
207                .expect("DataType/Array mismatch");
208            let val = arr.value(idx);
209            // RFC 4180: quote on delimiter, quote, LF — and CR, which readers
210            // in universal-newline mode treat as a record terminator. One pass
211            // instead of four `contains` scans per cell.
212            if val
213                .bytes()
214                .any(|b| matches!(b, b',' | b'"' | b'\n' | b'\r'))
215            {
216                writer.write_all(b"\"")?;
217                let mut rest = val;
218                while let Some(pos) = rest.find('"') {
219                    writer.write_all(&rest.as_bytes()[..pos])?;
220                    writer.write_all(b"\"\"")?;
221                    rest = &rest[pos + 1..];
222                }
223                writer.write_all(rest.as_bytes())?;
224                writer.write_all(b"\"")?;
225            } else {
226                writer.write_all(val.as_bytes())?;
227            }
228        }
229        DataType::Binary => {
230            let arr = array
231                .as_any()
232                .downcast_ref::<BinaryArray>()
233                .expect("DataType/Array mismatch");
234            write_lower_hex(writer, arr.value(idx))?;
235        }
236        // FixedSizeBinary today only carries 16-byte UUIDs (see
237        // `RivetType::Uuid` → `DataType::FixedSizeBinary(16)` in
238        // `src/types/mapping.rs`). CSV has no native binary cell; emit the
239        // canonical hyphenated lowercase form so downstream readers can
240        // recognise it as a UUID rather than 16 bytes of mojibake. Any
241        // future FixedSizeBinary use that is not a UUID should branch on
242        // the size argument before reaching this arm.
243        DataType::FixedSizeBinary(16) => {
244            let arr = array
245                .as_any()
246                .downcast_ref::<FixedSizeBinaryArray>()
247                .expect("DataType/Array mismatch");
248            let val = arr.value(idx);
249            let mut bytes = [0u8; 16];
250            bytes.copy_from_slice(val);
251            write!(writer, "{}", uuid::Uuid::from_bytes(bytes).to_hyphenated())?;
252        }
253        DataType::Date32 => {
254            let arr = array
255                .as_any()
256                .downcast_ref::<Date32Array>()
257                .expect("DataType/Array mismatch");
258            let days = arr.value(idx);
259            // `Date32` is "days since 1970-01-01"; a pathological value near
260            // i32::MAX overflows `NaiveDate + Duration` and panics in chrono.
261            // Fall back to checked arithmetic and emit an empty cell on
262            // overflow — matches the null-cell convention for unserialisable
263            // values elsewhere in this writer.
264            let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).expect("epoch is valid");
265            let date =
266                chrono::Duration::try_days(days as i64).and_then(|d| epoch.checked_add_signed(d));
267            if let Some(date) = date {
268                write!(writer, "{}", date)?;
269            }
270        }
271        DataType::Time64(TimeUnit::Microsecond) => {
272            let arr = array
273                .as_any()
274                .downcast_ref::<Time64MicrosecondArray>()
275                .expect("DataType/Array mismatch");
276            let micros = arr.value(idx);
277            let secs = micros / 1_000_000;
278            let frac_us = micros % 1_000_000;
279            write!(
280                writer,
281                "{:02}:{:02}:{:02}.{:06}",
282                secs / 3600,
283                (secs % 3600) / 60,
284                secs % 60,
285                frac_us
286            )?;
287        }
288        DataType::Timestamp(TimeUnit::Microsecond, tz) => {
289            let arr = array
290                .as_any()
291                .downcast_ref::<TimestampMicrosecondArray>()
292                .expect("DataType/Array mismatch");
293            // rivet's type model keeps the naive-vs-instant distinction that
294            // every other output (Parquet `isAdjustedToUTC`, the warehouse
295            // target types) preserves: a source TIMESTAMPTZ maps to
296            // `Timestamp(_, Some("UTC"))`, a naive TIMESTAMP to `None`. Without
297            // a marker on the instant, CSV renders BOTH identically — a consumer
298            // loading the instant into a tz-aware column on a NON-UTC session
299            // then reads the UTC wall-clock as local and shifts every value (the
300            // session-state class, on the CSV encode leg). Emit a trailing `Z`
301            // for the instant (rivet normalises to UTC) so it is unambiguous;
302            // the naive value stays bare (its wall-clock IS the value).
303            let is_instant = tz.is_some();
304            let micros = arr.value(idx);
305            // FLOOR division (div_euclid/rem_euclid), NOT truncating `/`+`%`: for
306            // a pre-1970 (negative) micros with a sub-second fraction, `micros %
307            // 1_000_000` is NEGATIVE, so `* 1_000 as u32` wraps to a huge value,
308            // `from_timestamp` returns None, and the cell was emitted EMPTY —
309            // silent data loss for every sub-second timestamp before the epoch
310            // (Parquet keeps them). Euclidean rem is always in [0, 1_000_000), so
311            // nsecs stays in range: micros=-1 → secs=-1, nsecs=999_999_000 →
312            // 1969-12-31T23:59:59.999999.
313            let secs = micros.div_euclid(1_000_000);
314            let nsecs = (micros.rem_euclid(1_000_000) * 1_000) as u32;
315            if let Some(dt) = chrono::DateTime::from_timestamp(secs, nsecs) {
316                use chrono::{Datelike as _, Timelike as _};
317                let y = dt.year();
318                // Fast path: emit the `%Y-%m-%dT%H:%M:%S%.6f` form by hand to
319                // skip chrono's per-value strftime-string re-parse. `{:04}`
320                // matches `%Y`'s zero-pad-to-4 exactly for years 0..=9999;
321                // outside that range (pre-CE / 5-digit years from pathological
322                // micros) fall back to chrono so the output stays identical.
323                if (0..=9999).contains(&y) {
324                    write!(
325                        writer,
326                        "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:06}",
327                        y,
328                        dt.month(),
329                        dt.day(),
330                        dt.hour(),
331                        dt.minute(),
332                        dt.second(),
333                        dt.nanosecond() / 1_000
334                    )?;
335                } else {
336                    write!(writer, "{}", dt.format("%Y-%m-%dT%H:%M:%S%.6f"))?;
337                }
338                if is_instant {
339                    writer.write_all(b"Z")?;
340                }
341            }
342        }
343        other => {
344            // Defensive: `create_writer` rejects unsupported types up front, so
345            // this should be unreachable. Bail rather than silently skip if a
346            // new type slips through.
347            anyhow::bail!(
348                "CSV: no serializer for Arrow type {other:?} (column should have been rejected at writer creation)"
349            );
350        }
351    }
352
353    Ok(())
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359    use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
360    use std::sync::Arc;
361
362    // Helper: render one cell to a String using write_csv_value.
363    fn cell<A: Array + 'static>(array: A, idx: usize) -> String {
364        let mut buf = Vec::new();
365        write_csv_value(&mut buf, &array, idx).unwrap();
366        String::from_utf8(buf).unwrap()
367    }
368
369    // Helper: render a null cell from any typed array.
370    fn null_cell(dt: DataType) -> String {
371        use arrow::array::new_null_array;
372        let arr = new_null_array(&dt, 1);
373        let mut buf = Vec::new();
374        write_csv_value(&mut buf, arr.as_ref(), 0).unwrap();
375        String::from_utf8(buf).unwrap()
376    }
377
378    // ── mutation-W5 gap closures ─────────────────────────────────────────────
379
380    #[test]
381    fn time64_renders_exact_wall_clock() {
382        // W5: the µs → hh:mm:ss.ffffff arithmetic (/, %) had no golden — a
383        // mutated divisor silently corrupts every TIME cell in CSV output.
384        use arrow::array::Time64MicrosecondArray;
385        assert_eq!(
386            cell(Time64MicrosecondArray::from(vec![86_399_999_999_i64]), 0),
387            "23:59:59.999999"
388        );
389        assert_eq!(
390            cell(Time64MicrosecondArray::from(vec![3_661_000_001_i64]), 0),
391            "01:01:01.000001"
392        );
393        assert_eq!(
394            cell(Time64MicrosecondArray::from(vec![0_i64]), 0),
395            "00:00:00.000000"
396        );
397    }
398
399    #[test]
400    fn write_batch_layout_and_byte_accounting_are_exact() {
401        // W5: `col_idx > 0` mutating to `>=` puts a LEADING comma on every row
402        // (corrupt CSV); the header `len + 1` and `bytes_written +=` mutants
403        // desync the rotation accounting. Pin the exact output AND the count.
404        use arrow::array::Int64Array;
405        use std::sync::{Arc as SArc, Mutex};
406
407        #[derive(Clone)]
408        struct SharedBuf(SArc<Mutex<Vec<u8>>>);
409        impl std::io::Write for SharedBuf {
410            fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
411                self.0.lock().unwrap().extend_from_slice(b);
412                Ok(b.len())
413            }
414            fn flush(&mut self) -> std::io::Result<()> {
415                Ok(())
416            }
417        }
418
419        let schema: SchemaRef = Arc::new(Schema::new(vec![
420            Field::new("a", DataType::Int64, false),
421            Field::new("b", DataType::Int64, false),
422        ]));
423        let sink = SharedBuf(SArc::new(Mutex::new(Vec::new())));
424        use crate::format::Format as _;
425        let mut w = CsvFormat
426            .create_writer(&schema, Box::new(sink.clone()))
427            .unwrap();
428        let batch = RecordBatch::try_new(
429            schema,
430            vec![
431                Arc::new(Int64Array::from(vec![1, 2])),
432                Arc::new(Int64Array::from(vec![10, 20])),
433            ],
434        )
435        .unwrap();
436        w.write_batch(&batch).unwrap();
437
438        let text = String::from_utf8(sink.0.lock().unwrap().clone()).unwrap();
439        assert_eq!(
440            text, "a,b\n1,10\n2,20\n",
441            "exact CSV layout (no leading commas)"
442        );
443        assert_eq!(
444            w.bytes_written(),
445            text.len() as u64,
446            "byte accounting must equal the physical output"
447        );
448    }
449
450    // ── null handling ────────────────────────────────────────────────────────
451
452    #[test]
453    fn null_value_writes_empty_string() {
454        assert_eq!(null_cell(DataType::Int64), "");
455        assert_eq!(null_cell(DataType::Utf8), "");
456        assert_eq!(null_cell(DataType::Boolean), "");
457    }
458
459    // ── scalars ─────────────────────────────────────────────────────────────
460
461    #[test]
462    fn bool_true_writes_true() {
463        assert_eq!(cell(BooleanArray::from(vec![true]), 0), "true");
464    }
465
466    #[test]
467    fn bool_false_writes_false() {
468        assert_eq!(cell(BooleanArray::from(vec![false]), 0), "false");
469    }
470
471    #[test]
472    fn int16_value() {
473        assert_eq!(cell(Int16Array::from(vec![42i16]), 0), "42");
474    }
475
476    #[test]
477    fn int32_negative() {
478        assert_eq!(cell(Int32Array::from(vec![-7i32]), 0), "-7");
479    }
480
481    #[test]
482    fn decimal128_writes_exact_text() {
483        let arr = Decimal128Array::from(vec![10i128])
484            .with_precision_and_scale(18, 2)
485            .unwrap();
486        assert_eq!(cell(arr, 0), "0.10");
487        let scaled =
488            crate::types::decimal::decimal_str_to_scaled_i128("999999999999.99", 2).unwrap();
489        let arr = Decimal128Array::from(vec![scaled])
490            .with_precision_and_scale(18, 2)
491            .unwrap();
492        assert_eq!(cell(arr, 0), "999999999999.99");
493    }
494
495    #[test]
496    fn int64_large() {
497        assert_eq!(
498            cell(Int64Array::from(vec![9_999_999_999i64]), 0),
499            "9999999999"
500        );
501    }
502
503    #[test]
504    fn float32_value() {
505        let result = cell(Float32Array::from(vec![1.5f32]), 0);
506        assert!(result.starts_with("1.5"), "got: {result}");
507    }
508
509    #[test]
510    fn float64_value() {
511        let result = cell(Float64Array::from(vec![std::f64::consts::PI]), 0);
512        assert!(result.starts_with("3.14"), "got: {result}");
513    }
514
515    // Characterization: float NaN/±Infinity are IEEE-754 values a float column
516    // can legitimately hold (unlike `decimal`, whose Arrow `Decimal128` has no
517    // NaN bit pattern — see the NUMERIC NaN/infinity reject in
518    // `postgres::arrow_convert::build_array`). They are preserved natively in
519    // Parquet; in CSV we emit the Rust float literal (`NaN` / `inf` / `-inf`)
520    // rather than an empty cell, because writing empty would silently conflate
521    // a real NaN/Inf with NULL — corruption — whereas a recognizable literal
522    // round-trips into every major loader's float parser. This pins that
523    // contract so a future arrow/std change can't silently alter it. The CSV
524    // literal is documented in docs/type-mapping.md.
525    #[test]
526    fn float_special_values_emit_literals_not_empty() {
527        assert_eq!(cell(Float64Array::from(vec![f64::NAN]), 0), "NaN");
528        assert_eq!(cell(Float64Array::from(vec![f64::INFINITY]), 0), "inf");
529        assert_eq!(cell(Float64Array::from(vec![f64::NEG_INFINITY]), 0), "-inf");
530        assert_eq!(cell(Float32Array::from(vec![f32::NAN]), 0), "NaN");
531        assert_eq!(cell(Float32Array::from(vec![f32::INFINITY]), 0), "inf");
532        // -0.0 keeps its sign (a real IEEE-754 distinction), never becomes "0".
533        assert_eq!(cell(Float64Array::from(vec![-0.0f64]), 0), "-0");
534    }
535
536    // ── string escaping ──────────────────────────────────────────────────────
537
538    #[test]
539    fn plain_string_no_quoting() {
540        assert_eq!(cell(StringArray::from(vec!["hello"]), 0), "hello");
541    }
542
543    #[test]
544    fn string_with_comma_is_quoted() {
545        assert_eq!(cell(StringArray::from(vec!["a,b"]), 0), "\"a,b\"");
546    }
547
548    #[test]
549    fn string_with_double_quote_is_escaped() {
550        // say "hi" → opening " + say  + "" + hi + "" + closing " = "say ""hi"""
551        let result = cell(StringArray::from(vec![r#"say "hi""#]), 0);
552        assert_eq!(result, r#""say ""hi""""#);
553    }
554
555    #[test]
556    fn string_with_newline_is_quoted() {
557        let result = cell(StringArray::from(vec!["line1\nline2"]), 0);
558        assert!(
559            result.starts_with('"') && result.ends_with('"'),
560            "got: {result}"
561        );
562        assert!(result.contains("line1\nline2"), "got: {result}");
563    }
564
565    // ROAST-RED csv-cr-quote: the quote predicate checks ',', '"' and '\n' but
566    // not '\r', so a value containing a lone CR is emitted unquoted — RFC 4180
567    // requires quoting CR, and Excel/Python universal-newline readers split the
568    // row on the bare CR.
569    // Asserts CORRECT behavior; expected to FAIL until the fix lands.
570    #[test]
571    fn roast_string_with_carriage_return_is_quoted() {
572        let result = cell(StringArray::from(vec!["a\rb"]), 0);
573        assert_eq!(
574            result, "\"a\rb\"",
575            "lone CR must force quoting per RFC 4180, but got unquoted cell {result:?}"
576        );
577    }
578
579    // ── binary ───────────────────────────────────────────────────────────────
580
581    #[test]
582    fn binary_is_written_as_hex() {
583        let arr = BinaryArray::from_vec(vec![&[0xDE, 0xAD, 0xBE, 0xEF][..]]);
584        assert_eq!(cell(arr, 0), "deadbeef");
585    }
586
587    #[test]
588    fn binary_empty_writes_empty() {
589        let arr = BinaryArray::from_vec(vec![&[][..]]);
590        assert_eq!(cell(arr, 0), "");
591    }
592
593    // ── Date32 ───────────────────────────────────────────────────────────────
594
595    #[test]
596    fn date32_epoch_is_1970_01_01() {
597        assert_eq!(cell(Date32Array::from(vec![0i32]), 0), "1970-01-01");
598    }
599
600    #[test]
601    fn date32_positive_offset() {
602        // 365 days after epoch = 1971-01-01
603        assert_eq!(cell(Date32Array::from(vec![365i32]), 0), "1971-01-01");
604    }
605
606    // ── Timestamp(Microsecond) ───────────────────────────────────────────────
607
608    #[test]
609    fn timestamp_micros_formats_as_iso() {
610        // 2023-01-01T00:00:00.000000 = 1672531200_000000 micros since epoch
611        let micros: i64 = 1_672_531_200 * 1_000_000;
612        let _schema = Arc::new(Schema::new(vec![Field::new(
613            "ts",
614            DataType::Timestamp(TimeUnit::Microsecond, None),
615            true,
616        )]));
617        let arr = TimestampMicrosecondArray::from(vec![micros]);
618        let result = cell(arr, 0);
619        assert!(result.starts_with("2023-01-01T"), "got: {result}");
620        assert!(result.contains("00:00:00"), "got: {result}");
621    }
622
623    // rivet keeps the naive-vs-instant distinction in every other output
624    // (Parquet isAdjustedToUTC, the warehouse target types). CSV must too: a
625    // naive TIMESTAMP renders bare (its wall-clock IS the value), an instant
626    // TIMESTAMPTZ (Timestamp(_, Some("UTC"))) gets a trailing `Z` — else a
627    // consumer loading the instant on a NON-UTC session reads the UTC wall-clock
628    // as local and shifts every value.
629    #[test]
630    fn timestamp_marks_instant_utc_but_leaves_naive_bare() {
631        // 2024-06-15T03:00:00.000000 UTC.
632        let micros: i64 = 1_718_420_400 * 1_000_000;
633        let naive = TimestampMicrosecondArray::from(vec![micros]);
634        assert_eq!(
635            cell(naive, 0),
636            "2024-06-15T03:00:00.000000",
637            "a naive TIMESTAMP renders bare — no tz marker"
638        );
639        let instant = TimestampMicrosecondArray::from(vec![micros]).with_timezone("UTC");
640        assert_eq!(
641            cell(instant, 0),
642            "2024-06-15T03:00:00.000000Z",
643            "an instant TIMESTAMPTZ renders an explicit UTC `Z`"
644        );
645    }
646
647    // ── write_batch via CsvFormat ────────────────────────────────────────────
648
649    #[test]
650    fn csv_format_write_batch_tracks_bytes_and_succeeds() {
651        use crate::format::Format;
652
653        let schema = Arc::new(Schema::new(vec![
654            Field::new("id", DataType::Int64, false),
655            Field::new("name", DataType::Utf8, true),
656        ]));
657        let batch = arrow::record_batch::RecordBatch::try_new(
658            schema.clone(),
659            vec![
660                Arc::new(Int64Array::from(vec![1i64, 2])),
661                Arc::new(StringArray::from(vec![Some("alice"), None])),
662            ],
663        )
664        .unwrap();
665
666        // Pass Vec by value — avoids the &mut T 'static lifetime requirement.
667        let fmt = CsvFormat;
668        let mut writer = fmt
669            .create_writer(&schema, Box::new(Vec::<u8>::new()))
670            .unwrap();
671        writer.write_batch(&batch).unwrap();
672        // Header "id,name\n" + rows "1,alice\n" + "2,\n" = at least 18 bytes
673        assert!(
674            writer.bytes_written() > 10,
675            "expected >10 bytes, got {}",
676            writer.bytes_written()
677        );
678        writer.finish().unwrap();
679    }
680
681    // Regression: a column NAME containing a comma/quote must be RFC-4180 quoted
682    // in the header, exactly as data cells are — else a curated-query alias like
683    // `... AS "Amount, USD"` split the header and silently mis-aligned every
684    // downstream reader from the data columns. RED before the header quoter.
685    #[test]
686    fn csv_header_quotes_a_column_name_with_a_comma_or_quote() {
687        use crate::format::Format;
688        let schema = Arc::new(Schema::new(vec![
689            Field::new("Amount, USD", DataType::Int64, false), // comma in the name
690            Field::new("a\"b", DataType::Utf8, true),          // quote in the name
691            Field::new("plain", DataType::Utf8, true),         // untouched
692        ]));
693        let dir = tempfile::tempdir().unwrap();
694        let path = dir.path().join("out.csv");
695        let w = CsvFormat
696            .create_writer(&schema, Box::new(std::fs::File::create(&path).unwrap()))
697            .unwrap();
698        w.finish().unwrap();
699        let out = std::fs::read_to_string(&path).unwrap();
700        let header = out.lines().next().unwrap();
701        // comma-name quoted; quote-name quoted with the `"` doubled; plain bare —
702        // and crucially still exactly THREE header columns, the comma not a split.
703        assert_eq!(header, r#""Amount, USD","a""b",plain"#);
704    }
705
706    // ── fail loud on types CSV can't represent ───────────────────────────────
707
708    #[test]
709    fn csv_rejects_array_columns_loudly() {
710        use crate::format::Format;
711        let schema = Arc::new(Schema::new(vec![
712            Field::new("id", DataType::Int64, false),
713            Field::new(
714                "tags",
715                DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
716                true,
717            ),
718        ]));
719        let Err(err) = CsvFormat.create_writer(&schema, Box::new(Vec::<u8>::new())) else {
720            panic!("CSV must reject array columns, not silently drop them");
721        };
722        let msg = format!("{err:#}");
723        assert!(msg.contains("tags"), "error must name the column: {msg}");
724        assert!(msg.to_lowercase().contains("csv"), "{msg}");
725    }
726
727    /// Consistency guard: every type `csv_serializable` admits must actually be
728    /// handled by `write_csv_value` (not hit its `other => bail` fallthrough).
729    /// Keeps the whitelist and the writer in lock-step so one can't drift.
730    #[test]
731    fn every_serializable_type_is_actually_written() {
732        use crate::format::Format;
733        let cols: Vec<(&str, ArrayRef)> = vec![
734            ("b", Arc::new(BooleanArray::from(vec![true]))),
735            ("i16", Arc::new(Int16Array::from(vec![1i16]))),
736            ("i32", Arc::new(Int32Array::from(vec![1i32]))),
737            ("i64", Arc::new(Int64Array::from(vec![1i64]))),
738            ("u64", Arc::new(UInt64Array::from(vec![1u64]))),
739            (
740                "dec",
741                Arc::new(
742                    Decimal128Array::from(vec![100i128])
743                        .with_precision_and_scale(18, 2)
744                        .unwrap(),
745                ),
746            ),
747            ("f32", Arc::new(Float32Array::from(vec![1.0f32]))),
748            ("f64", Arc::new(Float64Array::from(vec![1.0f64]))),
749            ("s", Arc::new(StringArray::from(vec!["x"]))),
750            ("bin", Arc::new(BinaryArray::from_vec(vec![&[1u8][..]]))),
751            (
752                "uuid",
753                Arc::new(
754                    FixedSizeBinaryArray::try_from_iter(std::iter::once(vec![0u8; 16])).unwrap(),
755                ),
756            ),
757            ("d", Arc::new(Date32Array::from(vec![0i32]))),
758            ("t", Arc::new(Time64MicrosecondArray::from(vec![0i64]))),
759            ("ts", Arc::new(TimestampMicrosecondArray::from(vec![0i64]))),
760        ];
761        let fields: Vec<Field> = cols
762            .iter()
763            .map(|(n, a)| Field::new(*n, a.data_type().clone(), true))
764            .collect();
765        // Sanity: each column's type is on the whitelist.
766        for f in &fields {
767            assert!(
768                csv_serializable(f.data_type()),
769                "test type {:?} not in csv_serializable",
770                f.data_type()
771            );
772        }
773        let schema = Arc::new(Schema::new(fields));
774        let arrays: Vec<ArrayRef> = cols.into_iter().map(|(_, a)| a).collect();
775        let batch = RecordBatch::try_new(schema.clone(), arrays).unwrap();
776        let mut w = CsvFormat
777            .create_writer(&schema, Box::new(Vec::<u8>::new()))
778            .unwrap();
779        w.write_batch(&batch)
780            .expect("every serializable type must write without hitting the fallthrough");
781    }
782
783    // ── perf-wave byte-identity guards ───────────────────────────────────────
784    // The hex + timestamp fast paths must stay byte-for-byte identical to the
785    // `{:02x}`-per-byte / `dt.format(...)` forms they replaced.
786
787    #[test]
788    fn binary_hex_matches_per_byte_format_for_all_byte_values() {
789        // Every byte 0..=255 plus the empty slice — the table+chunk encoder
790        // must equal the old per-byte `{:02x}`.
791        let all: Vec<u8> = (0..=255u8).collect();
792        for case in [&all[..], &[][..], &[0x00, 0xff, 0x10, 0x0a]] {
793            let expected: String = case.iter().map(|b| format!("{b:02x}")).collect();
794            let got = cell(BinaryArray::from_vec(vec![case]), 0);
795            assert_eq!(got, expected, "hex mismatch for {case:?}");
796        }
797    }
798
799    #[test]
800    fn binary_hex_spans_chunk_boundary() {
801        // Larger than the 512-byte slab so the chunked write_all path is hit.
802        let big: Vec<u8> = (0..2000u32).map(|i| (i % 256) as u8).collect();
803        let expected: String = big.iter().map(|b| format!("{b:02x}")).collect();
804        let got = cell(BinaryArray::from_vec(vec![&big[..]]), 0);
805        assert_eq!(got, expected);
806    }
807
808    #[test]
809    fn timestamp_fast_path_matches_chrono_format() {
810        // Representative micros: epoch, sub-second, end-of-day, a far-future
811        // in-range year, and a value whose year leaves the 0..=9999 fast path
812        // (must hit the chrono fallback and still match). The recompute uses the
813        // SAME euclidean split the product now uses, so it is a self-consistency
814        // check of the fast-path-vs-chrono rendering — the pre-1970 correctness
815        // is pinned separately below with a HARD-CODED oracle.
816        let cases: [i64; 6] = [
817            0,
818            1_700_000_000_123_456,        // 2023-… with micros
819            1_000_000 * 86_399 + 999_999, // 1970-01-01T23:59:59.999999
820            -62_135_596_800_000_000,      // 0001-01-01T00:00:00 (negative, fast path)
821            253_402_300_799_000_000,      // 9999-12-31T23:59:59 (still fast path)
822            300_000_000_000_000_000,      // year > 9999 → chrono fallback
823        ];
824        for micros in cases {
825            let got = cell(TimestampMicrosecondArray::from(vec![micros]), 0);
826            let secs = micros.div_euclid(1_000_000);
827            let nsecs = (micros.rem_euclid(1_000_000) * 1_000) as u32;
828            let expected = match chrono::DateTime::from_timestamp(secs, nsecs) {
829                Some(dt) => format!("{}", dt.format("%Y-%m-%dT%H:%M:%S%.6f")),
830                None => String::new(),
831            };
832            assert_eq!(got, expected, "timestamp mismatch for micros={micros}");
833        }
834    }
835
836    // Regression (true oracle, NOT the recomputed self-oracle above): a pre-1970
837    // sub-second timestamp must render its real value, never an empty cell. With
838    // the old truncating `/`+`%` split the negative remainder wrapped `nsecs`,
839    // `from_timestamp` returned None, and the cell was emitted EMPTY — 100%
840    // silent loss for this whole class while Parquet kept the value. Hard-coded
841    // expected strings so the test can never re-derive the bug.
842    #[test]
843    fn pre_1970_subsecond_timestamp_is_not_dropped_to_an_empty_cell() {
844        let expect: [(i64, &str); 4] = [
845            (-1, "1969-12-31T23:59:59.999999"), // one micro before the epoch
846            (-500_000, "1969-12-31T23:59:59.500000"), // half a second before
847            (-86_400_000_000, "1969-12-31T00:00:00.000000"), // whole day before (whole second)
848            (-125_125_000, "1969-12-31T23:57:54.875000"), // multi-second, sub-second frac
849        ];
850        for (micros, want) in expect {
851            let got = cell(TimestampMicrosecondArray::from(vec![micros]), 0);
852            assert_eq!(
853                got, want,
854                "pre-1970 micros={micros} must render {want:?}, not an empty/other cell"
855            );
856        }
857    }
858}