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::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
996fn 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 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
1030fn 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 "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
1118fn 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 "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#[derive(Debug)]
1152pub struct ReaderBuilder {
1153 schema: SchemaRef,
1155 format: Format,
1157 batch_size: usize,
1161 bounds: Bounds,
1163 projection: Option<Vec<usize>>,
1165}
1166
1167impl ReaderBuilder {
1168 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 pub fn with_header(mut self, has_header: bool) -> Self {
1202 self.format.header = has_header;
1203 self
1204 }
1205
1206 pub fn with_header_validation(mut self, validate_header: bool) -> Self {
1210 self.format.header_validation = validate_header;
1211 self
1212 }
1213
1214 pub fn with_format(mut self, format: Format) -> Self {
1216 self.format = format;
1217 self
1218 }
1219
1220 pub fn with_delimiter(mut self, delimiter: u8) -> Self {
1222 self.format.delimiter = Some(delimiter);
1223 self
1224 }
1225
1226 pub fn with_escape(mut self, escape: u8) -> Self {
1228 self.format.escape = Some(escape);
1229 self
1230 }
1231
1232 pub fn with_quote(mut self, quote: u8) -> Self {
1234 self.format.quote = Some(quote);
1235 self
1236 }
1237
1238 pub fn with_terminator(mut self, terminator: u8) -> Self {
1240 self.format.terminator = Some(terminator);
1241 self
1242 }
1243
1244 pub fn with_comment(mut self, comment: u8) -> Self {
1246 self.format.comment = Some(comment);
1247 self
1248 }
1249
1250 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 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1258 self.batch_size = batch_size;
1259 self
1260 }
1261
1262 pub fn with_bounds(mut self, start: usize, end: usize) -> Self {
1265 self.bounds = Some((start, end));
1266 self
1267 }
1268
1269 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1271 self.projection = Some(projection);
1272 self
1273 }
1274
1275 pub fn with_truncated_rows(mut self, allow: bool) -> Self {
1282 self.format.truncated_rows = allow;
1283 self
1284 }
1285
1286 pub fn build<R: Read>(self, reader: R) -> Result<Reader<R>, ArrowError> {
1291 self.build_buffered(StdBufReader::new(reader))
1292 }
1293
1294 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 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 let lat = batch.column(1).as_primitive::<Float64Type>();
1358 assert_eq!(57.653484, lat.value(0));
1359
1360 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 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 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 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 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 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 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 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 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 let city = batch
1603 .column(0)
1604 .as_any()
1605 .downcast_ref::<StringArray>()
1606 .unwrap();
1607
1608 assert_eq!("Elgin, Scotland, the UK", city.value(0));
1610
1611 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 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 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(); 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 writeln!(csv2, "c1,c2,c3,c4").unwrap();
2150 writeln!(csv2, "10,,3.14,true").unwrap();
2151 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), 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 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'~'); 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'\\'); 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'); 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 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 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 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}