Skip to main content

arrow_csv/reader/
mod.rs

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