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