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.iter() {
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    arrays.and_then(|arr| {
1021        RecordBatch::try_new_with_options(
1022            projected_schema,
1023            arr,
1024            &RecordBatchOptions::new()
1025                .with_match_field_names(true)
1026                .with_row_count(Some(rows.len())),
1027        )
1028    })
1029}
1030
1031fn parse_bool(string: &str) -> Option<bool> {
1032    if string.eq_ignore_ascii_case("false") {
1033        Some(false)
1034    } else if string.eq_ignore_ascii_case("true") {
1035        Some(true)
1036    } else {
1037        None
1038    }
1039}
1040
1041// parse the column string to an Arrow Array
1042fn build_decimal_array<T: DecimalType>(
1043    _line_number: usize,
1044    rows: &StringRecords<'_>,
1045    col_idx: usize,
1046    precision: u8,
1047    scale: i8,
1048    null_regex: &NullRegex,
1049) -> Result<ArrayRef, ArrowError> {
1050    let mut decimal_builder = PrimitiveBuilder::<T>::with_capacity(rows.len());
1051    for row in rows.iter() {
1052        let s = row.get(col_idx);
1053        if null_regex.is_null(s) {
1054            // append null
1055            decimal_builder.append_null();
1056        } else {
1057            let decimal_value: Result<T::Native, _> = parse_decimal::<T>(s, precision, scale);
1058            match decimal_value {
1059                Ok(v) => {
1060                    decimal_builder.append_value(v);
1061                }
1062                Err(e) => {
1063                    return Err(e);
1064                }
1065            }
1066        }
1067    }
1068    Ok(Arc::new(
1069        decimal_builder
1070            .finish()
1071            .with_precision_and_scale(precision, scale)?,
1072    ))
1073}
1074
1075// parses a specific column (col_idx) into an Arrow Array.
1076fn build_primitive_array<T: ArrowPrimitiveType + Parser>(
1077    line_number: usize,
1078    rows: &StringRecords<'_>,
1079    col_idx: usize,
1080    null_regex: &NullRegex,
1081) -> Result<ArrayRef, ArrowError> {
1082    rows.iter()
1083        .enumerate()
1084        .map(|(row_index, row)| {
1085            let s = row.get(col_idx);
1086            if null_regex.is_null(s) {
1087                return Ok(None);
1088            }
1089
1090            match T::parse(s) {
1091                Some(e) => Ok(Some(e)),
1092                None => Err(ArrowError::ParseError(format!(
1093                    // TODO: we should surface the underlying error here.
1094                    "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1095                    s,
1096                    T::DATA_TYPE,
1097                    col_idx,
1098                    line_number + row_index,
1099                    row
1100                ))),
1101            }
1102        })
1103        .collect::<Result<PrimitiveArray<T>, ArrowError>>()
1104        .map(|e| Arc::new(e) as ArrayRef)
1105}
1106
1107fn build_timestamp_array<T: ArrowTimestampType>(
1108    line_number: usize,
1109    rows: &StringRecords<'_>,
1110    col_idx: usize,
1111    timezone: Option<&str>,
1112    null_regex: &NullRegex,
1113) -> Result<ArrayRef, ArrowError> {
1114    Ok(Arc::new(match timezone {
1115        Some(timezone) => {
1116            let tz: Tz = timezone.parse()?;
1117            build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &tz, null_regex)?
1118                .with_timezone(timezone)
1119        }
1120        None => build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &Utc, null_regex)?,
1121    }))
1122}
1123
1124fn build_timestamp_array_impl<T: ArrowTimestampType, Tz: TimeZone>(
1125    line_number: usize,
1126    rows: &StringRecords<'_>,
1127    col_idx: usize,
1128    timezone: &Tz,
1129    null_regex: &NullRegex,
1130) -> Result<PrimitiveArray<T>, ArrowError> {
1131    rows.iter()
1132        .enumerate()
1133        .map(|(row_index, row)| {
1134            let s = row.get(col_idx);
1135            if null_regex.is_null(s) {
1136                return Ok(None);
1137            }
1138
1139            let date = string_to_datetime(timezone, s)
1140                .and_then(|date| match T::UNIT {
1141                    TimeUnit::Second => Ok(date.timestamp()),
1142                    TimeUnit::Millisecond => Ok(date.timestamp_millis()),
1143                    TimeUnit::Microsecond => Ok(date.timestamp_micros()),
1144                    TimeUnit::Nanosecond => date.timestamp_nanos_opt().ok_or_else(|| {
1145                        ArrowError::ParseError(format!(
1146                            "{} would overflow 64-bit signed nanoseconds",
1147                            date.to_rfc3339(),
1148                        ))
1149                    }),
1150                })
1151                .map_err(|e| {
1152                    ArrowError::ParseError(format!(
1153                        "Error parsing column {col_idx} at line {}: {}",
1154                        line_number + row_index,
1155                        e
1156                    ))
1157                })?;
1158            Ok(Some(date))
1159        })
1160        .collect()
1161}
1162
1163// parses a specific column (col_idx) into an Arrow Array.
1164fn build_boolean_array(
1165    line_number: usize,
1166    rows: &StringRecords<'_>,
1167    col_idx: usize,
1168    null_regex: &NullRegex,
1169) -> Result<ArrayRef, ArrowError> {
1170    rows.iter()
1171        .enumerate()
1172        .map(|(row_index, row)| {
1173            let s = row.get(col_idx);
1174            if null_regex.is_null(s) {
1175                return Ok(None);
1176            }
1177            let parsed = parse_bool(s);
1178            match parsed {
1179                Some(e) => Ok(Some(e)),
1180                None => Err(ArrowError::ParseError(format!(
1181                    // TODO: we should surface the underlying error here.
1182                    "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1183                    s,
1184                    "Boolean",
1185                    col_idx,
1186                    line_number + row_index,
1187                    row
1188                ))),
1189            }
1190        })
1191        .collect::<Result<BooleanArray, _>>()
1192        .map(|e| Arc::new(e) as ArrayRef)
1193}
1194
1195/// Builder for CSV [`Reader`]s
1196#[derive(Debug)]
1197pub struct ReaderBuilder {
1198    /// Schema of the CSV file
1199    schema: SchemaRef,
1200    /// Format of the CSV file
1201    format: Format,
1202    /// Batch size (number of records to load each time)
1203    ///
1204    /// The default batch size when using the `ReaderBuilder` is 1024 records
1205    batch_size: usize,
1206    /// The bounds over which to scan the reader. `None` starts from 0 and runs until EOF.
1207    bounds: Bounds,
1208    /// Optional projection for which columns to load (zero-based column indices)
1209    projection: Option<Vec<usize>>,
1210}
1211
1212impl ReaderBuilder {
1213    /// Create a new builder for configuring [`Reader`] CSV parsing options.
1214    ///
1215    /// To convert a builder into a reader, call [`ReaderBuilder::build`]. See
1216    /// the [module-level documentation](crate::reader) for more details and examples.
1217    ///
1218    /// # Example
1219    ///
1220    /// ```
1221    /// # use arrow_csv::{Reader, ReaderBuilder};
1222    /// # use std::fs::File;
1223    /// # use std::io::Seek;
1224    /// # use std::sync::Arc;
1225    /// # use arrow_csv::reader::Format;
1226    /// #
1227    /// let mut file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1228    /// // Infer the schema with the first 100 records
1229    /// let (schema, _) = Format::default().infer_schema(&mut file, Some(100)).unwrap();
1230    /// file.rewind().unwrap();
1231    ///
1232    /// // create a builder
1233    /// ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
1234    /// ```
1235    pub fn new(schema: SchemaRef) -> ReaderBuilder {
1236        Self {
1237            schema,
1238            format: Format::default(),
1239            batch_size: 1024,
1240            bounds: None,
1241            projection: None,
1242        }
1243    }
1244
1245    /// Set whether the CSV file has a header
1246    pub fn with_header(mut self, has_header: bool) -> Self {
1247        self.format.header = has_header;
1248        self
1249    }
1250
1251    /// Set whether to validate the CSV header against the schema
1252    ///
1253    /// This option only applies when [`Self::with_header`] is set to `true`, and defaults to `false`
1254    pub fn with_header_validation(mut self, validate_header: bool) -> Self {
1255        self.format.header_validation = validate_header;
1256        self
1257    }
1258
1259    /// Overrides the [Format] of this [ReaderBuilder]
1260    pub fn with_format(mut self, format: Format) -> Self {
1261        self.format = format;
1262        self
1263    }
1264
1265    /// Set the CSV file's column delimiter as a byte character
1266    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
1267        self.format.delimiter = Some(delimiter);
1268        self
1269    }
1270
1271    /// Set the given character as the CSV file's escape character
1272    pub fn with_escape(mut self, escape: u8) -> Self {
1273        self.format.escape = Some(escape);
1274        self
1275    }
1276
1277    /// Set the given character as the CSV file's quote character, by default it is double quote
1278    pub fn with_quote(mut self, quote: u8) -> Self {
1279        self.format.quote = Some(quote);
1280        self
1281    }
1282
1283    /// Provide a custom terminator character, defaults to CRLF
1284    pub fn with_terminator(mut self, terminator: u8) -> Self {
1285        self.format.terminator = Some(terminator);
1286        self
1287    }
1288
1289    /// Provide a comment character, lines starting with this character will be ignored
1290    pub fn with_comment(mut self, comment: u8) -> Self {
1291        self.format.comment = Some(comment);
1292        self
1293    }
1294
1295    /// Provide a regex to match null values, defaults to `^$`
1296    pub fn with_null_regex(mut self, null_regex: Regex) -> Self {
1297        self.format.null_regex = NullRegex(Some(null_regex));
1298        self
1299    }
1300
1301    /// Set the batch size (number of records to load at one time)
1302    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1303        self.batch_size = batch_size;
1304        self
1305    }
1306
1307    /// Set the bounds over which to scan the reader.
1308    /// `start` and `end` are line numbers.
1309    pub fn with_bounds(mut self, start: usize, end: usize) -> Self {
1310        self.bounds = Some((start, end));
1311        self
1312    }
1313
1314    /// Set the reader's column projection
1315    pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1316        self.projection = Some(projection);
1317        self
1318    }
1319
1320    /// Whether to allow truncated rows when parsing.
1321    ///
1322    /// By default this is set to `false` and will error if the CSV rows have different lengths.
1323    /// When set to true then it will allow records with less than the expected number of columns
1324    /// and fill the missing columns with nulls. If the record's schema is not nullable, then it
1325    /// will still return an error.
1326    pub fn with_truncated_rows(mut self, allow: bool) -> Self {
1327        self.format.truncated_rows = allow;
1328        self
1329    }
1330
1331    /// Create a new `Reader` from a non-buffered reader
1332    ///
1333    /// If `R: BufRead` consider using [`Self::build_buffered`] to avoid unnecessary additional
1334    /// buffering, as internally this method wraps `reader` in [`std::io::BufReader`]
1335    pub fn build<R: Read>(self, reader: R) -> Result<Reader<R>, ArrowError> {
1336        self.build_buffered(StdBufReader::new(reader))
1337    }
1338
1339    /// Create a new `BufReader` from a buffered reader
1340    pub fn build_buffered<R: BufRead>(self, reader: R) -> Result<BufReader<R>, ArrowError> {
1341        let schema = match &self.projection {
1342            Some(projection) => Arc::new(self.schema.project(projection)?),
1343            None => self.schema.clone(),
1344        };
1345
1346        Ok(BufReader {
1347            reader,
1348            decoder: self.build_decoder(),
1349            schema,
1350        })
1351    }
1352
1353    /// Builds a decoder that can be used to decode CSV from an arbitrary byte stream
1354    pub fn build_decoder(self) -> Decoder {
1355        let delimiter = self.format.build_parser();
1356        let record_decoder = RecordDecoder::new(
1357            delimiter,
1358            self.schema.fields().len(),
1359            self.format.truncated_rows,
1360        );
1361
1362        let header = self.format.header as usize;
1363
1364        let (start, end) = match self.bounds {
1365            Some((start, end)) => (start + header, end + header),
1366            None => (header, usize::MAX),
1367        };
1368
1369        Decoder {
1370            schema: self.schema,
1371            to_skip: start,
1372            header_validation: self.format.header && self.format.header_validation,
1373            record_decoder,
1374            line_number: start,
1375            end,
1376            projection: self.projection,
1377            batch_size: self.batch_size,
1378            null_regex: self.format.null_regex,
1379        }
1380    }
1381}
1382
1383#[cfg(test)]
1384mod tests {
1385    use super::*;
1386
1387    use std::io::{Cursor, Seek, SeekFrom, Write};
1388    use tempfile::NamedTempFile;
1389
1390    use arrow_array::cast::AsArray;
1391
1392    #[test]
1393    fn test_csv() {
1394        let schema = Arc::new(Schema::new(vec![
1395            Field::new("city", DataType::Utf8, false),
1396            Field::new("lat", DataType::Float64, false),
1397            Field::new("lng", DataType::Float64, false),
1398        ]));
1399
1400        let file = File::open("test/data/uk_cities.csv").unwrap();
1401        let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1402        assert_eq!(schema, csv.schema());
1403        let batch = csv.next().unwrap().unwrap();
1404        assert_eq!(37, batch.num_rows());
1405        assert_eq!(3, batch.num_columns());
1406
1407        // access data from a primitive array
1408        let lat = batch.column(1).as_primitive::<Float64Type>();
1409        assert_eq!(57.653484, lat.value(0));
1410
1411        // access data from a string array (ListArray<u8>)
1412        let city = batch.column(0).as_string::<i32>();
1413
1414        assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1415    }
1416
1417    #[test]
1418    fn test_csv_schema_metadata() {
1419        let mut metadata = std::collections::HashMap::new();
1420        metadata.insert("foo".to_owned(), "bar".to_owned());
1421        let schema = Arc::new(Schema::new_with_metadata(
1422            vec![
1423                Field::new("city", DataType::Utf8, false),
1424                Field::new("lat", DataType::Float64, false),
1425                Field::new("lng", DataType::Float64, false),
1426            ],
1427            metadata.clone(),
1428        ));
1429
1430        let file = File::open("test/data/uk_cities.csv").unwrap();
1431
1432        let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1433        assert_eq!(schema, csv.schema());
1434        let batch = csv.next().unwrap().unwrap();
1435        assert_eq!(37, batch.num_rows());
1436        assert_eq!(3, batch.num_columns());
1437
1438        assert_eq!(batch.schema().metadata(), &metadata);
1439    }
1440
1441    #[test]
1442    fn test_csv_reader_with_decimal() {
1443        let schema = Arc::new(Schema::new(vec![
1444            Field::new("city", DataType::Utf8, false),
1445            Field::new("lat", DataType::Decimal128(38, 6), false),
1446            Field::new("lng", DataType::Decimal256(76, 6), false),
1447        ]));
1448
1449        let file = File::open("test/data/decimal_test.csv").unwrap();
1450
1451        let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1452        let batch = csv.next().unwrap().unwrap();
1453        // access data from a primitive array
1454        let lat = batch
1455            .column(1)
1456            .as_any()
1457            .downcast_ref::<Decimal128Array>()
1458            .unwrap();
1459
1460        assert_eq!("57.653484", lat.value_as_string(0));
1461        assert_eq!("53.002666", lat.value_as_string(1));
1462        assert_eq!("52.412811", lat.value_as_string(2));
1463        assert_eq!("51.481583", lat.value_as_string(3));
1464        assert_eq!("12.123456", lat.value_as_string(4));
1465        assert_eq!("50.760000", lat.value_as_string(5));
1466        assert_eq!("0.123000", lat.value_as_string(6));
1467        assert_eq!("123.000000", lat.value_as_string(7));
1468        assert_eq!("123.000000", lat.value_as_string(8));
1469        assert_eq!("-50.760000", lat.value_as_string(9));
1470
1471        let lng = batch
1472            .column(2)
1473            .as_any()
1474            .downcast_ref::<Decimal256Array>()
1475            .unwrap();
1476
1477        assert_eq!("-3.335724", lng.value_as_string(0));
1478        assert_eq!("-2.179404", lng.value_as_string(1));
1479        assert_eq!("-1.778197", lng.value_as_string(2));
1480        assert_eq!("-3.179090", lng.value_as_string(3));
1481        assert_eq!("-3.179090", lng.value_as_string(4));
1482        assert_eq!("0.290472", lng.value_as_string(5));
1483        assert_eq!("0.290472", lng.value_as_string(6));
1484        assert_eq!("0.290472", lng.value_as_string(7));
1485        assert_eq!("0.290472", lng.value_as_string(8));
1486        assert_eq!("0.290472", lng.value_as_string(9));
1487    }
1488
1489    #[test]
1490    fn test_csv_reader_with_decimal_3264() {
1491        let schema = Arc::new(Schema::new(vec![
1492            Field::new("city", DataType::Utf8, false),
1493            Field::new("lat", DataType::Decimal32(9, 6), false),
1494            Field::new("lng", DataType::Decimal64(16, 6), false),
1495        ]));
1496
1497        let file = File::open("test/data/decimal_test.csv").unwrap();
1498
1499        let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1500        let batch = csv.next().unwrap().unwrap();
1501        // access data from a primitive array
1502        let lat = batch
1503            .column(1)
1504            .as_any()
1505            .downcast_ref::<Decimal32Array>()
1506            .unwrap();
1507
1508        assert_eq!("57.653484", lat.value_as_string(0));
1509        assert_eq!("53.002666", lat.value_as_string(1));
1510        assert_eq!("52.412811", lat.value_as_string(2));
1511        assert_eq!("51.481583", lat.value_as_string(3));
1512        assert_eq!("12.123456", lat.value_as_string(4));
1513        assert_eq!("50.760000", lat.value_as_string(5));
1514        assert_eq!("0.123000", lat.value_as_string(6));
1515        assert_eq!("123.000000", lat.value_as_string(7));
1516        assert_eq!("123.000000", lat.value_as_string(8));
1517        assert_eq!("-50.760000", lat.value_as_string(9));
1518
1519        let lng = batch
1520            .column(2)
1521            .as_any()
1522            .downcast_ref::<Decimal64Array>()
1523            .unwrap();
1524
1525        assert_eq!("-3.335724", lng.value_as_string(0));
1526        assert_eq!("-2.179404", lng.value_as_string(1));
1527        assert_eq!("-1.778197", lng.value_as_string(2));
1528        assert_eq!("-3.179090", lng.value_as_string(3));
1529        assert_eq!("-3.179090", lng.value_as_string(4));
1530        assert_eq!("0.290472", lng.value_as_string(5));
1531        assert_eq!("0.290472", lng.value_as_string(6));
1532        assert_eq!("0.290472", lng.value_as_string(7));
1533        assert_eq!("0.290472", lng.value_as_string(8));
1534        assert_eq!("0.290472", lng.value_as_string(9));
1535    }
1536
1537    #[test]
1538    fn test_csv_from_buf_reader() {
1539        let schema = Schema::new(vec![
1540            Field::new("city", DataType::Utf8, false),
1541            Field::new("lat", DataType::Float64, false),
1542            Field::new("lng", DataType::Float64, false),
1543        ]);
1544
1545        let file_with_headers = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1546        let file_without_headers = File::open("test/data/uk_cities.csv").unwrap();
1547        let both_files = file_with_headers
1548            .chain(Cursor::new("\n".to_string()))
1549            .chain(file_without_headers);
1550        let mut csv = ReaderBuilder::new(Arc::new(schema))
1551            .with_header(true)
1552            .build(both_files)
1553            .unwrap();
1554        let batch = csv.next().unwrap().unwrap();
1555        assert_eq!(74, batch.num_rows());
1556        assert_eq!(3, batch.num_columns());
1557    }
1558
1559    #[test]
1560    #[cfg_attr(miri, ignore)] // Takes too long
1561    fn test_csv_with_schema_inference() {
1562        let mut file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1563
1564        let (schema, _) = Format::default()
1565            .with_header(true)
1566            .infer_schema(&mut file, None)
1567            .unwrap();
1568
1569        file.rewind().unwrap();
1570        let builder = ReaderBuilder::new(Arc::new(schema)).with_header(true);
1571
1572        let mut csv = builder.build(file).unwrap();
1573        let expected_schema = Schema::new(vec![
1574            Field::new("city", DataType::Utf8, true),
1575            Field::new("lat", DataType::Float64, true),
1576            Field::new("lng", DataType::Float64, true),
1577        ]);
1578        assert_eq!(Arc::new(expected_schema), csv.schema());
1579        let batch = csv.next().unwrap().unwrap();
1580        assert_eq!(37, batch.num_rows());
1581        assert_eq!(3, batch.num_columns());
1582
1583        // access data from a primitive array
1584        let lat = batch
1585            .column(1)
1586            .as_any()
1587            .downcast_ref::<Float64Array>()
1588            .unwrap();
1589        assert_eq!(57.653484, lat.value(0));
1590
1591        // access data from a string array (ListArray<u8>)
1592        let city = batch
1593            .column(0)
1594            .as_any()
1595            .downcast_ref::<StringArray>()
1596            .unwrap();
1597
1598        assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1599    }
1600
1601    #[test]
1602    #[cfg_attr(miri, ignore)] // Takes too long
1603    fn test_csv_with_schema_inference_no_headers() {
1604        let mut file = File::open("test/data/uk_cities.csv").unwrap();
1605
1606        let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1607        file.rewind().unwrap();
1608
1609        let mut csv = ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
1610
1611        // csv field names should be 'column_{number}'
1612        let schema = csv.schema();
1613        assert_eq!("column_1", schema.field(0).name());
1614        assert_eq!("column_2", schema.field(1).name());
1615        assert_eq!("column_3", schema.field(2).name());
1616        let batch = csv.next().unwrap().unwrap();
1617        let batch_schema = batch.schema();
1618
1619        assert_eq!(schema, batch_schema);
1620        assert_eq!(37, batch.num_rows());
1621        assert_eq!(3, batch.num_columns());
1622
1623        // access data from a primitive array
1624        let lat = batch
1625            .column(1)
1626            .as_any()
1627            .downcast_ref::<Float64Array>()
1628            .unwrap();
1629        assert_eq!(57.653484, lat.value(0));
1630
1631        // access data from a string array (ListArray<u8>)
1632        let city = batch
1633            .column(0)
1634            .as_any()
1635            .downcast_ref::<StringArray>()
1636            .unwrap();
1637
1638        assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1639    }
1640
1641    #[test]
1642    #[cfg_attr(miri, ignore)] // Takes too long
1643    fn test_csv_builder_with_bounds() {
1644        let mut file = File::open("test/data/uk_cities.csv").unwrap();
1645
1646        // Set the bounds to the lines 0, 1 and 2.
1647        let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1648        file.rewind().unwrap();
1649        let mut csv = ReaderBuilder::new(Arc::new(schema))
1650            .with_bounds(0, 2)
1651            .build(file)
1652            .unwrap();
1653        let batch = csv.next().unwrap().unwrap();
1654
1655        // access data from a string array (ListArray<u8>)
1656        let city = batch
1657            .column(0)
1658            .as_any()
1659            .downcast_ref::<StringArray>()
1660            .unwrap();
1661
1662        // The value on line 0 is within the bounds
1663        assert_eq!("Elgin, Scotland, the UK", city.value(0));
1664
1665        // The value on line 13 is outside of the bounds. Therefore
1666        // the call to .value() will panic.
1667        let result = std::panic::catch_unwind(|| city.value(13));
1668        assert!(result.is_err());
1669    }
1670
1671    #[test]
1672    fn test_csv_with_projection() {
1673        let schema = Arc::new(Schema::new(vec![
1674            Field::new("city", DataType::Utf8, false),
1675            Field::new("lat", DataType::Float64, false),
1676            Field::new("lng", DataType::Float64, false),
1677        ]));
1678
1679        let file = File::open("test/data/uk_cities.csv").unwrap();
1680
1681        let mut csv = ReaderBuilder::new(schema)
1682            .with_projection(vec![0, 1])
1683            .build(file)
1684            .unwrap();
1685
1686        let projected_schema = Arc::new(Schema::new(vec![
1687            Field::new("city", DataType::Utf8, false),
1688            Field::new("lat", DataType::Float64, false),
1689        ]));
1690        assert_eq!(projected_schema, csv.schema());
1691        let batch = csv.next().unwrap().unwrap();
1692        assert_eq!(projected_schema, batch.schema());
1693        assert_eq!(37, batch.num_rows());
1694        assert_eq!(2, batch.num_columns());
1695    }
1696
1697    #[test]
1698    fn test_csv_record_batch_reader_schema() {
1699        let schema = Arc::new(Schema::new(vec![
1700            Field::new("a", DataType::Int32, false),
1701            Field::new("b", DataType::Int32, false),
1702        ]));
1703
1704        let cases = [
1705            None,
1706            Some(vec![]),
1707            Some(vec![1]),
1708            Some(vec![1, 0]),
1709            Some(vec![1, 1]),
1710        ];
1711        for projection in cases {
1712            let builder = ReaderBuilder::new(schema.clone());
1713            let builder = match projection {
1714                Some(projection) => builder.with_projection(projection),
1715                None => builder,
1716            };
1717            let mut reader = builder.build(Cursor::new(b"1,2\n")).unwrap();
1718
1719            let reader_schema = RecordBatchReader::schema(&reader);
1720            let batch = reader.next().unwrap().unwrap();
1721
1722            assert_eq!(reader_schema, batch.schema());
1723        }
1724    }
1725
1726    #[test]
1727    fn test_csv_reader_rejects_invalid_projection() {
1728        let schema = Arc::new(Schema::new(vec![
1729            Field::new("a", DataType::Int32, false),
1730            Field::new("b", DataType::Int32, false),
1731        ]));
1732
1733        let result = ReaderBuilder::new(schema)
1734            .with_projection(vec![2])
1735            .build(Cursor::new(b"1,2\n"));
1736
1737        assert!(matches!(
1738            result,
1739            Err(ArrowError::SchemaError(message))
1740                if message == "project index 2 out of bounds, max field 2"
1741        ));
1742    }
1743
1744    #[test]
1745    fn test_csv_decoder_rejects_invalid_projection() {
1746        let schema = Arc::new(Schema::new(vec![
1747            Field::new("a", DataType::Int32, false),
1748            Field::new("b", DataType::Int32, false),
1749        ]));
1750        let mut decoder = ReaderBuilder::new(schema)
1751            .with_projection(vec![2])
1752            .build_decoder();
1753
1754        decoder.decode(b"1,2\n").unwrap();
1755        let result = decoder.flush();
1756
1757        assert!(matches!(
1758            result,
1759            Err(ArrowError::SchemaError(message))
1760                if message == "project index 2 out of bounds, max field 2"
1761        ));
1762    }
1763
1764    #[test]
1765    fn test_csv_with_dictionary() {
1766        let schema = Arc::new(Schema::new(vec![
1767            Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
1768            Field::new("lat", DataType::Float64, false),
1769            Field::new("lng", DataType::Float64, false),
1770        ]));
1771
1772        let file = File::open("test/data/uk_cities.csv").unwrap();
1773
1774        let mut csv = ReaderBuilder::new(schema)
1775            .with_projection(vec![0, 1])
1776            .build(file)
1777            .unwrap();
1778
1779        let projected_schema = Arc::new(Schema::new(vec![
1780            Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
1781            Field::new("lat", DataType::Float64, false),
1782        ]));
1783        assert_eq!(projected_schema, csv.schema());
1784        let batch = csv.next().unwrap().unwrap();
1785        assert_eq!(projected_schema, batch.schema());
1786        assert_eq!(37, batch.num_rows());
1787        assert_eq!(2, batch.num_columns());
1788
1789        let strings = arrow_cast::cast(batch.column(0), &DataType::Utf8).unwrap();
1790        let strings = strings.as_string::<i32>();
1791
1792        assert_eq!(strings.value(0), "Elgin, Scotland, the UK");
1793        assert_eq!(strings.value(4), "Eastbourne, East Sussex, UK");
1794        assert_eq!(strings.value(29), "Uckfield, East Sussex, UK");
1795    }
1796
1797    #[test]
1798    fn test_csv_with_nullable_dictionary() {
1799        let offset_type = vec![
1800            DataType::Int8,
1801            DataType::Int16,
1802            DataType::Int32,
1803            DataType::Int64,
1804            DataType::UInt8,
1805            DataType::UInt16,
1806            DataType::UInt32,
1807            DataType::UInt64,
1808        ];
1809        for data_type in offset_type {
1810            let file = File::open("test/data/dictionary_nullable_test.csv").unwrap();
1811            let dictionary_type =
1812                DataType::Dictionary(Box::new(data_type), Box::new(DataType::Utf8));
1813            let schema = Arc::new(Schema::new(vec![
1814                Field::new("id", DataType::Utf8, false),
1815                Field::new("name", dictionary_type.clone(), true),
1816            ]));
1817
1818            let mut csv = ReaderBuilder::new(schema)
1819                .build(file.try_clone().unwrap())
1820                .unwrap();
1821
1822            let batch = csv.next().unwrap().unwrap();
1823            assert_eq!(3, batch.num_rows());
1824            assert_eq!(2, batch.num_columns());
1825
1826            let names = arrow_cast::cast(batch.column(1), &dictionary_type).unwrap();
1827            assert!(!names.is_null(2));
1828            assert!(names.is_null(1));
1829        }
1830    }
1831    #[test]
1832    fn test_nulls() {
1833        let schema = Arc::new(Schema::new(vec![
1834            Field::new("c_int", DataType::UInt64, false),
1835            Field::new("c_float", DataType::Float32, true),
1836            Field::new("c_string", DataType::Utf8, true),
1837            Field::new("c_bool", DataType::Boolean, false),
1838        ]));
1839
1840        let file = File::open("test/data/null_test.csv").unwrap();
1841
1842        let mut csv = ReaderBuilder::new(schema)
1843            .with_header(true)
1844            .build(file)
1845            .unwrap();
1846
1847        let batch = csv.next().unwrap().unwrap();
1848
1849        assert!(!batch.column(1).is_null(0));
1850        assert!(!batch.column(1).is_null(1));
1851        assert!(batch.column(1).is_null(2));
1852        assert!(!batch.column(1).is_null(3));
1853        assert!(!batch.column(1).is_null(4));
1854    }
1855
1856    #[test]
1857    fn test_init_nulls() {
1858        let schema = Arc::new(Schema::new(vec![
1859            Field::new("c_int", DataType::UInt64, true),
1860            Field::new("c_float", DataType::Float32, true),
1861            Field::new("c_string", DataType::Utf8, true),
1862            Field::new("c_bool", DataType::Boolean, true),
1863            Field::new("c_null", DataType::Null, true),
1864        ]));
1865        let file = File::open("test/data/init_null_test.csv").unwrap();
1866
1867        let mut csv = ReaderBuilder::new(schema)
1868            .with_header(true)
1869            .build(file)
1870            .unwrap();
1871
1872        let batch = csv.next().unwrap().unwrap();
1873
1874        assert!(batch.column(1).is_null(0));
1875        assert!(!batch.column(1).is_null(1));
1876        assert!(batch.column(1).is_null(2));
1877        assert!(!batch.column(1).is_null(3));
1878        assert!(!batch.column(1).is_null(4));
1879    }
1880
1881    #[test]
1882    #[cfg_attr(miri, ignore)] // Takes too long
1883    fn test_init_nulls_with_inference() {
1884        let format = Format::default().with_header(true).with_delimiter(b',');
1885
1886        let mut file = File::open("test/data/init_null_test.csv").unwrap();
1887        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
1888        file.rewind().unwrap();
1889
1890        let expected_schema = Schema::new(vec![
1891            Field::new("c_int", DataType::Int64, true),
1892            Field::new("c_float", DataType::Float64, true),
1893            Field::new("c_string", DataType::Utf8, true),
1894            Field::new("c_bool", DataType::Boolean, true),
1895            Field::new("c_null", DataType::Null, true),
1896        ]);
1897        assert_eq!(schema, expected_schema);
1898
1899        let mut csv = ReaderBuilder::new(Arc::new(schema))
1900            .with_format(format)
1901            .build(file)
1902            .unwrap();
1903
1904        let batch = csv.next().unwrap().unwrap();
1905
1906        assert!(batch.column(1).is_null(0));
1907        assert!(!batch.column(1).is_null(1));
1908        assert!(batch.column(1).is_null(2));
1909        assert!(!batch.column(1).is_null(3));
1910        assert!(!batch.column(1).is_null(4));
1911    }
1912
1913    #[test]
1914    fn test_custom_nulls() {
1915        let schema = Arc::new(Schema::new(vec![
1916            Field::new("c_int", DataType::UInt64, true),
1917            Field::new("c_float", DataType::Float32, true),
1918            Field::new("c_string", DataType::Utf8, true),
1919            Field::new("c_bool", DataType::Boolean, true),
1920        ]));
1921
1922        let file = File::open("test/data/custom_null_test.csv").unwrap();
1923
1924        let null_regex = Regex::new("^nil$").unwrap();
1925
1926        let mut csv = ReaderBuilder::new(schema)
1927            .with_header(true)
1928            .with_null_regex(null_regex)
1929            .build(file)
1930            .unwrap();
1931
1932        let batch = csv.next().unwrap().unwrap();
1933
1934        // "nil"s should be NULL
1935        assert!(batch.column(0).is_null(1));
1936        assert!(batch.column(1).is_null(2));
1937        assert!(batch.column(3).is_null(4));
1938        assert!(batch.column(2).is_null(3));
1939        assert!(!batch.column(2).is_null(4));
1940    }
1941
1942    #[test]
1943    #[cfg_attr(miri, ignore)] // Takes too long
1944    fn test_nulls_with_inference() {
1945        let mut file = File::open("test/data/various_types.csv").unwrap();
1946        let format = Format::default().with_header(true).with_delimiter(b'|');
1947
1948        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
1949        file.rewind().unwrap();
1950
1951        let builder = ReaderBuilder::new(Arc::new(schema))
1952            .with_format(format)
1953            .with_batch_size(512)
1954            .with_projection(vec![0, 1, 2, 3, 4, 5]);
1955
1956        let mut csv = builder.build(file).unwrap();
1957        let batch = csv.next().unwrap().unwrap();
1958
1959        assert_eq!(10, batch.num_rows());
1960        assert_eq!(6, batch.num_columns());
1961
1962        let schema = batch.schema();
1963
1964        assert_eq!(&DataType::Int64, schema.field(0).data_type());
1965        assert_eq!(&DataType::Float64, schema.field(1).data_type());
1966        assert_eq!(&DataType::Float64, schema.field(2).data_type());
1967        assert_eq!(&DataType::Boolean, schema.field(3).data_type());
1968        assert_eq!(&DataType::Date32, schema.field(4).data_type());
1969        assert_eq!(
1970            &DataType::Timestamp(TimeUnit::Second, None),
1971            schema.field(5).data_type()
1972        );
1973
1974        let names: Vec<&str> = schema.fields().iter().map(|x| x.name().as_str()).collect();
1975        assert_eq!(
1976            names,
1977            vec![
1978                "c_int",
1979                "c_float",
1980                "c_string",
1981                "c_bool",
1982                "c_date",
1983                "c_datetime"
1984            ]
1985        );
1986
1987        assert!(schema.field(0).is_nullable());
1988        assert!(schema.field(1).is_nullable());
1989        assert!(schema.field(2).is_nullable());
1990        assert!(schema.field(3).is_nullable());
1991        assert!(schema.field(4).is_nullable());
1992        assert!(schema.field(5).is_nullable());
1993
1994        assert!(!batch.column(1).is_null(0));
1995        assert!(!batch.column(1).is_null(1));
1996        assert!(batch.column(1).is_null(2));
1997        assert!(!batch.column(1).is_null(3));
1998        assert!(!batch.column(1).is_null(4));
1999    }
2000
2001    #[test]
2002    #[cfg_attr(miri, ignore)] // Takes too long
2003    fn test_custom_nulls_with_inference() {
2004        let mut file = File::open("test/data/custom_null_test.csv").unwrap();
2005
2006        let null_regex = Regex::new("^nil$").unwrap();
2007
2008        let format = Format::default()
2009            .with_header(true)
2010            .with_null_regex(null_regex);
2011
2012        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2013        file.rewind().unwrap();
2014
2015        let expected_schema = Schema::new(vec![
2016            Field::new("c_int", DataType::Int64, true),
2017            Field::new("c_float", DataType::Float64, true),
2018            Field::new("c_string", DataType::Utf8, true),
2019            Field::new("c_bool", DataType::Boolean, true),
2020        ]);
2021
2022        assert_eq!(schema, expected_schema);
2023
2024        let builder = ReaderBuilder::new(Arc::new(schema))
2025            .with_format(format)
2026            .with_batch_size(512)
2027            .with_projection(vec![0, 1, 2, 3]);
2028
2029        let mut csv = builder.build(file).unwrap();
2030        let batch = csv.next().unwrap().unwrap();
2031
2032        assert_eq!(5, batch.num_rows());
2033        assert_eq!(4, batch.num_columns());
2034
2035        assert_eq!(batch.schema().as_ref(), &expected_schema);
2036    }
2037
2038    #[test]
2039    #[cfg_attr(miri, ignore)] // Takes too long
2040    fn test_scientific_notation_with_inference() {
2041        let mut file = File::open("test/data/scientific_notation_test.csv").unwrap();
2042        let format = Format::default().with_header(false).with_delimiter(b',');
2043
2044        let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2045        file.rewind().unwrap();
2046
2047        let builder = ReaderBuilder::new(Arc::new(schema))
2048            .with_format(format)
2049            .with_batch_size(512)
2050            .with_projection(vec![0, 1]);
2051
2052        let mut csv = builder.build(file).unwrap();
2053        let batch = csv.next().unwrap().unwrap();
2054
2055        let schema = batch.schema();
2056
2057        assert_eq!(&DataType::Float64, schema.field(0).data_type());
2058    }
2059
2060    fn invalid_csv_helper(file_name: &str) -> String {
2061        let file = File::open(file_name).unwrap();
2062        let schema = Schema::new(vec![
2063            Field::new("c_int", DataType::UInt64, false),
2064            Field::new("c_float", DataType::Float32, false),
2065            Field::new("c_string", DataType::Utf8, false),
2066            Field::new("c_bool", DataType::Boolean, false),
2067        ]);
2068
2069        let builder = ReaderBuilder::new(Arc::new(schema))
2070            .with_header(true)
2071            .with_delimiter(b'|')
2072            .with_batch_size(512)
2073            .with_projection(vec![0, 1, 2, 3]);
2074
2075        let mut csv = builder.build(file).unwrap();
2076
2077        csv.next().unwrap().unwrap_err().to_string()
2078    }
2079
2080    #[test]
2081    fn test_parse_invalid_csv_float() {
2082        let file_name = "test/data/various_invalid_types/invalid_float.csv";
2083
2084        let error = invalid_csv_helper(file_name);
2085        assert_eq!(
2086            "Parser error: Error while parsing value '4.x4' as type 'Float32' for column 1 at line 4. Row data: '[4,4.x4,,false]'",
2087            error
2088        );
2089    }
2090
2091    #[test]
2092    fn test_parse_invalid_csv_int() {
2093        let file_name = "test/data/various_invalid_types/invalid_int.csv";
2094
2095        let error = invalid_csv_helper(file_name);
2096        assert_eq!(
2097            "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]'",
2098            error
2099        );
2100    }
2101
2102    #[test]
2103    fn test_parse_invalid_csv_bool() {
2104        let file_name = "test/data/various_invalid_types/invalid_bool.csv";
2105
2106        let error = invalid_csv_helper(file_name);
2107        assert_eq!(
2108            "Parser error: Error while parsing value 'none' as type 'Boolean' for column 3 at line 2. Row data: '[2,2.2,2.22,none]'",
2109            error
2110        );
2111    }
2112
2113    /// Infer the data type of a record
2114    fn infer_field_schema(string: &str) -> DataType {
2115        let mut v = InferredDataType::default();
2116        v.update(string);
2117        v.get()
2118    }
2119
2120    #[test]
2121    #[cfg_attr(miri, ignore)] // Takes too long
2122    fn test_infer_field_schema() {
2123        assert_eq!(infer_field_schema("A"), DataType::Utf8);
2124        assert_eq!(infer_field_schema("\"123\""), DataType::Utf8);
2125        assert_eq!(infer_field_schema("10"), DataType::Int64);
2126        assert_eq!(infer_field_schema("10.2"), DataType::Float64);
2127        assert_eq!(infer_field_schema(".2"), DataType::Float64);
2128        assert_eq!(infer_field_schema("2."), DataType::Float64);
2129        assert_eq!(infer_field_schema("NaN"), DataType::Float64);
2130        assert_eq!(infer_field_schema("nan"), DataType::Float64);
2131        assert_eq!(infer_field_schema("inf"), DataType::Float64);
2132        assert_eq!(infer_field_schema("-inf"), DataType::Float64);
2133        assert_eq!(infer_field_schema("true"), DataType::Boolean);
2134        assert_eq!(infer_field_schema("trUe"), DataType::Boolean);
2135        assert_eq!(infer_field_schema("false"), DataType::Boolean);
2136        assert_eq!(infer_field_schema("2020-11-08"), DataType::Date32);
2137        assert_eq!(
2138            infer_field_schema("2020-11-08T14:20:01"),
2139            DataType::Timestamp(TimeUnit::Second, None)
2140        );
2141        assert_eq!(
2142            infer_field_schema("2020-11-08 14:20:01"),
2143            DataType::Timestamp(TimeUnit::Second, None)
2144        );
2145        assert_eq!(
2146            infer_field_schema("2020-11-08 14:20:01"),
2147            DataType::Timestamp(TimeUnit::Second, None)
2148        );
2149        assert_eq!(infer_field_schema("-5.13"), DataType::Float64);
2150        assert_eq!(infer_field_schema("0.1300"), DataType::Float64);
2151        assert_eq!(
2152            infer_field_schema("2021-12-19 13:12:30.921"),
2153            DataType::Timestamp(TimeUnit::Millisecond, None)
2154        );
2155        assert_eq!(
2156            infer_field_schema("2021-12-19T13:12:30.123456789"),
2157            DataType::Timestamp(TimeUnit::Nanosecond, None)
2158        );
2159        assert_eq!(infer_field_schema("–9223372036854775809"), DataType::Utf8);
2160        assert_eq!(infer_field_schema("9223372036854775808"), DataType::Utf8);
2161    }
2162
2163    #[test]
2164    fn parse_date32() {
2165        assert_eq!(Date32Type::parse("1970-01-01").unwrap(), 0);
2166        assert_eq!(Date32Type::parse("2020-03-15").unwrap(), 18336);
2167        assert_eq!(Date32Type::parse("1945-05-08").unwrap(), -9004);
2168    }
2169
2170    #[test]
2171    fn parse_time() {
2172        assert_eq!(
2173            Time64NanosecondType::parse("12:10:01.123456789 AM"),
2174            Some(601_123_456_789)
2175        );
2176        assert_eq!(
2177            Time64MicrosecondType::parse("12:10:01.123456 am"),
2178            Some(601_123_456)
2179        );
2180        assert_eq!(
2181            Time32MillisecondType::parse("2:10:01.12 PM"),
2182            Some(51_001_120)
2183        );
2184        assert_eq!(Time32SecondType::parse("2:10:01 pm"), Some(51_001));
2185    }
2186
2187    #[test]
2188    fn parse_date64() {
2189        assert_eq!(Date64Type::parse("1970-01-01T00:00:00").unwrap(), 0);
2190        assert_eq!(
2191            Date64Type::parse("2018-11-13T17:11:10").unwrap(),
2192            1542129070000
2193        );
2194        assert_eq!(
2195            Date64Type::parse("2018-11-13T17:11:10.011").unwrap(),
2196            1542129070011
2197        );
2198        assert_eq!(
2199            Date64Type::parse("1900-02-28T12:34:56").unwrap(),
2200            -2203932304000
2201        );
2202        assert_eq!(
2203            Date64Type::parse_formatted("1900-02-28 12:34:56", "%Y-%m-%d %H:%M:%S").unwrap(),
2204            -2203932304000
2205        );
2206        assert_eq!(
2207            Date64Type::parse_formatted("1900-02-28 12:34:56+0030", "%Y-%m-%d %H:%M:%S%z").unwrap(),
2208            -2203932304000 - (30 * 60 * 1000)
2209        );
2210    }
2211
2212    fn test_parse_timestamp_impl<T: ArrowTimestampType>(
2213        timezone: Option<Arc<str>>,
2214        expected: &[i64],
2215    ) {
2216        let csv = [
2217            "1970-01-01T00:00:00",
2218            "1970-01-01T00:00:00Z",
2219            "1970-01-01T00:00:00+02:00",
2220        ]
2221        .join("\n");
2222        let schema = Arc::new(Schema::new(vec![Field::new(
2223            "field",
2224            DataType::Timestamp(T::UNIT, timezone.clone()),
2225            true,
2226        )]));
2227
2228        let mut decoder = ReaderBuilder::new(schema).build_decoder();
2229
2230        let decoded = decoder.decode(csv.as_bytes()).unwrap();
2231        assert_eq!(decoded, csv.len());
2232        decoder.decode(&[]).unwrap();
2233
2234        let batch = decoder.flush().unwrap().unwrap();
2235        assert_eq!(batch.num_columns(), 1);
2236        assert_eq!(batch.num_rows(), 3);
2237        let col = batch.column(0).as_primitive::<T>();
2238        assert_eq!(col.values(), expected);
2239        assert_eq!(col.data_type(), &DataType::Timestamp(T::UNIT, timezone));
2240    }
2241
2242    #[test]
2243    fn test_parse_timestamp() {
2244        test_parse_timestamp_impl::<TimestampNanosecondType>(None, &[0, 0, -7_200_000_000_000]);
2245        test_parse_timestamp_impl::<TimestampNanosecondType>(
2246            Some("+00:00".into()),
2247            &[0, 0, -7_200_000_000_000],
2248        );
2249        test_parse_timestamp_impl::<TimestampNanosecondType>(
2250            Some("-05:00".into()),
2251            &[18_000_000_000_000, 0, -7_200_000_000_000],
2252        );
2253        test_parse_timestamp_impl::<TimestampMicrosecondType>(
2254            Some("-03".into()),
2255            &[10_800_000_000, 0, -7_200_000_000],
2256        );
2257        test_parse_timestamp_impl::<TimestampMillisecondType>(
2258            Some("-03".into()),
2259            &[10_800_000, 0, -7_200_000],
2260        );
2261        test_parse_timestamp_impl::<TimestampSecondType>(Some("-03".into()), &[10_800, 0, -7_200]);
2262    }
2263
2264    #[test]
2265    #[cfg_attr(miri, ignore)] // Takes too long
2266    fn test_infer_schema_from_multiple_files() {
2267        let mut csv1 = NamedTempFile::new().unwrap();
2268        let mut csv2 = NamedTempFile::new().unwrap();
2269        let csv3 = NamedTempFile::new().unwrap(); // empty csv file should be skipped
2270        let mut csv4 = NamedTempFile::new().unwrap();
2271        writeln!(csv1, "c1,c2,c3").unwrap();
2272        writeln!(csv1, "1,\"foo\",0.5").unwrap();
2273        writeln!(csv1, "3,\"bar\",1").unwrap();
2274        writeln!(csv1, "3,\"bar\",2e-06").unwrap();
2275        // reading csv2 will set c2 to optional
2276        writeln!(csv2, "c1,c2,c3,c4").unwrap();
2277        writeln!(csv2, "10,,3.14,true").unwrap();
2278        // reading csv4 will set c3 to optional
2279        writeln!(csv4, "c1,c2,c3").unwrap();
2280        writeln!(csv4, "10,\"foo\",").unwrap();
2281
2282        let schema = infer_schema_from_files(
2283            &[
2284                csv3.path().to_str().unwrap().to_string(),
2285                csv1.path().to_str().unwrap().to_string(),
2286                csv2.path().to_str().unwrap().to_string(),
2287                csv4.path().to_str().unwrap().to_string(),
2288            ],
2289            b',',
2290            Some(4), // only csv1 and csv2 should be read
2291            true,
2292        )
2293        .unwrap();
2294
2295        assert_eq!(schema.fields().len(), 4);
2296        assert!(schema.field(0).is_nullable());
2297        assert!(schema.field(1).is_nullable());
2298        assert!(schema.field(2).is_nullable());
2299        assert!(schema.field(3).is_nullable());
2300
2301        assert_eq!(&DataType::Int64, schema.field(0).data_type());
2302        assert_eq!(&DataType::Utf8, schema.field(1).data_type());
2303        assert_eq!(&DataType::Float64, schema.field(2).data_type());
2304        assert_eq!(&DataType::Boolean, schema.field(3).data_type());
2305    }
2306
2307    #[test]
2308    fn test_bounded() {
2309        let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2310        let data = [
2311            vec!["0"],
2312            vec!["1"],
2313            vec!["2"],
2314            vec!["3"],
2315            vec!["4"],
2316            vec!["5"],
2317            vec!["6"],
2318        ];
2319
2320        let data = data
2321            .iter()
2322            .map(|x| x.join(","))
2323            .collect::<Vec<_>>()
2324            .join("\n");
2325        let data = data.as_bytes();
2326
2327        let reader = std::io::Cursor::new(data);
2328
2329        let mut csv = ReaderBuilder::new(Arc::new(schema))
2330            .with_batch_size(2)
2331            .with_projection(vec![0])
2332            .with_bounds(2, 6)
2333            .build_buffered(reader)
2334            .unwrap();
2335
2336        let batch = csv.next().unwrap().unwrap();
2337        let a = batch.column(0);
2338        let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2339        assert_eq!(a, &UInt32Array::from(vec![2, 3]));
2340
2341        let batch = csv.next().unwrap().unwrap();
2342        let a = batch.column(0);
2343        let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2344        assert_eq!(a, &UInt32Array::from(vec![4, 5]));
2345
2346        assert!(csv.next().is_none());
2347    }
2348
2349    #[test]
2350    fn test_empty_projection() {
2351        let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2352        let data = [vec!["0"], vec!["1"]];
2353
2354        let data = data
2355            .iter()
2356            .map(|x| x.join(","))
2357            .collect::<Vec<_>>()
2358            .join("\n");
2359
2360        let mut csv = ReaderBuilder::new(Arc::new(schema))
2361            .with_batch_size(2)
2362            .with_projection(vec![])
2363            .build_buffered(Cursor::new(data.as_bytes()))
2364            .unwrap();
2365
2366        let batch = csv.next().unwrap().unwrap();
2367        assert_eq!(batch.columns().len(), 0);
2368        assert_eq!(batch.num_rows(), 2);
2369
2370        assert!(csv.next().is_none());
2371    }
2372
2373    #[test]
2374    fn test_parsing_bool() {
2375        // Encode the expected behavior of boolean parsing
2376        assert_eq!(Some(true), parse_bool("true"));
2377        assert_eq!(Some(true), parse_bool("tRUe"));
2378        assert_eq!(Some(true), parse_bool("True"));
2379        assert_eq!(Some(true), parse_bool("TRUE"));
2380        assert_eq!(None, parse_bool("t"));
2381        assert_eq!(None, parse_bool("T"));
2382        assert_eq!(None, parse_bool(""));
2383
2384        assert_eq!(Some(false), parse_bool("false"));
2385        assert_eq!(Some(false), parse_bool("fALse"));
2386        assert_eq!(Some(false), parse_bool("False"));
2387        assert_eq!(Some(false), parse_bool("FALSE"));
2388        assert_eq!(None, parse_bool("f"));
2389        assert_eq!(None, parse_bool("F"));
2390        assert_eq!(None, parse_bool(""));
2391    }
2392
2393    #[test]
2394    fn test_parsing_float() {
2395        assert_eq!(Some(12.34), Float64Type::parse("12.34"));
2396        assert_eq!(Some(-12.34), Float64Type::parse("-12.34"));
2397        assert_eq!(Some(12.0), Float64Type::parse("12"));
2398        assert_eq!(Some(0.0), Float64Type::parse("0"));
2399        assert_eq!(Some(2.0), Float64Type::parse("2."));
2400        assert_eq!(Some(0.2), Float64Type::parse(".2"));
2401        assert!(Float64Type::parse("nan").unwrap().is_nan());
2402        assert!(Float64Type::parse("NaN").unwrap().is_nan());
2403        assert!(Float64Type::parse("inf").unwrap().is_infinite());
2404        assert!(Float64Type::parse("inf").unwrap().is_sign_positive());
2405        assert!(Float64Type::parse("-inf").unwrap().is_infinite());
2406        assert!(Float64Type::parse("-inf").unwrap().is_sign_negative());
2407        assert_eq!(None, Float64Type::parse(""));
2408        assert_eq!(None, Float64Type::parse("dd"));
2409        assert_eq!(None, Float64Type::parse("12.34.56"));
2410    }
2411
2412    #[test]
2413    fn test_non_std_quote() {
2414        let schema = Schema::new(vec![
2415            Field::new("text1", DataType::Utf8, false),
2416            Field::new("text2", DataType::Utf8, false),
2417        ]);
2418        let builder = ReaderBuilder::new(Arc::new(schema))
2419            .with_header(false)
2420            .with_quote(b'~'); // default is ", change to ~
2421
2422        let mut csv_text = Vec::new();
2423        let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2424        for index in 0..10 {
2425            let text1 = format!("id{index:}");
2426            let text2 = format!("value{index:}");
2427            csv_writer
2428                .write_fmt(format_args!("~{text1}~,~{text2}~\r\n"))
2429                .unwrap();
2430        }
2431        let mut csv_reader = std::io::Cursor::new(&csv_text);
2432        let mut reader = builder.build(&mut csv_reader).unwrap();
2433        let batch = reader.next().unwrap().unwrap();
2434        let col0 = batch.column(0);
2435        assert_eq!(col0.len(), 10);
2436        let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2437        assert_eq!(col0_arr.value(0), "id0");
2438        let col1 = batch.column(1);
2439        assert_eq!(col1.len(), 10);
2440        let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2441        assert_eq!(col1_arr.value(5), "value5");
2442    }
2443
2444    #[test]
2445    fn test_non_std_escape() {
2446        let schema = Schema::new(vec![
2447            Field::new("text1", DataType::Utf8, false),
2448            Field::new("text2", DataType::Utf8, false),
2449        ]);
2450        let builder = ReaderBuilder::new(Arc::new(schema))
2451            .with_header(false)
2452            .with_escape(b'\\'); // default is None, change to \
2453
2454        let mut csv_text = Vec::new();
2455        let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2456        for index in 0..10 {
2457            let text1 = format!("id{index:}");
2458            let text2 = format!("value\\\"{index:}");
2459            csv_writer
2460                .write_fmt(format_args!("\"{text1}\",\"{text2}\"\r\n"))
2461                .unwrap();
2462        }
2463        let mut csv_reader = std::io::Cursor::new(&csv_text);
2464        let mut reader = builder.build(&mut csv_reader).unwrap();
2465        let batch = reader.next().unwrap().unwrap();
2466        let col0 = batch.column(0);
2467        assert_eq!(col0.len(), 10);
2468        let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2469        assert_eq!(col0_arr.value(0), "id0");
2470        let col1 = batch.column(1);
2471        assert_eq!(col1.len(), 10);
2472        let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2473        assert_eq!(col1_arr.value(5), "value\"5");
2474    }
2475
2476    #[test]
2477    fn test_non_std_terminator() {
2478        let schema = Schema::new(vec![
2479            Field::new("text1", DataType::Utf8, false),
2480            Field::new("text2", DataType::Utf8, false),
2481        ]);
2482        let builder = ReaderBuilder::new(Arc::new(schema))
2483            .with_header(false)
2484            .with_terminator(b'\n'); // default is CRLF, change to LF
2485
2486        let mut csv_text = Vec::new();
2487        let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2488        for index in 0..10 {
2489            let text1 = format!("id{index:}");
2490            let text2 = format!("value{index:}");
2491            csv_writer
2492                .write_fmt(format_args!("\"{text1}\",\"{text2}\"\n"))
2493                .unwrap();
2494        }
2495        let mut csv_reader = std::io::Cursor::new(&csv_text);
2496        let mut reader = builder.build(&mut csv_reader).unwrap();
2497        let batch = reader.next().unwrap().unwrap();
2498        let col0 = batch.column(0);
2499        assert_eq!(col0.len(), 10);
2500        let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2501        assert_eq!(col0_arr.value(0), "id0");
2502        let col1 = batch.column(1);
2503        assert_eq!(col1.len(), 10);
2504        let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2505        assert_eq!(col1_arr.value(5), "value5");
2506    }
2507
2508    #[test]
2509    fn test_header_bounds() {
2510        let csv = "a,b\na,b\na,b\na,b\na,b\n";
2511        let tests = [
2512            (None, false, 5),
2513            (None, true, 4),
2514            (Some((0, 4)), false, 4),
2515            (Some((1, 4)), false, 3),
2516            (Some((0, 4)), true, 4),
2517            (Some((1, 4)), true, 3),
2518        ];
2519        let schema = Arc::new(Schema::new(vec![
2520            Field::new("a", DataType::Utf8, false),
2521            Field::new("a", DataType::Utf8, false),
2522        ]));
2523
2524        for (idx, (bounds, has_header, expected)) in tests.into_iter().enumerate() {
2525            let mut reader = ReaderBuilder::new(schema.clone()).with_header(has_header);
2526            if let Some((start, end)) = bounds {
2527                reader = reader.with_bounds(start, end);
2528            }
2529            let b = reader
2530                .build_buffered(Cursor::new(csv.as_bytes()))
2531                .unwrap()
2532                .next()
2533                .unwrap()
2534                .unwrap();
2535            assert_eq!(b.num_rows(), expected, "{idx}");
2536        }
2537    }
2538
2539    #[test]
2540    fn test_header_validation() {
2541        let schema = Arc::new(Schema::new(vec![
2542            Field::new("a", DataType::Int32, false),
2543            Field::new("b", DataType::Int32, false),
2544        ]));
2545
2546        let csv = "a,c\n1,2\n";
2547        let err = ReaderBuilder::new(schema.clone())
2548            .with_header(true)
2549            .with_header_validation(true)
2550            .build_buffered(Cursor::new(csv.as_bytes()))
2551            .unwrap()
2552            .next()
2553            .unwrap()
2554            .unwrap_err()
2555            .to_string();
2556        assert_eq!(
2557            err,
2558            "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"c\""
2559        );
2560
2561        let batch = ReaderBuilder::new(schema)
2562            .with_header(true)
2563            .with_header_validation(false)
2564            .build_buffered(Cursor::new(csv.as_bytes()))
2565            .unwrap()
2566            .next()
2567            .unwrap()
2568            .unwrap();
2569        assert_eq!(batch.num_rows(), 1);
2570    }
2571
2572    #[test]
2573    fn test_header_validation_with_buffered_reader() {
2574        let schema = Arc::new(Schema::new(vec![
2575            Field::new("a", DataType::Int32, false),
2576            Field::new("b", DataType::Int32, false),
2577        ]));
2578
2579        let csv = "a,b\n1,2\n";
2580        let buffered = std::io::BufReader::with_capacity(1, Cursor::new(csv.as_bytes()));
2581        let batch = ReaderBuilder::new(schema)
2582            .with_header(true)
2583            .with_header_validation(true)
2584            .build_buffered(buffered)
2585            .unwrap()
2586            .next()
2587            .unwrap()
2588            .unwrap();
2589
2590        assert_eq!(batch.num_rows(), 1);
2591        let a = batch.column(0).as_primitive::<Int32Type>();
2592        assert_eq!(a.value(0), 1);
2593    }
2594
2595    #[test]
2596    fn test_header_validation_with_truncated_rows() {
2597        let schema = Arc::new(Schema::new(vec![
2598            Field::new("a", DataType::Int32, true),
2599            Field::new("b", DataType::Int32, true),
2600        ]));
2601
2602        let csv = "a\n1\n";
2603        let err = ReaderBuilder::new(schema.clone())
2604            .with_header(true)
2605            .with_header_validation(true)
2606            .with_truncated_rows(true)
2607            .build_buffered(Cursor::new(csv.as_bytes()))
2608            .unwrap()
2609            .next()
2610            .unwrap()
2611            .unwrap_err()
2612            .to_string();
2613        assert_eq!(
2614            err,
2615            "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"\"",
2616        )
2617    }
2618
2619    #[test]
2620    fn test_null_boolean() {
2621        let csv = "true,false\nFalse,True\n,True\nFalse,";
2622        let schema = Arc::new(Schema::new(vec![
2623            Field::new("a", DataType::Boolean, true),
2624            Field::new("a", DataType::Boolean, true),
2625        ]));
2626
2627        let b = ReaderBuilder::new(schema)
2628            .build_buffered(Cursor::new(csv.as_bytes()))
2629            .unwrap()
2630            .next()
2631            .unwrap()
2632            .unwrap();
2633
2634        assert_eq!(b.num_rows(), 4);
2635        assert_eq!(b.num_columns(), 2);
2636
2637        let c = b.column(0).as_boolean();
2638        assert_eq!(c.null_count(), 1);
2639        assert!(c.value(0));
2640        assert!(!c.value(1));
2641        assert!(c.is_null(2));
2642        assert!(!c.value(3));
2643
2644        let c = b.column(1).as_boolean();
2645        assert_eq!(c.null_count(), 1);
2646        assert!(!c.value(0));
2647        assert!(c.value(1));
2648        assert!(c.value(2));
2649        assert!(c.is_null(3));
2650    }
2651
2652    #[test]
2653    fn test_truncated_rows() {
2654        let data = "a,b,c\n1,2,3\n4,5\n\n6,7,8";
2655        let schema = Arc::new(Schema::new(vec![
2656            Field::new("a", DataType::Int32, true),
2657            Field::new("b", DataType::Int32, true),
2658            Field::new("c", DataType::Int32, true),
2659        ]));
2660
2661        let reader = ReaderBuilder::new(schema.clone())
2662            .with_header(true)
2663            .with_truncated_rows(true)
2664            .build(Cursor::new(data))
2665            .unwrap();
2666
2667        let batches = reader.collect::<Result<Vec<_>, _>>();
2668        assert!(batches.is_ok());
2669        let batch = batches.unwrap().into_iter().next().unwrap();
2670        // Empty rows are skipped by the underlying csv parser
2671        assert_eq!(batch.num_rows(), 3);
2672
2673        let reader = ReaderBuilder::new(schema.clone())
2674            .with_header(true)
2675            .with_truncated_rows(false)
2676            .build(Cursor::new(data))
2677            .unwrap();
2678
2679        let batches = reader.collect::<Result<Vec<_>, _>>();
2680        assert!(match batches {
2681            Err(ArrowError::CsvError(e)) => e.to_string().contains("incorrect number of fields"),
2682            _ => false,
2683        });
2684    }
2685
2686    #[test]
2687    fn test_truncated_rows_csv() {
2688        let file = File::open("test/data/truncated_rows.csv").unwrap();
2689        let schema = Arc::new(Schema::new(vec![
2690            Field::new("Name", DataType::Utf8, true),
2691            Field::new("Age", DataType::UInt32, true),
2692            Field::new("Occupation", DataType::Utf8, true),
2693            Field::new("DOB", DataType::Date32, true),
2694        ]));
2695        let reader = ReaderBuilder::new(schema.clone())
2696            .with_header(true)
2697            .with_batch_size(24)
2698            .with_truncated_rows(true);
2699        let csv = reader.build(file).unwrap();
2700        let batches = csv.collect::<Result<Vec<_>, _>>().unwrap();
2701
2702        assert_eq!(batches.len(), 1);
2703        let batch = &batches[0];
2704        assert_eq!(batch.num_rows(), 6);
2705        assert_eq!(batch.num_columns(), 4);
2706        let name = batch
2707            .column(0)
2708            .as_any()
2709            .downcast_ref::<StringArray>()
2710            .unwrap();
2711        let age = batch
2712            .column(1)
2713            .as_any()
2714            .downcast_ref::<UInt32Array>()
2715            .unwrap();
2716        let occupation = batch
2717            .column(2)
2718            .as_any()
2719            .downcast_ref::<StringArray>()
2720            .unwrap();
2721        let dob = batch
2722            .column(3)
2723            .as_any()
2724            .downcast_ref::<Date32Array>()
2725            .unwrap();
2726
2727        assert_eq!(name.value(0), "A1");
2728        assert_eq!(name.value(1), "B2");
2729        assert!(name.is_null(2));
2730        assert_eq!(name.value(3), "C3");
2731        assert_eq!(name.value(4), "D4");
2732        assert_eq!(name.value(5), "E5");
2733
2734        assert_eq!(age.value(0), 34);
2735        assert_eq!(age.value(1), 29);
2736        assert!(age.is_null(2));
2737        assert_eq!(age.value(3), 45);
2738        assert!(age.is_null(4));
2739        assert_eq!(age.value(5), 31);
2740
2741        assert_eq!(occupation.value(0), "Engineer");
2742        assert_eq!(occupation.value(1), "Doctor");
2743        assert!(occupation.is_null(2));
2744        assert_eq!(occupation.value(3), "Artist");
2745        assert!(occupation.is_null(4));
2746        assert!(occupation.is_null(5));
2747
2748        assert_eq!(dob.value(0), 5675);
2749        assert!(dob.is_null(1));
2750        assert!(dob.is_null(2));
2751        assert_eq!(dob.value(3), -1858);
2752        assert!(dob.is_null(4));
2753        assert!(dob.is_null(5));
2754    }
2755
2756    /// Schema used by the `truncated_row_count` tests below
2757    fn truncated_row_count_schema() -> SchemaRef {
2758        Arc::new(Schema::new(vec![
2759            Field::new("name", DataType::Utf8, true),
2760            Field::new("age", DataType::Int32, true),
2761            Field::new("city", DataType::Utf8, true),
2762        ]))
2763    }
2764
2765    #[test]
2766    fn test_truncated_row_count_counts_padded_rows() {
2767        let data = "name,age,city\nAlice,25,Rome\nBob,30\n";
2768
2769        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
2770            .with_header(true)
2771            .with_truncated_rows(true)
2772            .build(Cursor::new(data))
2773            .unwrap();
2774
2775        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
2776        assert_eq!(batches[0].num_rows(), 2);
2777        assert_eq!(reader.truncated_row_count(), 1);
2778    }
2779
2780    #[test]
2781    fn test_truncated_row_count_ignores_empty_trailing_field() {
2782        // "Carol,35," has all three fields, the last one just happens to be empty, so it
2783        // parses to the same null as a padded row would. The count must not be inferred
2784        // from the nulls in the batch
2785        let data = "name,age,city\nAlice,25,Rome\nCarol,35,\n";
2786
2787        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
2788            .with_header(true)
2789            .with_truncated_rows(true)
2790            .build(Cursor::new(data))
2791            .unwrap();
2792
2793        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
2794        let batch = &batches[0];
2795        assert_eq!(batch.num_rows(), 2);
2796        assert!(batch.column(2).is_null(1));
2797        assert_eq!(reader.truncated_row_count(), 0);
2798    }
2799
2800    #[test]
2801    fn test_truncated_row_count_clean_file() {
2802        let data = "name,age,city\nAlice,25,Rome\nBob,30,Milan\n";
2803
2804        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
2805            .with_header(true)
2806            .with_truncated_rows(true)
2807            .build(Cursor::new(data))
2808            .unwrap();
2809
2810        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
2811        assert_eq!(batches[0].num_rows(), 2);
2812        assert_eq!(reader.truncated_row_count(), 0);
2813    }
2814
2815    #[test]
2816    fn test_truncated_row_count_without_truncated_rows() {
2817        let data = "name,age,city\nAlice,25,Rome\nBob,30\n";
2818
2819        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
2820            .with_header(true)
2821            .with_truncated_rows(false)
2822            .build(Cursor::new(data))
2823            .unwrap();
2824
2825        // The short row is an error rather than something to count
2826        let err = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap_err();
2827        assert!(
2828            err.to_string().contains("incorrect number of fields"),
2829            "{err}"
2830        );
2831        assert_eq!(reader.truncated_row_count(), 0);
2832    }
2833
2834    #[test]
2835    fn test_truncated_row_count_accumulates_across_batches() {
2836        // Six short rows read two at a time
2837        let data = "name,age,city\nn0,0\nn1,1\nn2,2\nn3,3\nn4,4\nn5,5\n";
2838
2839        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
2840            .with_header(true)
2841            .with_truncated_rows(true)
2842            .with_batch_size(2)
2843            .build(Cursor::new(data))
2844            .unwrap();
2845
2846        let mut running = vec![];
2847        while let Some(batch) = reader.next().transpose().unwrap() {
2848            assert_eq!(batch.num_rows(), 2);
2849            running.push(reader.truncated_row_count());
2850        }
2851
2852        // A running total, not a per batch count
2853        assert_eq!(running, vec![2, 4, 6]);
2854        assert_eq!(reader.truncated_row_count(), 6);
2855    }
2856
2857    #[test]
2858    fn test_truncated_row_count_excludes_skipped_rows() {
2859        // The header is one field short of the schema, so skipping it pads it. Skipped
2860        // rows never reach a batch and must not be counted
2861        let data = "name,age\nAlice,25,Rome\nBob,30,Milan\n";
2862
2863        let mut reader = ReaderBuilder::new(truncated_row_count_schema())
2864            .with_header(true)
2865            .with_truncated_rows(true)
2866            .build(Cursor::new(data))
2867            .unwrap();
2868
2869        let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
2870        assert_eq!(batches[0].num_rows(), 2);
2871        assert_eq!(reader.truncated_row_count(), 0);
2872    }
2873
2874    #[test]
2875    fn test_truncated_row_count_on_decoder() {
2876        let data = "1,2\n3\n";
2877        let schema = Arc::new(Schema::new(vec![
2878            Field::new("a", DataType::Int32, true),
2879            Field::new("b", DataType::Int32, true),
2880        ]));
2881
2882        let mut decoder = ReaderBuilder::new(schema)
2883            .with_truncated_rows(true)
2884            .build_decoder();
2885
2886        assert_eq!(decoder.truncated_row_count(), 0);
2887        let decoded = decoder.decode(data.as_bytes()).unwrap();
2888        assert_eq!(decoded, data.len());
2889        decoder.flush().unwrap().unwrap();
2890        assert_eq!(decoder.truncated_row_count(), 1);
2891    }
2892
2893    #[test]
2894    fn test_truncated_rows_not_nullable_error() {
2895        let data = "a,b,c\n1,2,3\n4,5";
2896        let schema = Arc::new(Schema::new(vec![
2897            Field::new("a", DataType::Int32, false),
2898            Field::new("b", DataType::Int32, false),
2899            Field::new("c", DataType::Int32, false),
2900        ]));
2901
2902        let reader = ReaderBuilder::new(schema.clone())
2903            .with_header(true)
2904            .with_truncated_rows(true)
2905            .build(Cursor::new(data))
2906            .unwrap();
2907
2908        let batches = reader.collect::<Result<Vec<_>, _>>();
2909        assert!(match batches {
2910            Err(ArrowError::InvalidArgumentError(e)) =>
2911                e.to_string().contains("contains null values"),
2912            _ => false,
2913        });
2914    }
2915
2916    #[test]
2917    #[cfg_attr(miri, ignore)] // Takes too long
2918    fn test_buffered() {
2919        let tests = [
2920            ("test/data/uk_cities.csv", false, 37),
2921            ("test/data/various_types.csv", true, 10),
2922            ("test/data/decimal_test.csv", false, 10),
2923        ];
2924
2925        for (path, has_header, expected_rows) in tests {
2926            let (schema, _) = Format::default()
2927                .infer_schema(File::open(path).unwrap(), None)
2928                .unwrap();
2929            let schema = Arc::new(schema);
2930
2931            for batch_size in [1, 4] {
2932                for capacity in [1, 3, 7, 100] {
2933                    let reader = ReaderBuilder::new(schema.clone())
2934                        .with_batch_size(batch_size)
2935                        .with_header(has_header)
2936                        .build(File::open(path).unwrap())
2937                        .unwrap();
2938
2939                    let expected = reader.collect::<Result<Vec<_>, _>>().unwrap();
2940
2941                    assert_eq!(
2942                        expected.iter().map(|x| x.num_rows()).sum::<usize>(),
2943                        expected_rows
2944                    );
2945
2946                    let buffered =
2947                        std::io::BufReader::with_capacity(capacity, File::open(path).unwrap());
2948
2949                    let reader = ReaderBuilder::new(schema.clone())
2950                        .with_batch_size(batch_size)
2951                        .with_header(has_header)
2952                        .build_buffered(buffered)
2953                        .unwrap();
2954
2955                    let actual = reader.collect::<Result<Vec<_>, _>>().unwrap();
2956                    assert_eq!(expected, actual)
2957                }
2958            }
2959        }
2960    }
2961
2962    fn err_test(csv: &[u8], expected: &str) {
2963        fn err_test_with_schema(csv: &[u8], expected: &str, schema: Arc<Schema>) {
2964            let buffer = std::io::BufReader::with_capacity(2, Cursor::new(csv));
2965            let b = ReaderBuilder::new(schema)
2966                .with_batch_size(2)
2967                .build_buffered(buffer)
2968                .unwrap();
2969            let err = b.collect::<Result<Vec<_>, _>>().unwrap_err().to_string();
2970            assert_eq!(err, expected)
2971        }
2972
2973        let schema_utf8 = Arc::new(Schema::new(vec![
2974            Field::new("text1", DataType::Utf8, true),
2975            Field::new("text2", DataType::Utf8, true),
2976        ]));
2977        err_test_with_schema(csv, expected, schema_utf8);
2978
2979        let schema_utf8view = Arc::new(Schema::new(vec![
2980            Field::new("text1", DataType::Utf8View, true),
2981            Field::new("text2", DataType::Utf8View, true),
2982        ]));
2983        err_test_with_schema(csv, expected, schema_utf8view);
2984    }
2985
2986    #[test]
2987    fn test_invalid_utf8() {
2988        err_test(
2989            b"sdf,dsfg\ndfd,hgh\xFFue\n,sds\nFalhghse,",
2990            "Csv error: Encountered invalid UTF-8 data for line 2 and field 2",
2991        );
2992
2993        err_test(
2994            b"sdf,dsfg\ndksdk,jf\nd\xFFfd,hghue\n,sds\nFalhghse,",
2995            "Csv error: Encountered invalid UTF-8 data for line 3 and field 1",
2996        );
2997
2998        err_test(
2999            b"sdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
3000            "Csv error: Encountered invalid UTF-8 data for line 5 and field 2",
3001        );
3002
3003        err_test(
3004            b"\xFFsdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
3005            "Csv error: Encountered invalid UTF-8 data for line 1 and field 1",
3006        );
3007    }
3008
3009    struct InstrumentedRead<R> {
3010        r: R,
3011        fill_count: usize,
3012        fill_sizes: Vec<usize>,
3013    }
3014
3015    impl<R> InstrumentedRead<R> {
3016        fn new(r: R) -> Self {
3017            Self {
3018                r,
3019                fill_count: 0,
3020                fill_sizes: vec![],
3021            }
3022        }
3023    }
3024
3025    impl<R: Seek> Seek for InstrumentedRead<R> {
3026        fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
3027            self.r.seek(pos)
3028        }
3029    }
3030
3031    impl<R: BufRead> Read for InstrumentedRead<R> {
3032        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
3033            self.r.read(buf)
3034        }
3035    }
3036
3037    impl<R: BufRead> BufRead for InstrumentedRead<R> {
3038        fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
3039            self.fill_count += 1;
3040            let buf = self.r.fill_buf()?;
3041            self.fill_sizes.push(buf.len());
3042            Ok(buf)
3043        }
3044
3045        fn consume(&mut self, amt: usize) {
3046            self.r.consume(amt)
3047        }
3048    }
3049
3050    #[test]
3051    fn test_io() {
3052        let schema = Arc::new(Schema::new(vec![
3053            Field::new("a", DataType::Utf8, false),
3054            Field::new("b", DataType::Utf8, false),
3055        ]));
3056        let csv = "foo,bar\nbaz,foo\na,b\nc,d";
3057        let mut read = InstrumentedRead::new(Cursor::new(csv.as_bytes()));
3058        let reader = ReaderBuilder::new(schema)
3059            .with_batch_size(3)
3060            .build_buffered(&mut read)
3061            .unwrap();
3062
3063        let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3064        assert_eq!(batches.len(), 2);
3065        assert_eq!(batches[0].num_rows(), 3);
3066        assert_eq!(batches[1].num_rows(), 1);
3067
3068        // Expect 4 calls to fill_buf
3069        // 1. Read first 3 rows
3070        // 2. Read final row
3071        // 3. Delimit and flush final row
3072        // 4. Iterator finished
3073        assert_eq!(&read.fill_sizes, &[23, 3, 0, 0]);
3074        assert_eq!(read.fill_count, 4);
3075    }
3076
3077    #[test]
3078    #[cfg_attr(miri, ignore)] // Takes too long
3079    fn test_inference() {
3080        let cases: &[(&[&str], DataType)] = &[
3081            (&[], DataType::Null),
3082            (&["false", "12"], DataType::Utf8),
3083            (&["12", "cupcakes"], DataType::Utf8),
3084            (&["12", "12.4"], DataType::Float64),
3085            (&["14050", "24332"], DataType::Int64),
3086            (&["14050.0", "true"], DataType::Utf8),
3087            (&["14050", "2020-03-19 00:00:00"], DataType::Utf8),
3088            (&["14050", "2340.0", "2020-03-19 00:00:00"], DataType::Utf8),
3089            (
3090                &["2020-03-19 02:00:00", "2020-03-19 00:00:00"],
3091                DataType::Timestamp(TimeUnit::Second, None),
3092            ),
3093            (&["2020-03-19", "2020-03-20"], DataType::Date32),
3094            (
3095                &["2020-03-19", "2020-03-19 02:00:00", "2020-03-19 00:00:00"],
3096                DataType::Timestamp(TimeUnit::Second, None),
3097            ),
3098            (
3099                &[
3100                    "2020-03-19",
3101                    "2020-03-19 02:00:00",
3102                    "2020-03-19 00:00:00.000",
3103                ],
3104                DataType::Timestamp(TimeUnit::Millisecond, None),
3105            ),
3106            (
3107                &[
3108                    "2020-03-19",
3109                    "2020-03-19 02:00:00",
3110                    "2020-03-19 00:00:00.000000",
3111                ],
3112                DataType::Timestamp(TimeUnit::Microsecond, None),
3113            ),
3114            (
3115                &["2020-03-19 02:00:00+02:00", "2020-03-19 02:00:00Z"],
3116                DataType::Timestamp(TimeUnit::Second, None),
3117            ),
3118            (
3119                &[
3120                    "2020-03-19",
3121                    "2020-03-19 02:00:00+02:00",
3122                    "2020-03-19 02:00:00Z",
3123                    "2020-03-19 02:00:00.12Z",
3124                ],
3125                DataType::Timestamp(TimeUnit::Millisecond, None),
3126            ),
3127            (
3128                &[
3129                    "2020-03-19",
3130                    "2020-03-19 02:00:00.000000000",
3131                    "2020-03-19 00:00:00.000000",
3132                ],
3133                DataType::Timestamp(TimeUnit::Nanosecond, None),
3134            ),
3135        ];
3136
3137        for (values, expected) in cases {
3138            let mut t = InferredDataType::default();
3139            for v in *values {
3140                t.update(v)
3141            }
3142            assert_eq!(&t.get(), expected, "{values:?}")
3143        }
3144    }
3145
3146    #[test]
3147    #[cfg_attr(miri, ignore)] // Takes too long
3148    fn test_record_length_mismatch() {
3149        let csv = "\
3150        a,b,c\n\
3151        1,2,3\n\
3152        4,5\n\
3153        6,7,8";
3154        let mut read = Cursor::new(csv.as_bytes());
3155        let result = Format::default()
3156            .with_header(true)
3157            .infer_schema(&mut read, None);
3158        assert!(result.is_err());
3159        // Include line number in the error message to help locate and fix the issue
3160        assert_eq!(
3161            result.err().unwrap().to_string(),
3162            "Csv error: Encountered unequal lengths between records on CSV file. Expected 3 records, found 2 records at line 3"
3163        );
3164    }
3165
3166    #[test]
3167    fn test_comment() {
3168        let schema = Schema::new(vec![
3169            Field::new("a", DataType::Int8, false),
3170            Field::new("b", DataType::Int8, false),
3171        ]);
3172
3173        let csv = "# comment1 \n1,2\n#comment2\n11,22";
3174        let mut read = Cursor::new(csv.as_bytes());
3175        let reader = ReaderBuilder::new(Arc::new(schema))
3176            .with_comment(b'#')
3177            .build(&mut read)
3178            .unwrap();
3179
3180        let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3181        assert_eq!(batches.len(), 1);
3182        let b = batches.first().unwrap();
3183        assert_eq!(b.num_columns(), 2);
3184        assert_eq!(
3185            b.column(0)
3186                .as_any()
3187                .downcast_ref::<Int8Array>()
3188                .unwrap()
3189                .values(),
3190            &vec![1, 11]
3191        );
3192        assert_eq!(
3193            b.column(1)
3194                .as_any()
3195                .downcast_ref::<Int8Array>()
3196                .unwrap()
3197                .values(),
3198            &vec![2, 22]
3199        );
3200    }
3201
3202    #[test]
3203    fn test_parse_string_view_single_column() {
3204        let csv = ["foo", "something_cannot_be_inlined", "foobar"].join("\n");
3205        let schema = Arc::new(Schema::new(vec![Field::new(
3206            "c1",
3207            DataType::Utf8View,
3208            true,
3209        )]));
3210
3211        let mut decoder = ReaderBuilder::new(schema).build_decoder();
3212
3213        let decoded = decoder.decode(csv.as_bytes()).unwrap();
3214        assert_eq!(decoded, csv.len());
3215        decoder.decode(&[]).unwrap();
3216
3217        let batch = decoder.flush().unwrap().unwrap();
3218        assert_eq!(batch.num_columns(), 1);
3219        assert_eq!(batch.num_rows(), 3);
3220        let col = batch.column(0).as_string_view();
3221        assert_eq!(col.data_type(), &DataType::Utf8View);
3222        assert_eq!(col.value(0), "foo");
3223        assert_eq!(col.value(1), "something_cannot_be_inlined");
3224        assert_eq!(col.value(2), "foobar");
3225    }
3226
3227    #[test]
3228    fn test_parse_string_view_multi_column() {
3229        let csv = ["foo,", ",something_cannot_be_inlined", "foobarfoobar,bar"].join("\n");
3230        let schema = Arc::new(Schema::new(vec![
3231            Field::new("c1", DataType::Utf8View, true),
3232            Field::new("c2", DataType::Utf8View, true),
3233        ]));
3234
3235        let mut decoder = ReaderBuilder::new(schema).build_decoder();
3236
3237        let decoded = decoder.decode(csv.as_bytes()).unwrap();
3238        assert_eq!(decoded, csv.len());
3239        decoder.decode(&[]).unwrap();
3240
3241        let batch = decoder.flush().unwrap().unwrap();
3242        assert_eq!(batch.num_columns(), 2);
3243        assert_eq!(batch.num_rows(), 3);
3244        let c1 = batch.column(0).as_string_view();
3245        let c2 = batch.column(1).as_string_view();
3246        assert_eq!(c1.data_type(), &DataType::Utf8View);
3247        assert_eq!(c2.data_type(), &DataType::Utf8View);
3248
3249        assert!(!c1.is_null(0));
3250        assert!(c1.is_null(1));
3251        assert!(!c1.is_null(2));
3252        assert_eq!(c1.value(0), "foo");
3253        assert_eq!(c1.value(2), "foobarfoobar");
3254
3255        assert!(c2.is_null(0));
3256        assert!(!c2.is_null(1));
3257        assert!(!c2.is_null(2));
3258        assert_eq!(c2.value(1), "something_cannot_be_inlined");
3259        assert_eq!(c2.value(2), "bar");
3260    }
3261
3262    #[test]
3263    #[cfg_attr(miri, ignore)] // Unsupported inline assembly
3264    fn test_float_precision() {
3265        let data = [
3266            "f16,f32,f64",
3267            "1.5,1.5,1.5",
3268            "0.25,0.25,0.25",
3269            "1.23456789,1.23456789,1.23456789",
3270            "1.234567890123456,1.234567890123456,1.234567890123456",
3271            "-2.5,-2.5,-2.5",
3272            "0,0,0",
3273            ",,",
3274        ]
3275        .join("\n");
3276
3277        let schema = Schema::new(vec![
3278            Field::new("f16", DataType::Float16, true),
3279            Field::new("f32", DataType::Float32, true),
3280            Field::new("f64", DataType::Float64, true),
3281        ]);
3282
3283        let mut reader = ReaderBuilder::new(Arc::new(schema))
3284            .with_header(true)
3285            .build(Cursor::new(data))
3286            .unwrap();
3287
3288        let batch = reader.next().unwrap().unwrap();
3289        assert_eq!(batch.num_rows(), 7);
3290
3291        let f16_col = batch.column(0).as_primitive::<Float16Type>();
3292        let f32_col = batch.column(1).as_primitive::<Float32Type>();
3293        let f64_col = batch.column(2).as_primitive::<Float64Type>();
3294
3295        assert_eq!(f16_col.value(0), half::f16::from_f32(1.5));
3296        assert_eq!(f32_col.value(0), 1.5f32);
3297        assert_eq!(f64_col.value(0), 1.5f64);
3298
3299        assert_eq!(f16_col.value(1), half::f16::from_f32(0.25));
3300        assert_eq!(f32_col.value(1), 0.25f32);
3301        assert_eq!(f64_col.value(1), 0.25f64);
3302
3303        assert_eq!(f16_col.value(2), half::f16::from_f32(1.234_567_9));
3304        assert_eq!(f32_col.value(2), 1.234_567_9_f32);
3305        assert_eq!(f64_col.value(2), 1.23456789f64);
3306
3307        assert_eq!(f16_col.value(3), half::f16::from_f64(1.234567890123456f64));
3308        assert_eq!(f32_col.value(3), 1.234_567_9_f32);
3309        assert_eq!(f64_col.value(3), 1.234567890123456f64);
3310
3311        assert_eq!(f16_col.value(4), half::f16::from_f32(-2.5));
3312        assert_eq!(f32_col.value(4), -2.5f32);
3313        assert_eq!(f64_col.value(4), -2.5f64);
3314
3315        assert_eq!(f16_col.value(5), half::f16::from_f32(0.0));
3316        assert_eq!(f32_col.value(5), 0.0f32);
3317        assert_eq!(f64_col.value(5), 0.0f64);
3318
3319        assert!(f16_col.is_null(6));
3320        assert!(f32_col.is_null(6));
3321        assert!(f64_col.is_null(6));
3322    }
3323}