1mod 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
185static REGEX_SET: LazyLock<RegexSet> = LazyLock::new(|| {
187 RegexSet::new([
188 r"(?i)^(true)$|^(false)$(?-i)", r"^-?(\d+)$", r"^-?((\d*\.\d+|\d+\.\d*)([eE][-+]?\d+)?|\d+([eE][-+]?\d+))$", r"^\d{4}-\d\d-\d\d$", r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d(?:[^\d\.].*)?$", r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,3}(?:[^\d].*)?$", r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,6}(?:[^\d].*)?$", r"^\d{4}-\d\d-\d\d[T ]\d\d:\d\d:\d\d\.\d{1,9}(?:[^\d].*)?$", ])
197 .unwrap()
198});
199
200#[derive(Debug, Clone, Default)]
202struct NullRegex(Option<Regex>);
203
204impl NullRegex {
205 #[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: u16,
230}
231
232impl InferredDataType {
233 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, b if b != 0 && (b & !0b11111000) == 0 => match b.leading_zeros() {
241 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 fn update(&mut self, string: &str) {
255 self.packed |= if string.starts_with('"') {
256 1 << 8 } 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 1 << 8
261 } else {
262 1 << m
263 }
264 } else if string == "NaN" || string == "nan" || string == "inf" || string == "-inf" {
265 1 << 2 } else {
267 1 << 8 }
269 }
270}
271
272#[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 pub fn with_header(mut self, has_header: bool) -> Self {
291 self.header = has_header;
292 self
293 }
294
295 pub fn with_header_validation(mut self, validate_header: bool) -> Self {
301 self.header_validation = validate_header;
302 self
303 }
304
305 pub fn with_delimiter(mut self, delimiter: u8) -> Self {
307 self.delimiter = Some(delimiter);
308 self
309 }
310
311 pub fn with_escape(mut self, escape: u8) -> Self {
313 self.escape = Some(escape);
314 self
315 }
316
317 pub fn with_quote(mut self, quote: u8) -> Self {
319 self.quote = Some(quote);
320 self
321 }
322
323 pub fn with_terminator(mut self, terminator: u8) -> Self {
325 self.terminator = Some(terminator);
326 self
327 }
328
329 pub fn with_comment(mut self, comment: u8) -> Self {
333 self.comment = Some(comment);
334 self
335 }
336
337 pub fn with_null_regex(mut self, null_regex: Regex) -> Self {
339 self.null_regex = NullRegex(Some(null_regex));
340 self
341 }
342
343 pub fn with_truncated_rows(mut self, allow: bool) -> Self {
350 self.truncated_rows = allow;
351 self
352 }
353
354 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 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 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 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 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 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 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
455pub 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
491type Bounds = Option<(usize, usize)>;
493
494pub type Reader<R> = BufReader<StdBufReader<R>>;
499
500pub struct BufReader<R> {
505 reader: R,
507 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 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 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#[derive(Debug)]
611pub struct Decoder {
612 schema: SchemaRef,
614
615 projection: Option<Vec<usize>>,
617
618 batch_size: usize,
620
621 to_skip: usize,
623
624 header_validation: bool,
626
627 line_number: usize,
629
630 end: usize,
632
633 record_decoder: RecordDecoder,
635
636 null_regex: NullRegex,
638}
639
640impl Decoder {
641 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 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 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 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
727fn 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::Float32 => {
805 build_primitive_array::<Float32Type>(line_number, rows, i, null_regex)
806 }
807 DataType::Float64 => {
808 build_primitive_array::<Float64Type>(line_number, rows, i, null_regex)
809 }
810 DataType::Date32 => {
811 build_primitive_array::<Date32Type>(line_number, rows, i, null_regex)
812 }
813 DataType::Date64 => {
814 build_primitive_array::<Date64Type>(line_number, rows, i, null_regex)
815 }
816 DataType::Time32(TimeUnit::Second) => {
817 build_primitive_array::<Time32SecondType>(line_number, rows, i, null_regex)
818 }
819 DataType::Time32(TimeUnit::Millisecond) => {
820 build_primitive_array::<Time32MillisecondType>(line_number, rows, i, null_regex)
821 }
822 DataType::Time64(TimeUnit::Microsecond) => {
823 build_primitive_array::<Time64MicrosecondType>(line_number, rows, i, null_regex)
824 }
825 DataType::Time64(TimeUnit::Nanosecond) => {
826 build_primitive_array::<Time64NanosecondType>(line_number, rows, i, null_regex)
827 }
828 DataType::Timestamp(TimeUnit::Second, tz) => {
829 build_timestamp_array::<TimestampSecondType>(
830 line_number,
831 rows,
832 i,
833 tz.as_deref(),
834 null_regex,
835 )
836 }
837 DataType::Timestamp(TimeUnit::Millisecond, tz) => {
838 build_timestamp_array::<TimestampMillisecondType>(
839 line_number,
840 rows,
841 i,
842 tz.as_deref(),
843 null_regex,
844 )
845 }
846 DataType::Timestamp(TimeUnit::Microsecond, tz) => {
847 build_timestamp_array::<TimestampMicrosecondType>(
848 line_number,
849 rows,
850 i,
851 tz.as_deref(),
852 null_regex,
853 )
854 }
855 DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
856 build_timestamp_array::<TimestampNanosecondType>(
857 line_number,
858 rows,
859 i,
860 tz.as_deref(),
861 null_regex,
862 )
863 }
864 DataType::Null => Ok(Arc::new({
865 let mut builder = NullBuilder::new();
866 builder.append_nulls(rows.len());
867 builder.finish()
868 }) as ArrayRef),
869 DataType::Utf8 => Ok(Arc::new(
870 rows.iter()
871 .map(|row| {
872 let s = row.get(i);
873 (!null_regex.is_null(s)).then_some(s)
874 })
875 .collect::<StringArray>(),
876 ) as ArrayRef),
877 DataType::Utf8View => Ok(Arc::new(
878 rows.iter()
879 .map(|row| {
880 let s = row.get(i);
881 (!null_regex.is_null(s)).then_some(s)
882 })
883 .collect::<StringViewArray>(),
884 ) as ArrayRef),
885 DataType::Dictionary(key_type, value_type)
886 if value_type.as_ref() == &DataType::Utf8 =>
887 {
888 match key_type.as_ref() {
889 DataType::Int8 => Ok(Arc::new(
890 rows.iter()
891 .map(|row| {
892 let s = row.get(i);
893 (!null_regex.is_null(s)).then_some(s)
894 })
895 .collect::<DictionaryArray<Int8Type>>(),
896 ) as ArrayRef),
897 DataType::Int16 => Ok(Arc::new(
898 rows.iter()
899 .map(|row| {
900 let s = row.get(i);
901 (!null_regex.is_null(s)).then_some(s)
902 })
903 .collect::<DictionaryArray<Int16Type>>(),
904 ) as ArrayRef),
905 DataType::Int32 => Ok(Arc::new(
906 rows.iter()
907 .map(|row| {
908 let s = row.get(i);
909 (!null_regex.is_null(s)).then_some(s)
910 })
911 .collect::<DictionaryArray<Int32Type>>(),
912 ) as ArrayRef),
913 DataType::Int64 => Ok(Arc::new(
914 rows.iter()
915 .map(|row| {
916 let s = row.get(i);
917 (!null_regex.is_null(s)).then_some(s)
918 })
919 .collect::<DictionaryArray<Int64Type>>(),
920 ) as ArrayRef),
921 DataType::UInt8 => Ok(Arc::new(
922 rows.iter()
923 .map(|row| {
924 let s = row.get(i);
925 (!null_regex.is_null(s)).then_some(s)
926 })
927 .collect::<DictionaryArray<UInt8Type>>(),
928 ) as ArrayRef),
929 DataType::UInt16 => Ok(Arc::new(
930 rows.iter()
931 .map(|row| {
932 let s = row.get(i);
933 (!null_regex.is_null(s)).then_some(s)
934 })
935 .collect::<DictionaryArray<UInt16Type>>(),
936 ) as ArrayRef),
937 DataType::UInt32 => Ok(Arc::new(
938 rows.iter()
939 .map(|row| {
940 let s = row.get(i);
941 (!null_regex.is_null(s)).then_some(s)
942 })
943 .collect::<DictionaryArray<UInt32Type>>(),
944 ) as ArrayRef),
945 DataType::UInt64 => Ok(Arc::new(
946 rows.iter()
947 .map(|row| {
948 let s = row.get(i);
949 (!null_regex.is_null(s)).then_some(s)
950 })
951 .collect::<DictionaryArray<UInt64Type>>(),
952 ) as ArrayRef),
953 _ => Err(ArrowError::ParseError(format!(
954 "Unsupported dictionary key type {key_type}"
955 ))),
956 }
957 }
958 other => Err(ArrowError::ParseError(format!(
959 "Unsupported data type {other:?}"
960 ))),
961 }
962 })
963 .collect();
964
965 let projected_fields: Fields = projection.iter().map(|i| fields[*i].clone()).collect();
966
967 let projected_schema = Arc::new(match metadata {
968 None => Schema::new(projected_fields),
969 Some(metadata) => Schema::new_with_metadata(projected_fields, metadata),
970 });
971
972 arrays.and_then(|arr| {
973 RecordBatch::try_new_with_options(
974 projected_schema,
975 arr,
976 &RecordBatchOptions::new()
977 .with_match_field_names(true)
978 .with_row_count(Some(rows.len())),
979 )
980 })
981}
982
983fn parse_bool(string: &str) -> Option<bool> {
984 if string.eq_ignore_ascii_case("false") {
985 Some(false)
986 } else if string.eq_ignore_ascii_case("true") {
987 Some(true)
988 } else {
989 None
990 }
991}
992
993fn build_decimal_array<T: DecimalType>(
995 _line_number: usize,
996 rows: &StringRecords<'_>,
997 col_idx: usize,
998 precision: u8,
999 scale: i8,
1000 null_regex: &NullRegex,
1001) -> Result<ArrayRef, ArrowError> {
1002 let mut decimal_builder = PrimitiveBuilder::<T>::with_capacity(rows.len());
1003 for row in rows.iter() {
1004 let s = row.get(col_idx);
1005 if null_regex.is_null(s) {
1006 decimal_builder.append_null();
1008 } else {
1009 let decimal_value: Result<T::Native, _> = parse_decimal::<T>(s, precision, scale);
1010 match decimal_value {
1011 Ok(v) => {
1012 decimal_builder.append_value(v);
1013 }
1014 Err(e) => {
1015 return Err(e);
1016 }
1017 }
1018 }
1019 }
1020 Ok(Arc::new(
1021 decimal_builder
1022 .finish()
1023 .with_precision_and_scale(precision, scale)?,
1024 ))
1025}
1026
1027fn build_primitive_array<T: ArrowPrimitiveType + Parser>(
1029 line_number: usize,
1030 rows: &StringRecords<'_>,
1031 col_idx: usize,
1032 null_regex: &NullRegex,
1033) -> Result<ArrayRef, ArrowError> {
1034 rows.iter()
1035 .enumerate()
1036 .map(|(row_index, row)| {
1037 let s = row.get(col_idx);
1038 if null_regex.is_null(s) {
1039 return Ok(None);
1040 }
1041
1042 match T::parse(s) {
1043 Some(e) => Ok(Some(e)),
1044 None => Err(ArrowError::ParseError(format!(
1045 "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1047 s,
1048 T::DATA_TYPE,
1049 col_idx,
1050 line_number + row_index,
1051 row
1052 ))),
1053 }
1054 })
1055 .collect::<Result<PrimitiveArray<T>, ArrowError>>()
1056 .map(|e| Arc::new(e) as ArrayRef)
1057}
1058
1059fn build_timestamp_array<T: ArrowTimestampType>(
1060 line_number: usize,
1061 rows: &StringRecords<'_>,
1062 col_idx: usize,
1063 timezone: Option<&str>,
1064 null_regex: &NullRegex,
1065) -> Result<ArrayRef, ArrowError> {
1066 Ok(Arc::new(match timezone {
1067 Some(timezone) => {
1068 let tz: Tz = timezone.parse()?;
1069 build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &tz, null_regex)?
1070 .with_timezone(timezone)
1071 }
1072 None => build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &Utc, null_regex)?,
1073 }))
1074}
1075
1076fn build_timestamp_array_impl<T: ArrowTimestampType, Tz: TimeZone>(
1077 line_number: usize,
1078 rows: &StringRecords<'_>,
1079 col_idx: usize,
1080 timezone: &Tz,
1081 null_regex: &NullRegex,
1082) -> Result<PrimitiveArray<T>, ArrowError> {
1083 rows.iter()
1084 .enumerate()
1085 .map(|(row_index, row)| {
1086 let s = row.get(col_idx);
1087 if null_regex.is_null(s) {
1088 return Ok(None);
1089 }
1090
1091 let date = string_to_datetime(timezone, s)
1092 .and_then(|date| match T::UNIT {
1093 TimeUnit::Second => Ok(date.timestamp()),
1094 TimeUnit::Millisecond => Ok(date.timestamp_millis()),
1095 TimeUnit::Microsecond => Ok(date.timestamp_micros()),
1096 TimeUnit::Nanosecond => date.timestamp_nanos_opt().ok_or_else(|| {
1097 ArrowError::ParseError(format!(
1098 "{} would overflow 64-bit signed nanoseconds",
1099 date.to_rfc3339(),
1100 ))
1101 }),
1102 })
1103 .map_err(|e| {
1104 ArrowError::ParseError(format!(
1105 "Error parsing column {col_idx} at line {}: {}",
1106 line_number + row_index,
1107 e
1108 ))
1109 })?;
1110 Ok(Some(date))
1111 })
1112 .collect()
1113}
1114
1115fn build_boolean_array(
1117 line_number: usize,
1118 rows: &StringRecords<'_>,
1119 col_idx: usize,
1120 null_regex: &NullRegex,
1121) -> Result<ArrayRef, ArrowError> {
1122 rows.iter()
1123 .enumerate()
1124 .map(|(row_index, row)| {
1125 let s = row.get(col_idx);
1126 if null_regex.is_null(s) {
1127 return Ok(None);
1128 }
1129 let parsed = parse_bool(s);
1130 match parsed {
1131 Some(e) => Ok(Some(e)),
1132 None => Err(ArrowError::ParseError(format!(
1133 "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1135 s,
1136 "Boolean",
1137 col_idx,
1138 line_number + row_index,
1139 row
1140 ))),
1141 }
1142 })
1143 .collect::<Result<BooleanArray, _>>()
1144 .map(|e| Arc::new(e) as ArrayRef)
1145}
1146
1147#[derive(Debug)]
1149pub struct ReaderBuilder {
1150 schema: SchemaRef,
1152 format: Format,
1154 batch_size: usize,
1158 bounds: Bounds,
1160 projection: Option<Vec<usize>>,
1162}
1163
1164impl ReaderBuilder {
1165 pub fn new(schema: SchemaRef) -> ReaderBuilder {
1188 Self {
1189 schema,
1190 format: Format::default(),
1191 batch_size: 1024,
1192 bounds: None,
1193 projection: None,
1194 }
1195 }
1196
1197 pub fn with_header(mut self, has_header: bool) -> Self {
1199 self.format.header = has_header;
1200 self
1201 }
1202
1203 pub fn with_header_validation(mut self, validate_header: bool) -> Self {
1207 self.format.header_validation = validate_header;
1208 self
1209 }
1210
1211 pub fn with_format(mut self, format: Format) -> Self {
1213 self.format = format;
1214 self
1215 }
1216
1217 pub fn with_delimiter(mut self, delimiter: u8) -> Self {
1219 self.format.delimiter = Some(delimiter);
1220 self
1221 }
1222
1223 pub fn with_escape(mut self, escape: u8) -> Self {
1225 self.format.escape = Some(escape);
1226 self
1227 }
1228
1229 pub fn with_quote(mut self, quote: u8) -> Self {
1231 self.format.quote = Some(quote);
1232 self
1233 }
1234
1235 pub fn with_terminator(mut self, terminator: u8) -> Self {
1237 self.format.terminator = Some(terminator);
1238 self
1239 }
1240
1241 pub fn with_comment(mut self, comment: u8) -> Self {
1243 self.format.comment = Some(comment);
1244 self
1245 }
1246
1247 pub fn with_null_regex(mut self, null_regex: Regex) -> Self {
1249 self.format.null_regex = NullRegex(Some(null_regex));
1250 self
1251 }
1252
1253 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1255 self.batch_size = batch_size;
1256 self
1257 }
1258
1259 pub fn with_bounds(mut self, start: usize, end: usize) -> Self {
1262 self.bounds = Some((start, end));
1263 self
1264 }
1265
1266 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1268 self.projection = Some(projection);
1269 self
1270 }
1271
1272 pub fn with_truncated_rows(mut self, allow: bool) -> Self {
1279 self.format.truncated_rows = allow;
1280 self
1281 }
1282
1283 pub fn build<R: Read>(self, reader: R) -> Result<Reader<R>, ArrowError> {
1288 self.build_buffered(StdBufReader::new(reader))
1289 }
1290
1291 pub fn build_buffered<R: BufRead>(self, reader: R) -> Result<BufReader<R>, ArrowError> {
1293 Ok(BufReader {
1294 reader,
1295 decoder: self.build_decoder(),
1296 })
1297 }
1298
1299 pub fn build_decoder(self) -> Decoder {
1301 let delimiter = self.format.build_parser();
1302 let record_decoder = RecordDecoder::new(
1303 delimiter,
1304 self.schema.fields().len(),
1305 self.format.truncated_rows,
1306 );
1307
1308 let header = self.format.header as usize;
1309
1310 let (start, end) = match self.bounds {
1311 Some((start, end)) => (start + header, end + header),
1312 None => (header, usize::MAX),
1313 };
1314
1315 Decoder {
1316 schema: self.schema,
1317 to_skip: start,
1318 header_validation: self.format.header && self.format.header_validation,
1319 record_decoder,
1320 line_number: start,
1321 end,
1322 projection: self.projection,
1323 batch_size: self.batch_size,
1324 null_regex: self.format.null_regex,
1325 }
1326 }
1327}
1328
1329#[cfg(test)]
1330mod tests {
1331 use super::*;
1332
1333 use std::io::{Cursor, Seek, SeekFrom, Write};
1334 use tempfile::NamedTempFile;
1335
1336 use arrow_array::cast::AsArray;
1337
1338 #[test]
1339 fn test_csv() {
1340 let schema = Arc::new(Schema::new(vec![
1341 Field::new("city", DataType::Utf8, false),
1342 Field::new("lat", DataType::Float64, false),
1343 Field::new("lng", DataType::Float64, false),
1344 ]));
1345
1346 let file = File::open("test/data/uk_cities.csv").unwrap();
1347 let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1348 assert_eq!(schema, csv.schema());
1349 let batch = csv.next().unwrap().unwrap();
1350 assert_eq!(37, batch.num_rows());
1351 assert_eq!(3, batch.num_columns());
1352
1353 let lat = batch.column(1).as_primitive::<Float64Type>();
1355 assert_eq!(57.653484, lat.value(0));
1356
1357 let city = batch.column(0).as_string::<i32>();
1359
1360 assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1361 }
1362
1363 #[test]
1364 fn test_csv_schema_metadata() {
1365 let mut metadata = std::collections::HashMap::new();
1366 metadata.insert("foo".to_owned(), "bar".to_owned());
1367 let schema = Arc::new(Schema::new_with_metadata(
1368 vec![
1369 Field::new("city", DataType::Utf8, false),
1370 Field::new("lat", DataType::Float64, false),
1371 Field::new("lng", DataType::Float64, false),
1372 ],
1373 metadata.clone(),
1374 ));
1375
1376 let file = File::open("test/data/uk_cities.csv").unwrap();
1377
1378 let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1379 assert_eq!(schema, csv.schema());
1380 let batch = csv.next().unwrap().unwrap();
1381 assert_eq!(37, batch.num_rows());
1382 assert_eq!(3, batch.num_columns());
1383
1384 assert_eq!(&metadata, batch.schema().metadata());
1385 }
1386
1387 #[test]
1388 fn test_csv_reader_with_decimal() {
1389 let schema = Arc::new(Schema::new(vec![
1390 Field::new("city", DataType::Utf8, false),
1391 Field::new("lat", DataType::Decimal128(38, 6), false),
1392 Field::new("lng", DataType::Decimal256(76, 6), false),
1393 ]));
1394
1395 let file = File::open("test/data/decimal_test.csv").unwrap();
1396
1397 let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1398 let batch = csv.next().unwrap().unwrap();
1399 let lat = batch
1401 .column(1)
1402 .as_any()
1403 .downcast_ref::<Decimal128Array>()
1404 .unwrap();
1405
1406 assert_eq!("57.653484", lat.value_as_string(0));
1407 assert_eq!("53.002666", lat.value_as_string(1));
1408 assert_eq!("52.412811", lat.value_as_string(2));
1409 assert_eq!("51.481583", lat.value_as_string(3));
1410 assert_eq!("12.123456", lat.value_as_string(4));
1411 assert_eq!("50.760000", lat.value_as_string(5));
1412 assert_eq!("0.123000", lat.value_as_string(6));
1413 assert_eq!("123.000000", lat.value_as_string(7));
1414 assert_eq!("123.000000", lat.value_as_string(8));
1415 assert_eq!("-50.760000", lat.value_as_string(9));
1416
1417 let lng = batch
1418 .column(2)
1419 .as_any()
1420 .downcast_ref::<Decimal256Array>()
1421 .unwrap();
1422
1423 assert_eq!("-3.335724", lng.value_as_string(0));
1424 assert_eq!("-2.179404", lng.value_as_string(1));
1425 assert_eq!("-1.778197", lng.value_as_string(2));
1426 assert_eq!("-3.179090", lng.value_as_string(3));
1427 assert_eq!("-3.179090", lng.value_as_string(4));
1428 assert_eq!("0.290472", lng.value_as_string(5));
1429 assert_eq!("0.290472", lng.value_as_string(6));
1430 assert_eq!("0.290472", lng.value_as_string(7));
1431 assert_eq!("0.290472", lng.value_as_string(8));
1432 assert_eq!("0.290472", lng.value_as_string(9));
1433 }
1434
1435 #[test]
1436 fn test_csv_reader_with_decimal_3264() {
1437 let schema = Arc::new(Schema::new(vec![
1438 Field::new("city", DataType::Utf8, false),
1439 Field::new("lat", DataType::Decimal32(9, 6), false),
1440 Field::new("lng", DataType::Decimal64(16, 6), false),
1441 ]));
1442
1443 let file = File::open("test/data/decimal_test.csv").unwrap();
1444
1445 let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1446 let batch = csv.next().unwrap().unwrap();
1447 let lat = batch
1449 .column(1)
1450 .as_any()
1451 .downcast_ref::<Decimal32Array>()
1452 .unwrap();
1453
1454 assert_eq!("57.653484", lat.value_as_string(0));
1455 assert_eq!("53.002666", lat.value_as_string(1));
1456 assert_eq!("52.412811", lat.value_as_string(2));
1457 assert_eq!("51.481583", lat.value_as_string(3));
1458 assert_eq!("12.123456", lat.value_as_string(4));
1459 assert_eq!("50.760000", lat.value_as_string(5));
1460 assert_eq!("0.123000", lat.value_as_string(6));
1461 assert_eq!("123.000000", lat.value_as_string(7));
1462 assert_eq!("123.000000", lat.value_as_string(8));
1463 assert_eq!("-50.760000", lat.value_as_string(9));
1464
1465 let lng = batch
1466 .column(2)
1467 .as_any()
1468 .downcast_ref::<Decimal64Array>()
1469 .unwrap();
1470
1471 assert_eq!("-3.335724", lng.value_as_string(0));
1472 assert_eq!("-2.179404", lng.value_as_string(1));
1473 assert_eq!("-1.778197", lng.value_as_string(2));
1474 assert_eq!("-3.179090", lng.value_as_string(3));
1475 assert_eq!("-3.179090", lng.value_as_string(4));
1476 assert_eq!("0.290472", lng.value_as_string(5));
1477 assert_eq!("0.290472", lng.value_as_string(6));
1478 assert_eq!("0.290472", lng.value_as_string(7));
1479 assert_eq!("0.290472", lng.value_as_string(8));
1480 assert_eq!("0.290472", lng.value_as_string(9));
1481 }
1482
1483 #[test]
1484 fn test_csv_from_buf_reader() {
1485 let schema = Schema::new(vec![
1486 Field::new("city", DataType::Utf8, false),
1487 Field::new("lat", DataType::Float64, false),
1488 Field::new("lng", DataType::Float64, false),
1489 ]);
1490
1491 let file_with_headers = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1492 let file_without_headers = File::open("test/data/uk_cities.csv").unwrap();
1493 let both_files = file_with_headers
1494 .chain(Cursor::new("\n".to_string()))
1495 .chain(file_without_headers);
1496 let mut csv = ReaderBuilder::new(Arc::new(schema))
1497 .with_header(true)
1498 .build(both_files)
1499 .unwrap();
1500 let batch = csv.next().unwrap().unwrap();
1501 assert_eq!(74, batch.num_rows());
1502 assert_eq!(3, batch.num_columns());
1503 }
1504
1505 #[test]
1506 fn test_csv_with_schema_inference() {
1507 let mut file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1508
1509 let (schema, _) = Format::default()
1510 .with_header(true)
1511 .infer_schema(&mut file, None)
1512 .unwrap();
1513
1514 file.rewind().unwrap();
1515 let builder = ReaderBuilder::new(Arc::new(schema)).with_header(true);
1516
1517 let mut csv = builder.build(file).unwrap();
1518 let expected_schema = Schema::new(vec![
1519 Field::new("city", DataType::Utf8, true),
1520 Field::new("lat", DataType::Float64, true),
1521 Field::new("lng", DataType::Float64, true),
1522 ]);
1523 assert_eq!(Arc::new(expected_schema), csv.schema());
1524 let batch = csv.next().unwrap().unwrap();
1525 assert_eq!(37, batch.num_rows());
1526 assert_eq!(3, batch.num_columns());
1527
1528 let lat = batch
1530 .column(1)
1531 .as_any()
1532 .downcast_ref::<Float64Array>()
1533 .unwrap();
1534 assert_eq!(57.653484, lat.value(0));
1535
1536 let city = batch
1538 .column(0)
1539 .as_any()
1540 .downcast_ref::<StringArray>()
1541 .unwrap();
1542
1543 assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1544 }
1545
1546 #[test]
1547 fn test_csv_with_schema_inference_no_headers() {
1548 let mut file = File::open("test/data/uk_cities.csv").unwrap();
1549
1550 let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1551 file.rewind().unwrap();
1552
1553 let mut csv = ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
1554
1555 let schema = csv.schema();
1557 assert_eq!("column_1", schema.field(0).name());
1558 assert_eq!("column_2", schema.field(1).name());
1559 assert_eq!("column_3", schema.field(2).name());
1560 let batch = csv.next().unwrap().unwrap();
1561 let batch_schema = batch.schema();
1562
1563 assert_eq!(schema, batch_schema);
1564 assert_eq!(37, batch.num_rows());
1565 assert_eq!(3, batch.num_columns());
1566
1567 let lat = batch
1569 .column(1)
1570 .as_any()
1571 .downcast_ref::<Float64Array>()
1572 .unwrap();
1573 assert_eq!(57.653484, lat.value(0));
1574
1575 let city = batch
1577 .column(0)
1578 .as_any()
1579 .downcast_ref::<StringArray>()
1580 .unwrap();
1581
1582 assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1583 }
1584
1585 #[test]
1586 fn test_csv_builder_with_bounds() {
1587 let mut file = File::open("test/data/uk_cities.csv").unwrap();
1588
1589 let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1591 file.rewind().unwrap();
1592 let mut csv = ReaderBuilder::new(Arc::new(schema))
1593 .with_bounds(0, 2)
1594 .build(file)
1595 .unwrap();
1596 let batch = csv.next().unwrap().unwrap();
1597
1598 let city = batch
1600 .column(0)
1601 .as_any()
1602 .downcast_ref::<StringArray>()
1603 .unwrap();
1604
1605 assert_eq!("Elgin, Scotland, the UK", city.value(0));
1607
1608 let result = std::panic::catch_unwind(|| city.value(13));
1611 assert!(result.is_err());
1612 }
1613
1614 #[test]
1615 fn test_csv_with_projection() {
1616 let schema = Arc::new(Schema::new(vec![
1617 Field::new("city", DataType::Utf8, false),
1618 Field::new("lat", DataType::Float64, false),
1619 Field::new("lng", DataType::Float64, false),
1620 ]));
1621
1622 let file = File::open("test/data/uk_cities.csv").unwrap();
1623
1624 let mut csv = ReaderBuilder::new(schema)
1625 .with_projection(vec![0, 1])
1626 .build(file)
1627 .unwrap();
1628
1629 let projected_schema = Arc::new(Schema::new(vec![
1630 Field::new("city", DataType::Utf8, false),
1631 Field::new("lat", DataType::Float64, false),
1632 ]));
1633 assert_eq!(projected_schema, csv.schema());
1634 let batch = csv.next().unwrap().unwrap();
1635 assert_eq!(projected_schema, batch.schema());
1636 assert_eq!(37, batch.num_rows());
1637 assert_eq!(2, batch.num_columns());
1638 }
1639
1640 #[test]
1641 fn test_csv_with_dictionary() {
1642 let schema = Arc::new(Schema::new(vec![
1643 Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
1644 Field::new("lat", DataType::Float64, false),
1645 Field::new("lng", DataType::Float64, false),
1646 ]));
1647
1648 let file = File::open("test/data/uk_cities.csv").unwrap();
1649
1650 let mut csv = ReaderBuilder::new(schema)
1651 .with_projection(vec![0, 1])
1652 .build(file)
1653 .unwrap();
1654
1655 let projected_schema = Arc::new(Schema::new(vec![
1656 Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
1657 Field::new("lat", DataType::Float64, false),
1658 ]));
1659 assert_eq!(projected_schema, csv.schema());
1660 let batch = csv.next().unwrap().unwrap();
1661 assert_eq!(projected_schema, batch.schema());
1662 assert_eq!(37, batch.num_rows());
1663 assert_eq!(2, batch.num_columns());
1664
1665 let strings = arrow_cast::cast(batch.column(0), &DataType::Utf8).unwrap();
1666 let strings = strings.as_string::<i32>();
1667
1668 assert_eq!(strings.value(0), "Elgin, Scotland, the UK");
1669 assert_eq!(strings.value(4), "Eastbourne, East Sussex, UK");
1670 assert_eq!(strings.value(29), "Uckfield, East Sussex, UK");
1671 }
1672
1673 #[test]
1674 fn test_csv_with_nullable_dictionary() {
1675 let offset_type = vec![
1676 DataType::Int8,
1677 DataType::Int16,
1678 DataType::Int32,
1679 DataType::Int64,
1680 DataType::UInt8,
1681 DataType::UInt16,
1682 DataType::UInt32,
1683 DataType::UInt64,
1684 ];
1685 for data_type in offset_type {
1686 let file = File::open("test/data/dictionary_nullable_test.csv").unwrap();
1687 let dictionary_type =
1688 DataType::Dictionary(Box::new(data_type), Box::new(DataType::Utf8));
1689 let schema = Arc::new(Schema::new(vec![
1690 Field::new("id", DataType::Utf8, false),
1691 Field::new("name", dictionary_type.clone(), true),
1692 ]));
1693
1694 let mut csv = ReaderBuilder::new(schema)
1695 .build(file.try_clone().unwrap())
1696 .unwrap();
1697
1698 let batch = csv.next().unwrap().unwrap();
1699 assert_eq!(3, batch.num_rows());
1700 assert_eq!(2, batch.num_columns());
1701
1702 let names = arrow_cast::cast(batch.column(1), &dictionary_type).unwrap();
1703 assert!(!names.is_null(2));
1704 assert!(names.is_null(1));
1705 }
1706 }
1707 #[test]
1708 fn test_nulls() {
1709 let schema = Arc::new(Schema::new(vec![
1710 Field::new("c_int", DataType::UInt64, false),
1711 Field::new("c_float", DataType::Float32, true),
1712 Field::new("c_string", DataType::Utf8, true),
1713 Field::new("c_bool", DataType::Boolean, false),
1714 ]));
1715
1716 let file = File::open("test/data/null_test.csv").unwrap();
1717
1718 let mut csv = ReaderBuilder::new(schema)
1719 .with_header(true)
1720 .build(file)
1721 .unwrap();
1722
1723 let batch = csv.next().unwrap().unwrap();
1724
1725 assert!(!batch.column(1).is_null(0));
1726 assert!(!batch.column(1).is_null(1));
1727 assert!(batch.column(1).is_null(2));
1728 assert!(!batch.column(1).is_null(3));
1729 assert!(!batch.column(1).is_null(4));
1730 }
1731
1732 #[test]
1733 fn test_init_nulls() {
1734 let schema = Arc::new(Schema::new(vec![
1735 Field::new("c_int", DataType::UInt64, true),
1736 Field::new("c_float", DataType::Float32, true),
1737 Field::new("c_string", DataType::Utf8, true),
1738 Field::new("c_bool", DataType::Boolean, true),
1739 Field::new("c_null", DataType::Null, true),
1740 ]));
1741 let file = File::open("test/data/init_null_test.csv").unwrap();
1742
1743 let mut csv = ReaderBuilder::new(schema)
1744 .with_header(true)
1745 .build(file)
1746 .unwrap();
1747
1748 let batch = csv.next().unwrap().unwrap();
1749
1750 assert!(batch.column(1).is_null(0));
1751 assert!(!batch.column(1).is_null(1));
1752 assert!(batch.column(1).is_null(2));
1753 assert!(!batch.column(1).is_null(3));
1754 assert!(!batch.column(1).is_null(4));
1755 }
1756
1757 #[test]
1758 fn test_init_nulls_with_inference() {
1759 let format = Format::default().with_header(true).with_delimiter(b',');
1760
1761 let mut file = File::open("test/data/init_null_test.csv").unwrap();
1762 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
1763 file.rewind().unwrap();
1764
1765 let expected_schema = Schema::new(vec![
1766 Field::new("c_int", DataType::Int64, true),
1767 Field::new("c_float", DataType::Float64, true),
1768 Field::new("c_string", DataType::Utf8, true),
1769 Field::new("c_bool", DataType::Boolean, true),
1770 Field::new("c_null", DataType::Null, true),
1771 ]);
1772 assert_eq!(schema, expected_schema);
1773
1774 let mut csv = ReaderBuilder::new(Arc::new(schema))
1775 .with_format(format)
1776 .build(file)
1777 .unwrap();
1778
1779 let batch = csv.next().unwrap().unwrap();
1780
1781 assert!(batch.column(1).is_null(0));
1782 assert!(!batch.column(1).is_null(1));
1783 assert!(batch.column(1).is_null(2));
1784 assert!(!batch.column(1).is_null(3));
1785 assert!(!batch.column(1).is_null(4));
1786 }
1787
1788 #[test]
1789 fn test_custom_nulls() {
1790 let schema = Arc::new(Schema::new(vec![
1791 Field::new("c_int", DataType::UInt64, true),
1792 Field::new("c_float", DataType::Float32, true),
1793 Field::new("c_string", DataType::Utf8, true),
1794 Field::new("c_bool", DataType::Boolean, true),
1795 ]));
1796
1797 let file = File::open("test/data/custom_null_test.csv").unwrap();
1798
1799 let null_regex = Regex::new("^nil$").unwrap();
1800
1801 let mut csv = ReaderBuilder::new(schema)
1802 .with_header(true)
1803 .with_null_regex(null_regex)
1804 .build(file)
1805 .unwrap();
1806
1807 let batch = csv.next().unwrap().unwrap();
1808
1809 assert!(batch.column(0).is_null(1));
1811 assert!(batch.column(1).is_null(2));
1812 assert!(batch.column(3).is_null(4));
1813 assert!(batch.column(2).is_null(3));
1814 assert!(!batch.column(2).is_null(4));
1815 }
1816
1817 #[test]
1818 fn test_nulls_with_inference() {
1819 let mut file = File::open("test/data/various_types.csv").unwrap();
1820 let format = Format::default().with_header(true).with_delimiter(b'|');
1821
1822 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
1823 file.rewind().unwrap();
1824
1825 let builder = ReaderBuilder::new(Arc::new(schema))
1826 .with_format(format)
1827 .with_batch_size(512)
1828 .with_projection(vec![0, 1, 2, 3, 4, 5]);
1829
1830 let mut csv = builder.build(file).unwrap();
1831 let batch = csv.next().unwrap().unwrap();
1832
1833 assert_eq!(10, batch.num_rows());
1834 assert_eq!(6, batch.num_columns());
1835
1836 let schema = batch.schema();
1837
1838 assert_eq!(&DataType::Int64, schema.field(0).data_type());
1839 assert_eq!(&DataType::Float64, schema.field(1).data_type());
1840 assert_eq!(&DataType::Float64, schema.field(2).data_type());
1841 assert_eq!(&DataType::Boolean, schema.field(3).data_type());
1842 assert_eq!(&DataType::Date32, schema.field(4).data_type());
1843 assert_eq!(
1844 &DataType::Timestamp(TimeUnit::Second, None),
1845 schema.field(5).data_type()
1846 );
1847
1848 let names: Vec<&str> = schema.fields().iter().map(|x| x.name().as_str()).collect();
1849 assert_eq!(
1850 names,
1851 vec![
1852 "c_int",
1853 "c_float",
1854 "c_string",
1855 "c_bool",
1856 "c_date",
1857 "c_datetime"
1858 ]
1859 );
1860
1861 assert!(schema.field(0).is_nullable());
1862 assert!(schema.field(1).is_nullable());
1863 assert!(schema.field(2).is_nullable());
1864 assert!(schema.field(3).is_nullable());
1865 assert!(schema.field(4).is_nullable());
1866 assert!(schema.field(5).is_nullable());
1867
1868 assert!(!batch.column(1).is_null(0));
1869 assert!(!batch.column(1).is_null(1));
1870 assert!(batch.column(1).is_null(2));
1871 assert!(!batch.column(1).is_null(3));
1872 assert!(!batch.column(1).is_null(4));
1873 }
1874
1875 #[test]
1876 fn test_custom_nulls_with_inference() {
1877 let mut file = File::open("test/data/custom_null_test.csv").unwrap();
1878
1879 let null_regex = Regex::new("^nil$").unwrap();
1880
1881 let format = Format::default()
1882 .with_header(true)
1883 .with_null_regex(null_regex);
1884
1885 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
1886 file.rewind().unwrap();
1887
1888 let expected_schema = Schema::new(vec![
1889 Field::new("c_int", DataType::Int64, true),
1890 Field::new("c_float", DataType::Float64, true),
1891 Field::new("c_string", DataType::Utf8, true),
1892 Field::new("c_bool", DataType::Boolean, true),
1893 ]);
1894
1895 assert_eq!(schema, expected_schema);
1896
1897 let builder = ReaderBuilder::new(Arc::new(schema))
1898 .with_format(format)
1899 .with_batch_size(512)
1900 .with_projection(vec![0, 1, 2, 3]);
1901
1902 let mut csv = builder.build(file).unwrap();
1903 let batch = csv.next().unwrap().unwrap();
1904
1905 assert_eq!(5, batch.num_rows());
1906 assert_eq!(4, batch.num_columns());
1907
1908 assert_eq!(batch.schema().as_ref(), &expected_schema);
1909 }
1910
1911 #[test]
1912 fn test_scientific_notation_with_inference() {
1913 let mut file = File::open("test/data/scientific_notation_test.csv").unwrap();
1914 let format = Format::default().with_header(false).with_delimiter(b',');
1915
1916 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
1917 file.rewind().unwrap();
1918
1919 let builder = ReaderBuilder::new(Arc::new(schema))
1920 .with_format(format)
1921 .with_batch_size(512)
1922 .with_projection(vec![0, 1]);
1923
1924 let mut csv = builder.build(file).unwrap();
1925 let batch = csv.next().unwrap().unwrap();
1926
1927 let schema = batch.schema();
1928
1929 assert_eq!(&DataType::Float64, schema.field(0).data_type());
1930 }
1931
1932 fn invalid_csv_helper(file_name: &str) -> String {
1933 let file = File::open(file_name).unwrap();
1934 let schema = Schema::new(vec![
1935 Field::new("c_int", DataType::UInt64, false),
1936 Field::new("c_float", DataType::Float32, false),
1937 Field::new("c_string", DataType::Utf8, false),
1938 Field::new("c_bool", DataType::Boolean, false),
1939 ]);
1940
1941 let builder = ReaderBuilder::new(Arc::new(schema))
1942 .with_header(true)
1943 .with_delimiter(b'|')
1944 .with_batch_size(512)
1945 .with_projection(vec![0, 1, 2, 3]);
1946
1947 let mut csv = builder.build(file).unwrap();
1948
1949 csv.next().unwrap().unwrap_err().to_string()
1950 }
1951
1952 #[test]
1953 fn test_parse_invalid_csv_float() {
1954 let file_name = "test/data/various_invalid_types/invalid_float.csv";
1955
1956 let error = invalid_csv_helper(file_name);
1957 assert_eq!(
1958 "Parser error: Error while parsing value '4.x4' as type 'Float32' for column 1 at line 4. Row data: '[4,4.x4,,false]'",
1959 error
1960 );
1961 }
1962
1963 #[test]
1964 fn test_parse_invalid_csv_int() {
1965 let file_name = "test/data/various_invalid_types/invalid_int.csv";
1966
1967 let error = invalid_csv_helper(file_name);
1968 assert_eq!(
1969 "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]'",
1970 error
1971 );
1972 }
1973
1974 #[test]
1975 fn test_parse_invalid_csv_bool() {
1976 let file_name = "test/data/various_invalid_types/invalid_bool.csv";
1977
1978 let error = invalid_csv_helper(file_name);
1979 assert_eq!(
1980 "Parser error: Error while parsing value 'none' as type 'Boolean' for column 3 at line 2. Row data: '[2,2.2,2.22,none]'",
1981 error
1982 );
1983 }
1984
1985 fn infer_field_schema(string: &str) -> DataType {
1987 let mut v = InferredDataType::default();
1988 v.update(string);
1989 v.get()
1990 }
1991
1992 #[test]
1993 fn test_infer_field_schema() {
1994 assert_eq!(infer_field_schema("A"), DataType::Utf8);
1995 assert_eq!(infer_field_schema("\"123\""), DataType::Utf8);
1996 assert_eq!(infer_field_schema("10"), DataType::Int64);
1997 assert_eq!(infer_field_schema("10.2"), DataType::Float64);
1998 assert_eq!(infer_field_schema(".2"), DataType::Float64);
1999 assert_eq!(infer_field_schema("2."), DataType::Float64);
2000 assert_eq!(infer_field_schema("NaN"), DataType::Float64);
2001 assert_eq!(infer_field_schema("nan"), DataType::Float64);
2002 assert_eq!(infer_field_schema("inf"), DataType::Float64);
2003 assert_eq!(infer_field_schema("-inf"), DataType::Float64);
2004 assert_eq!(infer_field_schema("true"), DataType::Boolean);
2005 assert_eq!(infer_field_schema("trUe"), DataType::Boolean);
2006 assert_eq!(infer_field_schema("false"), DataType::Boolean);
2007 assert_eq!(infer_field_schema("2020-11-08"), DataType::Date32);
2008 assert_eq!(
2009 infer_field_schema("2020-11-08T14:20:01"),
2010 DataType::Timestamp(TimeUnit::Second, None)
2011 );
2012 assert_eq!(
2013 infer_field_schema("2020-11-08 14:20:01"),
2014 DataType::Timestamp(TimeUnit::Second, None)
2015 );
2016 assert_eq!(
2017 infer_field_schema("2020-11-08 14:20:01"),
2018 DataType::Timestamp(TimeUnit::Second, None)
2019 );
2020 assert_eq!(infer_field_schema("-5.13"), DataType::Float64);
2021 assert_eq!(infer_field_schema("0.1300"), DataType::Float64);
2022 assert_eq!(
2023 infer_field_schema("2021-12-19 13:12:30.921"),
2024 DataType::Timestamp(TimeUnit::Millisecond, None)
2025 );
2026 assert_eq!(
2027 infer_field_schema("2021-12-19T13:12:30.123456789"),
2028 DataType::Timestamp(TimeUnit::Nanosecond, None)
2029 );
2030 assert_eq!(infer_field_schema("–9223372036854775809"), DataType::Utf8);
2031 assert_eq!(infer_field_schema("9223372036854775808"), DataType::Utf8);
2032 }
2033
2034 #[test]
2035 fn parse_date32() {
2036 assert_eq!(Date32Type::parse("1970-01-01").unwrap(), 0);
2037 assert_eq!(Date32Type::parse("2020-03-15").unwrap(), 18336);
2038 assert_eq!(Date32Type::parse("1945-05-08").unwrap(), -9004);
2039 }
2040
2041 #[test]
2042 fn parse_time() {
2043 assert_eq!(
2044 Time64NanosecondType::parse("12:10:01.123456789 AM"),
2045 Some(601_123_456_789)
2046 );
2047 assert_eq!(
2048 Time64MicrosecondType::parse("12:10:01.123456 am"),
2049 Some(601_123_456)
2050 );
2051 assert_eq!(
2052 Time32MillisecondType::parse("2:10:01.12 PM"),
2053 Some(51_001_120)
2054 );
2055 assert_eq!(Time32SecondType::parse("2:10:01 pm"), Some(51_001));
2056 }
2057
2058 #[test]
2059 fn parse_date64() {
2060 assert_eq!(Date64Type::parse("1970-01-01T00:00:00").unwrap(), 0);
2061 assert_eq!(
2062 Date64Type::parse("2018-11-13T17:11:10").unwrap(),
2063 1542129070000
2064 );
2065 assert_eq!(
2066 Date64Type::parse("2018-11-13T17:11:10.011").unwrap(),
2067 1542129070011
2068 );
2069 assert_eq!(
2070 Date64Type::parse("1900-02-28T12:34:56").unwrap(),
2071 -2203932304000
2072 );
2073 assert_eq!(
2074 Date64Type::parse_formatted("1900-02-28 12:34:56", "%Y-%m-%d %H:%M:%S").unwrap(),
2075 -2203932304000
2076 );
2077 assert_eq!(
2078 Date64Type::parse_formatted("1900-02-28 12:34:56+0030", "%Y-%m-%d %H:%M:%S%z").unwrap(),
2079 -2203932304000 - (30 * 60 * 1000)
2080 );
2081 }
2082
2083 fn test_parse_timestamp_impl<T: ArrowTimestampType>(
2084 timezone: Option<Arc<str>>,
2085 expected: &[i64],
2086 ) {
2087 let csv = [
2088 "1970-01-01T00:00:00",
2089 "1970-01-01T00:00:00Z",
2090 "1970-01-01T00:00:00+02:00",
2091 ]
2092 .join("\n");
2093 let schema = Arc::new(Schema::new(vec![Field::new(
2094 "field",
2095 DataType::Timestamp(T::UNIT, timezone.clone()),
2096 true,
2097 )]));
2098
2099 let mut decoder = ReaderBuilder::new(schema).build_decoder();
2100
2101 let decoded = decoder.decode(csv.as_bytes()).unwrap();
2102 assert_eq!(decoded, csv.len());
2103 decoder.decode(&[]).unwrap();
2104
2105 let batch = decoder.flush().unwrap().unwrap();
2106 assert_eq!(batch.num_columns(), 1);
2107 assert_eq!(batch.num_rows(), 3);
2108 let col = batch.column(0).as_primitive::<T>();
2109 assert_eq!(col.values(), expected);
2110 assert_eq!(col.data_type(), &DataType::Timestamp(T::UNIT, timezone));
2111 }
2112
2113 #[test]
2114 fn test_parse_timestamp() {
2115 test_parse_timestamp_impl::<TimestampNanosecondType>(None, &[0, 0, -7_200_000_000_000]);
2116 test_parse_timestamp_impl::<TimestampNanosecondType>(
2117 Some("+00:00".into()),
2118 &[0, 0, -7_200_000_000_000],
2119 );
2120 test_parse_timestamp_impl::<TimestampNanosecondType>(
2121 Some("-05:00".into()),
2122 &[18_000_000_000_000, 0, -7_200_000_000_000],
2123 );
2124 test_parse_timestamp_impl::<TimestampMicrosecondType>(
2125 Some("-03".into()),
2126 &[10_800_000_000, 0, -7_200_000_000],
2127 );
2128 test_parse_timestamp_impl::<TimestampMillisecondType>(
2129 Some("-03".into()),
2130 &[10_800_000, 0, -7_200_000],
2131 );
2132 test_parse_timestamp_impl::<TimestampSecondType>(Some("-03".into()), &[10_800, 0, -7_200]);
2133 }
2134
2135 #[test]
2136 fn test_infer_schema_from_multiple_files() {
2137 let mut csv1 = NamedTempFile::new().unwrap();
2138 let mut csv2 = NamedTempFile::new().unwrap();
2139 let csv3 = NamedTempFile::new().unwrap(); let mut csv4 = NamedTempFile::new().unwrap();
2141 writeln!(csv1, "c1,c2,c3").unwrap();
2142 writeln!(csv1, "1,\"foo\",0.5").unwrap();
2143 writeln!(csv1, "3,\"bar\",1").unwrap();
2144 writeln!(csv1, "3,\"bar\",2e-06").unwrap();
2145 writeln!(csv2, "c1,c2,c3,c4").unwrap();
2147 writeln!(csv2, "10,,3.14,true").unwrap();
2148 writeln!(csv4, "c1,c2,c3").unwrap();
2150 writeln!(csv4, "10,\"foo\",").unwrap();
2151
2152 let schema = infer_schema_from_files(
2153 &[
2154 csv3.path().to_str().unwrap().to_string(),
2155 csv1.path().to_str().unwrap().to_string(),
2156 csv2.path().to_str().unwrap().to_string(),
2157 csv4.path().to_str().unwrap().to_string(),
2158 ],
2159 b',',
2160 Some(4), true,
2162 )
2163 .unwrap();
2164
2165 assert_eq!(schema.fields().len(), 4);
2166 assert!(schema.field(0).is_nullable());
2167 assert!(schema.field(1).is_nullable());
2168 assert!(schema.field(2).is_nullable());
2169 assert!(schema.field(3).is_nullable());
2170
2171 assert_eq!(&DataType::Int64, schema.field(0).data_type());
2172 assert_eq!(&DataType::Utf8, schema.field(1).data_type());
2173 assert_eq!(&DataType::Float64, schema.field(2).data_type());
2174 assert_eq!(&DataType::Boolean, schema.field(3).data_type());
2175 }
2176
2177 #[test]
2178 fn test_bounded() {
2179 let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2180 let data = [
2181 vec!["0"],
2182 vec!["1"],
2183 vec!["2"],
2184 vec!["3"],
2185 vec!["4"],
2186 vec!["5"],
2187 vec!["6"],
2188 ];
2189
2190 let data = data
2191 .iter()
2192 .map(|x| x.join(","))
2193 .collect::<Vec<_>>()
2194 .join("\n");
2195 let data = data.as_bytes();
2196
2197 let reader = std::io::Cursor::new(data);
2198
2199 let mut csv = ReaderBuilder::new(Arc::new(schema))
2200 .with_batch_size(2)
2201 .with_projection(vec![0])
2202 .with_bounds(2, 6)
2203 .build_buffered(reader)
2204 .unwrap();
2205
2206 let batch = csv.next().unwrap().unwrap();
2207 let a = batch.column(0);
2208 let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2209 assert_eq!(a, &UInt32Array::from(vec![2, 3]));
2210
2211 let batch = csv.next().unwrap().unwrap();
2212 let a = batch.column(0);
2213 let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2214 assert_eq!(a, &UInt32Array::from(vec![4, 5]));
2215
2216 assert!(csv.next().is_none());
2217 }
2218
2219 #[test]
2220 fn test_empty_projection() {
2221 let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2222 let data = [vec!["0"], vec!["1"]];
2223
2224 let data = data
2225 .iter()
2226 .map(|x| x.join(","))
2227 .collect::<Vec<_>>()
2228 .join("\n");
2229
2230 let mut csv = ReaderBuilder::new(Arc::new(schema))
2231 .with_batch_size(2)
2232 .with_projection(vec![])
2233 .build_buffered(Cursor::new(data.as_bytes()))
2234 .unwrap();
2235
2236 let batch = csv.next().unwrap().unwrap();
2237 assert_eq!(batch.columns().len(), 0);
2238 assert_eq!(batch.num_rows(), 2);
2239
2240 assert!(csv.next().is_none());
2241 }
2242
2243 #[test]
2244 fn test_parsing_bool() {
2245 assert_eq!(Some(true), parse_bool("true"));
2247 assert_eq!(Some(true), parse_bool("tRUe"));
2248 assert_eq!(Some(true), parse_bool("True"));
2249 assert_eq!(Some(true), parse_bool("TRUE"));
2250 assert_eq!(None, parse_bool("t"));
2251 assert_eq!(None, parse_bool("T"));
2252 assert_eq!(None, parse_bool(""));
2253
2254 assert_eq!(Some(false), parse_bool("false"));
2255 assert_eq!(Some(false), parse_bool("fALse"));
2256 assert_eq!(Some(false), parse_bool("False"));
2257 assert_eq!(Some(false), parse_bool("FALSE"));
2258 assert_eq!(None, parse_bool("f"));
2259 assert_eq!(None, parse_bool("F"));
2260 assert_eq!(None, parse_bool(""));
2261 }
2262
2263 #[test]
2264 fn test_parsing_float() {
2265 assert_eq!(Some(12.34), Float64Type::parse("12.34"));
2266 assert_eq!(Some(-12.34), Float64Type::parse("-12.34"));
2267 assert_eq!(Some(12.0), Float64Type::parse("12"));
2268 assert_eq!(Some(0.0), Float64Type::parse("0"));
2269 assert_eq!(Some(2.0), Float64Type::parse("2."));
2270 assert_eq!(Some(0.2), Float64Type::parse(".2"));
2271 assert!(Float64Type::parse("nan").unwrap().is_nan());
2272 assert!(Float64Type::parse("NaN").unwrap().is_nan());
2273 assert!(Float64Type::parse("inf").unwrap().is_infinite());
2274 assert!(Float64Type::parse("inf").unwrap().is_sign_positive());
2275 assert!(Float64Type::parse("-inf").unwrap().is_infinite());
2276 assert!(Float64Type::parse("-inf").unwrap().is_sign_negative());
2277 assert_eq!(None, Float64Type::parse(""));
2278 assert_eq!(None, Float64Type::parse("dd"));
2279 assert_eq!(None, Float64Type::parse("12.34.56"));
2280 }
2281
2282 #[test]
2283 fn test_non_std_quote() {
2284 let schema = Schema::new(vec![
2285 Field::new("text1", DataType::Utf8, false),
2286 Field::new("text2", DataType::Utf8, false),
2287 ]);
2288 let builder = ReaderBuilder::new(Arc::new(schema))
2289 .with_header(false)
2290 .with_quote(b'~'); let mut csv_text = Vec::new();
2293 let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2294 for index in 0..10 {
2295 let text1 = format!("id{index:}");
2296 let text2 = format!("value{index:}");
2297 csv_writer
2298 .write_fmt(format_args!("~{text1}~,~{text2}~\r\n"))
2299 .unwrap();
2300 }
2301 let mut csv_reader = std::io::Cursor::new(&csv_text);
2302 let mut reader = builder.build(&mut csv_reader).unwrap();
2303 let batch = reader.next().unwrap().unwrap();
2304 let col0 = batch.column(0);
2305 assert_eq!(col0.len(), 10);
2306 let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2307 assert_eq!(col0_arr.value(0), "id0");
2308 let col1 = batch.column(1);
2309 assert_eq!(col1.len(), 10);
2310 let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2311 assert_eq!(col1_arr.value(5), "value5");
2312 }
2313
2314 #[test]
2315 fn test_non_std_escape() {
2316 let schema = Schema::new(vec![
2317 Field::new("text1", DataType::Utf8, false),
2318 Field::new("text2", DataType::Utf8, false),
2319 ]);
2320 let builder = ReaderBuilder::new(Arc::new(schema))
2321 .with_header(false)
2322 .with_escape(b'\\'); let mut csv_text = Vec::new();
2325 let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2326 for index in 0..10 {
2327 let text1 = format!("id{index:}");
2328 let text2 = format!("value\\\"{index:}");
2329 csv_writer
2330 .write_fmt(format_args!("\"{text1}\",\"{text2}\"\r\n"))
2331 .unwrap();
2332 }
2333 let mut csv_reader = std::io::Cursor::new(&csv_text);
2334 let mut reader = builder.build(&mut csv_reader).unwrap();
2335 let batch = reader.next().unwrap().unwrap();
2336 let col0 = batch.column(0);
2337 assert_eq!(col0.len(), 10);
2338 let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2339 assert_eq!(col0_arr.value(0), "id0");
2340 let col1 = batch.column(1);
2341 assert_eq!(col1.len(), 10);
2342 let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2343 assert_eq!(col1_arr.value(5), "value\"5");
2344 }
2345
2346 #[test]
2347 fn test_non_std_terminator() {
2348 let schema = Schema::new(vec![
2349 Field::new("text1", DataType::Utf8, false),
2350 Field::new("text2", DataType::Utf8, false),
2351 ]);
2352 let builder = ReaderBuilder::new(Arc::new(schema))
2353 .with_header(false)
2354 .with_terminator(b'\n'); let mut csv_text = Vec::new();
2357 let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2358 for index in 0..10 {
2359 let text1 = format!("id{index:}");
2360 let text2 = format!("value{index:}");
2361 csv_writer
2362 .write_fmt(format_args!("\"{text1}\",\"{text2}\"\n"))
2363 .unwrap();
2364 }
2365 let mut csv_reader = std::io::Cursor::new(&csv_text);
2366 let mut reader = builder.build(&mut csv_reader).unwrap();
2367 let batch = reader.next().unwrap().unwrap();
2368 let col0 = batch.column(0);
2369 assert_eq!(col0.len(), 10);
2370 let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2371 assert_eq!(col0_arr.value(0), "id0");
2372 let col1 = batch.column(1);
2373 assert_eq!(col1.len(), 10);
2374 let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2375 assert_eq!(col1_arr.value(5), "value5");
2376 }
2377
2378 #[test]
2379 fn test_header_bounds() {
2380 let csv = "a,b\na,b\na,b\na,b\na,b\n";
2381 let tests = [
2382 (None, false, 5),
2383 (None, true, 4),
2384 (Some((0, 4)), false, 4),
2385 (Some((1, 4)), false, 3),
2386 (Some((0, 4)), true, 4),
2387 (Some((1, 4)), true, 3),
2388 ];
2389 let schema = Arc::new(Schema::new(vec![
2390 Field::new("a", DataType::Utf8, false),
2391 Field::new("a", DataType::Utf8, false),
2392 ]));
2393
2394 for (idx, (bounds, has_header, expected)) in tests.into_iter().enumerate() {
2395 let mut reader = ReaderBuilder::new(schema.clone()).with_header(has_header);
2396 if let Some((start, end)) = bounds {
2397 reader = reader.with_bounds(start, end);
2398 }
2399 let b = reader
2400 .build_buffered(Cursor::new(csv.as_bytes()))
2401 .unwrap()
2402 .next()
2403 .unwrap()
2404 .unwrap();
2405 assert_eq!(b.num_rows(), expected, "{idx}");
2406 }
2407 }
2408
2409 #[test]
2410 fn test_header_validation() {
2411 let schema = Arc::new(Schema::new(vec![
2412 Field::new("a", DataType::Int32, false),
2413 Field::new("b", DataType::Int32, false),
2414 ]));
2415
2416 let csv = "a,c\n1,2\n";
2417 let err = ReaderBuilder::new(schema.clone())
2418 .with_header(true)
2419 .with_header_validation(true)
2420 .build_buffered(Cursor::new(csv.as_bytes()))
2421 .unwrap()
2422 .next()
2423 .unwrap()
2424 .unwrap_err()
2425 .to_string();
2426 assert_eq!(
2427 err,
2428 "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"c\""
2429 );
2430
2431 let batch = ReaderBuilder::new(schema)
2432 .with_header(true)
2433 .with_header_validation(false)
2434 .build_buffered(Cursor::new(csv.as_bytes()))
2435 .unwrap()
2436 .next()
2437 .unwrap()
2438 .unwrap();
2439 assert_eq!(batch.num_rows(), 1);
2440 }
2441
2442 #[test]
2443 fn test_header_validation_with_buffered_reader() {
2444 let schema = Arc::new(Schema::new(vec![
2445 Field::new("a", DataType::Int32, false),
2446 Field::new("b", DataType::Int32, false),
2447 ]));
2448
2449 let csv = "a,b\n1,2\n";
2450 let buffered = std::io::BufReader::with_capacity(1, Cursor::new(csv.as_bytes()));
2451 let batch = ReaderBuilder::new(schema)
2452 .with_header(true)
2453 .with_header_validation(true)
2454 .build_buffered(buffered)
2455 .unwrap()
2456 .next()
2457 .unwrap()
2458 .unwrap();
2459
2460 assert_eq!(batch.num_rows(), 1);
2461 let a = batch.column(0).as_primitive::<Int32Type>();
2462 assert_eq!(a.value(0), 1);
2463 }
2464
2465 #[test]
2466 fn test_header_validation_with_truncated_rows() {
2467 let schema = Arc::new(Schema::new(vec![
2468 Field::new("a", DataType::Int32, true),
2469 Field::new("b", DataType::Int32, true),
2470 ]));
2471
2472 let csv = "a\n1\n";
2473 let err = ReaderBuilder::new(schema.clone())
2474 .with_header(true)
2475 .with_header_validation(true)
2476 .with_truncated_rows(true)
2477 .build_buffered(Cursor::new(csv.as_bytes()))
2478 .unwrap()
2479 .next()
2480 .unwrap()
2481 .unwrap_err()
2482 .to_string();
2483 assert_eq!(
2484 err,
2485 "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"\"",
2486 )
2487 }
2488
2489 #[test]
2490 fn test_null_boolean() {
2491 let csv = "true,false\nFalse,True\n,True\nFalse,";
2492 let schema = Arc::new(Schema::new(vec![
2493 Field::new("a", DataType::Boolean, true),
2494 Field::new("a", DataType::Boolean, true),
2495 ]));
2496
2497 let b = ReaderBuilder::new(schema)
2498 .build_buffered(Cursor::new(csv.as_bytes()))
2499 .unwrap()
2500 .next()
2501 .unwrap()
2502 .unwrap();
2503
2504 assert_eq!(b.num_rows(), 4);
2505 assert_eq!(b.num_columns(), 2);
2506
2507 let c = b.column(0).as_boolean();
2508 assert_eq!(c.null_count(), 1);
2509 assert!(c.value(0));
2510 assert!(!c.value(1));
2511 assert!(c.is_null(2));
2512 assert!(!c.value(3));
2513
2514 let c = b.column(1).as_boolean();
2515 assert_eq!(c.null_count(), 1);
2516 assert!(!c.value(0));
2517 assert!(c.value(1));
2518 assert!(c.value(2));
2519 assert!(c.is_null(3));
2520 }
2521
2522 #[test]
2523 fn test_truncated_rows() {
2524 let data = "a,b,c\n1,2,3\n4,5\n\n6,7,8";
2525 let schema = Arc::new(Schema::new(vec![
2526 Field::new("a", DataType::Int32, true),
2527 Field::new("b", DataType::Int32, true),
2528 Field::new("c", DataType::Int32, true),
2529 ]));
2530
2531 let reader = ReaderBuilder::new(schema.clone())
2532 .with_header(true)
2533 .with_truncated_rows(true)
2534 .build(Cursor::new(data))
2535 .unwrap();
2536
2537 let batches = reader.collect::<Result<Vec<_>, _>>();
2538 assert!(batches.is_ok());
2539 let batch = batches.unwrap().into_iter().next().unwrap();
2540 assert_eq!(batch.num_rows(), 3);
2542
2543 let reader = ReaderBuilder::new(schema.clone())
2544 .with_header(true)
2545 .with_truncated_rows(false)
2546 .build(Cursor::new(data))
2547 .unwrap();
2548
2549 let batches = reader.collect::<Result<Vec<_>, _>>();
2550 assert!(match batches {
2551 Err(ArrowError::CsvError(e)) => e.to_string().contains("incorrect number of fields"),
2552 _ => false,
2553 });
2554 }
2555
2556 #[test]
2557 fn test_truncated_rows_csv() {
2558 let file = File::open("test/data/truncated_rows.csv").unwrap();
2559 let schema = Arc::new(Schema::new(vec![
2560 Field::new("Name", DataType::Utf8, true),
2561 Field::new("Age", DataType::UInt32, true),
2562 Field::new("Occupation", DataType::Utf8, true),
2563 Field::new("DOB", DataType::Date32, true),
2564 ]));
2565 let reader = ReaderBuilder::new(schema.clone())
2566 .with_header(true)
2567 .with_batch_size(24)
2568 .with_truncated_rows(true);
2569 let csv = reader.build(file).unwrap();
2570 let batches = csv.collect::<Result<Vec<_>, _>>().unwrap();
2571
2572 assert_eq!(batches.len(), 1);
2573 let batch = &batches[0];
2574 assert_eq!(batch.num_rows(), 6);
2575 assert_eq!(batch.num_columns(), 4);
2576 let name = batch
2577 .column(0)
2578 .as_any()
2579 .downcast_ref::<StringArray>()
2580 .unwrap();
2581 let age = batch
2582 .column(1)
2583 .as_any()
2584 .downcast_ref::<UInt32Array>()
2585 .unwrap();
2586 let occupation = batch
2587 .column(2)
2588 .as_any()
2589 .downcast_ref::<StringArray>()
2590 .unwrap();
2591 let dob = batch
2592 .column(3)
2593 .as_any()
2594 .downcast_ref::<Date32Array>()
2595 .unwrap();
2596
2597 assert_eq!(name.value(0), "A1");
2598 assert_eq!(name.value(1), "B2");
2599 assert!(name.is_null(2));
2600 assert_eq!(name.value(3), "C3");
2601 assert_eq!(name.value(4), "D4");
2602 assert_eq!(name.value(5), "E5");
2603
2604 assert_eq!(age.value(0), 34);
2605 assert_eq!(age.value(1), 29);
2606 assert!(age.is_null(2));
2607 assert_eq!(age.value(3), 45);
2608 assert!(age.is_null(4));
2609 assert_eq!(age.value(5), 31);
2610
2611 assert_eq!(occupation.value(0), "Engineer");
2612 assert_eq!(occupation.value(1), "Doctor");
2613 assert!(occupation.is_null(2));
2614 assert_eq!(occupation.value(3), "Artist");
2615 assert!(occupation.is_null(4));
2616 assert!(occupation.is_null(5));
2617
2618 assert_eq!(dob.value(0), 5675);
2619 assert!(dob.is_null(1));
2620 assert!(dob.is_null(2));
2621 assert_eq!(dob.value(3), -1858);
2622 assert!(dob.is_null(4));
2623 assert!(dob.is_null(5));
2624 }
2625
2626 #[test]
2627 fn test_truncated_rows_not_nullable_error() {
2628 let data = "a,b,c\n1,2,3\n4,5";
2629 let schema = Arc::new(Schema::new(vec![
2630 Field::new("a", DataType::Int32, false),
2631 Field::new("b", DataType::Int32, false),
2632 Field::new("c", DataType::Int32, false),
2633 ]));
2634
2635 let reader = ReaderBuilder::new(schema.clone())
2636 .with_header(true)
2637 .with_truncated_rows(true)
2638 .build(Cursor::new(data))
2639 .unwrap();
2640
2641 let batches = reader.collect::<Result<Vec<_>, _>>();
2642 assert!(match batches {
2643 Err(ArrowError::InvalidArgumentError(e)) =>
2644 e.to_string().contains("contains null values"),
2645 _ => false,
2646 });
2647 }
2648
2649 #[test]
2650 fn test_buffered() {
2651 let tests = [
2652 ("test/data/uk_cities.csv", false, 37),
2653 ("test/data/various_types.csv", true, 10),
2654 ("test/data/decimal_test.csv", false, 10),
2655 ];
2656
2657 for (path, has_header, expected_rows) in tests {
2658 let (schema, _) = Format::default()
2659 .infer_schema(File::open(path).unwrap(), None)
2660 .unwrap();
2661 let schema = Arc::new(schema);
2662
2663 for batch_size in [1, 4] {
2664 for capacity in [1, 3, 7, 100] {
2665 let reader = ReaderBuilder::new(schema.clone())
2666 .with_batch_size(batch_size)
2667 .with_header(has_header)
2668 .build(File::open(path).unwrap())
2669 .unwrap();
2670
2671 let expected = reader.collect::<Result<Vec<_>, _>>().unwrap();
2672
2673 assert_eq!(
2674 expected.iter().map(|x| x.num_rows()).sum::<usize>(),
2675 expected_rows
2676 );
2677
2678 let buffered =
2679 std::io::BufReader::with_capacity(capacity, File::open(path).unwrap());
2680
2681 let reader = ReaderBuilder::new(schema.clone())
2682 .with_batch_size(batch_size)
2683 .with_header(has_header)
2684 .build_buffered(buffered)
2685 .unwrap();
2686
2687 let actual = reader.collect::<Result<Vec<_>, _>>().unwrap();
2688 assert_eq!(expected, actual)
2689 }
2690 }
2691 }
2692 }
2693
2694 fn err_test(csv: &[u8], expected: &str) {
2695 fn err_test_with_schema(csv: &[u8], expected: &str, schema: Arc<Schema>) {
2696 let buffer = std::io::BufReader::with_capacity(2, Cursor::new(csv));
2697 let b = ReaderBuilder::new(schema)
2698 .with_batch_size(2)
2699 .build_buffered(buffer)
2700 .unwrap();
2701 let err = b.collect::<Result<Vec<_>, _>>().unwrap_err().to_string();
2702 assert_eq!(err, expected)
2703 }
2704
2705 let schema_utf8 = Arc::new(Schema::new(vec![
2706 Field::new("text1", DataType::Utf8, true),
2707 Field::new("text2", DataType::Utf8, true),
2708 ]));
2709 err_test_with_schema(csv, expected, schema_utf8);
2710
2711 let schema_utf8view = Arc::new(Schema::new(vec![
2712 Field::new("text1", DataType::Utf8View, true),
2713 Field::new("text2", DataType::Utf8View, true),
2714 ]));
2715 err_test_with_schema(csv, expected, schema_utf8view);
2716 }
2717
2718 #[test]
2719 fn test_invalid_utf8() {
2720 err_test(
2721 b"sdf,dsfg\ndfd,hgh\xFFue\n,sds\nFalhghse,",
2722 "Csv error: Encountered invalid UTF-8 data for line 2 and field 2",
2723 );
2724
2725 err_test(
2726 b"sdf,dsfg\ndksdk,jf\nd\xFFfd,hghue\n,sds\nFalhghse,",
2727 "Csv error: Encountered invalid UTF-8 data for line 3 and field 1",
2728 );
2729
2730 err_test(
2731 b"sdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
2732 "Csv error: Encountered invalid UTF-8 data for line 5 and field 2",
2733 );
2734
2735 err_test(
2736 b"\xFFsdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
2737 "Csv error: Encountered invalid UTF-8 data for line 1 and field 1",
2738 );
2739 }
2740
2741 struct InstrumentedRead<R> {
2742 r: R,
2743 fill_count: usize,
2744 fill_sizes: Vec<usize>,
2745 }
2746
2747 impl<R> InstrumentedRead<R> {
2748 fn new(r: R) -> Self {
2749 Self {
2750 r,
2751 fill_count: 0,
2752 fill_sizes: vec![],
2753 }
2754 }
2755 }
2756
2757 impl<R: Seek> Seek for InstrumentedRead<R> {
2758 fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
2759 self.r.seek(pos)
2760 }
2761 }
2762
2763 impl<R: BufRead> Read for InstrumentedRead<R> {
2764 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
2765 self.r.read(buf)
2766 }
2767 }
2768
2769 impl<R: BufRead> BufRead for InstrumentedRead<R> {
2770 fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
2771 self.fill_count += 1;
2772 let buf = self.r.fill_buf()?;
2773 self.fill_sizes.push(buf.len());
2774 Ok(buf)
2775 }
2776
2777 fn consume(&mut self, amt: usize) {
2778 self.r.consume(amt)
2779 }
2780 }
2781
2782 #[test]
2783 fn test_io() {
2784 let schema = Arc::new(Schema::new(vec![
2785 Field::new("a", DataType::Utf8, false),
2786 Field::new("b", DataType::Utf8, false),
2787 ]));
2788 let csv = "foo,bar\nbaz,foo\na,b\nc,d";
2789 let mut read = InstrumentedRead::new(Cursor::new(csv.as_bytes()));
2790 let reader = ReaderBuilder::new(schema)
2791 .with_batch_size(3)
2792 .build_buffered(&mut read)
2793 .unwrap();
2794
2795 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
2796 assert_eq!(batches.len(), 2);
2797 assert_eq!(batches[0].num_rows(), 3);
2798 assert_eq!(batches[1].num_rows(), 1);
2799
2800 assert_eq!(&read.fill_sizes, &[23, 3, 0, 0]);
2806 assert_eq!(read.fill_count, 4);
2807 }
2808
2809 #[test]
2810 fn test_inference() {
2811 let cases: &[(&[&str], DataType)] = &[
2812 (&[], DataType::Null),
2813 (&["false", "12"], DataType::Utf8),
2814 (&["12", "cupcakes"], DataType::Utf8),
2815 (&["12", "12.4"], DataType::Float64),
2816 (&["14050", "24332"], DataType::Int64),
2817 (&["14050.0", "true"], DataType::Utf8),
2818 (&["14050", "2020-03-19 00:00:00"], DataType::Utf8),
2819 (&["14050", "2340.0", "2020-03-19 00:00:00"], DataType::Utf8),
2820 (
2821 &["2020-03-19 02:00:00", "2020-03-19 00:00:00"],
2822 DataType::Timestamp(TimeUnit::Second, None),
2823 ),
2824 (&["2020-03-19", "2020-03-20"], DataType::Date32),
2825 (
2826 &["2020-03-19", "2020-03-19 02:00:00", "2020-03-19 00:00:00"],
2827 DataType::Timestamp(TimeUnit::Second, None),
2828 ),
2829 (
2830 &[
2831 "2020-03-19",
2832 "2020-03-19 02:00:00",
2833 "2020-03-19 00:00:00.000",
2834 ],
2835 DataType::Timestamp(TimeUnit::Millisecond, None),
2836 ),
2837 (
2838 &[
2839 "2020-03-19",
2840 "2020-03-19 02:00:00",
2841 "2020-03-19 00:00:00.000000",
2842 ],
2843 DataType::Timestamp(TimeUnit::Microsecond, None),
2844 ),
2845 (
2846 &["2020-03-19 02:00:00+02:00", "2020-03-19 02:00:00Z"],
2847 DataType::Timestamp(TimeUnit::Second, None),
2848 ),
2849 (
2850 &[
2851 "2020-03-19",
2852 "2020-03-19 02:00:00+02:00",
2853 "2020-03-19 02:00:00Z",
2854 "2020-03-19 02:00:00.12Z",
2855 ],
2856 DataType::Timestamp(TimeUnit::Millisecond, None),
2857 ),
2858 (
2859 &[
2860 "2020-03-19",
2861 "2020-03-19 02:00:00.000000000",
2862 "2020-03-19 00:00:00.000000",
2863 ],
2864 DataType::Timestamp(TimeUnit::Nanosecond, None),
2865 ),
2866 ];
2867
2868 for (values, expected) in cases {
2869 let mut t = InferredDataType::default();
2870 for v in *values {
2871 t.update(v)
2872 }
2873 assert_eq!(&t.get(), expected, "{values:?}")
2874 }
2875 }
2876
2877 #[test]
2878 fn test_record_length_mismatch() {
2879 let csv = "\
2880 a,b,c\n\
2881 1,2,3\n\
2882 4,5\n\
2883 6,7,8";
2884 let mut read = Cursor::new(csv.as_bytes());
2885 let result = Format::default()
2886 .with_header(true)
2887 .infer_schema(&mut read, None);
2888 assert!(result.is_err());
2889 assert_eq!(
2891 result.err().unwrap().to_string(),
2892 "Csv error: Encountered unequal lengths between records on CSV file. Expected 3 records, found 2 records at line 3"
2893 );
2894 }
2895
2896 #[test]
2897 fn test_comment() {
2898 let schema = Schema::new(vec![
2899 Field::new("a", DataType::Int8, false),
2900 Field::new("b", DataType::Int8, false),
2901 ]);
2902
2903 let csv = "# comment1 \n1,2\n#comment2\n11,22";
2904 let mut read = Cursor::new(csv.as_bytes());
2905 let reader = ReaderBuilder::new(Arc::new(schema))
2906 .with_comment(b'#')
2907 .build(&mut read)
2908 .unwrap();
2909
2910 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
2911 assert_eq!(batches.len(), 1);
2912 let b = batches.first().unwrap();
2913 assert_eq!(b.num_columns(), 2);
2914 assert_eq!(
2915 b.column(0)
2916 .as_any()
2917 .downcast_ref::<Int8Array>()
2918 .unwrap()
2919 .values(),
2920 &vec![1, 11]
2921 );
2922 assert_eq!(
2923 b.column(1)
2924 .as_any()
2925 .downcast_ref::<Int8Array>()
2926 .unwrap()
2927 .values(),
2928 &vec![2, 22]
2929 );
2930 }
2931
2932 #[test]
2933 fn test_parse_string_view_single_column() {
2934 let csv = ["foo", "something_cannot_be_inlined", "foobar"].join("\n");
2935 let schema = Arc::new(Schema::new(vec![Field::new(
2936 "c1",
2937 DataType::Utf8View,
2938 true,
2939 )]));
2940
2941 let mut decoder = ReaderBuilder::new(schema).build_decoder();
2942
2943 let decoded = decoder.decode(csv.as_bytes()).unwrap();
2944 assert_eq!(decoded, csv.len());
2945 decoder.decode(&[]).unwrap();
2946
2947 let batch = decoder.flush().unwrap().unwrap();
2948 assert_eq!(batch.num_columns(), 1);
2949 assert_eq!(batch.num_rows(), 3);
2950 let col = batch.column(0).as_string_view();
2951 assert_eq!(col.data_type(), &DataType::Utf8View);
2952 assert_eq!(col.value(0), "foo");
2953 assert_eq!(col.value(1), "something_cannot_be_inlined");
2954 assert_eq!(col.value(2), "foobar");
2955 }
2956
2957 #[test]
2958 fn test_parse_string_view_multi_column() {
2959 let csv = ["foo,", ",something_cannot_be_inlined", "foobarfoobar,bar"].join("\n");
2960 let schema = Arc::new(Schema::new(vec![
2961 Field::new("c1", DataType::Utf8View, true),
2962 Field::new("c2", DataType::Utf8View, true),
2963 ]));
2964
2965 let mut decoder = ReaderBuilder::new(schema).build_decoder();
2966
2967 let decoded = decoder.decode(csv.as_bytes()).unwrap();
2968 assert_eq!(decoded, csv.len());
2969 decoder.decode(&[]).unwrap();
2970
2971 let batch = decoder.flush().unwrap().unwrap();
2972 assert_eq!(batch.num_columns(), 2);
2973 assert_eq!(batch.num_rows(), 3);
2974 let c1 = batch.column(0).as_string_view();
2975 let c2 = batch.column(1).as_string_view();
2976 assert_eq!(c1.data_type(), &DataType::Utf8View);
2977 assert_eq!(c2.data_type(), &DataType::Utf8View);
2978
2979 assert!(!c1.is_null(0));
2980 assert!(c1.is_null(1));
2981 assert!(!c1.is_null(2));
2982 assert_eq!(c1.value(0), "foo");
2983 assert_eq!(c1.value(2), "foobarfoobar");
2984
2985 assert!(c2.is_null(0));
2986 assert!(!c2.is_null(1));
2987 assert!(!c2.is_null(2));
2988 assert_eq!(c2.value(1), "something_cannot_be_inlined");
2989 assert_eq!(c2.value(2), "bar");
2990 }
2991}