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