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