laminar_connectors/schema/
csv.rs

1//! CSV format decoder implementing [`FormatDecoder`].
2//!
3//! Converts raw CSV byte payloads into Arrow `RecordBatch`es.
4//! Constructed once at `CREATE SOURCE` time with a frozen Arrow schema
5//! and CSV format configuration. The decoder is stateless after
6//! construction so the Ring 1 hot path has zero configuration lookups.
7//!
8//! Uses the `csv` crate's `ByteRecord` API for zero-copy field access
9//! where possible. Type coercion (string → int, string → timestamp, etc.)
10//! is performed during the Arrow builder append phase.
11
12use std::sync::atomic::{AtomicU64, Ordering};
13use std::sync::Arc;
14
15use arrow_array::builder::{
16    BooleanBuilder, Date32Builder, Float64Builder, Int64Builder, StringBuilder,
17    TimestampNanosecondBuilder,
18};
19use arrow_array::{ArrayRef, RecordBatch};
20use arrow_schema::{DataType, SchemaRef, TimeUnit};
21
22use crate::schema::error::{SchemaError, SchemaResult};
23use crate::schema::traits::{FormatDecoder, FormatEncoder};
24use crate::schema::types::RawRecord;
25
26/// Strategy for rows with incorrect field count.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum FieldCountMismatchStrategy {
29    /// Pad missing fields with null, ignore extra fields. Default.
30    Null,
31    /// Skip the malformed row entirely.
32    Skip,
33    /// Return a decode error on the first malformed row.
34    Reject,
35}
36
37/// CSV decoder configuration.
38///
39/// Maps directly to the SQL `FORMAT CSV (...)` options.
40/// All fields have sensible defaults matching RFC 4180.
41#[derive(Debug, Clone)]
42pub struct CsvDecoderConfig {
43    /// Field delimiter character. Default: `','` (comma).
44    /// Common alternatives: `'\t'` (tab), `'|'` (pipe), `';'` (semicolon).
45    pub delimiter: u8,
46
47    /// Quote character for fields containing delimiters or newlines.
48    /// Default: `'"'` (double quote). Set to `None` to disable quoting.
49    pub quote: Option<u8>,
50
51    /// Escape character within quoted fields.
52    /// Default: `None` (RFC 4180 uses doubled quote chars for escaping).
53    /// Set to `Some(b'\\')` for backslash-escaped CSVs.
54    pub escape: Option<u8>,
55
56    /// Whether the first row is a header row with column names.
57    /// Default: `true`.
58    pub has_header: bool,
59
60    /// String value to interpret as SQL NULL.
61    /// Default: `""` (empty string). Common alternatives: `"NA"`, `"null"`, `"\\N"`.
62    pub null_string: String,
63
64    /// Comment line prefix. Lines starting with this character are skipped.
65    /// Default: `None` (no comment support).
66    pub comment: Option<u8>,
67
68    /// Number of rows to skip at the beginning of the data (after header).
69    /// Default: `0`.
70    pub skip_rows: usize,
71
72    /// Timestamp format pattern for parsing timestamp columns.
73    /// Default: `"%Y-%m-%d %H:%M:%S%.f"`.
74    pub timestamp_format: String,
75
76    /// Date format pattern for parsing date columns.
77    /// Default: `"%Y-%m-%d"`.
78    pub date_format: String,
79
80    /// How to handle rows with wrong number of fields.
81    /// Default: `Null` (pad missing fields with null, truncate extra).
82    pub field_count_mismatch: FieldCountMismatchStrategy,
83}
84
85impl Default for CsvDecoderConfig {
86    fn default() -> Self {
87        Self {
88            delimiter: b',',
89            quote: Some(b'"'),
90            escape: None,
91            has_header: true,
92            null_string: String::new(),
93            comment: None,
94            skip_rows: 0,
95            timestamp_format: "%Y-%m-%d %H:%M:%S%.f".into(),
96            date_format: "%Y-%m-%d".into(),
97            field_count_mismatch: FieldCountMismatchStrategy::Null,
98        }
99    }
100}
101
102/// Pre-computed coercion strategy for a single CSV column.
103#[derive(Debug, Clone)]
104enum CsvCoercion {
105    /// Parse as boolean (`"true"`/`"false"`, `"1"`/`"0"`, `"yes"`/`"no"`).
106    Boolean,
107    /// Parse as i64.
108    Int64,
109    /// Parse as f64.
110    Float64,
111    /// Parse as `Timestamp(Nanosecond, UTC)` using the configured format.
112    Timestamp(String),
113    /// Parse as `Date32` using the configured format.
114    Date(String),
115    /// No coercion needed — keep as UTF-8 string.
116    Utf8,
117}
118
119/// Decodes CSV byte payloads into Arrow `RecordBatch`es.
120///
121/// # Ring Placement
122///
123/// - **Ring 1**: `decode_batch()` — parse CSV, build columnar Arrow output
124/// - **Ring 2**: Construction (`new` / `with_config`) — one-time setup
125pub struct CsvDecoder {
126    /// Frozen output schema.
127    schema: SchemaRef,
128    /// CSV format configuration.
129    config: CsvDecoderConfig,
130    /// Per-column type coercion functions, indexed by column position.
131    /// Pre-computed at construction time to avoid per-record dispatch.
132    coercions: Vec<CsvCoercion>,
133    /// Cumulative count of parse errors (for diagnostics).
134    parse_error_count: AtomicU64,
135}
136
137#[allow(clippy::missing_fields_in_debug)]
138impl std::fmt::Debug for CsvDecoder {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        f.debug_struct("CsvDecoder")
141            .field("schema", &self.schema)
142            .field("config", &self.config)
143            .field(
144                "parse_error_count",
145                &self.parse_error_count.load(Ordering::Relaxed),
146            )
147            .finish()
148    }
149}
150
151impl CsvDecoder {
152    /// Creates a new CSV decoder for the given Arrow schema with default config.
153    #[must_use]
154    pub fn new(schema: SchemaRef) -> Self {
155        Self::with_config(schema, CsvDecoderConfig::default())
156    }
157
158    /// Creates a new CSV decoder with custom configuration.
159    #[must_use]
160    pub fn with_config(schema: SchemaRef, config: CsvDecoderConfig) -> Self {
161        let coercions: Vec<CsvCoercion> = schema
162            .fields()
163            .iter()
164            .map(|field| Self::coercion_for_type(field.data_type(), &config))
165            .collect();
166
167        Self {
168            schema,
169            config,
170            coercions,
171            parse_error_count: AtomicU64::new(0),
172        }
173    }
174
175    /// Returns the cumulative parse error count.
176    pub fn parse_error_count(&self) -> u64 {
177        self.parse_error_count.load(Ordering::Relaxed)
178    }
179
180    /// Determines the coercion strategy for an Arrow data type.
181    fn coercion_for_type(data_type: &DataType, config: &CsvDecoderConfig) -> CsvCoercion {
182        match data_type {
183            DataType::Boolean => CsvCoercion::Boolean,
184            DataType::Int8
185            | DataType::Int16
186            | DataType::Int32
187            | DataType::Int64
188            | DataType::UInt8
189            | DataType::UInt16
190            | DataType::UInt32
191            | DataType::UInt64 => CsvCoercion::Int64,
192            DataType::Float16 | DataType::Float32 | DataType::Float64 => CsvCoercion::Float64,
193            DataType::Timestamp(_, _) => CsvCoercion::Timestamp(config.timestamp_format.clone()),
194            DataType::Date32 | DataType::Date64 => CsvCoercion::Date(config.date_format.clone()),
195            _ => CsvCoercion::Utf8,
196        }
197    }
198
199    /// Builds a `csv::ReaderBuilder` from the decoder config.
200    fn make_reader_builder(&self) -> csv::ReaderBuilder {
201        let mut rb = csv::ReaderBuilder::new();
202        rb.delimiter(self.config.delimiter)
203            .has_headers(false) // We handle headers ourselves
204            .flexible(true); // Allow variable field counts
205
206        if let Some(q) = self.config.quote {
207            rb.quote(q);
208        }
209        if let Some(e) = self.config.escape {
210            rb.escape(Some(e));
211        }
212        if let Some(c) = self.config.comment {
213            rb.comment(Some(c));
214        }
215
216        rb
217    }
218}
219
220impl FormatDecoder for CsvDecoder {
221    fn output_schema(&self) -> SchemaRef {
222        self.schema.clone()
223    }
224
225    fn decode_batch(&self, records: &[RawRecord]) -> SchemaResult<RecordBatch> {
226        let values: Vec<&[u8]> = records
227            .iter()
228            .map(|record| record.value.as_slice())
229            .collect();
230        self.decode_slices(&values)
231    }
232
233    fn format_name(&self) -> &'static str {
234        "csv"
235    }
236}
237
238impl CsvDecoder {
239    /// Decodes borrowed CSV payloads without first copying them into raw records.
240    ///
241    /// # Errors
242    ///
243    /// Returns a decode error for malformed CSV or values incompatible with the schema.
244    pub fn decode_slices(&self, records: &[&[u8]]) -> SchemaResult<RecordBatch> {
245        if records.is_empty() {
246            return Ok(RecordBatch::new_empty(self.schema.clone()));
247        }
248
249        let num_fields = self.schema.fields().len();
250        let capacity = records.len();
251
252        // Initialize one builder per schema column.
253        let mut builders = create_builders(&self.schema, capacity);
254
255        // Concatenate all raw record bytes, ensuring newline separation.
256        let mut combined = Vec::with_capacity(records.iter().map(|record| record.len() + 1).sum());
257        for record in records {
258            combined.extend_from_slice(record);
259            if !record.ends_with(b"\n") {
260                combined.push(b'\n');
261            }
262        }
263
264        let rb = self.make_reader_builder();
265        let mut reader = rb.from_reader(combined.as_slice());
266
267        let mut rows_skipped = 0usize;
268        let mut header_skipped = false;
269        let mut row_count = 0usize;
270
271        let mut byte_record = csv::ByteRecord::new();
272        while reader
273            .read_byte_record(&mut byte_record)
274            .map_err(|e| SchemaError::DecodeError(format!("CSV parse error: {e}")))?
275        {
276            // Skip header row if configured.
277            if self.config.has_header && !header_skipped {
278                header_skipped = true;
279                continue;
280            }
281
282            // Skip initial data rows per config.
283            if rows_skipped < self.config.skip_rows {
284                rows_skipped += 1;
285                continue;
286            }
287
288            let field_count = byte_record.len();
289
290            // Handle field count mismatch.
291            if field_count != num_fields {
292                match self.config.field_count_mismatch {
293                    FieldCountMismatchStrategy::Reject => {
294                        return Err(SchemaError::DecodeError(format!(
295                            "field count mismatch: expected {num_fields}, got {field_count}"
296                        )));
297                    }
298                    FieldCountMismatchStrategy::Skip => {
299                        self.parse_error_count.fetch_add(1, Ordering::Relaxed);
300                        continue;
301                    }
302                    FieldCountMismatchStrategy::Null => {
303                        // Will pad/truncate below.
304                    }
305                }
306            }
307
308            // Process each column.
309            for col_idx in 0..num_fields {
310                if col_idx >= field_count {
311                    // Missing field — append null.
312                    append_null(&mut builders[col_idx]);
313                    continue;
314                }
315
316                let raw_field = &byte_record[col_idx];
317                let field_str = std::str::from_utf8(raw_field).unwrap_or("");
318                let trimmed = field_str.trim();
319
320                // Check for null string.
321                if trimmed == self.config.null_string {
322                    append_null(&mut builders[col_idx]);
323                    continue;
324                }
325
326                // Apply coercion.
327                let ok = append_coerced(&mut builders[col_idx], &self.coercions[col_idx], trimmed);
328
329                if !ok {
330                    self.parse_error_count.fetch_add(1, Ordering::Relaxed);
331                    append_null(&mut builders[col_idx]);
332                }
333            }
334
335            row_count += 1;
336        }
337
338        // If no data rows were processed, return empty batch.
339        if row_count == 0 {
340            return Ok(RecordBatch::new_empty(self.schema.clone()));
341        }
342
343        // Finish all builders into arrays.
344        let columns: Vec<ArrayRef> = builders.into_iter().map(|mut b| b.finish()).collect();
345
346        RecordBatch::try_new(self.schema.clone(), columns)
347            .map_err(|e| SchemaError::DecodeError(format!("RecordBatch construction: {e}")))
348    }
349}
350
351// ── Builder helpers ────────────────────────────────────────────────
352
353/// Trait-object wrapper so we can store heterogeneous builders in a `Vec`.
354trait ColumnBuilder: Send {
355    fn finish(&mut self) -> ArrayRef;
356    fn append_null_value(&mut self);
357    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
358}
359
360macro_rules! impl_column_builder {
361    ($builder:ty) => {
362        impl ColumnBuilder for $builder {
363            fn finish(&mut self) -> ArrayRef {
364                Arc::new(<$builder>::finish(self))
365            }
366            fn append_null_value(&mut self) {
367                self.append_null();
368            }
369            fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
370                self
371            }
372        }
373    };
374}
375
376impl_column_builder!(BooleanBuilder);
377impl_column_builder!(Int64Builder);
378impl_column_builder!(Float64Builder);
379impl_column_builder!(StringBuilder);
380impl_column_builder!(TimestampNanosecondBuilder);
381impl_column_builder!(Date32Builder);
382
383fn create_builders(schema: &SchemaRef, capacity: usize) -> Vec<Box<dyn ColumnBuilder>> {
384    schema
385        .fields()
386        .iter()
387        .map(|f| create_builder(f.data_type(), capacity))
388        .collect()
389}
390
391fn create_builder(data_type: &DataType, capacity: usize) -> Box<dyn ColumnBuilder> {
392    match data_type {
393        DataType::Boolean => Box::new(BooleanBuilder::with_capacity(capacity)),
394        DataType::Int8
395        | DataType::Int16
396        | DataType::Int32
397        | DataType::Int64
398        | DataType::UInt8
399        | DataType::UInt16
400        | DataType::UInt32
401        | DataType::UInt64 => Box::new(Int64Builder::with_capacity(capacity)),
402        DataType::Float16 | DataType::Float32 | DataType::Float64 => {
403            Box::new(Float64Builder::with_capacity(capacity))
404        }
405        DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
406            let builder =
407                TimestampNanosecondBuilder::with_capacity(capacity).with_timezone_opt(tz.clone());
408            Box::new(builder)
409        }
410        DataType::Date32 | DataType::Date64 => Box::new(Date32Builder::with_capacity(capacity)),
411        // Fallback: store as UTF-8 string.
412        _ => Box::new(StringBuilder::with_capacity(capacity, capacity * 32)),
413    }
414}
415
416fn append_null(builder: &mut Box<dyn ColumnBuilder>) {
417    builder.append_null_value();
418}
419
420/// Appends a coerced value to the appropriate builder. Returns `true` on
421/// success, `false` if the value could not be parsed.
422fn append_coerced(
423    builder: &mut Box<dyn ColumnBuilder>,
424    coercion: &CsvCoercion,
425    value: &str,
426) -> bool {
427    match coercion {
428        CsvCoercion::Boolean => {
429            let b = builder
430                .as_any_mut()
431                .downcast_mut::<BooleanBuilder>()
432                .unwrap();
433            match value.to_ascii_lowercase().as_str() {
434                "true" | "1" | "yes" | "t" | "y" => {
435                    b.append_value(true);
436                    true
437                }
438                "false" | "0" | "no" | "f" | "n" => {
439                    b.append_value(false);
440                    true
441                }
442                _ => false,
443            }
444        }
445        CsvCoercion::Int64 => {
446            let b = builder.as_any_mut().downcast_mut::<Int64Builder>().unwrap();
447            match value.parse::<i64>() {
448                Ok(v) => {
449                    b.append_value(v);
450                    true
451                }
452                Err(_) => false,
453            }
454        }
455        CsvCoercion::Float64 => {
456            let b = builder
457                .as_any_mut()
458                .downcast_mut::<Float64Builder>()
459                .unwrap();
460            match value.parse::<f64>() {
461                Ok(v) => {
462                    b.append_value(v);
463                    true
464                }
465                Err(_) => false,
466            }
467        }
468        CsvCoercion::Timestamp(fmt) => {
469            let b = builder
470                .as_any_mut()
471                .downcast_mut::<TimestampNanosecondBuilder>()
472                .unwrap();
473            // Try the configured format first.
474            if let Ok(ndt) = chrono::NaiveDateTime::parse_from_str(value, fmt) {
475                let nanos = ndt.and_utc().timestamp_nanos_opt().unwrap_or(0);
476                b.append_value(nanos);
477                return true;
478            }
479            // Try ISO 8601 fallback.
480            if let Ok(nanos) = arrow_cast::parse::string_to_timestamp_nanos(value) {
481                b.append_value(nanos);
482                return true;
483            }
484            false
485        }
486        CsvCoercion::Date(fmt) => {
487            let b = builder
488                .as_any_mut()
489                .downcast_mut::<Date32Builder>()
490                .unwrap();
491            if let Ok(date) = chrono::NaiveDate::parse_from_str(value, fmt) {
492                // Date32 stores days since epoch (1970-01-01).
493                let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
494                let days = (date - epoch).num_days();
495                #[allow(clippy::cast_possible_truncation)]
496                {
497                    b.append_value(days as i32);
498                }
499                return true;
500            }
501            false
502        }
503        CsvCoercion::Utf8 => {
504            let b = builder
505                .as_any_mut()
506                .downcast_mut::<StringBuilder>()
507                .unwrap();
508            b.append_value(value);
509            true
510        }
511    }
512}
513
514/// Configuration for [`CsvEncoder`].
515#[derive(Debug, Clone)]
516pub struct CsvEncoderConfig {
517    /// Field delimiter. Default: `','`.
518    pub delimiter: u8,
519    /// Whether to include a header row. Default: `false`.
520    pub has_header: bool,
521}
522
523impl Default for CsvEncoderConfig {
524    fn default() -> Self {
525        Self {
526            delimiter: b',',
527            has_header: false,
528        }
529    }
530}
531
532/// Encodes Arrow `RecordBatch`es into CSV byte records via `arrow_csv::writer`.
533#[derive(Debug)]
534pub struct CsvEncoder {
535    schema: SchemaRef,
536    config: CsvEncoderConfig,
537}
538
539impl CsvEncoder {
540    /// Creates a new CSV encoder for the given schema with default config.
541    #[must_use]
542    pub fn new(schema: SchemaRef) -> Self {
543        Self::with_config(schema, CsvEncoderConfig::default())
544    }
545
546    /// Creates a new CSV encoder with custom configuration.
547    #[must_use]
548    pub fn with_config(schema: SchemaRef, config: CsvEncoderConfig) -> Self {
549        Self { schema, config }
550    }
551}
552
553impl FormatEncoder for CsvEncoder {
554    fn input_schema(&self) -> SchemaRef {
555        self.schema.clone()
556    }
557
558    fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>> {
559        if batch.num_rows() == 0 {
560            return Ok(Vec::new());
561        }
562
563        let mut buf = Vec::new();
564        {
565            let writer = arrow_csv::writer::WriterBuilder::new()
566                .with_header(self.config.has_header)
567                .with_delimiter(self.config.delimiter);
568            let mut csv_writer = writer.build(&mut buf);
569            csv_writer
570                .write(batch)
571                .map_err(|e| SchemaError::DecodeError(format!("CSV encode error: {e}")))?;
572        }
573
574        let output: Vec<Vec<u8>> = buf
575            .split(|&b| b == b'\n')
576            .filter(|line| !line.is_empty())
577            .map(<[u8]>::to_vec)
578            .collect();
579
580        Ok(output)
581    }
582
583    fn format_name(&self) -> &'static str {
584        "csv"
585    }
586}
587
588#[cfg(test)]
589mod tests {
590    use super::*;
591    use crate::schema::traits::FormatEncoder;
592    use arrow_array::cast::AsArray;
593    use arrow_schema::{Field, Schema};
594
595    fn make_schema(fields: Vec<(&str, DataType, bool)>) -> SchemaRef {
596        Arc::new(Schema::new(
597            fields
598                .into_iter()
599                .map(|(name, dt, nullable)| Field::new(name, dt, nullable))
600                .collect::<Vec<_>>(),
601        ))
602    }
603
604    fn csv_record(line: &str) -> RawRecord {
605        RawRecord::new(line.as_bytes().to_vec())
606    }
607
608    fn csv_block(lines: &str) -> RawRecord {
609        RawRecord::new(lines.as_bytes().to_vec())
610    }
611
612    // ── Basic decode tests ────────────────────────────────────
613
614    #[test]
615    fn test_decode_empty_batch() {
616        let schema = make_schema(vec![("id", DataType::Int64, false)]);
617        let decoder = CsvDecoder::new(schema.clone());
618        let batch = decoder.decode_batch(&[]).unwrap();
619        assert_eq!(batch.num_rows(), 0);
620        assert_eq!(batch.schema(), schema);
621    }
622
623    #[test]
624    fn test_decode_single_row_with_header() {
625        let schema = make_schema(vec![
626            ("id", DataType::Int64, false),
627            ("name", DataType::Utf8, true),
628        ]);
629        let decoder = CsvDecoder::new(schema);
630        let records = vec![csv_block("id,name\n42,Alice")];
631        let batch = decoder.decode_batch(&records).unwrap();
632
633        assert_eq!(batch.num_rows(), 1);
634        assert_eq!(
635            batch
636                .column(0)
637                .as_primitive::<arrow_array::types::Int64Type>()
638                .value(0),
639            42
640        );
641        assert_eq!(batch.column(1).as_string::<i32>().value(0), "Alice");
642    }
643
644    #[test]
645    fn test_decode_multiple_rows() {
646        let schema = make_schema(vec![
647            ("x", DataType::Int64, false),
648            ("y", DataType::Float64, false),
649        ]);
650        let decoder = CsvDecoder::new(schema);
651        let records = vec![csv_block("x,y\n1,1.5\n2,2.5\n3,3.5")];
652        let batch = decoder.decode_batch(&records).unwrap();
653
654        assert_eq!(batch.num_rows(), 3);
655        let x_col = batch
656            .column(0)
657            .as_primitive::<arrow_array::types::Int64Type>();
658        assert_eq!(x_col.value(0), 1);
659        assert_eq!(x_col.value(1), 2);
660        assert_eq!(x_col.value(2), 3);
661    }
662
663    #[test]
664    fn test_decode_all_types() {
665        let schema = make_schema(vec![
666            ("bool_col", DataType::Boolean, false),
667            ("int_col", DataType::Int64, false),
668            ("float_col", DataType::Float64, false),
669            ("str_col", DataType::Utf8, false),
670        ]);
671        let decoder = CsvDecoder::new(schema);
672        let records = vec![csv_block(
673            "bool_col,int_col,float_col,str_col\ntrue,42,3.14,hello",
674        )];
675        let batch = decoder.decode_batch(&records).unwrap();
676
677        assert_eq!(batch.num_rows(), 1);
678        assert!(batch.column(0).as_boolean().value(0));
679        assert_eq!(
680            batch
681                .column(1)
682                .as_primitive::<arrow_array::types::Int64Type>()
683                .value(0),
684            42
685        );
686        let f = batch
687            .column(2)
688            .as_primitive::<arrow_array::types::Float64Type>()
689            .value(0);
690        assert!((f - 3.14).abs() < f64::EPSILON);
691        assert_eq!(batch.column(3).as_string::<i32>().value(0), "hello");
692    }
693
694    // ── Null handling ─────────────────────────────────────────
695
696    #[test]
697    fn test_decode_null_string_default() {
698        // Default null_string is empty string.
699        let schema = make_schema(vec![
700            ("a", DataType::Int64, true),
701            ("b", DataType::Utf8, true),
702        ]);
703        let decoder = CsvDecoder::new(schema);
704        let records = vec![csv_block("a,b\n,")];
705        let batch = decoder.decode_batch(&records).unwrap();
706
707        assert!(batch.column(0).is_null(0));
708        assert!(batch.column(1).is_null(0));
709    }
710
711    #[test]
712    fn test_decode_null_string_custom() {
713        let schema = make_schema(vec![("val", DataType::Int64, true)]);
714        let config = CsvDecoderConfig {
715            null_string: "NA".into(),
716            ..Default::default()
717        };
718        let decoder = CsvDecoder::with_config(schema, config);
719        let records = vec![csv_block("val\nNA\n42")];
720        let batch = decoder.decode_batch(&records).unwrap();
721
722        assert_eq!(batch.num_rows(), 2);
723        assert!(batch.column(0).is_null(0));
724        assert_eq!(
725            batch
726                .column(0)
727                .as_primitive::<arrow_array::types::Int64Type>()
728                .value(1),
729            42
730        );
731    }
732
733    // ── Field count mismatch strategies ───────────────────────
734
735    #[test]
736    fn test_mismatch_null_strategy() {
737        let schema = make_schema(vec![
738            ("a", DataType::Int64, true),
739            ("b", DataType::Utf8, true),
740            ("c", DataType::Int64, true),
741        ]);
742        let decoder = CsvDecoder::new(schema);
743        // Row only has 2 fields, schema expects 3.
744        let records = vec![csv_block("a,b,c\n1,hello")];
745        let batch = decoder.decode_batch(&records).unwrap();
746
747        assert_eq!(batch.num_rows(), 1);
748        assert_eq!(
749            batch
750                .column(0)
751                .as_primitive::<arrow_array::types::Int64Type>()
752                .value(0),
753            1
754        );
755        assert_eq!(batch.column(1).as_string::<i32>().value(0), "hello");
756        assert!(batch.column(2).is_null(0)); // padded with null
757    }
758
759    #[test]
760    fn test_mismatch_skip_strategy() {
761        let schema = make_schema(vec![
762            ("a", DataType::Int64, false),
763            ("b", DataType::Int64, false),
764        ]);
765        let config = CsvDecoderConfig {
766            field_count_mismatch: FieldCountMismatchStrategy::Skip,
767            ..Default::default()
768        };
769        let decoder = CsvDecoder::with_config(schema, config);
770        // One good row, one bad row (too few fields).
771        let records = vec![csv_block("a,b\n1,2\n3")];
772        let batch = decoder.decode_batch(&records).unwrap();
773
774        assert_eq!(batch.num_rows(), 1); // bad row skipped
775        assert_eq!(
776            batch
777                .column(0)
778                .as_primitive::<arrow_array::types::Int64Type>()
779                .value(0),
780            1
781        );
782    }
783
784    #[test]
785    fn test_mismatch_reject_strategy() {
786        let schema = make_schema(vec![
787            ("a", DataType::Int64, false),
788            ("b", DataType::Int64, false),
789        ]);
790        let config = CsvDecoderConfig {
791            field_count_mismatch: FieldCountMismatchStrategy::Reject,
792            ..Default::default()
793        };
794        let decoder = CsvDecoder::with_config(schema, config);
795        let records = vec![csv_block("a,b\n1")]; // too few fields
796        let result = decoder.decode_batch(&records);
797
798        assert!(result.is_err());
799        assert!(result
800            .unwrap_err()
801            .to_string()
802            .contains("field count mismatch"));
803    }
804
805    // ── Delimiter options ─────────────────────────────────────
806
807    #[test]
808    fn test_pipe_delimiter() {
809        let schema = make_schema(vec![
810            ("a", DataType::Int64, false),
811            ("b", DataType::Utf8, false),
812        ]);
813        let config = CsvDecoderConfig {
814            delimiter: b'|',
815            ..Default::default()
816        };
817        let decoder = CsvDecoder::with_config(schema, config);
818        let records = vec![csv_block("a|b\n42|hello")];
819        let batch = decoder.decode_batch(&records).unwrap();
820
821        assert_eq!(
822            batch
823                .column(0)
824                .as_primitive::<arrow_array::types::Int64Type>()
825                .value(0),
826            42
827        );
828        assert_eq!(batch.column(1).as_string::<i32>().value(0), "hello");
829    }
830
831    #[test]
832    fn test_tab_delimiter() {
833        let schema = make_schema(vec![
834            ("a", DataType::Int64, false),
835            ("b", DataType::Utf8, false),
836        ]);
837        let config = CsvDecoderConfig {
838            delimiter: b'\t',
839            ..Default::default()
840        };
841        let decoder = CsvDecoder::with_config(schema, config);
842        let records = vec![csv_block("a\tb\n42\thello")];
843        let batch = decoder.decode_batch(&records).unwrap();
844
845        assert_eq!(
846            batch
847                .column(0)
848                .as_primitive::<arrow_array::types::Int64Type>()
849                .value(0),
850            42
851        );
852        assert_eq!(batch.column(1).as_string::<i32>().value(0), "hello");
853    }
854
855    #[test]
856    fn test_semicolon_delimiter() {
857        let schema = make_schema(vec![
858            ("a", DataType::Int64, false),
859            ("b", DataType::Utf8, false),
860        ]);
861        let config = CsvDecoderConfig {
862            delimiter: b';',
863            ..Default::default()
864        };
865        let decoder = CsvDecoder::with_config(schema, config);
866        let records = vec![csv_block("a;b\n99;world")];
867        let batch = decoder.decode_batch(&records).unwrap();
868
869        assert_eq!(
870            batch
871                .column(0)
872                .as_primitive::<arrow_array::types::Int64Type>()
873                .value(0),
874            99
875        );
876        assert_eq!(batch.column(1).as_string::<i32>().value(0), "world");
877    }
878
879    // ── Comment lines ─────────────────────────────────────────
880
881    #[test]
882    fn test_comment_lines() {
883        let schema = make_schema(vec![("val", DataType::Int64, false)]);
884        let config = CsvDecoderConfig {
885            comment: Some(b'#'),
886            ..Default::default()
887        };
888        let decoder = CsvDecoder::with_config(schema, config);
889        let records = vec![csv_block("val\n# this is a comment\n42\n# another\n99")];
890        let batch = decoder.decode_batch(&records).unwrap();
891
892        assert_eq!(batch.num_rows(), 2);
893        let col = batch
894            .column(0)
895            .as_primitive::<arrow_array::types::Int64Type>();
896        assert_eq!(col.value(0), 42);
897        assert_eq!(col.value(1), 99);
898    }
899
900    // ── Skip rows ─────────────────────────────────────────────
901
902    #[test]
903    fn test_skip_rows() {
904        let schema = make_schema(vec![("val", DataType::Int64, false)]);
905        let config = CsvDecoderConfig {
906            skip_rows: 2,
907            ..Default::default()
908        };
909        let decoder = CsvDecoder::with_config(schema, config);
910        let records = vec![csv_block("val\nskip1\nskip2\n42\n99")];
911        let batch = decoder.decode_batch(&records).unwrap();
912
913        // "skip1" and "skip2" are skipped (parse errors counted), then 42 and 99.
914        // Actually skip_rows skips first N data rows; skip1/skip2 aren't valid i64
915        // so they'd be parse errors. But the skip_rows logic skips before type
916        // coercion, so they won't generate errors.
917        assert_eq!(batch.num_rows(), 2);
918    }
919
920    // ── No header mode ────────────────────────────────────────
921
922    #[test]
923    fn test_no_header() {
924        let schema = make_schema(vec![
925            ("col0", DataType::Int64, false),
926            ("col1", DataType::Utf8, false),
927        ]);
928        let config = CsvDecoderConfig {
929            has_header: false,
930            ..Default::default()
931        };
932        let decoder = CsvDecoder::with_config(schema, config);
933        let records = vec![csv_block("1,alpha\n2,beta")];
934        let batch = decoder.decode_batch(&records).unwrap();
935
936        assert_eq!(batch.num_rows(), 2);
937        let col0 = batch
938            .column(0)
939            .as_primitive::<arrow_array::types::Int64Type>();
940        assert_eq!(col0.value(0), 1);
941        assert_eq!(col0.value(1), 2);
942    }
943
944    // ── Multiple records (streaming) ──────────────────────────
945
946    #[test]
947    fn test_multiple_raw_records() {
948        // Simulate streaming: each RawRecord is one CSV line (no header).
949        let schema = make_schema(vec![
950            ("id", DataType::Int64, false),
951            ("val", DataType::Float64, false),
952        ]);
953        let config = CsvDecoderConfig {
954            has_header: false,
955            ..Default::default()
956        };
957        let decoder = CsvDecoder::with_config(schema, config);
958        let records = vec![
959            csv_record("1,1.5"),
960            csv_record("2,2.5"),
961            csv_record("3,3.5"),
962        ];
963        let batch = decoder.decode_batch(&records).unwrap();
964
965        assert_eq!(batch.num_rows(), 3);
966        let id_col = batch
967            .column(0)
968            .as_primitive::<arrow_array::types::Int64Type>();
969        let val_col = batch
970            .column(1)
971            .as_primitive::<arrow_array::types::Float64Type>();
972        assert_eq!(id_col.value(0), 1);
973        assert_eq!(id_col.value(2), 3);
974        assert!((val_col.value(1) - 2.5).abs() < f64::EPSILON);
975    }
976
977    // ── Quoted fields ─────────────────────────────────────────
978
979    #[test]
980    fn test_quoted_fields_with_delimiter() {
981        let schema = make_schema(vec![
982            ("name", DataType::Utf8, false),
983            ("desc", DataType::Utf8, false),
984        ]);
985        let decoder = CsvDecoder::new(schema);
986        let records = vec![csv_block("name,desc\n\"Smith, John\",\"A, B\"")];
987        let batch = decoder.decode_batch(&records).unwrap();
988
989        assert_eq!(batch.num_rows(), 1);
990        assert_eq!(batch.column(0).as_string::<i32>().value(0), "Smith, John");
991        assert_eq!(batch.column(1).as_string::<i32>().value(0), "A, B");
992    }
993
994    #[test]
995    fn test_quoted_fields_with_newline() {
996        let schema = make_schema(vec![
997            ("id", DataType::Int64, false),
998            ("text", DataType::Utf8, false),
999        ]);
1000        let decoder = CsvDecoder::new(schema);
1001        let records = vec![csv_block("id,text\n1,\"line1\nline2\"")];
1002        let batch = decoder.decode_batch(&records).unwrap();
1003
1004        assert_eq!(batch.num_rows(), 1);
1005        assert_eq!(batch.column(1).as_string::<i32>().value(0), "line1\nline2");
1006    }
1007
1008    #[test]
1009    fn test_escaped_quotes_rfc4180() {
1010        // RFC 4180: doubled quotes within quoted field.
1011        let schema = make_schema(vec![("val", DataType::Utf8, false)]);
1012        let decoder = CsvDecoder::new(schema);
1013        let records = vec![csv_block("val\n\"She said \"\"hello\"\"\"")];
1014        let batch = decoder.decode_batch(&records).unwrap();
1015
1016        assert_eq!(batch.num_rows(), 1);
1017        assert_eq!(
1018            batch.column(0).as_string::<i32>().value(0),
1019            "She said \"hello\""
1020        );
1021    }
1022
1023    // ── Timestamp parsing ─────────────────────────────────────
1024
1025    #[test]
1026    fn test_decode_timestamp() {
1027        let schema = make_schema(vec![(
1028            "ts",
1029            DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())),
1030            false,
1031        )]);
1032        let decoder = CsvDecoder::new(schema);
1033        let records = vec![csv_block("ts\n2025-01-15 10:30:00.000")];
1034        let batch = decoder.decode_batch(&records).unwrap();
1035
1036        assert_eq!(batch.num_rows(), 1);
1037        assert!(!batch.column(0).is_null(0));
1038    }
1039
1040    #[test]
1041    fn test_decode_timestamp_iso8601_fallback() {
1042        let schema = make_schema(vec![(
1043            "ts",
1044            DataType::Timestamp(TimeUnit::Nanosecond, None),
1045            false,
1046        )]);
1047        let decoder = CsvDecoder::new(schema);
1048        let records = vec![csv_block("ts\n2025-01-15T10:30:00Z")];
1049        let batch = decoder.decode_batch(&records).unwrap();
1050
1051        assert_eq!(batch.num_rows(), 1);
1052        assert!(!batch.column(0).is_null(0));
1053    }
1054
1055    // ── Date parsing ──────────────────────────────────────────
1056
1057    #[test]
1058    fn test_decode_date() {
1059        let schema = make_schema(vec![("d", DataType::Date32, false)]);
1060        let decoder = CsvDecoder::new(schema);
1061        let records = vec![csv_block("d\n2025-06-15")];
1062        let batch = decoder.decode_batch(&records).unwrap();
1063
1064        assert_eq!(batch.num_rows(), 1);
1065        assert!(!batch.column(0).is_null(0));
1066        // 2025-06-15 is day 20254 since epoch.
1067        let days = batch
1068            .column(0)
1069            .as_primitive::<arrow_array::types::Date32Type>()
1070            .value(0);
1071        let expected = chrono::NaiveDate::from_ymd_opt(2025, 6, 15)
1072            .unwrap()
1073            .signed_duration_since(chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap())
1074            .num_days();
1075        #[allow(clippy::cast_possible_truncation)]
1076        {
1077            assert_eq!(days, expected as i32);
1078        }
1079    }
1080
1081    // ── Boolean parsing ───────────────────────────────────────
1082
1083    #[test]
1084    fn test_decode_boolean_variants() {
1085        let schema = make_schema(vec![("b", DataType::Boolean, false)]);
1086        let config = CsvDecoderConfig {
1087            has_header: false,
1088            ..Default::default()
1089        };
1090        let decoder = CsvDecoder::with_config(schema, config);
1091        let records = vec![csv_block("true\nfalse\n1\n0\nyes\nno\nt\nf\ny\nn")];
1092        let batch = decoder.decode_batch(&records).unwrap();
1093
1094        assert_eq!(batch.num_rows(), 10);
1095        let col = batch.column(0).as_boolean();
1096        assert!(col.value(0)); // true
1097        assert!(!col.value(1)); // false
1098        assert!(col.value(2)); // 1
1099        assert!(!col.value(3)); // 0
1100        assert!(col.value(4)); // yes
1101        assert!(!col.value(5)); // no
1102        assert!(col.value(6)); // t
1103        assert!(!col.value(7)); // f
1104        assert!(col.value(8)); // y
1105        assert!(!col.value(9)); // n
1106    }
1107
1108    // ── Parse error counting ──────────────────────────────────
1109
1110    #[test]
1111    fn test_parse_error_count() {
1112        let schema = make_schema(vec![("val", DataType::Int64, true)]);
1113        let decoder = CsvDecoder::new(schema);
1114        let records = vec![csv_block("val\nnot_a_number\n42\nalso_bad")];
1115        let batch = decoder.decode_batch(&records).unwrap();
1116
1117        assert_eq!(batch.num_rows(), 3);
1118        assert!(batch.column(0).is_null(0));
1119        assert_eq!(
1120            batch
1121                .column(0)
1122                .as_primitive::<arrow_array::types::Int64Type>()
1123                .value(1),
1124            42
1125        );
1126        assert!(batch.column(0).is_null(2));
1127        assert_eq!(decoder.parse_error_count(), 2);
1128    }
1129
1130    // ── Extra fields ignored ──────────────────────────────────
1131
1132    #[test]
1133    fn test_extra_fields_truncated() {
1134        let schema = make_schema(vec![("a", DataType::Int64, false)]);
1135        let decoder = CsvDecoder::new(schema);
1136        // Row has 3 fields but schema only has 1.
1137        let records = vec![csv_block("a\n42,extra1,extra2")];
1138        let batch = decoder.decode_batch(&records).unwrap();
1139
1140        // Extra fields silently ignored (flexible mode).
1141        // field_count (3) != num_fields (1), but Null strategy just pads/truncates.
1142        assert_eq!(batch.num_rows(), 1);
1143        assert_eq!(
1144            batch
1145                .column(0)
1146                .as_primitive::<arrow_array::types::Int64Type>()
1147                .value(0),
1148            42
1149        );
1150    }
1151
1152    // ── FormatDecoder trait ───────────────────────────────────
1153
1154    #[test]
1155    fn test_format_name() {
1156        let schema = make_schema(vec![("a", DataType::Int64, false)]);
1157        let decoder = CsvDecoder::new(schema);
1158        assert_eq!(decoder.format_name(), "csv");
1159    }
1160
1161    #[test]
1162    fn test_output_schema() {
1163        let schema = make_schema(vec![
1164            ("a", DataType::Int64, false),
1165            ("b", DataType::Utf8, true),
1166        ]);
1167        let decoder = CsvDecoder::new(schema.clone());
1168        assert_eq!(decoder.output_schema(), schema);
1169    }
1170
1171    #[test]
1172    fn test_decode_one() {
1173        let schema = make_schema(vec![("x", DataType::Int64, false)]);
1174        let config = CsvDecoderConfig {
1175            has_header: false,
1176            ..Default::default()
1177        };
1178        let decoder = CsvDecoder::with_config(schema, config);
1179        let record = csv_record("99");
1180        let batch = decoder.decode_one(&record).unwrap();
1181        assert_eq!(batch.num_rows(), 1);
1182        assert_eq!(
1183            batch
1184                .column(0)
1185                .as_primitive::<arrow_array::types::Int64Type>()
1186                .value(0),
1187            99
1188        );
1189    }
1190
1191    // ── Edge cases ────────────────────────────────────────────
1192
1193    #[test]
1194    fn test_mixed_line_endings() {
1195        let schema = make_schema(vec![("val", DataType::Int64, false)]);
1196        let config = CsvDecoderConfig {
1197            has_header: false,
1198            ..Default::default()
1199        };
1200        let decoder = CsvDecoder::with_config(schema, config);
1201        let records = vec![csv_block("1\r\n2\n3\r\n")];
1202        let batch = decoder.decode_batch(&records).unwrap();
1203        assert_eq!(batch.num_rows(), 3);
1204    }
1205
1206    #[test]
1207    fn test_unicode_values() {
1208        let schema = make_schema(vec![("name", DataType::Utf8, false)]);
1209        let decoder = CsvDecoder::new(schema);
1210        let records = vec![csv_block("name\nこんにちは\nüber\nnaïve")];
1211        let batch = decoder.decode_batch(&records).unwrap();
1212
1213        assert_eq!(batch.num_rows(), 3);
1214        assert_eq!(batch.column(0).as_string::<i32>().value(0), "こんにちは");
1215        assert_eq!(batch.column(0).as_string::<i32>().value(1), "über");
1216        assert_eq!(batch.column(0).as_string::<i32>().value(2), "naïve");
1217    }
1218
1219    #[test]
1220    fn test_trailing_comma() {
1221        // Trailing comma creates an extra empty field.
1222        let schema = make_schema(vec![
1223            ("a", DataType::Int64, false),
1224            ("b", DataType::Int64, true),
1225        ]);
1226        let decoder = CsvDecoder::new(schema);
1227        let records = vec![csv_block("a,b\n1,")];
1228        let batch = decoder.decode_batch(&records).unwrap();
1229
1230        assert_eq!(batch.num_rows(), 1);
1231        assert_eq!(
1232            batch
1233                .column(0)
1234                .as_primitive::<arrow_array::types::Int64Type>()
1235                .value(0),
1236            1
1237        );
1238        // Empty string matches default null_string → null.
1239        assert!(batch.column(1).is_null(0));
1240    }
1241
1242    #[test]
1243    fn test_backslash_escape() {
1244        let schema = make_schema(vec![("val", DataType::Utf8, false)]);
1245        let config = CsvDecoderConfig {
1246            escape: Some(b'\\'),
1247            ..Default::default()
1248        };
1249        let decoder = CsvDecoder::with_config(schema, config);
1250        let records = vec![csv_block("val\n\"hello \\\"world\\\"\"")];
1251        let batch = decoder.decode_batch(&records).unwrap();
1252
1253        assert_eq!(batch.num_rows(), 1);
1254        assert_eq!(
1255            batch.column(0).as_string::<i32>().value(0),
1256            "hello \"world\""
1257        );
1258    }
1259
1260    // ── CsvEncoder tests ──────────────────────────────────────
1261
1262    #[test]
1263    fn test_csv_encode_basic() {
1264        let schema = make_schema(vec![
1265            ("id", DataType::Int64, false),
1266            ("name", DataType::Utf8, false),
1267        ]);
1268
1269        let batch = RecordBatch::try_new(
1270            schema.clone(),
1271            vec![
1272                Arc::new(arrow_array::Int64Array::from(vec![1, 2])),
1273                Arc::new(arrow_array::StringArray::from(vec!["Alice", "Bob"])),
1274            ],
1275        )
1276        .unwrap();
1277
1278        let encoder = CsvEncoder::new(schema);
1279        let records = encoder.encode_batch(&batch).unwrap();
1280
1281        assert_eq!(records.len(), 2);
1282        assert_eq!(std::str::from_utf8(&records[0]).unwrap(), "1,Alice");
1283        assert_eq!(std::str::from_utf8(&records[1]).unwrap(), "2,Bob");
1284    }
1285
1286    #[test]
1287    fn test_csv_encode_with_header() {
1288        let schema = make_schema(vec![
1289            ("id", DataType::Int64, false),
1290            ("name", DataType::Utf8, false),
1291        ]);
1292
1293        let batch = RecordBatch::try_new(
1294            schema.clone(),
1295            vec![
1296                Arc::new(arrow_array::Int64Array::from(vec![1])),
1297                Arc::new(arrow_array::StringArray::from(vec!["Alice"])),
1298            ],
1299        )
1300        .unwrap();
1301
1302        let config = CsvEncoderConfig {
1303            has_header: true,
1304            ..Default::default()
1305        };
1306        let encoder = CsvEncoder::with_config(schema, config);
1307        let records = encoder.encode_batch(&batch).unwrap();
1308
1309        assert_eq!(records.len(), 2); // header + 1 data row
1310        assert_eq!(std::str::from_utf8(&records[0]).unwrap(), "id,name");
1311        assert_eq!(std::str::from_utf8(&records[1]).unwrap(), "1,Alice");
1312    }
1313
1314    #[test]
1315    fn test_csv_encode_tab_delimiter() {
1316        let schema = make_schema(vec![
1317            ("a", DataType::Int64, false),
1318            ("b", DataType::Utf8, false),
1319        ]);
1320
1321        let batch = RecordBatch::try_new(
1322            schema.clone(),
1323            vec![
1324                Arc::new(arrow_array::Int64Array::from(vec![42])),
1325                Arc::new(arrow_array::StringArray::from(vec!["hello"])),
1326            ],
1327        )
1328        .unwrap();
1329
1330        let config = CsvEncoderConfig {
1331            delimiter: b'\t',
1332            ..Default::default()
1333        };
1334        let encoder = CsvEncoder::with_config(schema, config);
1335        let records = encoder.encode_batch(&batch).unwrap();
1336
1337        assert_eq!(records.len(), 1);
1338        assert_eq!(std::str::from_utf8(&records[0]).unwrap(), "42\thello");
1339    }
1340
1341    #[test]
1342    fn test_csv_encode_empty_batch() {
1343        let schema = make_schema(vec![("x", DataType::Int64, false)]);
1344        let batch = RecordBatch::new_empty(schema.clone());
1345        let encoder = CsvEncoder::new(schema);
1346        let records = encoder.encode_batch(&batch).unwrap();
1347        assert!(records.is_empty());
1348    }
1349
1350    #[test]
1351    fn test_csv_encode_nulls() {
1352        let schema = make_schema(vec![
1353            ("id", DataType::Int64, false),
1354            ("value", DataType::Int64, true),
1355        ]);
1356
1357        let batch = RecordBatch::try_new(
1358            schema.clone(),
1359            vec![
1360                Arc::new(arrow_array::Int64Array::from(vec![1, 2])),
1361                Arc::new(arrow_array::Int64Array::from(vec![Some(10), None])),
1362            ],
1363        )
1364        .unwrap();
1365
1366        let encoder = CsvEncoder::new(schema);
1367        let records = encoder.encode_batch(&batch).unwrap();
1368
1369        assert_eq!(records.len(), 2);
1370        assert_eq!(std::str::from_utf8(&records[0]).unwrap(), "1,10");
1371        assert_eq!(std::str::from_utf8(&records[1]).unwrap(), "2,");
1372    }
1373}
laminar_connectors/schema/csv.rs

laminar_connectors/schema/
csv.rs
