Skip to main content

arrow_csv/reader/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! CSV Reading: [`Reader`] and [`ReaderBuilder`]
6//!
7//! # Basic Usage
8//!
9//! This CSV reader allows CSV files to be read into the Arrow memory model. Records are
10//! loaded in batches and are then converted from row-based data to columnar data.
11//!
12//! Example:
13//!
14//! ```
15//! # use arrow_schema::*;
16//! # use arrow_csv::{Reader, ReaderBuilder};
17//! # use std::fs::File;
18//! # use std::sync::Arc;
19//!
20//! let schema = Schema::new(vec![
21//!     Field::new("city", DataType::Utf8, false),
22//!     Field::new("lat", DataType::Float64, false),
23//!     Field::new("lng", DataType::Float64, false),
24//! ]);
25//!
26//! let file = File::open("test/data/uk_cities.csv").unwrap();
27//!
28//! let mut csv = ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
29//! let batch = csv.next().unwrap().unwrap();
30//! ```
31//!
32//! # Timestamps
33//!
34//! Timestamps are parsed with [`arrow_cast::parse::string_to_datetime`]. A
35//! timestamp without an explicit offset that falls in a daylight savings gap or
36//! overlap of the target timezone is resolved as described there, rather than
37//! returning an error.
38//!
39//! # Example: Numeric calculations on CSV
40//! This code finds the maximum value in column 0 of a CSV file containing
41//! ```csv
42//! c1,c2,c3,c4
43//! 1,1.1,"hong kong",true
44//! 3,323.12,"XiAn",false
45//! 10,131323.12,"cheng du",false
46//! ```
47//!
48//! ```
49//! # use arrow_array::cast::AsArray;
50//! # use arrow_array::types::Int16Type;
51//! # use arrow_csv::ReaderBuilder;
52//! # use arrow_schema::{DataType, Field, Schema};
53//! # use std::fs::File;
54//! # use std::sync::Arc;
55//! // Open the example file
56//! let file = File::open("test/data/example.csv").unwrap();
57//! let csv_schema = Schema::new(vec![
58//!     Field::new("c1", DataType::Int16, true),
59//!     Field::new("c2", DataType::Float32, true),
60//!     Field::new("c3", DataType::Utf8, true),
61//!     Field::new("c4", DataType::Boolean, true),
62//! ]);
63//! let mut reader = ReaderBuilder::new(Arc::new(csv_schema))
64//!     .with_header(true)
65//!     .build(file)
66//!     .unwrap();
67//! // find the maximum value in column 0 across all batches
68//! let mut max_c0 = 0;
69//! while let Some(r) = reader.next() {
70//!   let r = r.unwrap(); // handle error
71//!   // get the max value in column(0) for this batch
72//!   let col = r.column(0).as_primitive::<Int16Type>();
73//!   let batch_max = col.iter().max().flatten().unwrap_or_default();
74//!   max_c0 = max_c0.max(batch_max);
75//! }
76//! assert_eq!(max_c0, 10);
77//!```
78//!
79//! # Async Usage
80//!
81//! The lower-level [`Decoder`] can be integrated with various forms of async data streams,
82//! and is designed to be agnostic to the various different kinds of async IO primitives found
83//! within the Rust ecosystem.
84//!
85//! For example, see below for how it can be used with an arbitrary `Stream` of `Bytes`
86//!
87//! ```
88//! # use std::task::{Poll, ready};
89//! # use bytes::{Buf, Bytes};
90//! # use arrow_schema::ArrowError;
91//! # use futures::stream::{Stream, StreamExt};
92//! # use arrow_array::RecordBatch;
93//! # use arrow_csv::reader::Decoder;
94//! #
95//! fn decode_stream<S: Stream<Item = Bytes> + Unpin>(
96//!     mut decoder: Decoder,
97//!     mut input: S,
98//! ) -> impl Stream<Item = Result<RecordBatch, ArrowError>> {
99//!     let mut buffered = Bytes::new();
100//!     futures::stream::poll_fn(move |cx| {
101//!         loop {
102//!             if buffered.is_empty() {
103//!                 if let Some(b) = ready!(input.poll_next_unpin(cx)) {
104//!                     buffered = b;
105//!                 }
106//!                 // Note: don't break on `None` as the decoder needs
107//!                 // to be called with an empty array to delimit the
108//!                 // final record
109//!             }
110//!             let decoded = match decoder.decode(buffered.as_ref()) {
111//!                 Ok(0) => break,
112//!                 Ok(decoded) => decoded,
113//!                 Err(e) => return Poll::Ready(Some(Err(e))),
114//!             };
115//!             buffered.advance(decoded);
116//!         }
117//!
118//!         Poll::Ready(decoder.flush().transpose())
119//!     })
120//! }
121//!
122//! ```
123//!
124//! In a similar vein, it can also be used with tokio-based IO primitives
125//!
126//! ```
127//! # use std::pin::Pin;
128//! # use std::task::{Poll, ready};
129//! # use futures::Stream;
130//! # use tokio::io::AsyncBufRead;
131//! # use arrow_array::RecordBatch;
132//! # use arrow_csv::reader::Decoder;
133//! # use arrow_schema::ArrowError;
134//! fn decode_stream<R: AsyncBufRead + Unpin>(
135//!     mut decoder: Decoder,
136//!     mut reader: R,
137//! ) -> impl Stream<Item = Result<RecordBatch, ArrowError>> {
138//!     futures::stream::poll_fn(move |cx| {
139//!         loop {
140//!             let b = match ready!(Pin::new(&mut reader).poll_fill_buf(cx)) {
141//!                 Ok(b) => b,
142//!                 Err(e) => return Poll::Ready(Some(Err(e.into()))),
143//!             };
144//!             let decoded = match decoder.decode(b) {
145//!                 // Note: the decoder needs to be called with an empty
146//!                 // array to delimit the final record
147//!                 Ok(0) => break,
148//!                 Ok(decoded) => decoded,
149//!                 Err(e) => return Poll::Ready(Some(Err(e))),
150//!             };
151//!             Pin::new(&mut reader).consume(decoded);
152//!         }
153//!
154//!         Poll::Ready(decoder.flush().transpose())
155//!     })
156//! }
157//! ```
158//!
159
160mod records;
161
162use arrow_array::builder::{NullBuilder, PrimitiveBuilder, StringDictionaryBuilder};
163use arrow_array::types::*;
164use arrow_array::*;
165use arrow_cast::parse::{Parser, parse_decimal, string_to_datetime};
166use arrow_schema::*;
167use chrono::{TimeZone, Utc};
168use csv::StringRecord;
169use regex::{Regex, RegexSet};
170use std::fmt::{self, Debug};
171use std::fs::File;
172use std::io::{BufRead, BufReader as StdBufReader, Read};
173use std::sync::{Arc, LazyLock};
174
175use crate::map_csv_error;
176use crate::reader::records::{RecordDecoder, StringRecords};
177use arrow_array::timezone::Tz;
178
179/// Order should match [`InferredDataType`]
180static REGEX_SET: LazyLock<RegexSet> = LazyLock::new(|| {
181    RegexSet::new([
182        r"(?i)^(true)$|^(false)$(?-i)", //BOOLEAN
183        r"^[+-]?(\d+)$",                //INTEGER
184        r"^[+-]?((\d*\.\d+|\d+\.\d*)([eE][-+]?\d+)?|\d+([eE][-+]?\d+))$", //DECIMAL
185        r"^\d{4}-\d\d-\d\d$",           //DATE32
186        r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d(?:[^\d\.].*)?$", //Timestamp(Second)
187        r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,3}(?:[^\d].*)?$", //Timestamp(Millisecond)
188        r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,6}(?:[^\d].*)?$", //Timestamp(Microsecond)
189        r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,9}(?:[^\d].*)?$", //Timestamp(Nanosecond)
190    ])
191    .unwrap()
192});
193
194/// A wrapper over `Option<Regex>` to check if the value is `NULL`.
195#[derive(Debug, Clone, Default)]
196struct NullRegex(Option<Regex>);
197
198impl NullRegex {
199    /// Returns true if the value should be considered as `NULL` according to
200    /// the provided regular expression.
201    #[inline]
202    fn is_null(&self, s: &str) -> bool {
203        match &self.0 {
204            Some(r) => r.is_match(s),
205            None => s.is_empty(),
206        }
207    }
208}
209
210#[derive(Default, Copy, Clone)]
211struct InferredDataType {
212    /// Packed booleans indicating type
213    ///
214    /// 0 - Boolean
215    /// 1 - Integer
216    /// 2 - Float64
217    /// 3 - Date32
218    /// 4 - Timestamp(Second)
219    /// 5 - Timestamp(Millisecond)
220    /// 6 - Timestamp(Microsecond)
221    /// 7 - Timestamp(Nanosecond)
222    /// 8 - Utf8
223    packed: u16,
224}
225
226impl InferredDataType {
227    /// Returns the inferred data type
228    fn get(&self) -> DataType {
229        match self.packed {
230            0 => DataType::Null,
231            1 => DataType::Boolean,
232            2 => DataType::Int64,
233            4 | 6 => DataType::Float64, // Promote Int64 to Float64
234            b if b != 0 && (b & !0b11111000) == 0 => match b.leading_zeros() {
235                // Promote to highest precision temporal type
236                8 => DataType::Timestamp(TimeUnit::Nanosecond, None),
237                9 => DataType::Timestamp(TimeUnit::Microsecond, None),
238                10 => DataType::Timestamp(TimeUnit::Millisecond, None),
239                11 => DataType::Timestamp(TimeUnit::Second, None),
240                12 => DataType::Date32,
241                _ => unreachable!(),
242            },
243            _ => DataType::Utf8,
244        }
245    }
246
247    /// Updates the [`InferredDataType`] with the given string
248    fn update(&mut self, string: &str) {
249        self.packed |= if string.starts_with('"') {
250            1 << 8 // Utf8
251        } else if let Some(m) = REGEX_SET.matches(string).into_iter().next() {
252            if m == 1 && string.len() >= 19 && string.parse::<i64>().is_err() {
253                // if overflow i64, fallback to utf8
254                1 << 8
255            } else {
256                1 << m
257            }
258        } else if string == "NaN" || string == "nan" || string == "inf" || string == "-inf" {
259            1 << 2 // Float64
260        } else {
261            1 << 8 // Utf8
262        }
263    }
264}
265
266/// The format specification for the CSV file
267#[derive(Debug, Clone, Default)]
268pub struct Format {
269    header: bool,
270    header_validation: bool,
271    delimiter: Option<u8>,
272    escape: Option<u8>,
273    quote: Option<u8>,
274    terminator: Option<u8>,
275    comment: Option<u8>,
276    null_regex: NullRegex,
277    truncated_rows: bool,
278}
279
280impl Format {
281    /// Specify whether the CSV file has a header, defaults to `false`
282    ///
283    /// When `true`, the first row of the CSV file is treated as a header row
284    pub fn with_header(mut self, has_header: bool) -> Self {
285        self.header = has_header;
286        self
287    }
288
289    /// Specify whether to validate the CSV header against the schema, defaults to `false`
290    ///
291    /// When `true`, the first row gets validated against the schema before any data is read
292    ///
293    /// Only applies when [`Self::with_header`] is set to `true`
294    pub fn with_header_validation(mut self, validate_header: bool) -> Self {
295        self.header_validation = validate_header;
296        self
297    }
298
299    /// Specify a custom delimiter character, defaults to comma `','`
300    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
301        self.delimiter = Some(delimiter);
302        self
303    }
304
305    /// Specify an escape character, defaults to `None`
306    pub fn with_escape(mut self, escape: u8) -> Self {
307        self.escape = Some(escape);
308        self
309    }
310
311    /// Specify a custom quote character, defaults to double quote `'"'`
312    pub fn with_quote(mut self, quote: u8) -> Self {
313        self.quote = Some(quote);
314        self
315    }
316
317    /// Specify a custom terminator character, defaults to CRLF
318    pub fn with_terminator(mut self, terminator: u8) -> Self {
319        self.terminator = Some(terminator);
320        self
321    }
322
323    /// Specify a comment character, defaults to `None`
324    ///
325    /// Lines starting with this character will be ignored
326    pub fn with_comment(mut self, comment: u8) -> Self {
327        self.comment = Some(comment);
328        self
329    }
330
331    /// Provide a regex to match null values, defaults to `^$`
332    pub fn with_null_regex(mut self, null_regex: Regex) -> Self {
333        self.null_regex = NullRegex(Some(null_regex));
334        self
335    }
336
337    /// Whether to allow truncated rows when parsing.
338    ///
339    /// By default this is set to `false` and will error if the CSV rows have different lengths.
340    /// When set to true then it will allow records with less than the expected number of columns
341    /// and fill the missing columns with nulls. If the record's schema is not nullable, then it
342    /// will still return an error.
343    pub fn with_truncated_rows(mut self, allow: bool) -> Self {
344        self.truncated_rows = allow;
345        self
346    }
347
348    /// Infer format settings from the CSV records in `reader`
349    ///
350    /// This currently infers whether the first record is a header. Up to
351    /// `max_records` records after the first record are inspected; if `None`, all
352    /// records are read. Detection is conservative and returns no header when the
353    /// sampled records do not provide type evidence. Returns the updated format
354    /// and the number of records read, including the first header candidate.
355    ///
356    /// # Example
357    ///
358    /// ```
359    /// use arrow_csv::reader::Format;
360    /// use std::io::Cursor;
361    ///
362    /// let csv = "name,count\nalice,1\nbob,2\n";
363    /// let (format, format_records_read) =
364    ///     Format::default().infer_format(Cursor::new(csv), Some(10))?;
365    /// let (schema, records_read) = format.infer_schema(Cursor::new(csv), None)?;
366    ///
367    /// assert_eq!(schema.field(0).name(), "name");
368    /// assert_eq!(format_records_read, 3);
369    /// assert_eq!(records_read, 2);
370    /// # Ok::<_, arrow_schema::ArrowError>(())
371    /// ```
372    pub fn infer_format<R: Read>(
373        mut self,
374        reader: R,
375        max_records: Option<usize>,
376    ) -> Result<(Self, usize), ArrowError> {
377        let (header, records_read) = self.infer_header(reader, max_records)?;
378        self.header = header;
379        Ok((self, records_read))
380    }
381
382    /// Infer whether the first CSV record is a header
383    ///
384    /// Inspects up to `max_records` records after the first record. Returns `true`
385    /// when a value in the first record is text while the remaining values in the
386    /// same column have a consistent non-text type.
387    fn infer_header<R: Read>(
388        &self,
389        reader: R,
390        max_records: Option<usize>,
391    ) -> Result<(bool, usize), ArrowError> {
392        let mut format = self.clone();
393        format.header = false;
394        let mut csv_reader = format.build_reader(reader);
395
396        let mut first_record = StringRecord::new();
397        if !csv_reader
398            .read_record(&mut first_record)
399            .map_err(map_csv_error)?
400        {
401            return Ok((false, 0));
402        }
403
404        let mut first_types = vec![InferredDataType::default(); first_record.len()];
405        for (value, inferred) in first_record.iter().zip(&mut first_types) {
406            if !self.null_regex.is_null(value) {
407                inferred.update(value);
408            }
409        }
410
411        let mut column_types = vec![InferredDataType::default(); first_record.len()];
412        let mut record = StringRecord::new();
413        let mut records_count = 0;
414        let max_records = max_records.unwrap_or(usize::MAX);
415        while records_count < max_records
416            && csv_reader.read_record(&mut record).map_err(map_csv_error)?
417        {
418            records_count += 1;
419            for (value, inferred) in record.iter().zip(&mut column_types) {
420                if !self.null_regex.is_null(value) {
421                    inferred.update(value);
422                }
423            }
424        }
425
426        let has_header = first_types
427            .iter()
428            .zip(&column_types)
429            .zip(first_record.iter())
430            .any(|((first, rest), value)| {
431                // Numeric-looking values (e.g. +1) are not header evidence, even
432                // when ordinary schema inference conservatively treats them as text.
433                first.get() == DataType::Utf8
434                    && value.parse::<f64>().is_err()
435                    && !matches!(rest.get(), DataType::Utf8 | DataType::Null)
436            });
437        Ok((has_header, records_count + 1))
438    }
439
440    /// Infer schema of CSV records from the provided `reader`
441    ///
442    /// If `max_records` is `None`, all records will be read, otherwise up to `max_records`
443    /// records are read to infer the schema
444    ///
445    /// Returns inferred schema and number of records read
446    pub fn infer_schema<R: Read>(
447        &self,
448        reader: R,
449        max_records: Option<usize>,
450    ) -> Result<(Schema, usize), ArrowError> {
451        let mut csv_reader = self.build_reader(reader);
452
453        // get or create header names
454        // when has_header is false, creates default column names with column_ prefix
455        let headers: Vec<String> = if self.header {
456            let headers = &csv_reader.headers().map_err(map_csv_error)?.clone();
457            headers.iter().map(|s| s.to_string()).collect()
458        } else {
459            let first_record_count = &csv_reader.headers().map_err(map_csv_error)?.len();
460            (0..*first_record_count)
461                .map(|i| format!("column_{}", i + 1))
462                .collect()
463        };
464
465        let header_length = headers.len();
466        // keep track of inferred field types
467        let mut column_types: Vec<InferredDataType> = vec![Default::default(); header_length];
468
469        let mut records_count = 0;
470
471        let mut record = StringRecord::new();
472        let max_records = max_records.unwrap_or(usize::MAX);
473        while records_count < max_records {
474            if !csv_reader.read_record(&mut record).map_err(map_csv_error)? {
475                break;
476            }
477            records_count += 1;
478
479            // Note since we may be looking at a sample of the data, we make the safe assumption that
480            // they could be nullable
481            for (i, column_type) in column_types.iter_mut().enumerate().take(header_length) {
482                if let Some(string) = record.get(i)
483                    && !self.null_regex.is_null(string)
484                {
485                    column_type.update(string)
486                }
487            }
488        }
489
490        // build schema from inference results
491        let fields: Fields = column_types
492            .iter()
493            .zip(&headers)
494            .map(|(inferred, field_name)| Field::new(field_name, inferred.get(), true))
495            .collect();
496
497        Ok((Schema::new(fields), records_count))
498    }
499
500    /// Build a [`csv::Reader`] for this [`Format`]
501    fn build_reader<R: Read>(&self, reader: R) -> csv::Reader<R> {
502        let mut builder = csv::ReaderBuilder::new();
503        builder.has_headers(self.header);
504        builder.flexible(self.truncated_rows);
505
506        if let Some(c) = self.delimiter {
507            builder.delimiter(c);
508        }
509        builder.escape(self.escape);
510        if let Some(c) = self.quote {
511            builder.quote(c);
512        }
513        if let Some(t) = self.terminator {
514            builder.terminator(csv::Terminator::Any(t));
515        }
516        if let Some(comment) = self.comment {
517            builder.comment(Some(comment));
518        }
519        builder.from_reader(reader)
520    }
521
522    /// Build a [`csv_core::Reader`] for this [`Format`]
523    fn build_parser(&self) -> csv_core::Reader {
524        let mut builder = csv_core::ReaderBuilder::new();
525        builder.escape(self.escape);
526        builder.comment(self.comment);
527
528        if let Some(c) = self.delimiter {
529            builder.delimiter(c);
530        }
531        if let Some(c) = self.quote {
532            builder.quote(c);
533        }
534        if let Some(t) = self.terminator {
535            builder.terminator(csv_core::Terminator::Any(t));
536        }
537        builder.build()
538    }
539}
540
541/// Infer schema from a list of CSV files by reading through first n records
542/// with `max_read_records` controlling the maximum number of records to read.
543///
544/// Files will be read in the given order until n records have been reached.
545///
546/// If `max_read_records` is not set, all files will be read fully to infer the schema.
547pub fn infer_schema_from_files(
548    files: &[String],
549    delimiter: u8,
550    max_read_records: Option<usize>,
551    has_header: bool,
552) -> Result<Schema, ArrowError> {
553    let mut schemas = vec![];
554    let mut records_to_read = max_read_records.unwrap_or(usize::MAX);
555    let format = Format {
556        delimiter: Some(delimiter),
557        header: has_header,
558        ..Default::default()
559    };
560
561    for fname in files {
562        let f = File::open(fname)?;
563        let (schema, records_read) = format.infer_schema(f, Some(records_to_read))?;
564        if records_read == 0 {
565            continue;
566        }
567        schemas.push(schema.clone());
568        records_to_read -= records_read;
569        if records_to_read == 0 {
570            break;
571        }
572    }
573
574    Schema::try_merge(schemas)
575}
576
577// optional bounds of the reader, of the form (min line, max line).
578type Bounds = Option<(usize, usize)>;
579
580/// CSV file reader using [`std::io::BufReader`]
581///
582/// See [`ReaderBuilder`] to construct a CSV reader with options and  the
583/// [module-level documentation](crate::reader) for more details and examples
584pub type Reader<R> = BufReader<StdBufReader<R>>;
585
586/// CSV file reader implementation. See [`Reader`] for usage
587///
588/// Despite having the same name as [`std::io::BufReader`, this structure does
589/// not buffer reads itself
590pub struct BufReader<R> {
591    /// File reader
592    reader: R,
593    /// The decoder
594    decoder: Decoder,
595    /// Schema of the record batches produced by this reader
596    schema: SchemaRef,
597}
598
599impl<R> fmt::Debug for BufReader<R>
600where
601    R: BufRead,
602{
603    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
604        f.debug_struct("Reader")
605            .field("decoder", &self.decoder)
606            .finish()
607    }
608}
609
610impl<R> BufReader<R> {
611    /// The number of rows padded because they had fewer fields than the schema
612    ///
613    /// Always 0 unless [`ReaderBuilder::with_truncated_rows`] was set to `true`.
614    ///
615    /// The count is cumulative over the lifetime of this reader, so reading it
616    /// between batches yields a running total of the rows read so far, and reading it
617    /// once the reader is exhausted yields the total for the whole input. Rows that
618    /// are skipped rather than read into a batch, such as a header row or rows before
619    /// the start bound, do not contribute.
620    ///
621    /// A padded row is indistinguishable from a row with genuinely empty trailing
622    /// fields once it has been read, so this counter is the only way to tell the two
623    /// apart.
624    ///
625    /// ```
626    /// # use std::io::Cursor;
627    /// # use std::sync::Arc;
628    /// # use arrow_csv::ReaderBuilder;
629    /// # use arrow_schema::{DataType, Field, Schema};
630    /// #
631    /// let schema = Arc::new(Schema::new(vec![
632    ///     Field::new("a", DataType::Int32, true),
633    ///     Field::new("b", DataType::Int32, true),
634    /// ]));
635    ///
636    /// let mut reader = ReaderBuilder::new(schema)
637    ///     .with_truncated_rows(true)
638    ///     .build(Cursor::new("1,2\n3\n"))
639    ///     .unwrap();
640    ///
641    /// let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
642    /// assert_eq!(batches[0].num_rows(), 2);
643    /// assert_eq!(reader.truncated_row_count(), 1);
644    /// ```
645    pub fn truncated_row_count(&self) -> usize {
646        self.decoder.truncated_row_count()
647    }
648}
649
650impl<R: Read> Reader<R> {
651    /// Returns the schema of the reader, useful for getting the schema without reading
652    /// record batches
653    pub fn schema(&self) -> SchemaRef {
654        self.schema.clone()
655    }
656}
657
658impl<R: BufRead> BufReader<R> {
659    fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
660        loop {
661            let buf = self.reader.fill_buf()?;
662            let decoded = self.decoder.decode(buf)?;
663            self.reader.consume(decoded);
664            // Yield if decoded no bytes or the decoder is full
665            //
666            // The capacity check avoids looping around and potentially
667            // blocking reading data in fill_buf that isn't needed
668            // to flush the next batch
669            if decoded == 0 || self.decoder.capacity() == 0 {
670                break;
671            }
672        }
673
674        self.decoder.flush()
675    }
676}
677
678impl<R: BufRead> Iterator for BufReader<R> {
679    type Item = Result<RecordBatch, ArrowError>;
680
681    fn next(&mut self) -> Option<Self::Item> {
682        self.read().transpose()
683    }
684}
685
686impl<R: BufRead> RecordBatchReader for BufReader<R> {
687    fn schema(&self) -> SchemaRef {
688        self.schema.clone()
689    }
690}
691
692/// A push-based interface for decoding CSV data from an arbitrary byte stream
693///
694/// See [`Reader`] for a higher-level interface for interface with [`Read`]
695///
696/// The push-based interface facilitates integration with sources that yield arbitrarily
697/// delimited bytes ranges, such as [`BufRead`], or a chunked byte stream received from
698/// object storage
699///
700/// ```
701/// # use std::io::BufRead;
702/// # use arrow_array::RecordBatch;
703/// # use arrow_csv::ReaderBuilder;
704/// # use arrow_schema::{ArrowError, SchemaRef};
705/// #
706/// fn read_from_csv<R: BufRead>(
707///     mut reader: R,
708///     schema: SchemaRef,
709///     batch_size: usize,
710/// ) -> Result<impl Iterator<Item = Result<RecordBatch, ArrowError>>, ArrowError> {
711///     let mut decoder = ReaderBuilder::new(schema)
712///         .with_batch_size(batch_size)
713///         .build_decoder();
714///
715///     let mut next = move || {
716///         loop {
717///             let buf = reader.fill_buf()?;
718///             let decoded = decoder.decode(buf)?;
719///             if decoded == 0 {
720///                 break;
721///             }
722///
723///             // Consume the number of bytes read
724///             reader.consume(decoded);
725///         }
726///         decoder.flush()
727///     };
728///     Ok(std::iter::from_fn(move || next().transpose()))
729/// }
730/// ```
731#[derive(Debug)]
732pub struct Decoder {
733    /// Explicit schema for the CSV file
734    schema: SchemaRef,
735
736    /// Optional projection for which columns to load (zero-based column indices)
737    projection: Option<Vec<usize>>,
738
739    /// Number of records per batch
740    batch_size: usize,
741
742    /// Rows to skip
743    to_skip: usize,
744
745    /// Whether to validate the first skipped row against the schema
746    header_validation: bool,
747
748    /// Current line number
749    line_number: usize,
750
751    /// End line number
752    end: usize,
753
754    /// A decoder for [`StringRecords`]
755    record_decoder: RecordDecoder,
756
757    /// Check if the string matches this pattern for `NULL`.
758    null_regex: NullRegex,
759}
760
761impl Decoder {
762    /// Decode records from `buf` returning the number of bytes read
763    ///
764    /// This method returns once `batch_size` objects have been parsed since the
765    /// last call to [`Self::flush`], or `buf` is exhausted. Any remaining bytes
766    /// should be included in the next call to [`Self::decode`]
767    ///
768    /// There is no requirement that `buf` contains a whole number of records, facilitating
769    /// integration with arbitrary byte streams, such as that yielded by [`BufRead`] or
770    /// network sources such as object storage
771    pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
772        if self.to_skip != 0 {
773            if self.header_validation {
774                let (skipped, bytes) = self.record_decoder.decode(buf, 1)?;
775
776                if skipped == 0 {
777                    return Ok(bytes);
778                }
779
780                let rows = self.record_decoder.flush()?;
781                validate_header(&rows, self.schema.fields())?;
782                self.header_validation = false;
783                self.to_skip -= 1;
784                return Ok(bytes);
785            }
786
787            // Skip in units of `to_read` to avoid over-allocating buffers
788            let to_skip = self.to_skip.min(self.batch_size);
789            let (skipped, bytes) = self.record_decoder.decode(buf, to_skip)?;
790            self.to_skip -= skipped;
791            self.record_decoder.clear();
792            return Ok(bytes);
793        }
794
795        let to_read = self.batch_size.min(self.end - self.line_number) - self.record_decoder.len();
796        let (_, bytes) = self.record_decoder.decode(buf, to_read)?;
797        Ok(bytes)
798    }
799
800    /// Flushes the currently buffered data to a [`RecordBatch`]
801    ///
802    /// This should only be called after [`Self::decode`] has returned `Ok(0)`,
803    /// otherwise may return an error if part way through decoding a record
804    ///
805    /// Returns `Ok(None)` if no buffered data
806    pub fn flush(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
807        if self.record_decoder.is_empty() {
808            return Ok(None);
809        }
810
811        let rows = self.record_decoder.flush()?;
812        let batch = parse(
813            &rows,
814            &self.schema,
815            self.projection.as_ref(),
816            self.line_number,
817            &self.null_regex,
818        )?;
819        self.line_number += rows.len();
820        Ok(Some(batch))
821    }
822
823    /// Returns the number of records that can be read before requiring a call to [`Self::flush`]
824    pub fn capacity(&self) -> usize {
825        self.batch_size - self.record_decoder.len()
826    }
827
828    /// The number of rows padded because they had fewer fields than the schema
829    ///
830    /// Always 0 unless [`ReaderBuilder::with_truncated_rows`] was set to `true`.
831    ///
832    /// The count is cumulative over the lifetime of this decoder and is not reset by
833    /// [`Self::flush`], so reading it between batches yields a running total of the
834    /// rows decoded so far, and reading it once the input is exhausted yields the
835    /// total for the whole stream. Rows that are skipped rather than decoded into a
836    /// batch, such as a header row or rows before the start bound, do not contribute.
837    ///
838    /// A padded row is indistinguishable from a row with genuinely empty trailing
839    /// fields once it has been decoded, so this counter is the only way to tell the
840    /// two apart.
841    pub fn truncated_row_count(&self) -> usize {
842        self.record_decoder.truncated_row_count()
843    }
844}
845
846fn validate_header(rows: &StringRecords<'_>, fields: &Fields) -> Result<(), ArrowError> {
847    let header = rows.iter().next().ok_or_else(|| {
848        ArrowError::CsvError("CSV header validation failed: no header row found".to_string())
849    })?;
850
851    for (idx, field) in fields.iter().enumerate() {
852        let actual = header.get(idx);
853        let expected = field.name();
854        if actual != expected {
855            return Err(ArrowError::CsvError(format!(
856                "CSV header does not match schema at column {idx}: expected {expected:?} but found {actual:?}"
857            )));
858        }
859    }
860
861    Ok(())
862}
863
864/// Parses a slice of [`StringRecords`] into a [RecordBatch]
865fn parse(
866    rows: &StringRecords<'_>,
867    schema: &Schema,
868    projection: Option<&Vec<usize>>,
869    line_number: usize,
870    null_regex: &NullRegex,
871) -> Result<RecordBatch, ArrowError> {
872    let fields = schema.fields();
873    let projection: Vec<usize> = match projection {
874        Some(v) => v.clone(),
875        None => fields.iter().enumerate().map(|(i, _)| i).collect(),
876    };
877    let projected_schema = Arc::new(schema.project(&projection)?);
878
879    let arrays: Result<Vec<ArrayRef>, _> = projection
880        .iter()
881        .map(|i| {
882            let i = *i;
883            let field = &fields[i];
884            match field.data_type() {
885                DataType::Boolean => build_boolean_array(line_number, rows, i, null_regex),
886                DataType::Decimal32(precision, scale) => build_decimal_array::<Decimal32Type>(
887                    line_number,
888                    rows,
889                    i,
890                    *precision,
891                    *scale,
892                    null_regex,
893                ),
894                DataType::Decimal64(precision, scale) => build_decimal_array::<Decimal64Type>(
895                    line_number,
896                    rows,
897                    i,
898                    *precision,
899                    *scale,
900                    null_regex,
901                ),
902                DataType::Decimal128(precision, scale) => build_decimal_array::<Decimal128Type>(
903                    line_number,
904                    rows,
905                    i,
906                    *precision,
907                    *scale,
908                    null_regex,
909                ),
910                DataType::Decimal256(precision, scale) => build_decimal_array::<Decimal256Type>(
911                    line_number,
912                    rows,
913                    i,
914                    *precision,
915                    *scale,
916                    null_regex,
917                ),
918                DataType::Int8 => {
919                    build_primitive_array::<Int8Type>(line_number, rows, i, null_regex)
920                }
921                DataType::Int16 => {
922                    build_primitive_array::<Int16Type>(line_number, rows, i, null_regex)
923                }
924                DataType::Int32 => {
925                    build_primitive_array::<Int32Type>(line_number, rows, i, null_regex)
926                }
927                DataType::Int64 => {
928                    build_primitive_array::<Int64Type>(line_number, rows, i, null_regex)
929                }
930                DataType::UInt8 => {
931                    build_primitive_array::<UInt8Type>(line_number, rows, i, null_regex)
932                }
933                DataType::UInt16 => {
934                    build_primitive_array::<UInt16Type>(line_number, rows, i, null_regex)
935                }
936                DataType::UInt32 => {
937                    build_primitive_array::<UInt32Type>(line_number, rows, i, null_regex)
938                }
939                DataType::UInt64 => {
940                    build_primitive_array::<UInt64Type>(line_number, rows, i, null_regex)
941                }
942                DataType::Float16 => {
943                    build_primitive_array::<Float16Type>(line_number, rows, i, null_regex)
944                }
945                DataType::Float32 => {
946                    build_primitive_array::<Float32Type>(line_number, rows, i, null_regex)
947                }
948                DataType::Float64 => {
949                    build_primitive_array::<Float64Type>(line_number, rows, i, null_regex)
950                }
951                DataType::Date32 => {
952                    build_primitive_array::<Date32Type>(line_number, rows, i, null_regex)
953                }
954                DataType::Date64 => {
955                    build_primitive_array::<Date64Type>(line_number, rows, i, null_regex)
956                }
957                DataType::Time32(TimeUnit::Second) => {
958                    build_primitive_array::<Time32SecondType>(line_number, rows, i, null_regex)
959                }
960                DataType::Time32(TimeUnit::Millisecond) => {
961                    build_primitive_array::<Time32MillisecondType>(line_number, rows, i, null_regex)
962                }
963                DataType::Time64(TimeUnit::Microsecond) => {
964                    build_primitive_array::<Time64MicrosecondType>(line_number, rows, i, null_regex)
965                }
966                DataType::Time64(TimeUnit::Nanosecond) => {
967                    build_primitive_array::<Time64NanosecondType>(line_number, rows, i, null_regex)
968                }
969                DataType::Timestamp(TimeUnit::Second, tz) => {
970                    build_timestamp_array::<TimestampSecondType>(
971                        line_number,
972                        rows,
973                        i,
974                        tz.as_deref(),
975                        null_regex,
976                    )
977                }
978                DataType::Timestamp(TimeUnit::Millisecond, tz) => {
979                    build_timestamp_array::<TimestampMillisecondType>(
980                        line_number,
981                        rows,
982                        i,
983                        tz.as_deref(),
984                        null_regex,
985                    )
986                }
987                DataType::Timestamp(TimeUnit::Microsecond, tz) => {
988                    build_timestamp_array::<TimestampMicrosecondType>(
989                        line_number,
990                        rows,
991                        i,
992                        tz.as_deref(),
993                        null_regex,
994                    )
995                }
996                DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
997                    build_timestamp_array::<TimestampNanosecondType>(
998                        line_number,
999                        rows,
1000                        i,
1001                        tz.as_deref(),
1002                        null_regex,
1003                    )
1004                }
1005                DataType::Null => Ok(Arc::new({
1006                    let mut builder = NullBuilder::new();
1007                    builder.append_nulls(rows.len());
1008                    builder.finish()
1009                }) as ArrayRef),
1010                DataType::Utf8 => Ok(Arc::new(
1011                    rows.iter()
1012                        .map(|row| {
1013                            let s = row.get(i);
1014                            (!null_regex.is_null(s)).then_some(s)
1015                        })
1016                        .collect::<StringArray>(),
1017                ) as ArrayRef),
1018                DataType::Utf8View => Ok(Arc::new(
1019                    rows.iter()
1020                        .map(|row| {
1021                            let s = row.get(i);
1022                            (!null_regex.is_null(s)).then_some(s)
1023                        })
1024                        .collect::<StringViewArray>(),
1025                ) as ArrayRef),
1026                DataType::Dictionary(key_type, value_type)
1027                    if value_type.as_ref() == &DataType::Utf8 =>
1028                {
1029                    match key_type.as_ref() {
1030                        DataType::Int8 => build_dictionary_array::<Int8Type>(rows, i, null_regex),
1031                        DataType::Int16 => build_dictionary_array::<Int16Type>(rows, i, null_regex),
1032                        DataType::Int32 => build_dictionary_array::<Int32Type>(rows, i, null_regex),
1033                        DataType::Int64 => build_dictionary_array::<Int64Type>(rows, i, null_regex),
1034                        DataType::UInt8 => build_dictionary_array::<UInt8Type>(rows, i, null_regex),
1035                        DataType::UInt16 => {
1036                            build_dictionary_array::<UInt16Type>(rows, i, null_regex)
1037                        }
1038                        DataType::UInt32 => {
1039                            build_dictionary_array::<UInt32Type>(rows, i, null_regex)
1040                        }
1041                        DataType::UInt64 => {
1042                            build_dictionary_array::<UInt64Type>(rows, i, null_regex)
1043                        }
1044                        _ => Err(ArrowError::ParseError(format!(
1045                            "Unsupported dictionary key type {key_type}"
1046                        ))),
1047                    }
1048                }
1049                other => Err(ArrowError::ParseError(format!(
1050                    "Unsupported data type {other:?}"
1051                ))),
1052            }
1053        })
1054        .collect();
1055
1056    RecordBatch::try_new_with_options(
1057        projected_schema,
1058        arrays?,
1059        &RecordBatchOptions::new()
1060            .with_match_field_names(true)
1061            .with_row_count(Some(rows.len())),
1062    )
1063}
1064
1065fn parse_bool(string: &str) -> Option<bool> {
1066    if string.eq_ignore_ascii_case("false") {
1067        Some(false)
1068    } else if string.eq_ignore_ascii_case("true") {
1069        Some(true)
1070    } else {
1071        None
1072    }
1073}
1074
1075// parse the column string to an Arrow Array
1076fn build_decimal_array<T: DecimalType>(
1077    _line_number: usize,
1078    rows: &StringRecords<'_>,
1079    col_idx: usize,
1080    precision: u8,
1081    scale: i8,
1082    null_regex: &NullRegex,
1083) -> Result<ArrayRef, ArrowError> {
1084    let mut decimal_builder = PrimitiveBuilder::<T>::with_capacity(rows.len());
1085    for row in rows.iter() {
1086        let s = row.get(col_idx);
1087        if null_regex.is_null(s) {
1088            // append null
1089            decimal_builder.append_null();
1090        } else {
1091            let decimal_value: Result<T::Native, _> = parse_decimal::<T>(s, precision, scale);
1092            match decimal_value {
1093                Ok(v) => {
1094                    decimal_builder.append_value(v);
1095                }
1096                Err(e) => {
1097                    return Err(e);
1098                }
1099            }
1100        }
1101    }
1102    Ok(Arc::new(
1103        decimal_builder
1104            .finish()
1105            .with_precision_and_scale(precision, scale)?,
1106    ))
1107}
1108
1109// parses a specific column (col_idx) into an Arrow Array.
1110fn build_primitive_array<T: ArrowPrimitiveType + Parser>(
1111    line_number: usize,
1112    rows: &StringRecords<'_>,
1113    col_idx: usize,
1114    null_regex: &NullRegex,
1115) -> Result<ArrayRef, ArrowError> {
1116    rows.iter()
1117        .enumerate()
1118        .map(|(row_index, row)| {
1119            let s = row.get(col_idx);
1120            if null_regex.is_null(s) {
1121                return Ok(None);
1122            }
1123
1124            match T::parse(s) {
1125                Some(e) => Ok(Some(e)),
1126                None => Err(ArrowError::ParseError(format!(
1127                    // TODO: we should surface the underlying error here.
1128                    "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1129                    s,
1130                    T::DATA_TYPE,
1131                    col_idx,
1132                    line_number + row_index,
1133                    row
1134                ))),
1135            }
1136        })
1137        .collect::<Result<PrimitiveArray<T>, ArrowError>>()
1138        .map(|e| Arc::new(e) as ArrayRef)
1139}
1140
1141fn build_timestamp_array<T: ArrowTimestampType>(
1142    line_number: usize,
1143    rows: &StringRecords<'_>,
1144    col_idx: usize,
1145    timezone: Option<&str>,
1146    null_regex: &NullRegex,
1147) -> Result<ArrayRef, ArrowError> {
1148    Ok(Arc::new(match timezone {
1149        Some(timezone) => {
1150            let tz: Tz = timezone.parse()?;
1151            build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &tz, null_regex)?
1152                .with_timezone(timezone)
1153        }
1154        None => build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &Utc, null_regex)?,
1155    }))
1156}
1157
1158fn build_timestamp_array_impl<T: ArrowTimestampType, Tz: TimeZone>(
1159    line_number: usize,
1160    rows: &StringRecords<'_>,
1161    col_idx: usize,
1162    timezone: &Tz,
1163    null_regex: &NullRegex,
1164) -> Result<PrimitiveArray<T>, ArrowError> {
1165    rows.iter()
1166        .enumerate()
1167        .map(|(row_index, row)| {
1168            let s = row.get(col_idx);
1169            if null_regex.is_null(s) {
1170                return Ok(None);
1171            }
1172
1173            let date = string_to_datetime(timezone, s)
1174                .and_then(|date| match T::UNIT {
1175                    TimeUnit::Second => Ok(date.timestamp()),
1176                    TimeUnit::Millisecond => Ok(date.timestamp_millis()),
1177                    TimeUnit::Microsecond => Ok(date.timestamp_micros()),
1178                    TimeUnit::Nanosecond => date.timestamp_nanos_opt().ok_or_else(|| {
1179                        ArrowError::ParseError(format!(
1180                            "{} would overflow 64-bit signed nanoseconds",
1181                            date.to_rfc3339(),
1182                        ))
1183                    }),
1184                })
1185                .map_err(|e| {
1186                    ArrowError::ParseError(format!(
1187                        "Error parsing column {col_idx} at line {}: {}",
1188                        line_number + row_index,
1189                        e
1190                    ))
1191                })?;
1192            Ok(Some(date))
1193        })
1194        .collect()
1195}
1196
1197// parses a specific column (col_idx) into an Arrow Array.
1198fn build_boolean_array(
1199    line_number: usize,
1200    rows: &StringRecords<'_>,
1201    col_idx: usize,
1202    null_regex: &NullRegex,
1203) -> Result<ArrayRef, ArrowError> {
1204    rows.iter()
1205        .enumerate()
1206        .map(|(row_index, row)| {
1207            let s = row.get(col_idx);
1208            if null_regex.is_null(s) {
1209                return Ok(None);
1210            }
1211            let parsed = parse_bool(s);
1212            match parsed {
1213                Some(e) => Ok(Some(e)),
1214                None => Err(ArrowError::ParseError(format!(
1215                    // TODO: we should surface the underlying error here.
1216                    "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1217                    s,
1218                    "Boolean",
1219                    col_idx,
1220                    line_number + row_index,
1221                    row
1222                ))),
1223            }
1224        })
1225        .collect::<Result<BooleanArray, _>>()
1226        .map(|e| Arc::new(e) as ArrayRef)
1227}
1228
1229// parses a specific column (col_idx) into an Arrow dictionary array of strings.
1230fn build_dictionary_array<K: ArrowDictionaryKeyType>(
1231    rows: &StringRecords<'_>,
1232    col_idx: usize,
1233    null_regex: &NullRegex,
1234) -> Result<ArrayRef, ArrowError> {
1235    let mut builder = StringDictionaryBuilder::<K>::with_capacity(rows.len(), 256, 1024);
1236    for row in rows.iter() {
1237        let s = row.get(col_idx);
1238        if null_regex.is_null(s) {
1239            builder.append_null();
1240        } else {
1241            builder.append(s)?;
1242        }
1243    }
1244    Ok(Arc::new(builder.finish()))
1245}
1246
1247/// Builder for CSV [`Reader`]s
1248#[derive(Debug)]
1249pub struct ReaderBuilder {
1250    /// Schema of the CSV file
1251    schema: SchemaRef,
1252    /// Format of the CSV file
1253    format: Format,
1254    /// Batch size (number of records to load each time)
1255    ///
1256    /// The default batch size when using the `ReaderBuilder` is 1024 records
1257    batch_size: usize,
1258    /// The bounds over which to scan the reader. `None` starts from 0 and runs until EOF.
1259    bounds: Bounds,
1260    /// Optional projection for which columns to load (zero-based column indices)
1261    projection: Option<Vec<usize>>,
1262}
1263
1264impl ReaderBuilder {
1265    /// Create a new builder for configuring [`Reader`] CSV parsing options.
1266    ///
1267    /// To convert a builder into a reader, call [`ReaderBuilder::build`]. See
1268    /// the [module-level documentation](crate::reader) for more details and examples.
1269    ///
1270    /// # Example
1271    ///
1272    /// ```
1273    /// # use arrow_csv::{Reader, ReaderBuilder};
1274    /// # use std::fs::File;
1275    /// # use std::io::Seek;
1276    /// # use std::sync::Arc;
1277    /// # use arrow_csv::reader::Format;
1278    /// #
1279    /// let mut file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1280    /// // Infer the schema with the first 100 records
1281    /// let (schema, _) = Format::default().infer_schema(&mut file, Some(100)).unwrap();
1282    /// file.rewind().unwrap();
1283    ///
1284    /// // create a builder
1285    /// ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
1286    /// ```
1287    pub fn new(schema: SchemaRef) -> ReaderBuilder {
1288        Self {
1289            schema,
1290            format: Format::default(),
1291            batch_size: 1024,
1292            bounds: None,
1293            projection: None,
1294        }
1295    }
1296
1297    /// Set whether the CSV file has a header
1298    pub fn with_header(mut self, has_header: bool) -> Self {
1299        self.format.header = has_header;
1300        self
1301    }
1302
1303    /// Set whether to validate the CSV header against the schema
1304    ///
1305    /// This option only applies when [`Self::with_header`] is set to `true`, and defaults to `false`
1306    pub fn with_header_validation(mut self, validate_header: bool) -> Self {
1307        self.format.header_validation = validate_header;
1308        self
1309    }
1310
1311    /// Overrides the [Format] of this [ReaderBuilder]
1312    pub fn with_format(mut self, format: Format) -> Self {
1313        self.format = format;
1314        self
1315    }
1316
1317    /// Set the CSV file's column delimiter as a byte character
1318    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
1319        self.format.delimiter = Some(delimiter);
1320        self
1321    }
1322
1323    /// Set the given character as the CSV file's escape character
1324    pub fn with_escape(mut self, escape: u8) -> Self {
1325        self.format.escape = Some(escape);
1326        self
1327    }
1328
1329    /// Set the given character as the CSV file's quote character, by default it is double quote
1330    pub fn with_quote(mut self, quote: u8) -> Self {
1331        self.format.quote = Some(quote);
1332        self
1333    }
1334
1335    /// Provide a custom terminator character, defaults to CRLF
1336    pub fn with_terminator(mut self, terminator: u8) -> Self {
1337        self.format.terminator = Some(terminator);
1338        self
1339    }
1340
1341    /// Provide a comment character, lines starting with this character will be ignored
1342    pub fn with_comment(mut self, comment: u8) -> Self {
1343        self.format.comment = Some(comment);
1344        self
1345    }
1346
1347    /// Provide a regex to match null values, defaults to `^$`
1348    pub fn with_null_regex(mut self, null_regex: Regex) -> Self {
1349        self.format.null_regex = NullRegex(Some(null_regex));
1350        self
1351    }
1352
1353    /// Set the batch size (number of records to load at one time)
1354    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1355        self.batch_size = batch_size;
1356        self
1357    }
1358
1359    /// Set the bounds over which to scan the reader.
1360    /// `start` and `end` are line numbers.
1361    pub fn with_bounds(mut self, start: usize, end: usize) -> Self {
1362        self.bounds = Some((start, end));
1363        self
1364    }
1365
1366    /// Set the reader's column projection
1367    pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1368        self.projection = Some(projection);
1369        self
1370    }
1371
1372    /// Whether to allow truncated rows when parsing.
1373    ///
1374    /// By default this is set to `false` and will error if the CSV rows have different lengths.
1375    /// When set to true then it will allow records with less than the expected number of columns
1376    /// and fill the missing columns with nulls. If the record's schema is not nullable, then it
1377    /// will still return an error.
1378    pub fn with_truncated_rows(mut self, allow: bool) -> Self {
1379        self.format.truncated_rows = allow;
1380        self
1381    }
1382
1383    /// Create a new `Reader` from a non-buffered reader
1384    ///
1385    /// If `R: BufRead` consider using [`Self::build_buffered`] to avoid unnecessary additional
1386    /// buffering, as internally this method wraps `reader` in [`std::io::BufReader`]
1387    pub fn build<R: Read>(self, reader: R) -> Result<Reader<R>, ArrowError> {
1388        self.build_buffered(StdBufReader::new(reader))
1389    }
1390
1391    /// Create a new `BufReader` from a buffered reader
1392    pub fn build_buffered<R: BufRead>(self, reader: R) -> Result<BufReader<R>, ArrowError> {
1393        let schema = match &self.projection {
1394            Some(projection) => Arc::new(self.schema.project(projection)?),
1395            None => self.schema.clone(),
1396        };
1397
1398        Ok(BufReader {
1399            reader,
1400            decoder: self.build_decoder(),
1401            schema,
1402        })
1403    }
1404
1405    /// Builds a decoder that can be used to decode CSV from an arbitrary byte stream
1406    pub fn build_decoder(self) -> Decoder {
1407        let delimiter = self.format.build_parser();
1408        let record_decoder = RecordDecoder::new(
1409            delimiter,
1410            self.schema.fields().len(),
1411            self.format.truncated_rows,
1412        );
1413
1414        let header = self.format.header as usize;
1415
1416        let (start, end) = match self.bounds {
1417            Some((start, end)) => (start + header, end + header),
1418            None => (header, usize::MAX),
1419        };
1420
1421        Decoder {
1422            schema: self.schema,
1423            to_skip: start,
1424            header_validation: self.format.header && self.format.header_validation,
1425            record_decoder,
1426            line_number: start,
1427            end,
1428            projection: self.projection,
1429            batch_size: self.batch_size,
1430            null_regex: self.format.null_regex,
1431        }
1432    }
1433}
1434
1435#[cfg(test)]
1436mod tests {
1437    use super::*;
1438
1439    use std::io::{Cursor, Seek, SeekFrom, Write};
1440    use tempfile::NamedTempFile;
1441
1442    use arrow_array::cast::AsArray;
1443    use arrow_cast::display::array_value_to_string;
1444
1445    #[test]
1446    fn test_infer_schema_leading_plus_numbers() {
1447        for (csv, expected_type) in [
1448            ("+1\n2\n-3\n", DataType::Int64),
1449            ("+1.5\n2.5\n-3.5\n", DataType::Float64),
1450            ("+1e3\n+2.5e-2\n-3E+2\n", DataType::Float64),
1451            ("+9223372036854775807\n0\n", DataType::Int64),
1452            ("+9223372036854775808\n0\n", DataType::Utf8),
1453            ("+-1\n2\n", DataType::Utf8),
1454            ("+\n2\n", DataType::Utf8),
1455        ] {
1456            let (schema, records_read) = Format::default()
1457                .infer_schema(Cursor::new(csv), None)
1458                .unwrap();
1459            assert_eq!(schema.field(0).data_type(), &expected_type, "CSV: {csv:?}");
1460            // Inferred numeric types must also be accepted by the CSV decoder.
1461            let reader = ReaderBuilder::new(Arc::new(schema))
1462                .build(Cursor::new(csv))
1463                .unwrap();
1464            let rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum();
1465            assert_eq!(rows, records_read, "CSV: {csv:?}");
1466        }
1467    }
1468
1469    #[test]
1470    fn test_csv() {
1471        let schema = Arc::new(Schema::new(vec![
1472            Field::new("city", DataType::Utf8, false),
1473            Field::new("lat", DataType::Float64, false),
1474            Field::new("lng", DataType::Float64, false),
1475        ]));
1476
1477        let file = File::open("test/data/uk_cities.csv").unwrap();
1478        let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1479        assert_eq!(schema, csv.schema());
1480        let batch = csv.next().unwrap().unwrap();
1481        assert_eq!(37, batch.num_rows());
1482        assert_eq!(3, batch.num_columns());
1483
1484        // access data from a primitive array
1485        let lat = batch.column(1).as_primitive::<Float64Type>();
1486        assert_eq!(57.653484, lat.value(0));
1487
1488        // access data from a string array (ListArray<u8>)
1489        let city = batch.column(0).as_string::<i32>();
1490
1491        assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1492    }
1493
1494    #[test]
1495    fn test_csv_schema_metadata() {
1496        let mut metadata = std::collections::HashMap::new();
1497        metadata.insert("foo".to_owned(), "bar".to_owned());
1498        let schema = Arc::new(Schema::new_with_metadata(
1499            vec![
1500                Field::new("city", DataType::Utf8, false),
1501                Field::new("lat", DataType::Float64, false),
1502                Field::new("lng", DataType::Float64, false),
1503            ],
1504            metadata.clone(),
1505        ));
1506
1507        let file = File::open("test/data/uk_cities.csv").unwrap();
1508
1509        let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1510        assert_eq!(schema, csv.schema());
1511        let batch = csv.next().unwrap().unwrap();
1512        assert_eq!(37, batch.num_rows());
1513        assert_eq!(3, batch.num_columns());
1514
1515        assert_eq!(batch.schema().metadata(), &metadata);
1516    }
1517
1518    #[test]
1519    fn test_csv_reader_with_decimal() {
1520        let schema = Arc::new(Schema::new(vec![
1521            Field::new("city", DataType::Utf8, false),
1522            Field::new("lat", DataType::Decimal128(38, 6), false),
1523            Field::new("lng", DataType::Decimal256(76, 6), false),
1524        ]));
1525
1526        let file = File::open("test/data/decimal_test.csv").unwrap();
1527
1528        let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1529        let batch = csv.next().unwrap().unwrap();
1530        // access data from a primitive array
1531        let lat = batch
1532            .column(1)
1533            .as_any()
1534            .downcast_ref::<Decimal128Array>()
1535            .unwrap();
1536
1537        assert_eq!("57.653484", lat.value_as_string(0));
1538        assert_eq!("53.002666", lat.value_as_string(1));
1539        assert_eq!("52.412811", lat.value_as_string(2));
1540        assert_eq!("51.481583", lat.value_as_string(3));
1541        assert_eq!("12.123457", lat.value_as_string(4));
1542        assert_eq!("50.760000", lat.value_as_string(5));
1543        assert_eq!("0.123000", lat.value_as_string(6));
1544        assert_eq!("123.000000", lat.value_as_string(7));
1545        assert_eq!("123.000000", lat.value_as_string(8));
1546        assert_eq!("-50.760000", lat.value_as_string(9));
1547
1548        let lng = batch
1549            .column(2)
1550            .as_any()
1551            .downcast_ref::<Decimal256Array>()
1552            .unwrap();
1553
1554        assert_eq!("-3.335724", lng.value_as_string(0));
1555        assert_eq!("-2.179404", lng.value_as_string(1));
1556        assert_eq!("-1.778197", lng.value_as_string(2));
1557        assert_eq!("-3.179090", lng.value_as_string(3));
1558        assert_eq!("-3.179090", lng.value_as_string(4));
1559        assert_eq!("0.290472", lng.value_as_string(5));
1560        assert_eq!("0.290472", lng.value_as_string(6));
1561        assert_eq!("0.290472", lng.value_as_string(7));
1562        assert_eq!("0.290472", lng.value_as_string(8));
1563        assert_eq!("0.290472", lng.value_as_string(9));
1564    }
1565
1566    #[test]
1567    fn test_csv_reader_decimal_parsing() {
1568        // Rounding half away from zero, surrounding whitespace, exponent
1569        // notation and negative scales are all accepted
1570        let data = " 1.995 ,1.5e2,1234.5,0e0\n-0.005,-1.5E-2,-150,1E+2\n123,+.5,5,-7\n";
1571        let schema = Arc::new(Schema::new(vec![
1572            Field::new("a", DataType::Decimal128(10, 2), false),
1573            Field::new("b", DataType::Decimal64(18, 2), false),
1574            Field::new("c", DataType::Decimal128(10, -2), false),
1575            Field::new("d", DataType::Decimal32(9, 0), false),
1576        ]));
1577        let mut csv = ReaderBuilder::new(schema).build(Cursor::new(data)).unwrap();
1578        let batch = csv.next().unwrap().unwrap();
1579        let column = |i: usize| {
1580            (0..batch.num_rows())
1581                .map(|row| array_value_to_string(batch.column(i), row).unwrap())
1582                .collect::<Vec<_>>()
1583        };
1584        assert_eq!(column(0), ["2.00", "-0.01", "123.00"]);
1585        assert_eq!(column(1), ["150.00", "-0.02", "0.50"]);
1586        assert_eq!(
1587            batch.column(2).as_primitive::<Decimal128Type>().values(),
1588            &[12, -2, 0]
1589        );
1590        assert_eq!(column(3), ["0", "100", "-7"]);
1591
1592        // Invalid and out-of-range values are errors, never panics
1593        for (data, expected) in [
1594            ("abc\n", "Invalid decimal format: \"abc\""),
1595            ("1.2.3\n", "Invalid decimal format: \"1.2.3\""),
1596            (
1597                "123456789\n",
1598                "\"123456789\" does not fit in Decimal128(5, 2)",
1599            ),
1600            ("1e99999\n", "does not fit in Decimal128(5, 2)"),
1601            (
1602                &format!("{}\n", "1".repeat(300)),
1603                "does not fit in Decimal128(5, 2)",
1604            ),
1605            (
1606                "4825037936439135476.2609835314269495255615E-14\n",
1607                "does not fit in Decimal128(5, 2)",
1608            ),
1609        ] {
1610            let schema = Arc::new(Schema::new(vec![Field::new(
1611                "a",
1612                DataType::Decimal128(5, 2),
1613                false,
1614            )]));
1615            let mut csv = ReaderBuilder::new(schema).build(Cursor::new(data)).unwrap();
1616            let err = csv.next().unwrap().unwrap_err().to_string();
1617            assert!(err.contains(expected), "{data:?}: {err}");
1618        }
1619    }
1620
1621    #[test]
1622    fn test_csv_reader_with_decimal_3264() {
1623        let schema = Arc::new(Schema::new(vec![
1624            Field::new("city", DataType::Utf8, false),
1625            Field::new("lat", DataType::Decimal32(9, 6), false),
1626            Field::new("lng", DataType::Decimal64(16, 6), false),
1627        ]));
1628
1629        let file = File::open("test/data/decimal_test.csv").unwrap();
1630
1631        let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1632        let batch = csv.next().unwrap().unwrap();
1633        // access data from a primitive array
1634        let lat = batch
1635            .column(1)
1636            .as_any()
1637            .downcast_ref::<Decimal32Array>()
1638            .unwrap();
1639
1640        assert_eq!("57.653484", lat.value_as_string(0));
1641        assert_eq!("53.002666", lat.value_as_string(1));
1642        assert_eq!("52.412811", lat.value_as_string(2));
1643        assert_eq!("51.481583", lat.value_as_string(3));
1644        assert_eq!("12.123457", lat.value_as_string(4));
1645        assert_eq!("50.760000", lat.value_as_string(5));
1646        assert_eq!("0.123000", lat.value_as_string(6));
1647        assert_eq!("123.000000", lat.value_as_string(7));
1648        assert_eq!("123.000000", lat.value_as_string(8));
1649        assert_eq!("-50.760000", lat.value_as_string(9));
1650
1651        let lng = batch
1652            .column(2)
1653            .as_any()
1654            .downcast_ref::<Decimal64Array>()
1655            .unwrap();
1656
1657        assert_eq!("-3.335724", lng.value_as_string(0));
1658        assert_eq!("-2.179404", lng.value_as_string(1));
1659        assert_eq!("-1.778197", lng.value_as_string(2));
1660        assert_eq!("-3.179090", lng.value_as_string(3));
1661        assert_eq!("-3.179090", lng.value_as_string(4));
1662        assert_eq!("0.290472", lng.value_as_string(5));
1663        assert_eq!("0.290472", lng.value_as_string(6));
1664        assert_eq!("0.290472", lng.value_as_string(7));
1665        assert_eq!("0.290472", lng.value_as_string(8));
1666        assert_eq!("0.290472", lng.value_as_string(9));
1667    }
1668
1669    #[test]
1670    fn test_csv_from_buf_reader() {
1671        let schema = Schema::new(vec![
1672            Field::new("city", DataType::Utf8, false),
1673            Field::new("lat", DataType::Float64, false),
1674            Field::new("lng", DataType::Float64, false),
1675        ]);
1676
1677        let file_with_headers = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1678        let file_without_headers = File::open("test/data/uk_cities.csv").unwrap();
1679        let both_files = file_with_headers
1680            .chain(Cursor::new("\n".to_string()))
1681            .chain(file_without_headers);
1682        let mut csv = ReaderBuilder::new(Arc::new(schema))
1683            .with_header(true)
1684            .build(both_files)
1685            .unwrap();
1686        let batch = csv.next().unwrap().unwrap();
1687        assert_eq!(74, batch.num_rows());
1688        assert_eq!(3, batch.num_columns());
1689    }
1690
1691    #[test]
1692    fn test_infer_format_with_typed_columns() {
1693        let csv = "name,count,active\nalice,1,true\nbob,2,false\n";
1694
1695        let (format, format_records_read) = Format::default()
1696            .infer_format(Cursor::new(csv), None)
1697            .unwrap();
1698        let (schema, records_read) = format.infer_schema(Cursor::new(csv), None).unwrap();
1699
1700        assert_eq!(schema.field(0).name(), "name");
1701        assert_eq!(schema.field(1).name(), "count");
1702        assert_eq!(schema.field(2).name(), "active");
1703        assert_eq!(format_records_read, 3);
1704        assert_eq!(records_read, 2);
1705    }
1706
1707    #[test]
1708    fn test_infer_format_without_header() {
1709        let csv = "1,true\n2,false\n";
1710
1711        let (format, format_records_read) = Format::default()
1712            .infer_format(Cursor::new(csv), None)
1713            .unwrap();
1714        let (schema, records_read) = format.infer_schema(Cursor::new(csv), None).unwrap();
1715
1716        assert_eq!(schema.field(0).name(), "column_1");
1717        assert_eq!(schema.field(1).name(), "column_2");
1718        assert_eq!(format_records_read, 2);
1719        assert_eq!(records_read, 2);
1720    }
1721
1722    #[test]
1723    fn test_infer_format_returns_no_header_when_ambiguous() {
1724        for csv in ["name,count\n", "alice,london\nbob,paris\n"] {
1725            let (format, _) = Format::default()
1726                .infer_format(Cursor::new(csv), None)
1727                .unwrap();
1728            let (schema, _) = format.infer_schema(Cursor::new(csv), None).unwrap();
1729            assert_eq!(schema.field(0).name(), "column_1", "CSV: {csv:?}");
1730        }
1731
1732        let (format, format_records_read) = Format::default()
1733            .infer_format(Cursor::new(""), None)
1734            .unwrap();
1735        let (schema, records_read) = format.infer_schema(Cursor::new(""), None).unwrap();
1736        assert!(schema.fields().is_empty());
1737        assert_eq!(format_records_read, 0);
1738        assert_eq!(records_read, 0);
1739    }
1740
1741    #[test]
1742    fn test_infer_format_honors_format_options() {
1743        let csv = "name;count\nalice;1\nbob;2\n";
1744        let (format, _) = Format::default()
1745            .with_delimiter(b';')
1746            .infer_format(Cursor::new(csv), None)
1747            .unwrap();
1748        let (schema, records_read) = format.infer_schema(Cursor::new(csv), None).unwrap();
1749
1750        assert_eq!(schema.field(0).name(), "name");
1751        assert_eq!(schema.field(1).name(), "count");
1752        assert_eq!(records_read, 2);
1753    }
1754
1755    #[test]
1756    fn test_infer_format_respects_max_records() {
1757        let csv = "name,count\nalice,1\nbob,unknown\n";
1758        let infer = |max_records| {
1759            let (format, records_read) = Format::default()
1760                .infer_format(Cursor::new(csv), max_records)
1761                .unwrap();
1762            let (schema, _) = format.infer_schema(Cursor::new(csv), None).unwrap();
1763            (schema.field(0).name().clone(), records_read)
1764        };
1765
1766        assert_eq!(infer(Some(1)), ("name".to_string(), 2));
1767        assert_eq!(infer(None), ("column_1".to_string(), 3));
1768        assert_eq!(infer(Some(0)), ("column_1".to_string(), 1));
1769    }
1770
1771    #[test]
1772    fn test_infer_format_numeric_text_is_not_header() {
1773        for csv in [
1774            "+1\n2\n3\n",
1775            "+1.5\n2.5\n3.5\n",
1776            "+1e3\n2e3\n3e3\n",
1777            "9223372036854775808\n2\n3\n",
1778        ] {
1779            let (format, records_read) = Format::default()
1780                .infer_format(Cursor::new(csv), None)
1781                .unwrap();
1782            let (schema, schema_records_read) =
1783                format.infer_schema(Cursor::new(csv), None).unwrap();
1784            let (ordinary_schema, _) = Format::default()
1785                .infer_schema(Cursor::new(csv), None)
1786                .unwrap();
1787
1788            assert_eq!(schema.field(0).name(), "column_1", "CSV: {csv:?}");
1789            assert_eq!(schema, ordinary_schema, "CSV: {csv:?}");
1790            assert_eq!(records_read, 3);
1791            assert_eq!(schema_records_read, 3);
1792        }
1793    }
1794
1795    #[test]
1796    #[cfg_attr(miri, ignore)] // Takes too long
1797    fn test_csv_with_schema_inference() {
1798        let mut file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1799
1800        let (schema, _) = Format::default()
1801            .with_header(true)
1802            .infer_schema(&mut file, None)
1803            .unwrap();
1804
1805        file.rewind().unwrap();
1806        let builder = ReaderBuilder::new(Arc::new(schema)).with_header(true);
1807
1808        let mut csv = builder.build(file).unwrap();
1809        let expected_schema = Schema::new(vec![
1810            Field::new("city", DataType::Utf8, true),
1811            Field::new("lat", DataType::Float64, true),
1812            Field::new("lng", DataType::Float64, true),
1813        ]);
1814        assert_eq!(Arc::new(expected_schema), csv.schema());
1815        let batch = csv.next().unwrap().unwrap();
1816        assert_eq!(37, batch.num_rows());
1817        assert_eq!(3, batch.num_columns());
1818
1819        // access data from a primitive array
1820        let lat = batch
1821            .column(1)
1822            .as_any()
1823            .downcast_ref::<Float64Array>()
1824            .unwrap();
1825        assert_eq!(57.653484, lat.value(0));
1826
1827        // access data from a string array (ListArray<u8>)
1828        let city = batch
1829            .column(0)
1830            .as_any()
1831            .downcast_ref::<StringArray>()
1832            .unwrap();
1833
1834        assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1835    }
1836
1837    #[test]
1838    #[cfg_attr(miri, ignore)] // Takes too long
1839    fn test_csv_with_schema_inference_no_headers() {
1840        let mut file = File::open("test/data/uk_cities.csv").unwrap();
1841
1842        let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1843        file.rewind().unwrap();
1844
1845        let mut csv = ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
1846
1847        // csv field names should be 'column_{number}'
1848        let schema = csv.schema();
1849        assert_eq!("column_1", schema.field(0).name());
1850        assert_eq!("column_2", schema.field(1).name());
1851        assert_eq!("column_3", schema.field(2).name());
1852        let batch = csv.next().unwrap().unwrap();
1853        let batch_schema = batch.schema();
1854
1855        assert_eq!(schema, batch_schema);
1856        assert_eq!(37, batch.num_rows());
1857        assert_eq!(3, batch.num_columns());
1858
1859        // access data from a primitive array
1860        let lat = batch
1861            .column(1)
1862            .as_any()
1863            .downcast_ref::<Float64Array>()
1864            .unwrap();
1865        assert_eq!(57.653484, lat.value(0));
1866
1867        // access data from a string array (ListArray<u8>)
1868        let city = batch
1869            .column(0)
1870            .as_any()
1871            .downcast_ref::<StringArray>()
1872            .unwrap();
1873
1874        assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1875    }
1876
1877    #[test]
1878    #[cfg_attr(miri, ignore)] // Takes too long
1879    fn test_csv_builder_with_bounds() {
1880        let mut file = File::open("test/data/uk_cities.csv").unwrap();
1881
1882        // Set the bounds to the lines 0, 1 and 2.
1883        let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1884        file.rewind().unwrap();
1885        let mut csv = ReaderBuilder::new(Arc::new(schema))
1886            .with_bounds(0, 2)
1887            .build(file)
1888            .unwrap();
1889        let batch = csv.next().unwrap().unwrap();
1890
1891        // access data from a string array (ListArray<u8>)
1892        let city = batch
1893            .column(0)
1894            .as_any()
1895            .downcast_ref::<StringArray>()
1896            .unwrap();
1897
1898        // The value on line 0 is within the bounds
1899        assert_eq!("Elgin, Scotland, the UK", city.value(0));
1900
1901        // The value on line 13 is outside of the bounds. Therefore
1902        // the call to .value() will panic.
1903        let result = std::panic::catch_unwind(|| city.value(13));
1904        assert!(result.is_err());
1905    }
1906
1907    #[test]
1908    fn test_csv_with_projection() {
1909        let schema = Arc::new(Schema::new(vec![
1910            Field::new("city", DataType::Utf8, false),
1911            Field::new("lat", DataType::Float64, false),
1912            Field::new("lng", DataType::Float64, false),
1913        ]));
1914
1915        let file = File::open("test/data/uk_cities.csv").unwrap();
1916
1917        let mut csv = ReaderBuilder::new(schema)
1918            .with_projection(vec![0, 1])
1919            .build(file)
1920            .unwrap();
1921
1922        let projected_schema = Arc::new(Schema::new(vec![
1923            Field::new("city", DataType::Utf8, false),
1924            Field::new("lat", DataType::Float64, false),
1925        ]));
1926        assert_eq!(projected_schema, csv.schema());
1927        let batch = csv.next().unwrap().unwrap();
1928        assert_eq!(projected_schema, batch.schema());
1929        assert_eq!(37, batch.num_rows());
1930        assert_eq!(2, batch.num_columns());
1931    }
1932
1933    #[test]
1934    fn test_csv_record_batch_reader_schema() {
1935        let schema = Arc::new(Schema::new(vec![
1936            Field::new("a", DataType::Int32, false),
1937            Field::new("b", DataType::Int32, false),
1938        ]));
1939
1940        let cases = [
1941            None,
1942            Some(vec![]),
1943            Some(vec![1]),
1944            Some(vec![1, 0]),
1945            Some(vec![1, 1]),
1946        ];
1947        for projection in cases {
1948            let builder = ReaderBuilder::new(schema.clone());
1949            let builder = match projection {
1950                Some(projection) => builder.with_projection(projection),
1951                None => builder,
1952            };
1953            let mut reader = builder.build(Cursor::new(b"1,2\n")).unwrap();
1954
1955            let reader_schema = RecordBatchReader::schema(&reader);
1956            let batch = reader.next().unwrap().unwrap();
1957
1958            assert_eq!(reader_schema, batch.schema());
1959        }
1960    }
1961
1962    #[test]
1963    fn test_csv_reader_rejects_invalid_projection() {
1964        let schema = Arc::new(Schema::new(vec![
1965            Field::new("a", DataType::Int32, false),
1966            Field::new("b", DataType::Int32, false),
1967        ]));
1968
1969        let result = ReaderBuilder::new(schema)
1970            .with_projection(vec![2])
1971            .build(Cursor::new(b"1,2\n"));
1972
1973        assert!(matches!(
1974            result,
1975            Err(ArrowError::SchemaError(message))
1976                if message == "project index 2 out of bounds, max field 2"
1977        ));
1978    }
1979
1980    #[test]
1981    fn test_csv_decoder_rejects_invalid_projection() {
1982        let schema = Arc::new(Schema::new(vec![
1983            Field::new("a", DataType::Int32, false),
1984            Field::new("b", DataType::Int32, false),
1985        ]));
1986        let mut decoder = ReaderBuilder::new(schema)
1987            .with_projection(vec![2])
1988            .build_decoder();
1989
1990        decoder.decode(b"1,2\n").unwrap();
1991        let result = decoder.flush();
1992
1993        assert!(matches!(
1994            result,
1995            Err(ArrowError::SchemaError(message))
1996                if message == "project index 2 out of bounds, max field 2"
1997        ));
1998    }
1999
2000    #[test]
2001    fn test_csv_with_dictionary() {
2002        let schema = Arc::new(Schema::new(vec![
2003            Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
2004            Field::new("lat", DataType::Float64, false),
2005            Field::new("lng", DataType::Float64, false),
2006        ]));
2007
2008        let file = File::open("test/data/uk_cities.csv").unwrap();
2009
2010        let mut csv = ReaderBuilder::new(schema)
2011            .with_projection(vec![0, 1])
2012            .build(file)
2013            .unwrap();
2014
2015        let projected_schema = Arc::new(Schema::new(vec![
2016            Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
2017            Field::new("lat", DataType::Float64, false),
2018        ]));
2019        assert_eq!(projected_schema, csv.schema());
2020        let batch = csv.next().unwrap().unwrap();
2021        assert_eq!(projected_schema, batch.schema());
2022        assert_eq!(37, batch.num_rows());
2023        assert_eq!(2, batch.num_columns());
2024
2025        let strings = arrow_cast::cast(batch.column(0), &DataType::Utf8).unwrap();
2026        let strings = strings.as_string::<i32>();
2027
2028        assert_eq!(strings.value(0), "Elgin, Scotland, the UK");
2029        assert_eq!(strings.value(4), "Eastbourne, East Sussex, UK");
2030        assert_eq!(strings.value(29), "Uckfield, East Sussex, UK");
2031    }
2032
2033    #[test]
2034    fn test_csv_dictionary_key_overflow() {
2035        let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
2036            "c",
2037            DataType::Int8,
2038            DataType::Utf8,
2039            false,
2040        )]));
2041        let read = |distinct: usize| {
2042            let rows = (0..distinct).map(|i| format!("v{i}")).collect::<Vec<_>>();
2043            let mut reader = ReaderBuilder::new(schema.clone())
2044                .build(Cursor::new(rows.join("\n")))
2045                .unwrap();
2046            reader.next().unwrap()
2047        };
2048
2049        assert_eq!(read(128).unwrap().num_rows(), 128);
2050        let err = read(129).unwrap_err();
2051        assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
2052    }
2053
2054    #[test]
2055    fn test_csv_with_nullable_dictionary() {
2056        let offset_type = vec![
2057            DataType::Int8,
2058            DataType::Int16,
2059            DataType::Int32,
2060            DataType::Int64,
2061            DataType::UInt8,
2062            DataType::UInt16,
2063            DataType::UInt32,
2064            DataType::UInt64,
2065        ];
2066        for data_type in offset_type {
2067            let file = File::open("test/data/dictionary_nullable_test.csv").unwrap();
2068            let dictionary_type =
2069                DataType::Dictionary(Box::new(data_type), Box::new(DataType::Utf8));
2070            let schema = Arc::new(Schema::new(vec![
2071                Field::new("id", DataType::Utf8, false),
2072                Field::new("name", dictionary_type.clone(), true),
2073            ]));
2074
2075            let mut csv = ReaderBuilder::new(schema)
2076                .build(file.try_clone().unwrap())
2077                .unwrap();
2078
2079            let batch = csv.next().unwrap().unwrap();
2080            assert_eq!(3, batch.num_rows());
2081            assert_eq!(2, batch.num_columns());
2082
2083            let names = arrow_cast::cast(batch.column(1), &dictionary_type).unwrap();
2084            assert!(!names.is_null(2));
2085            assert!(names.is_null(1));
2086        }
2087    }
2088    #[test]
2089    fn test_nulls() {
2090        let schema = Arc::new(Schema::new(vec![
2091            Field::new("c_int", DataType::UInt64, false),
2092            Field::new("c_float", DataType::Float32, true),
2093            Field::new("c_string", DataType::Utf8, true),
2094            Field::new("c_bool", DataType::Boolean, false),
2095        ]));
2096
2097        let file = File::open("test/data/null_test.csv").unwrap();
2098
2099        let mut csv = ReaderBuilder::new(schema)
2100            .with_header(true)
2101            .build(file)
2102            .unwrap();
2103
2104        let batch = csv.next().unwrap().unwrap();
2105
2106        assert!(!batch.column(1).is_null(0));
2107        assert!(!batch.column(1).is_null(1));
2108        assert!(batch.column(1).is_null(2));
2109        assert!(!batch.column(1).is_null(3));
2110        assert!(!batch.column(1).is_null(4));
2111    }
2112
2113    #[test]
2114    fn test_init_nulls() {
2115        let schema = Arc::new(Schema::new(vec![
2116            Field::new("c_int", DataType::UInt64, true),
2117            Field::new("c_float", DataType::Float32, true),
2118            Field::new("c_string", DataType::Utf8, true),
2119            Field::new("c_bool", DataType::Boolean, true),
2120            Field::new("c_null", DataType::Null, true),
2121        ]));
2122        let file = File::open("test/data/init_null_test.csv").unwrap();
2123
2124        let mut csv = ReaderBuilder::new(schema)
2125            .with_header(true)
2126            .build(file)
2127            .unwrap();
2128
2129        let batch = csv.next().unwrap().unwrap();
2130
2131        assert!(batch.column(1).is_null(0));
2132        assert!(!batch.column(1).is_null(1));
2133        assert!(batch.column(1).is_null(2));
2134        assert!(!batch.column(1).is_null(3));
2135        assert!(!batch.column(1).is_null(4));
2136    }
2137
2138    #[test]
2139    #[cfg_attr(miri, ignore)] // Takes too long
2140    fn test_init_nulls_with_inference() {
2141        let format = Format::default().with_header(true).with_delimiter(b',');
2142
2143        let mut file = File::open("test/data/init_null_test.csv").unwrap();
2144        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2145        file.rewind().unwrap();
2146
2147        let expected_schema = Schema::new(vec![
2148            Field::new("c_int", DataType::Int64, true),
2149            Field::new("c_float", DataType::Float64, true),
2150            Field::new("c_string", DataType::Utf8, true),
2151            Field::new("c_bool", DataType::Boolean, true),
2152            Field::new("c_null", DataType::Null, true),
2153        ]);
2154        assert_eq!(schema, expected_schema);
2155
2156        let mut csv = ReaderBuilder::new(Arc::new(schema))
2157            .with_format(format)
2158            .build(file)
2159            .unwrap();
2160
2161        let batch = csv.next().unwrap().unwrap();
2162
2163        assert!(batch.column(1).is_null(0));
2164        assert!(!batch.column(1).is_null(1));
2165        assert!(batch.column(1).is_null(2));
2166        assert!(!batch.column(1).is_null(3));
2167        assert!(!batch.column(1).is_null(4));
2168    }
2169
2170    #[test]
2171    fn test_custom_nulls() {
2172        let schema = Arc::new(Schema::new(vec![
2173            Field::new("c_int", DataType::UInt64, true),
2174            Field::new("c_float", DataType::Float32, true),
2175            Field::new("c_string", DataType::Utf8, true),
2176            Field::new("c_bool", DataType::Boolean, true),
2177        ]));
2178
2179        let file = File::open("test/data/custom_null_test.csv").unwrap();
2180
2181        let null_regex = Regex::new("^nil$").unwrap();
2182
2183        let mut csv = ReaderBuilder::new(schema)
2184            .with_header(true)
2185            .with_null_regex(null_regex)
2186            .build(file)
2187            .unwrap();
2188
2189        let batch = csv.next().unwrap().unwrap();
2190
2191        // "nil"s should be NULL
2192        assert!(batch.column(0).is_null(1));
2193        assert!(batch.column(1).is_null(2));
2194        assert!(batch.column(3).is_null(4));
2195        assert!(batch.column(2).is_null(3));
2196        assert!(!batch.column(2).is_null(4));
2197    }
2198
2199    #[test]
2200    #[cfg_attr(miri, ignore)] // Takes too long
2201    fn test_nulls_with_inference() {
2202        let mut file = File::open("test/data/various_types.csv").unwrap();
2203        let format = Format::default().with_header(true).with_delimiter(b'|');
2204
2205        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2206        file.rewind().unwrap();
2207
2208        let builder = ReaderBuilder::new(Arc::new(schema))
2209            .with_format(format)
2210            .with_batch_size(512)
2211            .with_projection(vec![0, 1, 2, 3, 4, 5]);
2212
2213        let mut csv = builder.build(file).unwrap();
2214        let batch = csv.next().unwrap().unwrap();
2215
2216        assert_eq!(10, batch.num_rows());
2217        assert_eq!(6, batch.num_columns());
2218
2219        let schema = batch.schema();
2220
2221        assert_eq!(&DataType::Int64, schema.field(0).data_type());
2222        assert_eq!(&DataType::Float64, schema.field(1).data_type());
2223        assert_eq!(&DataType::Float64, schema.field(2).data_type());
2224        assert_eq!(&DataType::Boolean, schema.field(3).data_type());
2225        assert_eq!(&DataType::Date32, schema.field(4).data_type());
2226        assert_eq!(
2227            &DataType::Timestamp(TimeUnit::Second, None),
2228            schema.field(5).data_type()
2229        );
2230
2231        let names: Vec<&str> = schema.fields().iter().map(|x| x.name().as_str()).collect();
2232        assert_eq!(
2233            names,
2234            vec![
2235                "c_int",
2236                "c_float",
2237                "c_string",
2238                "c_bool",
2239                "c_date",
2240                "c_datetime"
2241            ]
2242        );
2243
2244        assert!(schema.field(0).is_nullable());
2245        assert!(schema.field(1).is_nullable());
2246        assert!(schema.field(2).is_nullable());
2247        assert!(schema.field(3).is_nullable());
2248        assert!(schema.field(4).is_nullable());
2249        assert!(schema.field(5).is_nullable());
2250
2251        assert!(!batch.column(1).is_null(0));
2252        assert!(!batch.column(1).is_null(1));
2253        assert!(batch.column(1).is_null(2));
2254        assert!(!batch.column(1).is_null(3));
2255        assert!(!batch.column(1).is_null(4));
2256    }
2257
2258    #[test]
2259    #[cfg_attr(miri, ignore)] // Takes too long
2260    fn test_custom_nulls_with_inference() {
2261        let mut file = File::open("test/data/custom_null_test.csv").unwrap();
2262
2263        let null_regex = Regex::new("^nil$").unwrap();
2264
2265        let format = Format::default()
2266            .with_header(true)
2267            .with_null_regex(null_regex);
2268
2269        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2270        file.rewind().unwrap();
2271
2272        let expected_schema = Schema::new(vec![
2273            Field::new("c_int", DataType::Int64, true),
2274            Field::new("c_float", DataType::Float64, true),
2275            Field::new("c_string", DataType::Utf8, true),
2276            Field::new("c_bool", DataType::Boolean, true),
2277        ]);
2278
2279        assert_eq!(schema, expected_schema);
2280
2281        let builder = ReaderBuilder::new(Arc::new(schema))
2282            .with_format(format)
2283            .with_batch_size(512)
2284            .with_projection(vec![0, 1, 2, 3]);
2285
2286        let mut csv = builder.build(file).unwrap();
2287        let batch = csv.next().unwrap().unwrap();
2288
2289        assert_eq!(5, batch.num_rows());
2290        assert_eq!(4, batch.num_columns());
2291
2292        assert_eq!(batch.schema().as_ref(), &expected_schema);
2293    }
2294
2295    #[test]
2296    #[cfg_attr(miri, ignore)] // Takes too long
2297    fn test_scientific_notation_with_inference() {
2298        let mut file = File::open("test/data/scientific_notation_test.csv").unwrap();
2299        let format = Format::default().with_header(false).with_delimiter(b',');
2300
2301        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2302        file.rewind().unwrap();
2303
2304        let builder = ReaderBuilder::new(Arc::new(schema))
2305            .with_format(format)
2306            .with_batch_size(512)
2307            .with_projection(vec![0, 1]);
2308
2309        let mut csv = builder.build(file).unwrap();
2310        let batch = csv.next().unwrap().unwrap();
2311
2312        let schema = batch.schema();
2313
2314        assert_eq!(&DataType::Float64, schema.field(0).data_type());
2315    }
2316
2317    fn invalid_csv_helper(file_name: &str) -> String {
2318        let file = File::open(file_name).unwrap();
2319        let schema = Schema::new(vec![
2320            Field::new("c_int", DataType::UInt64, false),
2321            Field::new("c_float", DataType::Float32, false),
2322            Field::new("c_string", DataType::Utf8, false),
2323            Field::new("c_bool", DataType::Boolean, false),
2324        ]);
2325
2326        let builder = ReaderBuilder::new(Arc::new(schema))
2327            .with_header(true)
2328            .with_delimiter(b'|')
2329            .with_batch_size(512)
2330            .with_projection(vec![0, 1, 2, 3]);
2331
2332        let mut csv = builder.build(file).unwrap();
2333
2334        csv.next().unwrap().unwrap_err().to_string()
2335    }
2336
2337    #[test]
2338    fn test_parse_invalid_csv_float() {
2339        let file_name = "test/data/various_invalid_types/invalid_float.csv";
2340
2341        let error = invalid_csv_helper(file_name);
2342        assert_eq!(
2343            "Parser error: Error while parsing value '4.x4' as type 'Float32' for column 1 at line 4. Row data: '[4,4.x4,,false]'",
2344            error
2345        );
2346    }
2347
2348    #[test]
2349    fn test_parse_invalid_csv_int() {
2350        let file_name = "test/data/various_invalid_types/invalid_int.csv";
2351
2352        let error = invalid_csv_helper(file_name);
2353        assert_eq!(
2354            "Parser error: Error while parsing value '2.3' as type 'UInt64' for column 0 at line 2. Row data: '[2.3,2.2,2.22,false]'",
2355            error
2356        );
2357    }
2358
2359    #[test]
2360    fn test_parse_invalid_csv_bool() {
2361        let file_name = "test/data/various_invalid_types/invalid_bool.csv";
2362
2363        let error = invalid_csv_helper(file_name);
2364        assert_eq!(
2365            "Parser error: Error while parsing value 'none' as type 'Boolean' for column 3 at line 2. Row data: '[2,2.2,2.22,none]'",
2366            error
2367        );
2368    }
2369
2370    /// Infer the data type of a record
2371    fn infer_field_schema(string: &str) -> DataType {
2372        let mut v = InferredDataType::default();
2373        v.update(string);
2374        v.get()
2375    }
2376
2377    #[test]
2378    #[cfg_attr(miri, ignore)] // Takes too long
2379    fn test_infer_field_schema() {
2380        assert_eq!(infer_field_schema("A"), DataType::Utf8);
2381        assert_eq!(infer_field_schema("\"123\""), DataType::Utf8);
2382        assert_eq!(infer_field_schema("10"), DataType::Int64);
2383        assert_eq!(infer_field_schema("10.2"), DataType::Float64);
2384        assert_eq!(infer_field_schema(".2"), DataType::Float64);
2385        assert_eq!(infer_field_schema("2."), DataType::Float64);
2386        assert_eq!(infer_field_schema("NaN"), DataType::Float64);
2387        assert_eq!(infer_field_schema("nan"), DataType::Float64);
2388        assert_eq!(infer_field_schema("inf"), DataType::Float64);
2389        assert_eq!(infer_field_schema("-inf"), DataType::Float64);
2390        assert_eq!(infer_field_schema("true"), DataType::Boolean);
2391        assert_eq!(infer_field_schema("trUe"), DataType::Boolean);
2392        assert_eq!(infer_field_schema("false"), DataType::Boolean);
2393        assert_eq!(infer_field_schema("2020-11-08"), DataType::Date32);
2394        assert_eq!(
2395            infer_field_schema("2020-11-08T14:20:01"),
2396            DataType::Timestamp(TimeUnit::Second, None)
2397        );
2398        assert_eq!(
2399            infer_field_schema("2020-11-08 14:20:01"),
2400            DataType::Timestamp(TimeUnit::Second, None)
2401        );
2402        assert_eq!(
2403            infer_field_schema("2020-11-08 14:20:01"),
2404            DataType::Timestamp(TimeUnit::Second, None)
2405        );
2406        assert_eq!(infer_field_schema("-5.13"), DataType::Float64);
2407        assert_eq!(infer_field_schema("0.1300"), DataType::Float64);
2408        assert_eq!(
2409            infer_field_schema("2021-12-19 13:12:30.921"),
2410            DataType::Timestamp(TimeUnit::Millisecond, None)
2411        );
2412        assert_eq!(
2413            infer_field_schema("2021-12-19T13:12:30.123456789"),
2414            DataType::Timestamp(TimeUnit::Nanosecond, None)
2415        );
2416        assert_eq!(infer_field_schema("–9223372036854775809"), DataType::Utf8);
2417        assert_eq!(infer_field_schema("9223372036854775808"), DataType::Utf8);
2418    }
2419
2420    #[test]
2421    fn parse_date32() {
2422        assert_eq!(Date32Type::parse("1970-01-01").unwrap(), 0);
2423        assert_eq!(Date32Type::parse("2020-03-15").unwrap(), 18336);
2424        assert_eq!(Date32Type::parse("1945-05-08").unwrap(), -9004);
2425    }
2426
2427    #[test]
2428    fn parse_time() {
2429        assert_eq!(
2430            Time64NanosecondType::parse("12:10:01.123456789 AM"),
2431            Some(601_123_456_789)
2432        );
2433        assert_eq!(
2434            Time64MicrosecondType::parse("12:10:01.123456 am"),
2435            Some(601_123_456)
2436        );
2437        assert_eq!(
2438            Time32MillisecondType::parse("2:10:01.12 PM"),
2439            Some(51_001_120)
2440        );
2441        assert_eq!(Time32SecondType::parse("2:10:01 pm"), Some(51_001));
2442    }
2443
2444    #[test]
2445    fn parse_date64() {
2446        assert_eq!(Date64Type::parse("1970-01-01T00:00:00").unwrap(), 0);
2447        assert_eq!(
2448            Date64Type::parse("2018-11-13T17:11:10").unwrap(),
2449            1542129070000
2450        );
2451        assert_eq!(
2452            Date64Type::parse("2018-11-13T17:11:10.011").unwrap(),
2453            1542129070011
2454        );
2455        assert_eq!(
2456            Date64Type::parse("1900-02-28T12:34:56").unwrap(),
2457            -2203932304000
2458        );
2459        assert_eq!(
2460            Date64Type::parse_formatted("1900-02-28 12:34:56", "%Y-%m-%d %H:%M:%S").unwrap(),
2461            -2203932304000
2462        );
2463        assert_eq!(
2464            Date64Type::parse_formatted("1900-02-28 12:34:56+0030", "%Y-%m-%d %H:%M:%S%z").unwrap(),
2465            -2203932304000 - (30 * 60 * 1000)
2466        );
2467    }
2468
2469    fn test_parse_timestamp_impl<T: ArrowTimestampType>(
2470        timezone: Option<Arc<str>>,
2471        expected: &[i64],
2472    ) {
2473        let csv = [
2474            "1970-01-01T00:00:00",
2475            "1970-01-01T00:00:00Z",
2476            "1970-01-01T00:00:00+02:00",
2477        ]
2478        .join("\n");
2479        let schema = Arc::new(Schema::new(vec![Field::new(
2480            "field",
2481            DataType::Timestamp(T::UNIT, timezone.clone()),
2482            true,
2483        )]));
2484
2485        let mut decoder = ReaderBuilder::new(schema).build_decoder();
2486
2487        let decoded = decoder.decode(csv.as_bytes()).unwrap();
2488        assert_eq!(decoded, csv.len());
2489        decoder.decode(&[]).unwrap();
2490
2491        let batch = decoder.flush().unwrap().unwrap();
2492        assert_eq!(batch.num_columns(), 1);
2493        assert_eq!(batch.num_rows(), 3);
2494        let col = batch.column(0).as_primitive::<T>();
2495        assert_eq!(col.values(), expected);
2496        assert_eq!(col.data_type(), &DataType::Timestamp(T::UNIT, timezone));
2497    }
2498
2499    #[test]
2500    fn test_parse_timestamp() {
2501        test_parse_timestamp_impl::<TimestampNanosecondType>(None, &[0, 0, -7_200_000_000_000]);
2502        test_parse_timestamp_impl::<TimestampNanosecondType>(
2503            Some("+00:00".into()),
2504            &[0, 0, -7_200_000_000_000],
2505        );
2506        test_parse_timestamp_impl::<TimestampNanosecondType>(
2507            Some("-05:00".into()),
2508            &[18_000_000_000_000, 0, -7_200_000_000_000],
2509        );
2510        test_parse_timestamp_impl::<TimestampMicrosecondType>(
2511            Some("-03".into()),
2512            &[10_800_000_000, 0, -7_200_000_000],
2513        );
2514        test_parse_timestamp_impl::<TimestampMillisecondType>(
2515            Some("-03".into()),
2516            &[10_800_000, 0, -7_200_000],
2517        );
2518        test_parse_timestamp_impl::<TimestampSecondType>(Some("-03".into()), &[10_800, 0, -7_200]);
2519    }
2520
2521    // A local time in a DST gap or overlap of the target timezone is resolved
2522    // as described in `arrow_cast::parse::string_to_datetime`: an ambiguous
2523    // reading takes the later instant, and a nonexistent reading is shifted
2524    // forward by the length of the gap.
2525    #[test]
2526    fn test_parse_timestamp_dst_transitions() {
2527        let csv = [
2528            // Nonexistent: 02:30 is skipped, read as 03:30 EDT
2529            "2024-03-10 02:30:00",
2530            // Ambiguous: 01:30 happens twice, the later (EST) instant wins
2531            "2024-11-03 01:30:00",
2532        ]
2533        .join("\n");
2534        let timezone: Option<Arc<str>> = Some("America/New_York".into());
2535        let schema = Arc::new(Schema::new(vec![Field::new(
2536            "field",
2537            DataType::Timestamp(TimeUnit::Second, timezone.clone()),
2538            true,
2539        )]));
2540
2541        let mut decoder = ReaderBuilder::new(schema).build_decoder();
2542        let decoded = decoder.decode(csv.as_bytes()).unwrap();
2543        assert_eq!(decoded, csv.len());
2544        decoder.decode(&[]).unwrap();
2545
2546        let batch = decoder.flush().unwrap().unwrap();
2547        let col = batch.column(0).as_primitive::<TimestampSecondType>();
2548        // 2024-03-10T07:30:00Z and 2024-11-03T06:30:00Z
2549        assert_eq!(col.values(), &[1_710_055_800, 1_730_615_400]);
2550        assert_eq!(
2551            col.data_type(),
2552            &DataType::Timestamp(TimeUnit::Second, timezone)
2553        );
2554    }
2555
2556    #[test]
2557    #[cfg_attr(miri, ignore)] // Takes too long
2558    fn test_infer_schema_from_multiple_files() {
2559        let mut csv1 = NamedTempFile::new().unwrap();
2560        let mut csv2 = NamedTempFile::new().unwrap();
2561        let csv3 = NamedTempFile::new().unwrap(); // empty csv file should be skipped
2562        let mut csv4 = NamedTempFile::new().unwrap();
2563        writeln!(csv1, "c1,c2,c3").unwrap();
2564        writeln!(csv1, "1,\"foo\",0.5").unwrap();
2565        writeln!(csv1, "3,\"bar\",1").unwrap();
2566        writeln!(csv1, "3,\"bar\",2e-06").unwrap();
2567        // reading csv2 will set c2 to optional
2568        writeln!(csv2, "c1,c2,c3,c4").unwrap();
2569        writeln!(csv2, "10,,3.14,true").unwrap();
2570        // reading csv4 will set c3 to optional
2571        writeln!(csv4, "c1,c2,c3").unwrap();
2572        writeln!(csv4, "10,\"foo\",").unwrap();
2573
2574        let schema = infer_schema_from_files(
2575            &[
2576                csv3.path().to_str().unwrap().to_string(),
2577                csv1.path().to_str().unwrap().to_string(),
2578                csv2.path().to_str().unwrap().to_string(),
2579                csv4.path().to_str().unwrap().to_string(),
2580            ],
2581            b',',
2582            Some(4), // only csv1 and csv2 should be read
2583            true,
2584        )
2585        .unwrap();
2586
2587        assert_eq!(schema.fields().len(), 4);
2588        assert!(schema.field(0).is_nullable());
2589        assert!(schema.field(1).is_nullable());
2590        assert!(schema.field(2).is_nullable());
2591        assert!(schema.field(3).is_nullable());
2592
2593        assert_eq!(&DataType::Int64, schema.field(0).data_type());
2594        assert_eq!(&DataType::Utf8, schema.field(1).data_type());
2595        assert_eq!(&DataType::Float64, schema.field(2).data_type());
2596        assert_eq!(&DataType::Boolean, schema.field(3).data_type());
2597    }
2598
2599    #[test]
2600    fn test_bounded() {
2601        let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2602        let data = [
2603            vec!["0"],
2604            vec!["1"],
2605            vec!["2"],
2606            vec!["3"],
2607            vec!["4"],
2608            vec!["5"],
2609            vec!["6"],
2610        ];
2611
2612        let data = data
2613            .iter()
2614            .map(|x| x.join(","))
2615            .collect::<Vec<_>>()
2616            .join("\n");
2617        let data = data.as_bytes();
2618
2619        let reader = std::io::Cursor::new(data);
2620
2621        let mut csv = ReaderBuilder::new(Arc::new(schema))
2622            .with_batch_size(2)
2623            .with_projection(vec![0])
2624            .with_bounds(2, 6)
2625            .build_buffered(reader)
2626            .unwrap();
2627
2628        let batch = csv.next().unwrap().unwrap();
2629        let a = batch.column(0);
2630        let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2631        assert_eq!(a, &UInt32Array::from(vec![2, 3]));
2632
2633        let batch = csv.next().unwrap().unwrap();
2634        let a = batch.column(0);
2635        let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2636        assert_eq!(a, &UInt32Array::from(vec![4, 5]));
2637
2638        assert!(csv.next().is_none());
2639    }
2640
2641    #[test]
2642    fn test_empty_projection() {
2643        let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2644        let data = [vec!["0"], vec!["1"]];
2645
2646        let data = data
2647            .iter()
2648            .map(|x| x.join(","))
2649            .collect::<Vec<_>>()
2650            .join("\n");
2651
2652        let mut csv = ReaderBuilder::new(Arc::new(schema))
2653            .with_batch_size(2)
2654            .with_projection(vec![])
2655            .build_buffered(Cursor::new(data.as_bytes()))
2656            .unwrap();
2657
2658        let batch = csv.next().unwrap().unwrap();
2659        assert_eq!(batch.columns().len(), 0);
2660        assert_eq!(batch.num_rows(), 2);
2661
2662        assert!(csv.next().is_none());
2663    }
2664
2665    #[test]
2666    fn test_parsing_bool() {
2667        // Encode the expected behavior of boolean parsing
2668        assert_eq!(Some(true), parse_bool("true"));
2669        assert_eq!(Some(true), parse_bool("tRUe"));
2670        assert_eq!(Some(true), parse_bool("True"));
2671        assert_eq!(Some(true), parse_bool("TRUE"));
2672        assert_eq!(None, parse_bool("t"));
2673        assert_eq!(None, parse_bool("T"));
2674        assert_eq!(None, parse_bool(""));
2675
2676        assert_eq!(Some(false), parse_bool("false"));
2677        assert_eq!(Some(false), parse_bool("fALse"));
2678        assert_eq!(Some(false), parse_bool("False"));
2679        assert_eq!(Some(false), parse_bool("FALSE"));
2680        assert_eq!(None, parse_bool("f"));
2681        assert_eq!(None, parse_bool("F"));
2682        assert_eq!(None, parse_bool(""));
2683    }
2684
2685    #[test]
2686    fn test_parsing_float() {
2687        assert_eq!(Some(12.34), Float64Type::parse("12.34"));
2688        assert_eq!(Some(-12.34), Float64Type::parse("-12.34"));
2689        assert_eq!(Some(12.0), Float64Type::parse("12"));
2690        assert_eq!(Some(0.0), Float64Type::parse("0"));
2691        assert_eq!(Some(2.0), Float64Type::parse("2."));
2692        assert_eq!(Some(0.2), Float64Type::parse(".2"));
2693        assert!(Float64Type::parse("nan").unwrap().is_nan());
2694        assert!(Float64Type::parse("NaN").unwrap().is_nan());
2695        assert!(Float64Type::parse("inf").unwrap().is_infinite());
2696        assert!(Float64Type::parse("inf").unwrap().is_sign_positive());
2697        assert!(Float64Type::parse("-inf").unwrap().is_infinite());
2698        assert!(Float64Type::parse("-inf").unwrap().is_sign_negative());
2699        assert_eq!(None, Float64Type::parse(""));
2700        assert_eq!(None, Float64Type::parse("dd"));
2701        assert_eq!(None, Float64Type::parse("12.34.56"));
2702    }
2703
2704    #[test]
2705    fn test_non_std_quote() {
2706        let schema = Schema::new(vec![
2707            Field::new("text1", DataType::Utf8, false),
2708            Field::new("text2", DataType::Utf8, false),
2709        ]);
2710        let builder = ReaderBuilder::new(Arc::new(schema))
2711            .with_header(false)
2712            .with_quote(b'~'); // default is ", change to ~
2713
2714        let mut csv_text = Vec::new();
2715        let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2716        for index in 0..10 {
2717            let text1 = format!("id{index:}");
2718            let text2 = format!("value{index:}");
2719            csv_writer
2720                .write_fmt(format_args!("~{text1}~,~{text2}~\r\n"))
2721                .unwrap();
2722        }
2723        let mut csv_reader = std::io::Cursor::new(&csv_text);
2724        let mut reader = builder.build(&mut csv_reader).unwrap();
2725        let batch = reader.next().unwrap().unwrap();
2726        let col0 = batch.column(0);
2727        assert_eq!(col0.len(), 10);
2728        let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2729        assert_eq!(col0_arr.value(0), "id0");
2730        let col1 = batch.column(1);
2731        assert_eq!(col1.len(), 10);
2732        let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2733        assert_eq!(col1_arr.value(5), "value5");
2734    }
2735
2736    #[test]
2737    fn test_non_std_escape() {
2738        let schema = Schema::new(vec![
2739            Field::new("text1", DataType::Utf8, false),
2740            Field::new("text2", DataType::Utf8, false),
2741        ]);
2742        let builder = ReaderBuilder::new(Arc::new(schema))
2743            .with_header(false)
2744            .with_escape(b'\\'); // default is None, change to \
2745
2746        let mut csv_text = Vec::new();
2747        let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2748        for index in 0..10 {
2749            let text1 = format!("id{index:}");
2750            let text2 = format!("value\\\"{index:}");
2751            csv_writer
2752                .write_fmt(format_args!("\"{text1}\",\"{text2}\"\r\n"))
2753                .unwrap();
2754        }
2755        let mut csv_reader = std::io::Cursor::new(&csv_text);
2756        let mut reader = builder.build(&mut csv_reader).unwrap();
2757        let batch = reader.next().unwrap().unwrap();
2758        let col0 = batch.column(0);
2759        assert_eq!(col0.len(), 10);
2760        let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2761        assert_eq!(col0_arr.value(0), "id0");
2762        let col1 = batch.column(1);
2763        assert_eq!(col1.len(), 10);
2764        let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2765        assert_eq!(col1_arr.value(5), "value\"5");
2766    }
2767
2768    #[test]
2769    fn test_non_std_terminator() {
2770        let schema = Schema::new(vec![
2771            Field::new("text1", DataType::Utf8, false),
2772            Field::new("text2", DataType::Utf8, false),
2773        ]);
2774        let builder = ReaderBuilder::new(Arc::new(schema))
2775            .with_header(false)
2776            .with_terminator(b'\n'); // default is CRLF, change to LF
2777
2778        let mut csv_text = Vec::new();
2779        let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2780        for index in 0..10 {
2781            let text1 = format!("id{index:}");
2782            let text2 = format!("value{index:}");
2783            csv_writer
2784                .write_fmt(format_args!("\"{text1}\",\"{text2}\"\n"))
2785                .unwrap();
2786        }
2787        let mut csv_reader = std::io::Cursor::new(&csv_text);
2788        let mut reader = builder.build(&mut csv_reader).unwrap();
2789        let batch = reader.next().unwrap().unwrap();
2790        let col0 = batch.column(0);
2791        assert_eq!(col0.len(), 10);
2792        let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2793        assert_eq!(col0_arr.value(0), "id0");
2794        let col1 = batch.column(1);
2795        assert_eq!(col1.len(), 10);
2796        let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2797        assert_eq!(col1_arr.value(5), "value5");
2798    }
2799
2800    #[test]
2801    fn test_header_bounds() {
2802        let csv = "a,b\na,b\na,b\na,b\na,b\n";
2803        let tests = [
2804            (None, false, 5),
2805            (None, true, 4),
2806            (Some((0, 4)), false, 4),
2807            (Some((1, 4)), false, 3),
2808            (Some((0, 4)), true, 4),
2809            (Some((1, 4)), true, 3),
2810        ];
2811        let schema = Arc::new(Schema::new(vec![
2812            Field::new("a", DataType::Utf8, false),
2813            Field::new("a", DataType::Utf8, false),
2814        ]));
2815
2816        for (idx, (bounds, has_header, expected)) in tests.into_iter().enumerate() {
2817            let mut reader = ReaderBuilder::new(schema.clone()).with_header(has_header);
2818            if let Some((start, end)) = bounds {
2819                reader = reader.with_bounds(start, end);
2820            }
2821            let b = reader
2822                .build_buffered(Cursor::new(csv.as_bytes()))
2823                .unwrap()
2824                .next()
2825                .unwrap()
2826                .unwrap();
2827            assert_eq!(b.num_rows(), expected, "{idx}");
2828        }
2829    }
2830
2831    #[test]
2832    fn test_header_validation() {
2833        let schema = Arc::new(Schema::new(vec![
2834            Field::new("a", DataType::Int32, false),
2835            Field::new("b", DataType::Int32, false),
2836        ]));
2837
2838        let csv = "a,c\n1,2\n";
2839        let err = ReaderBuilder::new(schema.clone())
2840            .with_header(true)
2841            .with_header_validation(true)
2842            .build_buffered(Cursor::new(csv.as_bytes()))
2843            .unwrap()
2844            .next()
2845            .unwrap()
2846            .unwrap_err()
2847            .to_string();
2848        assert_eq!(
2849            err,
2850            "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"c\""
2851        );
2852
2853        let batch = ReaderBuilder::new(schema)
2854            .with_header(true)
2855            .with_header_validation(false)
2856            .build_buffered(Cursor::new(csv.as_bytes()))
2857            .unwrap()
2858            .next()
2859            .unwrap()
2860            .unwrap();
2861        assert_eq!(batch.num_rows(), 1);
2862    }
2863
2864    #[test]
2865    fn test_header_validation_with_buffered_reader() {
2866        let schema = Arc::new(Schema::new(vec![
2867            Field::new("a", DataType::Int32, false),
2868            Field::new("b", DataType::Int32, false),
2869        ]));
2870
2871        let csv = "a,b\n1,2\n";
2872        let buffered = std::io::BufReader::with_capacity(1, Cursor::new(csv.as_bytes()));
2873        let batch = ReaderBuilder::new(schema)
2874            .with_header(true)
2875            .with_header_validation(true)
2876            .build_buffered(buffered)
2877            .unwrap()
2878            .next()
2879            .unwrap()
2880            .unwrap();
2881
2882        assert_eq!(batch.num_rows(), 1);
2883        let a = batch.column(0).as_primitive::<Int32Type>();
2884        assert_eq!(a.value(0), 1);
2885    }
2886
2887    #[test]
2888    fn test_header_validation_with_truncated_rows() {
2889        let schema = Arc::new(Schema::new(vec![
2890            Field::new("a", DataType::Int32, true),
2891            Field::new("b", DataType::Int32, true),
2892        ]));
2893
2894        let csv = "a\n1\n";
2895        let err = ReaderBuilder::new(schema.clone())
2896            .with_header(true)
2897            .with_header_validation(true)
2898            .with_truncated_rows(true)
2899            .build_buffered(Cursor::new(csv.as_bytes()))
2900            .unwrap()
2901            .next()
2902            .unwrap()
2903            .unwrap_err()
2904            .to_string();
2905        assert_eq!(
2906            err,
2907            "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"\"",
2908        )
2909    }
2910
2911    #[test]
2912    fn test_null_boolean() {
2913        let csv = "true,false\nFalse,True\n,True\nFalse,";
2914        let schema = Arc::new(Schema::new(vec![
2915            Field::new("a", DataType::Boolean, true),
2916            Field::new("a", DataType::Boolean, true),
2917        ]));
2918
2919        let b = ReaderBuilder::new(schema)
2920            .build_buffered(Cursor::new(csv.as_bytes()))
2921            .unwrap()
2922            .next()
2923            .unwrap()
2924            .unwrap();
2925
2926        assert_eq!(b.num_rows(), 4);
2927        assert_eq!(b.num_columns(), 2);
2928
2929        let c = b.column(0).as_boolean();
2930        assert_eq!(c.null_count(), 1);
2931        assert!(c.value(0));
2932        assert!(!c.value(1));
2933        assert!(c.is_null(2));
2934        assert!(!c.value(3));
2935
2936        let c = b.column(1).as_boolean();
2937        assert_eq!(c.null_count(), 1);
2938        assert!(!c.value(0));
2939        assert!(c.value(1));
2940        assert!(c.value(2));
2941        assert!(c.is_null(3));
2942    }
2943
2944    #[test]
2945    fn test_truncated_rows() {
2946        let data = "a,b,c\n1,2,3\n4,5\n\n6,7,8";
2947        let schema = Arc::new(Schema::new(vec![
2948            Field::new("a", DataType::Int32, true),
2949            Field::new("b", DataType::Int32, true),
2950            Field::new("c", DataType::Int32, true),
2951        ]));
2952
2953        let reader = ReaderBuilder::new(schema.clone())
2954            .with_header(true)
2955            .with_truncated_rows(true)
2956            .build(Cursor::new(data))
2957            .unwrap();
2958
2959        let batches = reader.collect::<Result<Vec<_>, _>>();
2960        assert!(batches.is_ok());
2961        let batch = batches.unwrap().into_iter().next().unwrap();
2962        // Empty rows are skipped by the underlying csv parser
2963        assert_eq!(batch.num_rows(), 3);
2964
2965        let reader = ReaderBuilder::new(schema.clone())
2966            .with_header(true)
2967            .with_truncated_rows(false)
2968            .build(Cursor::new(data))
2969            .unwrap();
2970
2971        let batches = reader.collect::<Result<Vec<_>, _>>();
2972        assert!(match batches {
2973            Err(ArrowError::CsvError(e)) => e.contains("incorrect number of fields"),
2974            _ => false,
2975        });
2976    }
2977
2978    #[test]
2979    fn test_truncated_rows_csv() {
2980        let file = File::open("test/data/truncated_rows.csv").unwrap();
2981        let schema = Arc::new(Schema::new(vec![
2982            Field::new("Name", DataType::Utf8, true),
2983            Field::new("Age", DataType::UInt32, true),
2984            Field::new("Occupation", DataType::Utf8, true),
2985            Field::new("DOB", DataType::Date32, true),
2986        ]));
2987        let reader = ReaderBuilder::new(schema.clone())
2988            .with_header(true)
2989            .with_batch_size(24)
2990            .with_truncated_rows(true);
2991        let csv = reader.build(file).unwrap();
2992        let batches = csv.collect::<Result<Vec<_>, _>>().unwrap();
2993
2994        assert_eq!(batches.len(), 1);
2995        let batch = &batches[0];
2996        assert_eq!(batch.num_rows(), 6);
2997        assert_eq!(batch.num_columns(), 4);
2998        let name = batch
2999            .column(0)
3000            .as_any()
3001            .downcast_ref::<StringArray>()
3002            .unwrap();
3003        let age = batch
3004            .column(1)
3005            .as_any()
3006            .downcast_ref::<UInt32Array>()
3007            .unwrap();
3008        let occupation = batch
3009            .column(2)
3010            .as_any()
3011            .downcast_ref::<StringArray>()
3012            .unwrap();
3013        let dob = batch
3014            .column(3)
3015            .as_any()
3016            .downcast_ref::<Date32Array>()
3017            .unwrap();
3018
3019        assert_eq!(name.value(0), "A1");
3020        assert_eq!(name.value(1), "B2");
3021        assert!(name.is_null(2));
3022        assert_eq!(name.value(3), "C3");
3023        assert_eq!(name.value(4), "D4");
3024        assert_eq!(name.value(5), "E5");
3025
3026        assert_eq!(age.value(0), 34);
3027        assert_eq!(age.value(1), 29);
3028        assert!(age.is_null(2));
3029        assert_eq!(age.value(3), 45);
3030        assert!(age.is_null(4));
3031        assert_eq!(age.value(5), 31);
3032
3033        assert_eq!(occupation.value(0), "Engineer");
3034        assert_eq!(occupation.value(1), "Doctor");
3035        assert!(occupation.is_null(2));
3036        assert_eq!(occupation.value(3), "Artist");
3037        assert!(occupation.is_null(4));
3038        assert!(occupation.is_null(5));
3039
3040        assert_eq!(dob.value(0), 5675);
3041        assert!(dob.is_null(1));
3042        assert!(dob.is_null(2));
3043        assert_eq!(dob.value(3), -1858);
3044        assert!(dob.is_null(4));
3045        assert!(dob.is_null(5));
3046    }
3047
3048    /// Schema used by the `truncated_row_count` tests below
3049    fn truncated_row_count_schema() -> SchemaRef {
3050        Arc::new(Schema::new(vec![
3051            Field::new("name", DataType::Utf8, true),
3052            Field::new("age", DataType::Int32, true),
3053            Field::new("city", DataType::Utf8, true),
3054        ]))
3055    }
3056
3057    #[test]
3058    fn test_truncated_row_count_counts_padded_rows() {
3059        let data = "name,age,city\nAlice,25,Rome\nBob,30\n";
3060
3061        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3062            .with_header(true)
3063            .with_truncated_rows(true)
3064            .build(Cursor::new(data))
3065            .unwrap();
3066
3067        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3068        assert_eq!(batches[0].num_rows(), 2);
3069        assert_eq!(reader.truncated_row_count(), 1);
3070    }
3071
3072    #[test]
3073    fn test_truncated_row_count_ignores_empty_trailing_field() {
3074        // "Carol,35," has all three fields, the last one just happens to be empty, so it
3075        // parses to the same null as a padded row would. The count must not be inferred
3076        // from the nulls in the batch
3077        let data = "name,age,city\nAlice,25,Rome\nCarol,35,\n";
3078
3079        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3080            .with_header(true)
3081            .with_truncated_rows(true)
3082            .build(Cursor::new(data))
3083            .unwrap();
3084
3085        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3086        let batch = &batches[0];
3087        assert_eq!(batch.num_rows(), 2);
3088        assert!(batch.column(2).is_null(1));
3089        assert_eq!(reader.truncated_row_count(), 0);
3090    }
3091
3092    #[test]
3093    fn test_truncated_row_count_clean_file() {
3094        let data = "name,age,city\nAlice,25,Rome\nBob,30,Milan\n";
3095
3096        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3097            .with_header(true)
3098            .with_truncated_rows(true)
3099            .build(Cursor::new(data))
3100            .unwrap();
3101
3102        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3103        assert_eq!(batches[0].num_rows(), 2);
3104        assert_eq!(reader.truncated_row_count(), 0);
3105    }
3106
3107    #[test]
3108    fn test_truncated_row_count_without_truncated_rows() {
3109        let data = "name,age,city\nAlice,25,Rome\nBob,30\n";
3110
3111        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3112            .with_header(true)
3113            .with_truncated_rows(false)
3114            .build(Cursor::new(data))
3115            .unwrap();
3116
3117        // The short row is an error rather than something to count
3118        let err = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap_err();
3119        assert!(
3120            err.to_string().contains("incorrect number of fields"),
3121            "{err}"
3122        );
3123        assert_eq!(reader.truncated_row_count(), 0);
3124    }
3125
3126    #[test]
3127    fn test_truncated_row_count_accumulates_across_batches() {
3128        // Six short rows read two at a time
3129        let data = "name,age,city\nn0,0\nn1,1\nn2,2\nn3,3\nn4,4\nn5,5\n";
3130
3131        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3132            .with_header(true)
3133            .with_truncated_rows(true)
3134            .with_batch_size(2)
3135            .build(Cursor::new(data))
3136            .unwrap();
3137
3138        let mut running = vec![];
3139        while let Some(batch) = reader.next().transpose().unwrap() {
3140            assert_eq!(batch.num_rows(), 2);
3141            running.push(reader.truncated_row_count());
3142        }
3143
3144        // A running total, not a per batch count
3145        assert_eq!(running, vec![2, 4, 6]);
3146        assert_eq!(reader.truncated_row_count(), 6);
3147    }
3148
3149    #[test]
3150    fn test_truncated_row_count_excludes_skipped_rows() {
3151        // The header is one field short of the schema, so skipping it pads it. Skipped
3152        // rows never reach a batch and must not be counted
3153        let data = "name,age\nAlice,25,Rome\nBob,30,Milan\n";
3154
3155        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3156            .with_header(true)
3157            .with_truncated_rows(true)
3158            .build(Cursor::new(data))
3159            .unwrap();
3160
3161        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3162        assert_eq!(batches[0].num_rows(), 2);
3163        assert_eq!(reader.truncated_row_count(), 0);
3164    }
3165
3166    #[test]
3167    fn test_truncated_row_count_on_decoder() {
3168        let data = "1,2\n3\n";
3169        let schema = Arc::new(Schema::new(vec![
3170            Field::new("a", DataType::Int32, true),
3171            Field::new("b", DataType::Int32, true),
3172        ]));
3173
3174        let mut decoder = ReaderBuilder::new(schema)
3175            .with_truncated_rows(true)
3176            .build_decoder();
3177
3178        assert_eq!(decoder.truncated_row_count(), 0);
3179        let decoded = decoder.decode(data.as_bytes()).unwrap();
3180        assert_eq!(decoded, data.len());
3181        decoder.flush().unwrap().unwrap();
3182        assert_eq!(decoder.truncated_row_count(), 1);
3183    }
3184
3185    #[test]
3186    fn test_truncated_rows_not_nullable_error() {
3187        let data = "a,b,c\n1,2,3\n4,5";
3188        let schema = Arc::new(Schema::new(vec![
3189            Field::new("a", DataType::Int32, false),
3190            Field::new("b", DataType::Int32, false),
3191            Field::new("c", DataType::Int32, false),
3192        ]));
3193
3194        let reader = ReaderBuilder::new(schema.clone())
3195            .with_header(true)
3196            .with_truncated_rows(true)
3197            .build(Cursor::new(data))
3198            .unwrap();
3199
3200        let batches = reader.collect::<Result<Vec<_>, _>>();
3201        assert!(match batches {
3202            Err(ArrowError::InvalidArgumentError(e)) => e.contains("contains null values"),
3203            _ => false,
3204        });
3205    }
3206
3207    #[test]
3208    #[cfg_attr(miri, ignore)] // Takes too long
3209    fn test_buffered() {
3210        let tests = [
3211            ("test/data/uk_cities.csv", false, 37),
3212            ("test/data/various_types.csv", true, 10),
3213            ("test/data/decimal_test.csv", false, 10),
3214        ];
3215
3216        for (path, has_header, expected_rows) in tests {
3217            let (schema, _) = Format::default()
3218                .infer_schema(File::open(path).unwrap(), None)
3219                .unwrap();
3220            let schema = Arc::new(schema);
3221
3222            for batch_size in [1, 4] {
3223                for capacity in [1, 3, 7, 100] {
3224                    let reader = ReaderBuilder::new(schema.clone())
3225                        .with_batch_size(batch_size)
3226                        .with_header(has_header)
3227                        .build(File::open(path).unwrap())
3228                        .unwrap();
3229
3230                    let expected = reader.collect::<Result<Vec<_>, _>>().unwrap();
3231
3232                    assert_eq!(
3233                        expected.iter().map(|x| x.num_rows()).sum::<usize>(),
3234                        expected_rows
3235                    );
3236
3237                    let buffered =
3238                        std::io::BufReader::with_capacity(capacity, File::open(path).unwrap());
3239
3240                    let reader = ReaderBuilder::new(schema.clone())
3241                        .with_batch_size(batch_size)
3242                        .with_header(has_header)
3243                        .build_buffered(buffered)
3244                        .unwrap();
3245
3246                    let actual = reader.collect::<Result<Vec<_>, _>>().unwrap();
3247                    assert_eq!(expected, actual)
3248                }
3249            }
3250        }
3251    }
3252
3253    fn err_test(csv: &[u8], expected: &str) {
3254        fn err_test_with_schema(csv: &[u8], expected: &str, schema: Arc<Schema>) {
3255            let buffer = std::io::BufReader::with_capacity(2, Cursor::new(csv));
3256            let b = ReaderBuilder::new(schema)
3257                .with_batch_size(2)
3258                .build_buffered(buffer)
3259                .unwrap();
3260            let err = b.collect::<Result<Vec<_>, _>>().unwrap_err().to_string();
3261            assert_eq!(err, expected)
3262        }
3263
3264        let schema_utf8 = Arc::new(Schema::new(vec![
3265            Field::new("text1", DataType::Utf8, true),
3266            Field::new("text2", DataType::Utf8, true),
3267        ]));
3268        err_test_with_schema(csv, expected, schema_utf8);
3269
3270        let schema_utf8view = Arc::new(Schema::new(vec![
3271            Field::new("text1", DataType::Utf8View, true),
3272            Field::new("text2", DataType::Utf8View, true),
3273        ]));
3274        err_test_with_schema(csv, expected, schema_utf8view);
3275    }
3276
3277    #[test]
3278    fn test_invalid_utf8() {
3279        err_test(
3280            b"sdf,dsfg\ndfd,hgh\xFFue\n,sds\nFalhghse,",
3281            "Csv error: Encountered invalid UTF-8 data for line 2 and field 2",
3282        );
3283
3284        err_test(
3285            b"sdf,dsfg\ndksdk,jf\nd\xFFfd,hghue\n,sds\nFalhghse,",
3286            "Csv error: Encountered invalid UTF-8 data for line 3 and field 1",
3287        );
3288
3289        err_test(
3290            b"sdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
3291            "Csv error: Encountered invalid UTF-8 data for line 5 and field 2",
3292        );
3293
3294        err_test(
3295            b"\xFFsdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
3296            "Csv error: Encountered invalid UTF-8 data for line 1 and field 1",
3297        );
3298    }
3299
3300    struct InstrumentedRead<R> {
3301        r: R,
3302        fill_count: usize,
3303        fill_sizes: Vec<usize>,
3304    }
3305
3306    impl<R> InstrumentedRead<R> {
3307        fn new(r: R) -> Self {
3308            Self {
3309                r,
3310                fill_count: 0,
3311                fill_sizes: vec![],
3312            }
3313        }
3314    }
3315
3316    impl<R: Seek> Seek for InstrumentedRead<R> {
3317        fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
3318            self.r.seek(pos)
3319        }
3320    }
3321
3322    impl<R: BufRead> Read for InstrumentedRead<R> {
3323        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
3324            self.r.read(buf)
3325        }
3326    }
3327
3328    impl<R: BufRead> BufRead for InstrumentedRead<R> {
3329        fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
3330            self.fill_count += 1;
3331            let buf = self.r.fill_buf()?;
3332            self.fill_sizes.push(buf.len());
3333            Ok(buf)
3334        }
3335
3336        fn consume(&mut self, amt: usize) {
3337            self.r.consume(amt)
3338        }
3339    }
3340
3341    #[test]
3342    fn test_io() {
3343        let schema = Arc::new(Schema::new(vec![
3344            Field::new("a", DataType::Utf8, false),
3345            Field::new("b", DataType::Utf8, false),
3346        ]));
3347        let csv = "foo,bar\nbaz,foo\na,b\nc,d";
3348        let mut read = InstrumentedRead::new(Cursor::new(csv.as_bytes()));
3349        let reader = ReaderBuilder::new(schema)
3350            .with_batch_size(3)
3351            .build_buffered(&mut read)
3352            .unwrap();
3353
3354        let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3355        assert_eq!(batches.len(), 2);
3356        assert_eq!(batches[0].num_rows(), 3);
3357        assert_eq!(batches[1].num_rows(), 1);
3358
3359        // Expect 4 calls to fill_buf
3360        // 1. Read first 3 rows
3361        // 2. Read final row
3362        // 3. Delimit and flush final row
3363        // 4. Iterator finished
3364        assert_eq!(&read.fill_sizes, &[23, 3, 0, 0]);
3365        assert_eq!(read.fill_count, 4);
3366    }
3367
3368    #[test]
3369    #[cfg_attr(miri, ignore)] // Takes too long
3370    fn test_inference() {
3371        let cases: &[(&[&str], DataType)] = &[
3372            (&[], DataType::Null),
3373            (&["false", "12"], DataType::Utf8),
3374            (&["12", "cupcakes"], DataType::Utf8),
3375            (&["12", "12.4"], DataType::Float64),
3376            (&["14050", "24332"], DataType::Int64),
3377            (&["14050.0", "true"], DataType::Utf8),
3378            (&["14050", "2020-03-19 00:00:00"], DataType::Utf8),
3379            (&["14050", "2340.0", "2020-03-19 00:00:00"], DataType::Utf8),
3380            (
3381                &["2020-03-19 02:00:00", "2020-03-19 00:00:00"],
3382                DataType::Timestamp(TimeUnit::Second, None),
3383            ),
3384            (&["2020-03-19", "2020-03-20"], DataType::Date32),
3385            (
3386                &["2020-03-19", "2020-03-19 02:00:00", "2020-03-19 00:00:00"],
3387                DataType::Timestamp(TimeUnit::Second, None),
3388            ),
3389            (
3390                &[
3391                    "2020-03-19",
3392                    "2020-03-19 02:00:00",
3393                    "2020-03-19 00:00:00.000",
3394                ],
3395                DataType::Timestamp(TimeUnit::Millisecond, None),
3396            ),
3397            (
3398                &[
3399                    "2020-03-19",
3400                    "2020-03-19 02:00:00",
3401                    "2020-03-19 00:00:00.000000",
3402                ],
3403                DataType::Timestamp(TimeUnit::Microsecond, None),
3404            ),
3405            (
3406                &["2020-03-19 02:00:00+02:00", "2020-03-19 02:00:00Z"],
3407                DataType::Timestamp(TimeUnit::Second, None),
3408            ),
3409            (
3410                &[
3411                    "2020-03-19",
3412                    "2020-03-19 02:00:00+02:00",
3413                    "2020-03-19 02:00:00Z",
3414                    "2020-03-19 02:00:00.12Z",
3415                ],
3416                DataType::Timestamp(TimeUnit::Millisecond, None),
3417            ),
3418            (
3419                &[
3420                    "2020-03-19",
3421                    "2020-03-19 02:00:00.000000000",
3422                    "2020-03-19 00:00:00.000000",
3423                ],
3424                DataType::Timestamp(TimeUnit::Nanosecond, None),
3425            ),
3426        ];
3427
3428        for (values, expected) in cases {
3429            let mut t = InferredDataType::default();
3430            for v in *values {
3431                t.update(v)
3432            }
3433            assert_eq!(&t.get(), expected, "{values:?}")
3434        }
3435    }
3436
3437    #[test]
3438    #[cfg_attr(miri, ignore)] // Takes too long
3439    fn test_record_length_mismatch() {
3440        let csv = "\
3441        a,b,c\n\
3442        1,2,3\n\
3443        4,5\n\
3444        6,7,8";
3445        let mut read = Cursor::new(csv.as_bytes());
3446        let result = Format::default()
3447            .with_header(true)
3448            .infer_schema(&mut read, None);
3449        assert!(result.is_err());
3450        // Include line number in the error message to help locate and fix the issue
3451        assert_eq!(
3452            result.err().unwrap().to_string(),
3453            "Csv error: Encountered unequal lengths between records on CSV file. Expected 3 records, found 2 records at line 3"
3454        );
3455    }
3456
3457    #[test]
3458    fn test_comment() {
3459        let schema = Schema::new(vec![
3460            Field::new("a", DataType::Int8, false),
3461            Field::new("b", DataType::Int8, false),
3462        ]);
3463
3464        let csv = "# comment1 \n1,2\n#comment2\n11,22";
3465        let mut read = Cursor::new(csv.as_bytes());
3466        let reader = ReaderBuilder::new(Arc::new(schema))
3467            .with_comment(b'#')
3468            .build(&mut read)
3469            .unwrap();
3470
3471        let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3472        assert_eq!(batches.len(), 1);
3473        let b = batches.first().unwrap();
3474        assert_eq!(b.num_columns(), 2);
3475        assert_eq!(
3476            b.column(0)
3477                .as_any()
3478                .downcast_ref::<Int8Array>()
3479                .unwrap()
3480                .values(),
3481            &vec![1, 11]
3482        );
3483        assert_eq!(
3484            b.column(1)
3485                .as_any()
3486                .downcast_ref::<Int8Array>()
3487                .unwrap()
3488                .values(),
3489            &vec![2, 22]
3490        );
3491    }
3492
3493    #[test]
3494    fn test_parse_string_view_single_column() {
3495        let csv = ["foo", "something_cannot_be_inlined", "foobar"].join("\n");
3496        let schema = Arc::new(Schema::new(vec![Field::new(
3497            "c1",
3498            DataType::Utf8View,
3499            true,
3500        )]));
3501
3502        let mut decoder = ReaderBuilder::new(schema).build_decoder();
3503
3504        let decoded = decoder.decode(csv.as_bytes()).unwrap();
3505        assert_eq!(decoded, csv.len());
3506        decoder.decode(&[]).unwrap();
3507
3508        let batch = decoder.flush().unwrap().unwrap();
3509        assert_eq!(batch.num_columns(), 1);
3510        assert_eq!(batch.num_rows(), 3);
3511        let col = batch.column(0).as_string_view();
3512        assert_eq!(col.data_type(), &DataType::Utf8View);
3513        assert_eq!(col.value(0), "foo");
3514        assert_eq!(col.value(1), "something_cannot_be_inlined");
3515        assert_eq!(col.value(2), "foobar");
3516    }
3517
3518    #[test]
3519    fn test_parse_string_view_multi_column() {
3520        let csv = ["foo,", ",something_cannot_be_inlined", "foobarfoobar,bar"].join("\n");
3521        let schema = Arc::new(Schema::new(vec![
3522            Field::new("c1", DataType::Utf8View, true),
3523            Field::new("c2", DataType::Utf8View, true),
3524        ]));
3525
3526        let mut decoder = ReaderBuilder::new(schema).build_decoder();
3527
3528        let decoded = decoder.decode(csv.as_bytes()).unwrap();
3529        assert_eq!(decoded, csv.len());
3530        decoder.decode(&[]).unwrap();
3531
3532        let batch = decoder.flush().unwrap().unwrap();
3533        assert_eq!(batch.num_columns(), 2);
3534        assert_eq!(batch.num_rows(), 3);
3535        let c1 = batch.column(0).as_string_view();
3536        let c2 = batch.column(1).as_string_view();
3537        assert_eq!(c1.data_type(), &DataType::Utf8View);
3538        assert_eq!(c2.data_type(), &DataType::Utf8View);
3539
3540        assert!(!c1.is_null(0));
3541        assert!(c1.is_null(1));
3542        assert!(!c1.is_null(2));
3543        assert_eq!(c1.value(0), "foo");
3544        assert_eq!(c1.value(2), "foobarfoobar");
3545
3546        assert!(c2.is_null(0));
3547        assert!(!c2.is_null(1));
3548        assert!(!c2.is_null(2));
3549        assert_eq!(c2.value(1), "something_cannot_be_inlined");
3550        assert_eq!(c2.value(2), "bar");
3551    }
3552
3553    #[test]
3554    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
3555    fn test_float_precision() {
3556        let data = [
3557            "f16,f32,f64",
3558            "1.5,1.5,1.5",
3559            "0.25,0.25,0.25",
3560            "1.23456789,1.23456789,1.23456789",
3561            "1.234567890123456,1.234567890123456,1.234567890123456",
3562            "-2.5,-2.5,-2.5",
3563            "0,0,0",
3564            ",,",
3565        ]
3566        .join("\n");
3567
3568        let schema = Schema::new(vec![
3569            Field::new("f16", DataType::Float16, true),
3570            Field::new("f32", DataType::Float32, true),
3571            Field::new("f64", DataType::Float64, true),
3572        ]);
3573
3574        let mut reader = ReaderBuilder::new(Arc::new(schema))
3575            .with_header(true)
3576            .build(Cursor::new(data))
3577            .unwrap();
3578
3579        let batch = reader.next().unwrap().unwrap();
3580        assert_eq!(batch.num_rows(), 7);
3581
3582        let f16_col = batch.column(0).as_primitive::<Float16Type>();
3583        let f32_col = batch.column(1).as_primitive::<Float32Type>();
3584        let f64_col = batch.column(2).as_primitive::<Float64Type>();
3585
3586        assert_eq!(f16_col.value(0), half::f16::from_f32(1.5));
3587        assert_eq!(f32_col.value(0), 1.5f32);
3588        assert_eq!(f64_col.value(0), 1.5f64);
3589
3590        assert_eq!(f16_col.value(1), half::f16::from_f32(0.25));
3591        assert_eq!(f32_col.value(1), 0.25f32);
3592        assert_eq!(f64_col.value(1), 0.25f64);
3593
3594        assert_eq!(f16_col.value(2), half::f16::from_f32(1.234_567_9));
3595        assert_eq!(f32_col.value(2), 1.234_567_9_f32);
3596        assert_eq!(f64_col.value(2), 1.23456789f64);
3597
3598        assert_eq!(f16_col.value(3), half::f16::from_f64(1.234567890123456f64));
3599        assert_eq!(f32_col.value(3), 1.234_567_9_f32);
3600        assert_eq!(f64_col.value(3), 1.234567890123456f64);
3601
3602        assert_eq!(f16_col.value(4), half::f16::from_f32(-2.5));
3603        assert_eq!(f32_col.value(4), -2.5f32);
3604        assert_eq!(f64_col.value(4), -2.5f64);
3605
3606        assert_eq!(f16_col.value(5), half::f16::from_f32(0.0));
3607        assert_eq!(f32_col.value(5), 0.0f32);
3608        assert_eq!(f64_col.value(5), 0.0f64);
3609
3610        assert!(f16_col.is_null(6));
3611        assert!(f32_col.is_null(6));
3612        assert!(f64_col.is_null(6));
3613    }
3614}