1mod records;
167
168use arrow_array::builder::{NullBuilder, PrimitiveBuilder, StringDictionaryBuilder};
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_format<R: Read>(
379 mut self,
380 reader: R,
381 max_records: Option<usize>,
382 ) -> Result<(Self, usize), ArrowError> {
383 let (header, records_read) = self.infer_header(reader, max_records)?;
384 self.header = header;
385 Ok((self, records_read))
386 }
387
388 fn infer_header<R: Read>(
394 &self,
395 reader: R,
396 max_records: Option<usize>,
397 ) -> Result<(bool, usize), ArrowError> {
398 let mut format = self.clone();
399 format.header = false;
400 let mut csv_reader = format.build_reader(reader);
401
402 let mut first_record = StringRecord::new();
403 if !csv_reader
404 .read_record(&mut first_record)
405 .map_err(map_csv_error)?
406 {
407 return Ok((false, 0));
408 }
409
410 let mut first_types = vec![InferredDataType::default(); first_record.len()];
411 for (value, inferred) in first_record.iter().zip(&mut first_types) {
412 if !self.null_regex.is_null(value) {
413 inferred.update(value);
414 }
415 }
416
417 let mut column_types = vec![InferredDataType::default(); first_record.len()];
418 let mut record = StringRecord::new();
419 let mut records_count = 0;
420 let max_records = max_records.unwrap_or(usize::MAX);
421 while records_count < max_records
422 && csv_reader.read_record(&mut record).map_err(map_csv_error)?
423 {
424 records_count += 1;
425 for (value, inferred) in record.iter().zip(&mut column_types) {
426 if !self.null_regex.is_null(value) {
427 inferred.update(value);
428 }
429 }
430 }
431
432 let has_header = first_types
433 .iter()
434 .zip(&column_types)
435 .zip(first_record.iter())
436 .any(|((first, rest), value)| {
437 first.get() == DataType::Utf8
440 && value.parse::<f64>().is_err()
441 && !matches!(rest.get(), DataType::Utf8 | DataType::Null)
442 });
443 Ok((has_header, records_count + 1))
444 }
445
446 pub fn infer_schema<R: Read>(
453 &self,
454 reader: R,
455 max_records: Option<usize>,
456 ) -> Result<(Schema, usize), ArrowError> {
457 let mut csv_reader = self.build_reader(reader);
458
459 let headers: Vec<String> = if self.header {
462 let headers = &csv_reader.headers().map_err(map_csv_error)?.clone();
463 headers.iter().map(|s| s.to_string()).collect()
464 } else {
465 let first_record_count = &csv_reader.headers().map_err(map_csv_error)?.len();
466 (0..*first_record_count)
467 .map(|i| format!("column_{}", i + 1))
468 .collect()
469 };
470
471 let header_length = headers.len();
472 let mut column_types: Vec<InferredDataType> = vec![Default::default(); header_length];
474
475 let mut records_count = 0;
476
477 let mut record = StringRecord::new();
478 let max_records = max_records.unwrap_or(usize::MAX);
479 while records_count < max_records {
480 if !csv_reader.read_record(&mut record).map_err(map_csv_error)? {
481 break;
482 }
483 records_count += 1;
484
485 for (i, column_type) in column_types.iter_mut().enumerate().take(header_length) {
488 if let Some(string) = record.get(i)
489 && !self.null_regex.is_null(string)
490 {
491 column_type.update(string)
492 }
493 }
494 }
495
496 let fields: Fields = column_types
498 .iter()
499 .zip(&headers)
500 .map(|(inferred, field_name)| Field::new(field_name, inferred.get(), true))
501 .collect();
502
503 Ok((Schema::new(fields), records_count))
504 }
505
506 fn build_reader<R: Read>(&self, reader: R) -> csv::Reader<R> {
508 let mut builder = csv::ReaderBuilder::new();
509 builder.has_headers(self.header);
510 builder.flexible(self.truncated_rows);
511
512 if let Some(c) = self.delimiter {
513 builder.delimiter(c);
514 }
515 builder.escape(self.escape);
516 if let Some(c) = self.quote {
517 builder.quote(c);
518 }
519 if let Some(t) = self.terminator {
520 builder.terminator(csv::Terminator::Any(t));
521 }
522 if let Some(comment) = self.comment {
523 builder.comment(Some(comment));
524 }
525 builder.from_reader(reader)
526 }
527
528 fn build_parser(&self) -> csv_core::Reader {
530 let mut builder = csv_core::ReaderBuilder::new();
531 builder.escape(self.escape);
532 builder.comment(self.comment);
533
534 if let Some(c) = self.delimiter {
535 builder.delimiter(c);
536 }
537 if let Some(c) = self.quote {
538 builder.quote(c);
539 }
540 if let Some(t) = self.terminator {
541 builder.terminator(csv_core::Terminator::Any(t));
542 }
543 builder.build()
544 }
545}
546
547pub fn infer_schema_from_files(
554 files: &[String],
555 delimiter: u8,
556 max_read_records: Option<usize>,
557 has_header: bool,
558) -> Result<Schema, ArrowError> {
559 let mut schemas = vec![];
560 let mut records_to_read = max_read_records.unwrap_or(usize::MAX);
561 let format = Format {
562 delimiter: Some(delimiter),
563 header: has_header,
564 ..Default::default()
565 };
566
567 for fname in files {
568 let f = File::open(fname)?;
569 let (schema, records_read) = format.infer_schema(f, Some(records_to_read))?;
570 if records_read == 0 {
571 continue;
572 }
573 schemas.push(schema.clone());
574 records_to_read -= records_read;
575 if records_to_read == 0 {
576 break;
577 }
578 }
579
580 Schema::try_merge(schemas)
581}
582
583type Bounds = Option<(usize, usize)>;
585
586pub type Reader<R> = BufReader<StdBufReader<R>>;
591
592pub struct BufReader<R> {
597 reader: R,
599 decoder: Decoder,
601 schema: SchemaRef,
603}
604
605impl<R> fmt::Debug for BufReader<R>
606where
607 R: BufRead,
608{
609 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
610 f.debug_struct("Reader")
611 .field("decoder", &self.decoder)
612 .finish()
613 }
614}
615
616impl<R> BufReader<R> {
617 pub fn truncated_row_count(&self) -> usize {
652 self.decoder.truncated_row_count()
653 }
654}
655
656impl<R: Read> Reader<R> {
657 pub fn schema(&self) -> SchemaRef {
660 self.schema.clone()
661 }
662}
663
664impl<R: BufRead> BufReader<R> {
665 fn read(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
666 loop {
667 let buf = self.reader.fill_buf()?;
668 let decoded = self.decoder.decode(buf)?;
669 self.reader.consume(decoded);
670 if decoded == 0 || self.decoder.capacity() == 0 {
676 break;
677 }
678 }
679
680 self.decoder.flush()
681 }
682}
683
684impl<R: BufRead> Iterator for BufReader<R> {
685 type Item = Result<RecordBatch, ArrowError>;
686
687 fn next(&mut self) -> Option<Self::Item> {
688 self.read().transpose()
689 }
690}
691
692impl<R: BufRead> RecordBatchReader for BufReader<R> {
693 fn schema(&self) -> SchemaRef {
694 self.schema.clone()
695 }
696}
697
698#[derive(Debug)]
738pub struct Decoder {
739 schema: SchemaRef,
741
742 projection: Option<Vec<usize>>,
744
745 batch_size: usize,
747
748 to_skip: usize,
750
751 header_validation: bool,
753
754 line_number: usize,
756
757 end: usize,
759
760 record_decoder: RecordDecoder,
762
763 null_regex: NullRegex,
765}
766
767impl Decoder {
768 pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
778 if self.to_skip != 0 {
779 if self.header_validation {
780 let (skipped, bytes) = self.record_decoder.decode(buf, 1)?;
781
782 if skipped == 0 {
783 return Ok(bytes);
784 }
785
786 let rows = self.record_decoder.flush()?;
787 validate_header(&rows, self.schema.fields())?;
788 self.header_validation = false;
789 self.to_skip -= 1;
790 return Ok(bytes);
791 }
792
793 let to_skip = self.to_skip.min(self.batch_size);
795 let (skipped, bytes) = self.record_decoder.decode(buf, to_skip)?;
796 self.to_skip -= skipped;
797 self.record_decoder.clear();
798 return Ok(bytes);
799 }
800
801 let to_read = self.batch_size.min(self.end - self.line_number) - self.record_decoder.len();
802 let (_, bytes) = self.record_decoder.decode(buf, to_read)?;
803 Ok(bytes)
804 }
805
806 pub fn flush(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
813 if self.record_decoder.is_empty() {
814 return Ok(None);
815 }
816
817 let rows = self.record_decoder.flush()?;
818 let batch = parse(
819 &rows,
820 &self.schema,
821 self.projection.as_ref(),
822 self.line_number,
823 &self.null_regex,
824 )?;
825 self.line_number += rows.len();
826 Ok(Some(batch))
827 }
828
829 pub fn capacity(&self) -> usize {
831 self.batch_size - self.record_decoder.len()
832 }
833
834 pub fn truncated_row_count(&self) -> usize {
848 self.record_decoder.truncated_row_count()
849 }
850}
851
852fn validate_header(rows: &StringRecords<'_>, fields: &Fields) -> Result<(), ArrowError> {
853 let header = rows.iter().next().ok_or_else(|| {
854 ArrowError::CsvError("CSV header validation failed: no header row found".to_string())
855 })?;
856
857 for (idx, field) in fields.iter().enumerate() {
858 let actual = header.get(idx);
859 let expected = field.name();
860 if actual != expected {
861 return Err(ArrowError::CsvError(format!(
862 "CSV header does not match schema at column {idx}: expected {expected:?} but found {actual:?}"
863 )));
864 }
865 }
866
867 Ok(())
868}
869
870fn parse(
872 rows: &StringRecords<'_>,
873 schema: &Schema,
874 projection: Option<&Vec<usize>>,
875 line_number: usize,
876 null_regex: &NullRegex,
877) -> Result<RecordBatch, ArrowError> {
878 let fields = schema.fields();
879 let projection: Vec<usize> = match projection {
880 Some(v) => v.clone(),
881 None => fields.iter().enumerate().map(|(i, _)| i).collect(),
882 };
883 let projected_schema = Arc::new(schema.project(&projection)?);
884
885 let arrays: Result<Vec<ArrayRef>, _> = projection
886 .iter()
887 .map(|i| {
888 let i = *i;
889 let field = &fields[i];
890 match field.data_type() {
891 DataType::Boolean => build_boolean_array(line_number, rows, i, null_regex),
892 DataType::Decimal32(precision, scale) => build_decimal_array::<Decimal32Type>(
893 line_number,
894 rows,
895 i,
896 *precision,
897 *scale,
898 null_regex,
899 ),
900 DataType::Decimal64(precision, scale) => build_decimal_array::<Decimal64Type>(
901 line_number,
902 rows,
903 i,
904 *precision,
905 *scale,
906 null_regex,
907 ),
908 DataType::Decimal128(precision, scale) => build_decimal_array::<Decimal128Type>(
909 line_number,
910 rows,
911 i,
912 *precision,
913 *scale,
914 null_regex,
915 ),
916 DataType::Decimal256(precision, scale) => build_decimal_array::<Decimal256Type>(
917 line_number,
918 rows,
919 i,
920 *precision,
921 *scale,
922 null_regex,
923 ),
924 DataType::Int8 => {
925 build_primitive_array::<Int8Type>(line_number, rows, i, null_regex)
926 }
927 DataType::Int16 => {
928 build_primitive_array::<Int16Type>(line_number, rows, i, null_regex)
929 }
930 DataType::Int32 => {
931 build_primitive_array::<Int32Type>(line_number, rows, i, null_regex)
932 }
933 DataType::Int64 => {
934 build_primitive_array::<Int64Type>(line_number, rows, i, null_regex)
935 }
936 DataType::UInt8 => {
937 build_primitive_array::<UInt8Type>(line_number, rows, i, null_regex)
938 }
939 DataType::UInt16 => {
940 build_primitive_array::<UInt16Type>(line_number, rows, i, null_regex)
941 }
942 DataType::UInt32 => {
943 build_primitive_array::<UInt32Type>(line_number, rows, i, null_regex)
944 }
945 DataType::UInt64 => {
946 build_primitive_array::<UInt64Type>(line_number, rows, i, null_regex)
947 }
948 DataType::Float16 => {
949 build_primitive_array::<Float16Type>(line_number, rows, i, null_regex)
950 }
951 DataType::Float32 => {
952 build_primitive_array::<Float32Type>(line_number, rows, i, null_regex)
953 }
954 DataType::Float64 => {
955 build_primitive_array::<Float64Type>(line_number, rows, i, null_regex)
956 }
957 DataType::Date32 => {
958 build_primitive_array::<Date32Type>(line_number, rows, i, null_regex)
959 }
960 DataType::Date64 => {
961 build_primitive_array::<Date64Type>(line_number, rows, i, null_regex)
962 }
963 DataType::Time32(TimeUnit::Second) => {
964 build_primitive_array::<Time32SecondType>(line_number, rows, i, null_regex)
965 }
966 DataType::Time32(TimeUnit::Millisecond) => {
967 build_primitive_array::<Time32MillisecondType>(line_number, rows, i, null_regex)
968 }
969 DataType::Time64(TimeUnit::Microsecond) => {
970 build_primitive_array::<Time64MicrosecondType>(line_number, rows, i, null_regex)
971 }
972 DataType::Time64(TimeUnit::Nanosecond) => {
973 build_primitive_array::<Time64NanosecondType>(line_number, rows, i, null_regex)
974 }
975 DataType::Timestamp(TimeUnit::Second, tz) => {
976 build_timestamp_array::<TimestampSecondType>(
977 line_number,
978 rows,
979 i,
980 tz.as_deref(),
981 null_regex,
982 )
983 }
984 DataType::Timestamp(TimeUnit::Millisecond, tz) => {
985 build_timestamp_array::<TimestampMillisecondType>(
986 line_number,
987 rows,
988 i,
989 tz.as_deref(),
990 null_regex,
991 )
992 }
993 DataType::Timestamp(TimeUnit::Microsecond, tz) => {
994 build_timestamp_array::<TimestampMicrosecondType>(
995 line_number,
996 rows,
997 i,
998 tz.as_deref(),
999 null_regex,
1000 )
1001 }
1002 DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
1003 build_timestamp_array::<TimestampNanosecondType>(
1004 line_number,
1005 rows,
1006 i,
1007 tz.as_deref(),
1008 null_regex,
1009 )
1010 }
1011 DataType::Null => Ok(Arc::new({
1012 let mut builder = NullBuilder::new();
1013 builder.append_nulls(rows.len());
1014 builder.finish()
1015 }) as ArrayRef),
1016 DataType::Utf8 => Ok(Arc::new(
1017 rows.iter()
1018 .map(|row| {
1019 let s = row.get(i);
1020 (!null_regex.is_null(s)).then_some(s)
1021 })
1022 .collect::<StringArray>(),
1023 ) as ArrayRef),
1024 DataType::Utf8View => Ok(Arc::new(
1025 rows.iter()
1026 .map(|row| {
1027 let s = row.get(i);
1028 (!null_regex.is_null(s)).then_some(s)
1029 })
1030 .collect::<StringViewArray>(),
1031 ) as ArrayRef),
1032 DataType::Dictionary(key_type, value_type)
1033 if value_type.as_ref() == &DataType::Utf8 =>
1034 {
1035 match key_type.as_ref() {
1036 DataType::Int8 => build_dictionary_array::<Int8Type>(rows, i, null_regex),
1037 DataType::Int16 => build_dictionary_array::<Int16Type>(rows, i, null_regex),
1038 DataType::Int32 => build_dictionary_array::<Int32Type>(rows, i, null_regex),
1039 DataType::Int64 => build_dictionary_array::<Int64Type>(rows, i, null_regex),
1040 DataType::UInt8 => build_dictionary_array::<UInt8Type>(rows, i, null_regex),
1041 DataType::UInt16 => {
1042 build_dictionary_array::<UInt16Type>(rows, i, null_regex)
1043 }
1044 DataType::UInt32 => {
1045 build_dictionary_array::<UInt32Type>(rows, i, null_regex)
1046 }
1047 DataType::UInt64 => {
1048 build_dictionary_array::<UInt64Type>(rows, i, null_regex)
1049 }
1050 _ => Err(ArrowError::ParseError(format!(
1051 "Unsupported dictionary key type {key_type}"
1052 ))),
1053 }
1054 }
1055 other => Err(ArrowError::ParseError(format!(
1056 "Unsupported data type {other:?}"
1057 ))),
1058 }
1059 })
1060 .collect();
1061
1062 RecordBatch::try_new_with_options(
1063 projected_schema,
1064 arrays?,
1065 &RecordBatchOptions::new()
1066 .with_match_field_names(true)
1067 .with_row_count(Some(rows.len())),
1068 )
1069}
1070
1071fn parse_bool(string: &str) -> Option<bool> {
1072 if string.eq_ignore_ascii_case("false") {
1073 Some(false)
1074 } else if string.eq_ignore_ascii_case("true") {
1075 Some(true)
1076 } else {
1077 None
1078 }
1079}
1080
1081fn build_decimal_array<T: DecimalType>(
1083 _line_number: usize,
1084 rows: &StringRecords<'_>,
1085 col_idx: usize,
1086 precision: u8,
1087 scale: i8,
1088 null_regex: &NullRegex,
1089) -> Result<ArrayRef, ArrowError> {
1090 let mut decimal_builder = PrimitiveBuilder::<T>::with_capacity(rows.len());
1091 for row in rows.iter() {
1092 let s = row.get(col_idx);
1093 if null_regex.is_null(s) {
1094 decimal_builder.append_null();
1096 } else {
1097 let decimal_value: Result<T::Native, _> = parse_decimal::<T>(s, precision, scale);
1098 match decimal_value {
1099 Ok(v) => {
1100 decimal_builder.append_value(v);
1101 }
1102 Err(e) => {
1103 return Err(e);
1104 }
1105 }
1106 }
1107 }
1108 Ok(Arc::new(
1109 decimal_builder
1110 .finish()
1111 .with_precision_and_scale(precision, scale)?,
1112 ))
1113}
1114
1115fn build_primitive_array<T: ArrowPrimitiveType + Parser>(
1117 line_number: usize,
1118 rows: &StringRecords<'_>,
1119 col_idx: usize,
1120 null_regex: &NullRegex,
1121) -> Result<ArrayRef, ArrowError> {
1122 rows.iter()
1123 .enumerate()
1124 .map(|(row_index, row)| {
1125 let s = row.get(col_idx);
1126 if null_regex.is_null(s) {
1127 return Ok(None);
1128 }
1129
1130 match T::parse(s) {
1131 Some(e) => Ok(Some(e)),
1132 None => Err(ArrowError::ParseError(format!(
1133 "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1135 s,
1136 T::DATA_TYPE,
1137 col_idx,
1138 line_number + row_index,
1139 row
1140 ))),
1141 }
1142 })
1143 .collect::<Result<PrimitiveArray<T>, ArrowError>>()
1144 .map(|e| Arc::new(e) as ArrayRef)
1145}
1146
1147fn build_timestamp_array<T: ArrowTimestampType>(
1148 line_number: usize,
1149 rows: &StringRecords<'_>,
1150 col_idx: usize,
1151 timezone: Option<&str>,
1152 null_regex: &NullRegex,
1153) -> Result<ArrayRef, ArrowError> {
1154 Ok(Arc::new(match timezone {
1155 Some(timezone) => {
1156 let tz: Tz = timezone.parse()?;
1157 build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &tz, null_regex)?
1158 .with_timezone(timezone)
1159 }
1160 None => build_timestamp_array_impl::<T, _>(line_number, rows, col_idx, &Utc, null_regex)?,
1161 }))
1162}
1163
1164fn build_timestamp_array_impl<T: ArrowTimestampType, Tz: TimeZone>(
1165 line_number: usize,
1166 rows: &StringRecords<'_>,
1167 col_idx: usize,
1168 timezone: &Tz,
1169 null_regex: &NullRegex,
1170) -> Result<PrimitiveArray<T>, ArrowError> {
1171 rows.iter()
1172 .enumerate()
1173 .map(|(row_index, row)| {
1174 let s = row.get(col_idx);
1175 if null_regex.is_null(s) {
1176 return Ok(None);
1177 }
1178
1179 let date = string_to_datetime(timezone, s)
1180 .and_then(|date| match T::UNIT {
1181 TimeUnit::Second => Ok(date.timestamp()),
1182 TimeUnit::Millisecond => Ok(date.timestamp_millis()),
1183 TimeUnit::Microsecond => Ok(date.timestamp_micros()),
1184 TimeUnit::Nanosecond => date.timestamp_nanos_opt().ok_or_else(|| {
1185 ArrowError::ParseError(format!(
1186 "{} would overflow 64-bit signed nanoseconds",
1187 date.to_rfc3339(),
1188 ))
1189 }),
1190 })
1191 .map_err(|e| {
1192 ArrowError::ParseError(format!(
1193 "Error parsing column {col_idx} at line {}: {}",
1194 line_number + row_index,
1195 e
1196 ))
1197 })?;
1198 Ok(Some(date))
1199 })
1200 .collect()
1201}
1202
1203fn build_boolean_array(
1205 line_number: usize,
1206 rows: &StringRecords<'_>,
1207 col_idx: usize,
1208 null_regex: &NullRegex,
1209) -> Result<ArrayRef, ArrowError> {
1210 rows.iter()
1211 .enumerate()
1212 .map(|(row_index, row)| {
1213 let s = row.get(col_idx);
1214 if null_regex.is_null(s) {
1215 return Ok(None);
1216 }
1217 let parsed = parse_bool(s);
1218 match parsed {
1219 Some(e) => Ok(Some(e)),
1220 None => Err(ArrowError::ParseError(format!(
1221 "Error while parsing value '{}' as type '{}' for column {} at line {}. Row data: '{}'",
1223 s,
1224 "Boolean",
1225 col_idx,
1226 line_number + row_index,
1227 row
1228 ))),
1229 }
1230 })
1231 .collect::<Result<BooleanArray, _>>()
1232 .map(|e| Arc::new(e) as ArrayRef)
1233}
1234
1235fn build_dictionary_array<K: ArrowDictionaryKeyType>(
1237 rows: &StringRecords<'_>,
1238 col_idx: usize,
1239 null_regex: &NullRegex,
1240) -> Result<ArrayRef, ArrowError> {
1241 let mut builder = StringDictionaryBuilder::<K>::with_capacity(rows.len(), 256, 1024);
1242 for row in rows.iter() {
1243 let s = row.get(col_idx);
1244 if null_regex.is_null(s) {
1245 builder.append_null();
1246 } else {
1247 builder.append(s)?;
1248 }
1249 }
1250 Ok(Arc::new(builder.finish()))
1251}
1252
1253#[derive(Debug)]
1255pub struct ReaderBuilder {
1256 schema: SchemaRef,
1258 format: Format,
1260 batch_size: usize,
1264 bounds: Bounds,
1266 projection: Option<Vec<usize>>,
1268}
1269
1270impl ReaderBuilder {
1271 pub fn new(schema: SchemaRef) -> ReaderBuilder {
1294 Self {
1295 schema,
1296 format: Format::default(),
1297 batch_size: 1024,
1298 bounds: None,
1299 projection: None,
1300 }
1301 }
1302
1303 pub fn with_header(mut self, has_header: bool) -> Self {
1305 self.format.header = has_header;
1306 self
1307 }
1308
1309 pub fn with_header_validation(mut self, validate_header: bool) -> Self {
1313 self.format.header_validation = validate_header;
1314 self
1315 }
1316
1317 pub fn with_format(mut self, format: Format) -> Self {
1319 self.format = format;
1320 self
1321 }
1322
1323 pub fn with_delimiter(mut self, delimiter: u8) -> Self {
1325 self.format.delimiter = Some(delimiter);
1326 self
1327 }
1328
1329 pub fn with_escape(mut self, escape: u8) -> Self {
1331 self.format.escape = Some(escape);
1332 self
1333 }
1334
1335 pub fn with_quote(mut self, quote: u8) -> Self {
1337 self.format.quote = Some(quote);
1338 self
1339 }
1340
1341 pub fn with_terminator(mut self, terminator: u8) -> Self {
1343 self.format.terminator = Some(terminator);
1344 self
1345 }
1346
1347 pub fn with_comment(mut self, comment: u8) -> Self {
1349 self.format.comment = Some(comment);
1350 self
1351 }
1352
1353 pub fn with_null_regex(mut self, null_regex: Regex) -> Self {
1355 self.format.null_regex = NullRegex(Some(null_regex));
1356 self
1357 }
1358
1359 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1361 self.batch_size = batch_size;
1362 self
1363 }
1364
1365 pub fn with_bounds(mut self, start: usize, end: usize) -> Self {
1368 self.bounds = Some((start, end));
1369 self
1370 }
1371
1372 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1374 self.projection = Some(projection);
1375 self
1376 }
1377
1378 pub fn with_truncated_rows(mut self, allow: bool) -> Self {
1385 self.format.truncated_rows = allow;
1386 self
1387 }
1388
1389 pub fn build<R: Read>(self, reader: R) -> Result<Reader<R>, ArrowError> {
1394 self.build_buffered(StdBufReader::new(reader))
1395 }
1396
1397 pub fn build_buffered<R: BufRead>(self, reader: R) -> Result<BufReader<R>, ArrowError> {
1399 let schema = match &self.projection {
1400 Some(projection) => Arc::new(self.schema.project(projection)?),
1401 None => self.schema.clone(),
1402 };
1403
1404 Ok(BufReader {
1405 reader,
1406 decoder: self.build_decoder(),
1407 schema,
1408 })
1409 }
1410
1411 pub fn build_decoder(self) -> Decoder {
1413 let delimiter = self.format.build_parser();
1414 let record_decoder = RecordDecoder::new(
1415 delimiter,
1416 self.schema.fields().len(),
1417 self.format.truncated_rows,
1418 );
1419
1420 let header = self.format.header as usize;
1421
1422 let (start, end) = match self.bounds {
1423 Some((start, end)) => (start + header, end + header),
1424 None => (header, usize::MAX),
1425 };
1426
1427 Decoder {
1428 schema: self.schema,
1429 to_skip: start,
1430 header_validation: self.format.header && self.format.header_validation,
1431 record_decoder,
1432 line_number: start,
1433 end,
1434 projection: self.projection,
1435 batch_size: self.batch_size,
1436 null_regex: self.format.null_regex,
1437 }
1438 }
1439}
1440
1441#[cfg(test)]
1442mod tests {
1443 use super::*;
1444
1445 use std::io::{Cursor, Seek, SeekFrom, Write};
1446 use tempfile::NamedTempFile;
1447
1448 use arrow_array::cast::AsArray;
1449 use arrow_cast::display::array_value_to_string;
1450
1451 #[test]
1452 fn test_infer_schema_leading_plus_numbers() {
1453 for (csv, expected_type) in [
1454 ("+1\n2\n-3\n", DataType::Int64),
1455 ("+1.5\n2.5\n-3.5\n", DataType::Float64),
1456 ("+1e3\n+2.5e-2\n-3E+2\n", DataType::Float64),
1457 ("+9223372036854775807\n0\n", DataType::Int64),
1458 ("+9223372036854775808\n0\n", DataType::Utf8),
1459 ("+-1\n2\n", DataType::Utf8),
1460 ("+\n2\n", DataType::Utf8),
1461 ] {
1462 let (schema, records_read) = Format::default()
1463 .infer_schema(Cursor::new(csv), None)
1464 .unwrap();
1465 assert_eq!(schema.field(0).data_type(), &expected_type, "CSV: {csv:?}");
1466 let reader = ReaderBuilder::new(Arc::new(schema))
1468 .build(Cursor::new(csv))
1469 .unwrap();
1470 let rows: usize = reader.map(|batch| batch.unwrap().num_rows()).sum();
1471 assert_eq!(rows, records_read, "CSV: {csv:?}");
1472 }
1473 }
1474
1475 #[test]
1476 fn test_csv() {
1477 let schema = Arc::new(Schema::new(vec![
1478 Field::new("city", DataType::Utf8, false),
1479 Field::new("lat", DataType::Float64, false),
1480 Field::new("lng", DataType::Float64, false),
1481 ]));
1482
1483 let file = File::open("test/data/uk_cities.csv").unwrap();
1484 let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1485 assert_eq!(schema, csv.schema());
1486 let batch = csv.next().unwrap().unwrap();
1487 assert_eq!(37, batch.num_rows());
1488 assert_eq!(3, batch.num_columns());
1489
1490 let lat = batch.column(1).as_primitive::<Float64Type>();
1492 assert_eq!(57.653484, lat.value(0));
1493
1494 let city = batch.column(0).as_string::<i32>();
1496
1497 assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1498 }
1499
1500 #[test]
1501 fn test_csv_schema_metadata() {
1502 let mut metadata = std::collections::HashMap::new();
1503 metadata.insert("foo".to_owned(), "bar".to_owned());
1504 let schema = Arc::new(Schema::new_with_metadata(
1505 vec![
1506 Field::new("city", DataType::Utf8, false),
1507 Field::new("lat", DataType::Float64, false),
1508 Field::new("lng", DataType::Float64, false),
1509 ],
1510 metadata.clone(),
1511 ));
1512
1513 let file = File::open("test/data/uk_cities.csv").unwrap();
1514
1515 let mut csv = ReaderBuilder::new(schema.clone()).build(file).unwrap();
1516 assert_eq!(schema, csv.schema());
1517 let batch = csv.next().unwrap().unwrap();
1518 assert_eq!(37, batch.num_rows());
1519 assert_eq!(3, batch.num_columns());
1520
1521 assert_eq!(batch.schema().metadata(), &metadata);
1522 }
1523
1524 #[test]
1525 fn test_csv_reader_with_decimal() {
1526 let schema = Arc::new(Schema::new(vec![
1527 Field::new("city", DataType::Utf8, false),
1528 Field::new("lat", DataType::Decimal128(38, 6), false),
1529 Field::new("lng", DataType::Decimal256(76, 6), false),
1530 ]));
1531
1532 let file = File::open("test/data/decimal_test.csv").unwrap();
1533
1534 let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1535 let batch = csv.next().unwrap().unwrap();
1536 let lat = batch
1538 .column(1)
1539 .as_any()
1540 .downcast_ref::<Decimal128Array>()
1541 .unwrap();
1542
1543 assert_eq!("57.653484", lat.value_as_string(0));
1544 assert_eq!("53.002666", lat.value_as_string(1));
1545 assert_eq!("52.412811", lat.value_as_string(2));
1546 assert_eq!("51.481583", lat.value_as_string(3));
1547 assert_eq!("12.123457", lat.value_as_string(4));
1548 assert_eq!("50.760000", lat.value_as_string(5));
1549 assert_eq!("0.123000", lat.value_as_string(6));
1550 assert_eq!("123.000000", lat.value_as_string(7));
1551 assert_eq!("123.000000", lat.value_as_string(8));
1552 assert_eq!("-50.760000", lat.value_as_string(9));
1553
1554 let lng = batch
1555 .column(2)
1556 .as_any()
1557 .downcast_ref::<Decimal256Array>()
1558 .unwrap();
1559
1560 assert_eq!("-3.335724", lng.value_as_string(0));
1561 assert_eq!("-2.179404", lng.value_as_string(1));
1562 assert_eq!("-1.778197", lng.value_as_string(2));
1563 assert_eq!("-3.179090", lng.value_as_string(3));
1564 assert_eq!("-3.179090", lng.value_as_string(4));
1565 assert_eq!("0.290472", lng.value_as_string(5));
1566 assert_eq!("0.290472", lng.value_as_string(6));
1567 assert_eq!("0.290472", lng.value_as_string(7));
1568 assert_eq!("0.290472", lng.value_as_string(8));
1569 assert_eq!("0.290472", lng.value_as_string(9));
1570 }
1571
1572 #[test]
1573 fn test_csv_reader_decimal_parsing() {
1574 let data = " 1.995 ,1.5e2,1234.5,0e0\n-0.005,-1.5E-2,-150,1E+2\n123,+.5,5,-7\n";
1577 let schema = Arc::new(Schema::new(vec![
1578 Field::new("a", DataType::Decimal128(10, 2), false),
1579 Field::new("b", DataType::Decimal64(18, 2), false),
1580 Field::new("c", DataType::Decimal128(10, -2), false),
1581 Field::new("d", DataType::Decimal32(9, 0), false),
1582 ]));
1583 let mut csv = ReaderBuilder::new(schema).build(Cursor::new(data)).unwrap();
1584 let batch = csv.next().unwrap().unwrap();
1585 let column = |i: usize| {
1586 (0..batch.num_rows())
1587 .map(|row| array_value_to_string(batch.column(i), row).unwrap())
1588 .collect::<Vec<_>>()
1589 };
1590 assert_eq!(column(0), ["2.00", "-0.01", "123.00"]);
1591 assert_eq!(column(1), ["150.00", "-0.02", "0.50"]);
1592 assert_eq!(
1593 batch.column(2).as_primitive::<Decimal128Type>().values(),
1594 &[12, -2, 0]
1595 );
1596 assert_eq!(column(3), ["0", "100", "-7"]);
1597
1598 for (data, expected) in [
1600 ("abc\n", "Invalid decimal format: \"abc\""),
1601 ("1.2.3\n", "Invalid decimal format: \"1.2.3\""),
1602 (
1603 "123456789\n",
1604 "\"123456789\" does not fit in Decimal128(5, 2)",
1605 ),
1606 ("1e99999\n", "does not fit in Decimal128(5, 2)"),
1607 (
1608 &format!("{}\n", "1".repeat(300)),
1609 "does not fit in Decimal128(5, 2)",
1610 ),
1611 (
1612 "4825037936439135476.2609835314269495255615E-14\n",
1613 "does not fit in Decimal128(5, 2)",
1614 ),
1615 ] {
1616 let schema = Arc::new(Schema::new(vec![Field::new(
1617 "a",
1618 DataType::Decimal128(5, 2),
1619 false,
1620 )]));
1621 let mut csv = ReaderBuilder::new(schema).build(Cursor::new(data)).unwrap();
1622 let err = csv.next().unwrap().unwrap_err().to_string();
1623 assert!(err.contains(expected), "{data:?}: {err}");
1624 }
1625 }
1626
1627 #[test]
1628 fn test_csv_reader_with_decimal_3264() {
1629 let schema = Arc::new(Schema::new(vec![
1630 Field::new("city", DataType::Utf8, false),
1631 Field::new("lat", DataType::Decimal32(9, 6), false),
1632 Field::new("lng", DataType::Decimal64(16, 6), false),
1633 ]));
1634
1635 let file = File::open("test/data/decimal_test.csv").unwrap();
1636
1637 let mut csv = ReaderBuilder::new(schema).build(file).unwrap();
1638 let batch = csv.next().unwrap().unwrap();
1639 let lat = batch
1641 .column(1)
1642 .as_any()
1643 .downcast_ref::<Decimal32Array>()
1644 .unwrap();
1645
1646 assert_eq!("57.653484", lat.value_as_string(0));
1647 assert_eq!("53.002666", lat.value_as_string(1));
1648 assert_eq!("52.412811", lat.value_as_string(2));
1649 assert_eq!("51.481583", lat.value_as_string(3));
1650 assert_eq!("12.123457", lat.value_as_string(4));
1651 assert_eq!("50.760000", lat.value_as_string(5));
1652 assert_eq!("0.123000", lat.value_as_string(6));
1653 assert_eq!("123.000000", lat.value_as_string(7));
1654 assert_eq!("123.000000", lat.value_as_string(8));
1655 assert_eq!("-50.760000", lat.value_as_string(9));
1656
1657 let lng = batch
1658 .column(2)
1659 .as_any()
1660 .downcast_ref::<Decimal64Array>()
1661 .unwrap();
1662
1663 assert_eq!("-3.335724", lng.value_as_string(0));
1664 assert_eq!("-2.179404", lng.value_as_string(1));
1665 assert_eq!("-1.778197", lng.value_as_string(2));
1666 assert_eq!("-3.179090", lng.value_as_string(3));
1667 assert_eq!("-3.179090", lng.value_as_string(4));
1668 assert_eq!("0.290472", lng.value_as_string(5));
1669 assert_eq!("0.290472", lng.value_as_string(6));
1670 assert_eq!("0.290472", lng.value_as_string(7));
1671 assert_eq!("0.290472", lng.value_as_string(8));
1672 assert_eq!("0.290472", lng.value_as_string(9));
1673 }
1674
1675 #[test]
1676 fn test_csv_from_buf_reader() {
1677 let schema = Schema::new(vec![
1678 Field::new("city", DataType::Utf8, false),
1679 Field::new("lat", DataType::Float64, false),
1680 Field::new("lng", DataType::Float64, false),
1681 ]);
1682
1683 let file_with_headers = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1684 let file_without_headers = File::open("test/data/uk_cities.csv").unwrap();
1685 let both_files = file_with_headers
1686 .chain(Cursor::new("\n".to_string()))
1687 .chain(file_without_headers);
1688 let mut csv = ReaderBuilder::new(Arc::new(schema))
1689 .with_header(true)
1690 .build(both_files)
1691 .unwrap();
1692 let batch = csv.next().unwrap().unwrap();
1693 assert_eq!(74, batch.num_rows());
1694 assert_eq!(3, batch.num_columns());
1695 }
1696
1697 #[test]
1698 fn test_infer_format_with_typed_columns() {
1699 let csv = "name,count,active\nalice,1,true\nbob,2,false\n";
1700
1701 let (format, format_records_read) = Format::default()
1702 .infer_format(Cursor::new(csv), None)
1703 .unwrap();
1704 let (schema, records_read) = format.infer_schema(Cursor::new(csv), None).unwrap();
1705
1706 assert_eq!(schema.field(0).name(), "name");
1707 assert_eq!(schema.field(1).name(), "count");
1708 assert_eq!(schema.field(2).name(), "active");
1709 assert_eq!(format_records_read, 3);
1710 assert_eq!(records_read, 2);
1711 }
1712
1713 #[test]
1714 fn test_infer_format_without_header() {
1715 let csv = "1,true\n2,false\n";
1716
1717 let (format, format_records_read) = Format::default()
1718 .infer_format(Cursor::new(csv), None)
1719 .unwrap();
1720 let (schema, records_read) = format.infer_schema(Cursor::new(csv), None).unwrap();
1721
1722 assert_eq!(schema.field(0).name(), "column_1");
1723 assert_eq!(schema.field(1).name(), "column_2");
1724 assert_eq!(format_records_read, 2);
1725 assert_eq!(records_read, 2);
1726 }
1727
1728 #[test]
1729 fn test_infer_format_returns_no_header_when_ambiguous() {
1730 for csv in ["name,count\n", "alice,london\nbob,paris\n"] {
1731 let (format, _) = Format::default()
1732 .infer_format(Cursor::new(csv), None)
1733 .unwrap();
1734 let (schema, _) = format.infer_schema(Cursor::new(csv), None).unwrap();
1735 assert_eq!(schema.field(0).name(), "column_1", "CSV: {csv:?}");
1736 }
1737
1738 let (format, format_records_read) = Format::default()
1739 .infer_format(Cursor::new(""), None)
1740 .unwrap();
1741 let (schema, records_read) = format.infer_schema(Cursor::new(""), None).unwrap();
1742 assert!(schema.fields().is_empty());
1743 assert_eq!(format_records_read, 0);
1744 assert_eq!(records_read, 0);
1745 }
1746
1747 #[test]
1748 fn test_infer_format_honors_format_options() {
1749 let csv = "name;count\nalice;1\nbob;2\n";
1750 let (format, _) = Format::default()
1751 .with_delimiter(b';')
1752 .infer_format(Cursor::new(csv), None)
1753 .unwrap();
1754 let (schema, records_read) = format.infer_schema(Cursor::new(csv), None).unwrap();
1755
1756 assert_eq!(schema.field(0).name(), "name");
1757 assert_eq!(schema.field(1).name(), "count");
1758 assert_eq!(records_read, 2);
1759 }
1760
1761 #[test]
1762 fn test_infer_format_respects_max_records() {
1763 let csv = "name,count\nalice,1\nbob,unknown\n";
1764 let infer = |max_records| {
1765 let (format, records_read) = Format::default()
1766 .infer_format(Cursor::new(csv), max_records)
1767 .unwrap();
1768 let (schema, _) = format.infer_schema(Cursor::new(csv), None).unwrap();
1769 (schema.field(0).name().clone(), records_read)
1770 };
1771
1772 assert_eq!(infer(Some(1)), ("name".to_string(), 2));
1773 assert_eq!(infer(None), ("column_1".to_string(), 3));
1774 assert_eq!(infer(Some(0)), ("column_1".to_string(), 1));
1775 }
1776
1777 #[test]
1778 fn test_infer_format_numeric_text_is_not_header() {
1779 for csv in [
1780 "+1\n2\n3\n",
1781 "+1.5\n2.5\n3.5\n",
1782 "+1e3\n2e3\n3e3\n",
1783 "9223372036854775808\n2\n3\n",
1784 ] {
1785 let (format, records_read) = Format::default()
1786 .infer_format(Cursor::new(csv), None)
1787 .unwrap();
1788 let (schema, schema_records_read) =
1789 format.infer_schema(Cursor::new(csv), None).unwrap();
1790 let (ordinary_schema, _) = Format::default()
1791 .infer_schema(Cursor::new(csv), None)
1792 .unwrap();
1793
1794 assert_eq!(schema.field(0).name(), "column_1", "CSV: {csv:?}");
1795 assert_eq!(schema, ordinary_schema, "CSV: {csv:?}");
1796 assert_eq!(records_read, 3);
1797 assert_eq!(schema_records_read, 3);
1798 }
1799 }
1800
1801 #[test]
1802 #[cfg_attr(miri, ignore)] fn test_csv_with_schema_inference() {
1804 let mut file = File::open("test/data/uk_cities_with_headers.csv").unwrap();
1805
1806 let (schema, _) = Format::default()
1807 .with_header(true)
1808 .infer_schema(&mut file, None)
1809 .unwrap();
1810
1811 file.rewind().unwrap();
1812 let builder = ReaderBuilder::new(Arc::new(schema)).with_header(true);
1813
1814 let mut csv = builder.build(file).unwrap();
1815 let expected_schema = Schema::new(vec![
1816 Field::new("city", DataType::Utf8, true),
1817 Field::new("lat", DataType::Float64, true),
1818 Field::new("lng", DataType::Float64, true),
1819 ]);
1820 assert_eq!(Arc::new(expected_schema), csv.schema());
1821 let batch = csv.next().unwrap().unwrap();
1822 assert_eq!(37, batch.num_rows());
1823 assert_eq!(3, batch.num_columns());
1824
1825 let lat = batch
1827 .column(1)
1828 .as_any()
1829 .downcast_ref::<Float64Array>()
1830 .unwrap();
1831 assert_eq!(57.653484, lat.value(0));
1832
1833 let city = batch
1835 .column(0)
1836 .as_any()
1837 .downcast_ref::<StringArray>()
1838 .unwrap();
1839
1840 assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1841 }
1842
1843 #[test]
1844 #[cfg_attr(miri, ignore)] fn test_csv_with_schema_inference_no_headers() {
1846 let mut file = File::open("test/data/uk_cities.csv").unwrap();
1847
1848 let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1849 file.rewind().unwrap();
1850
1851 let mut csv = ReaderBuilder::new(Arc::new(schema)).build(file).unwrap();
1852
1853 let schema = csv.schema();
1855 assert_eq!("column_1", schema.field(0).name());
1856 assert_eq!("column_2", schema.field(1).name());
1857 assert_eq!("column_3", schema.field(2).name());
1858 let batch = csv.next().unwrap().unwrap();
1859 let batch_schema = batch.schema();
1860
1861 assert_eq!(schema, batch_schema);
1862 assert_eq!(37, batch.num_rows());
1863 assert_eq!(3, batch.num_columns());
1864
1865 let lat = batch
1867 .column(1)
1868 .as_any()
1869 .downcast_ref::<Float64Array>()
1870 .unwrap();
1871 assert_eq!(57.653484, lat.value(0));
1872
1873 let city = batch
1875 .column(0)
1876 .as_any()
1877 .downcast_ref::<StringArray>()
1878 .unwrap();
1879
1880 assert_eq!("Aberdeen, Aberdeen City, UK", city.value(13));
1881 }
1882
1883 #[test]
1884 #[cfg_attr(miri, ignore)] fn test_csv_builder_with_bounds() {
1886 let mut file = File::open("test/data/uk_cities.csv").unwrap();
1887
1888 let (schema, _) = Format::default().infer_schema(&mut file, None).unwrap();
1890 file.rewind().unwrap();
1891 let mut csv = ReaderBuilder::new(Arc::new(schema))
1892 .with_bounds(0, 2)
1893 .build(file)
1894 .unwrap();
1895 let batch = csv.next().unwrap().unwrap();
1896
1897 let city = batch
1899 .column(0)
1900 .as_any()
1901 .downcast_ref::<StringArray>()
1902 .unwrap();
1903
1904 assert_eq!("Elgin, Scotland, the UK", city.value(0));
1906
1907 let result = std::panic::catch_unwind(|| city.value(13));
1910 assert!(result.is_err());
1911 }
1912
1913 #[test]
1914 fn test_csv_with_projection() {
1915 let schema = Arc::new(Schema::new(vec![
1916 Field::new("city", DataType::Utf8, false),
1917 Field::new("lat", DataType::Float64, false),
1918 Field::new("lng", DataType::Float64, false),
1919 ]));
1920
1921 let file = File::open("test/data/uk_cities.csv").unwrap();
1922
1923 let mut csv = ReaderBuilder::new(schema)
1924 .with_projection(vec![0, 1])
1925 .build(file)
1926 .unwrap();
1927
1928 let projected_schema = Arc::new(Schema::new(vec![
1929 Field::new("city", DataType::Utf8, false),
1930 Field::new("lat", DataType::Float64, false),
1931 ]));
1932 assert_eq!(projected_schema, csv.schema());
1933 let batch = csv.next().unwrap().unwrap();
1934 assert_eq!(projected_schema, batch.schema());
1935 assert_eq!(37, batch.num_rows());
1936 assert_eq!(2, batch.num_columns());
1937 }
1938
1939 #[test]
1940 fn test_csv_record_batch_reader_schema() {
1941 let schema = Arc::new(Schema::new(vec![
1942 Field::new("a", DataType::Int32, false),
1943 Field::new("b", DataType::Int32, false),
1944 ]));
1945
1946 let cases = [
1947 None,
1948 Some(vec![]),
1949 Some(vec![1]),
1950 Some(vec![1, 0]),
1951 Some(vec![1, 1]),
1952 ];
1953 for projection in cases {
1954 let builder = ReaderBuilder::new(schema.clone());
1955 let builder = match projection {
1956 Some(projection) => builder.with_projection(projection),
1957 None => builder,
1958 };
1959 let mut reader = builder.build(Cursor::new(b"1,2\n")).unwrap();
1960
1961 let reader_schema = RecordBatchReader::schema(&reader);
1962 let batch = reader.next().unwrap().unwrap();
1963
1964 assert_eq!(reader_schema, batch.schema());
1965 }
1966 }
1967
1968 #[test]
1969 fn test_csv_reader_rejects_invalid_projection() {
1970 let schema = Arc::new(Schema::new(vec![
1971 Field::new("a", DataType::Int32, false),
1972 Field::new("b", DataType::Int32, false),
1973 ]));
1974
1975 let result = ReaderBuilder::new(schema)
1976 .with_projection(vec![2])
1977 .build(Cursor::new(b"1,2\n"));
1978
1979 assert!(matches!(
1980 result,
1981 Err(ArrowError::SchemaError(message))
1982 if message == "project index 2 out of bounds, max field 2"
1983 ));
1984 }
1985
1986 #[test]
1987 fn test_csv_decoder_rejects_invalid_projection() {
1988 let schema = Arc::new(Schema::new(vec![
1989 Field::new("a", DataType::Int32, false),
1990 Field::new("b", DataType::Int32, false),
1991 ]));
1992 let mut decoder = ReaderBuilder::new(schema)
1993 .with_projection(vec![2])
1994 .build_decoder();
1995
1996 decoder.decode(b"1,2\n").unwrap();
1997 let result = decoder.flush();
1998
1999 assert!(matches!(
2000 result,
2001 Err(ArrowError::SchemaError(message))
2002 if message == "project index 2 out of bounds, max field 2"
2003 ));
2004 }
2005
2006 #[test]
2007 fn test_csv_with_dictionary() {
2008 let schema = Arc::new(Schema::new(vec![
2009 Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
2010 Field::new("lat", DataType::Float64, false),
2011 Field::new("lng", DataType::Float64, false),
2012 ]));
2013
2014 let file = File::open("test/data/uk_cities.csv").unwrap();
2015
2016 let mut csv = ReaderBuilder::new(schema)
2017 .with_projection(vec![0, 1])
2018 .build(file)
2019 .unwrap();
2020
2021 let projected_schema = Arc::new(Schema::new(vec![
2022 Field::new_dictionary("city", DataType::Int32, DataType::Utf8, false),
2023 Field::new("lat", DataType::Float64, false),
2024 ]));
2025 assert_eq!(projected_schema, csv.schema());
2026 let batch = csv.next().unwrap().unwrap();
2027 assert_eq!(projected_schema, batch.schema());
2028 assert_eq!(37, batch.num_rows());
2029 assert_eq!(2, batch.num_columns());
2030
2031 let strings = arrow_cast::cast(batch.column(0), &DataType::Utf8).unwrap();
2032 let strings = strings.as_string::<i32>();
2033
2034 assert_eq!(strings.value(0), "Elgin, Scotland, the UK");
2035 assert_eq!(strings.value(4), "Eastbourne, East Sussex, UK");
2036 assert_eq!(strings.value(29), "Uckfield, East Sussex, UK");
2037 }
2038
2039 #[test]
2040 fn test_csv_dictionary_key_overflow() {
2041 let schema = Arc::new(Schema::new(vec![Field::new_dictionary(
2042 "c",
2043 DataType::Int8,
2044 DataType::Utf8,
2045 false,
2046 )]));
2047 let read = |distinct: usize| {
2048 let rows = (0..distinct).map(|i| format!("v{i}")).collect::<Vec<_>>();
2049 let mut reader = ReaderBuilder::new(schema.clone())
2050 .build(Cursor::new(rows.join("\n")))
2051 .unwrap();
2052 reader.next().unwrap()
2053 };
2054
2055 assert_eq!(read(128).unwrap().num_rows(), 128);
2056 let err = read(129).unwrap_err();
2057 assert!(matches!(err, ArrowError::DictionaryKeyOverflowError));
2058 }
2059
2060 #[test]
2061 fn test_csv_with_nullable_dictionary() {
2062 let offset_type = vec![
2063 DataType::Int8,
2064 DataType::Int16,
2065 DataType::Int32,
2066 DataType::Int64,
2067 DataType::UInt8,
2068 DataType::UInt16,
2069 DataType::UInt32,
2070 DataType::UInt64,
2071 ];
2072 for data_type in offset_type {
2073 let file = File::open("test/data/dictionary_nullable_test.csv").unwrap();
2074 let dictionary_type =
2075 DataType::Dictionary(Box::new(data_type), Box::new(DataType::Utf8));
2076 let schema = Arc::new(Schema::new(vec![
2077 Field::new("id", DataType::Utf8, false),
2078 Field::new("name", dictionary_type.clone(), true),
2079 ]));
2080
2081 let mut csv = ReaderBuilder::new(schema)
2082 .build(file.try_clone().unwrap())
2083 .unwrap();
2084
2085 let batch = csv.next().unwrap().unwrap();
2086 assert_eq!(3, batch.num_rows());
2087 assert_eq!(2, batch.num_columns());
2088
2089 let names = arrow_cast::cast(batch.column(1), &dictionary_type).unwrap();
2090 assert!(!names.is_null(2));
2091 assert!(names.is_null(1));
2092 }
2093 }
2094 #[test]
2095 fn test_nulls() {
2096 let schema = Arc::new(Schema::new(vec![
2097 Field::new("c_int", DataType::UInt64, false),
2098 Field::new("c_float", DataType::Float32, true),
2099 Field::new("c_string", DataType::Utf8, true),
2100 Field::new("c_bool", DataType::Boolean, false),
2101 ]));
2102
2103 let file = File::open("test/data/null_test.csv").unwrap();
2104
2105 let mut csv = ReaderBuilder::new(schema)
2106 .with_header(true)
2107 .build(file)
2108 .unwrap();
2109
2110 let batch = csv.next().unwrap().unwrap();
2111
2112 assert!(!batch.column(1).is_null(0));
2113 assert!(!batch.column(1).is_null(1));
2114 assert!(batch.column(1).is_null(2));
2115 assert!(!batch.column(1).is_null(3));
2116 assert!(!batch.column(1).is_null(4));
2117 }
2118
2119 #[test]
2120 fn test_init_nulls() {
2121 let schema = Arc::new(Schema::new(vec![
2122 Field::new("c_int", DataType::UInt64, true),
2123 Field::new("c_float", DataType::Float32, true),
2124 Field::new("c_string", DataType::Utf8, true),
2125 Field::new("c_bool", DataType::Boolean, true),
2126 Field::new("c_null", DataType::Null, true),
2127 ]));
2128 let file = File::open("test/data/init_null_test.csv").unwrap();
2129
2130 let mut csv = ReaderBuilder::new(schema)
2131 .with_header(true)
2132 .build(file)
2133 .unwrap();
2134
2135 let batch = csv.next().unwrap().unwrap();
2136
2137 assert!(batch.column(1).is_null(0));
2138 assert!(!batch.column(1).is_null(1));
2139 assert!(batch.column(1).is_null(2));
2140 assert!(!batch.column(1).is_null(3));
2141 assert!(!batch.column(1).is_null(4));
2142 }
2143
2144 #[test]
2145 #[cfg_attr(miri, ignore)] fn test_init_nulls_with_inference() {
2147 let format = Format::default().with_header(true).with_delimiter(b',');
2148
2149 let mut file = File::open("test/data/init_null_test.csv").unwrap();
2150 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2151 file.rewind().unwrap();
2152
2153 let expected_schema = Schema::new(vec![
2154 Field::new("c_int", DataType::Int64, true),
2155 Field::new("c_float", DataType::Float64, true),
2156 Field::new("c_string", DataType::Utf8, true),
2157 Field::new("c_bool", DataType::Boolean, true),
2158 Field::new("c_null", DataType::Null, true),
2159 ]);
2160 assert_eq!(schema, expected_schema);
2161
2162 let mut csv = ReaderBuilder::new(Arc::new(schema))
2163 .with_format(format)
2164 .build(file)
2165 .unwrap();
2166
2167 let batch = csv.next().unwrap().unwrap();
2168
2169 assert!(batch.column(1).is_null(0));
2170 assert!(!batch.column(1).is_null(1));
2171 assert!(batch.column(1).is_null(2));
2172 assert!(!batch.column(1).is_null(3));
2173 assert!(!batch.column(1).is_null(4));
2174 }
2175
2176 #[test]
2177 fn test_custom_nulls() {
2178 let schema = Arc::new(Schema::new(vec![
2179 Field::new("c_int", DataType::UInt64, true),
2180 Field::new("c_float", DataType::Float32, true),
2181 Field::new("c_string", DataType::Utf8, true),
2182 Field::new("c_bool", DataType::Boolean, true),
2183 ]));
2184
2185 let file = File::open("test/data/custom_null_test.csv").unwrap();
2186
2187 let null_regex = Regex::new("^nil$").unwrap();
2188
2189 let mut csv = ReaderBuilder::new(schema)
2190 .with_header(true)
2191 .with_null_regex(null_regex)
2192 .build(file)
2193 .unwrap();
2194
2195 let batch = csv.next().unwrap().unwrap();
2196
2197 assert!(batch.column(0).is_null(1));
2199 assert!(batch.column(1).is_null(2));
2200 assert!(batch.column(3).is_null(4));
2201 assert!(batch.column(2).is_null(3));
2202 assert!(!batch.column(2).is_null(4));
2203 }
2204
2205 #[test]
2206 #[cfg_attr(miri, ignore)] fn test_nulls_with_inference() {
2208 let mut file = File::open("test/data/various_types.csv").unwrap();
2209 let format = Format::default().with_header(true).with_delimiter(b'|');
2210
2211 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2212 file.rewind().unwrap();
2213
2214 let builder = ReaderBuilder::new(Arc::new(schema))
2215 .with_format(format)
2216 .with_batch_size(512)
2217 .with_projection(vec![0, 1, 2, 3, 4, 5]);
2218
2219 let mut csv = builder.build(file).unwrap();
2220 let batch = csv.next().unwrap().unwrap();
2221
2222 assert_eq!(10, batch.num_rows());
2223 assert_eq!(6, batch.num_columns());
2224
2225 let schema = batch.schema();
2226
2227 assert_eq!(&DataType::Int64, schema.field(0).data_type());
2228 assert_eq!(&DataType::Float64, schema.field(1).data_type());
2229 assert_eq!(&DataType::Float64, schema.field(2).data_type());
2230 assert_eq!(&DataType::Boolean, schema.field(3).data_type());
2231 assert_eq!(&DataType::Date32, schema.field(4).data_type());
2232 assert_eq!(
2233 &DataType::Timestamp(TimeUnit::Second, None),
2234 schema.field(5).data_type()
2235 );
2236
2237 let names: Vec<&str> = schema.fields().iter().map(|x| x.name().as_str()).collect();
2238 assert_eq!(
2239 names,
2240 vec![
2241 "c_int",
2242 "c_float",
2243 "c_string",
2244 "c_bool",
2245 "c_date",
2246 "c_datetime"
2247 ]
2248 );
2249
2250 assert!(schema.field(0).is_nullable());
2251 assert!(schema.field(1).is_nullable());
2252 assert!(schema.field(2).is_nullable());
2253 assert!(schema.field(3).is_nullable());
2254 assert!(schema.field(4).is_nullable());
2255 assert!(schema.field(5).is_nullable());
2256
2257 assert!(!batch.column(1).is_null(0));
2258 assert!(!batch.column(1).is_null(1));
2259 assert!(batch.column(1).is_null(2));
2260 assert!(!batch.column(1).is_null(3));
2261 assert!(!batch.column(1).is_null(4));
2262 }
2263
2264 #[test]
2265 #[cfg_attr(miri, ignore)] fn test_custom_nulls_with_inference() {
2267 let mut file = File::open("test/data/custom_null_test.csv").unwrap();
2268
2269 let null_regex = Regex::new("^nil$").unwrap();
2270
2271 let format = Format::default()
2272 .with_header(true)
2273 .with_null_regex(null_regex);
2274
2275 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2276 file.rewind().unwrap();
2277
2278 let expected_schema = Schema::new(vec![
2279 Field::new("c_int", DataType::Int64, true),
2280 Field::new("c_float", DataType::Float64, true),
2281 Field::new("c_string", DataType::Utf8, true),
2282 Field::new("c_bool", DataType::Boolean, true),
2283 ]);
2284
2285 assert_eq!(schema, expected_schema);
2286
2287 let builder = ReaderBuilder::new(Arc::new(schema))
2288 .with_format(format)
2289 .with_batch_size(512)
2290 .with_projection(vec![0, 1, 2, 3]);
2291
2292 let mut csv = builder.build(file).unwrap();
2293 let batch = csv.next().unwrap().unwrap();
2294
2295 assert_eq!(5, batch.num_rows());
2296 assert_eq!(4, batch.num_columns());
2297
2298 assert_eq!(batch.schema().as_ref(), &expected_schema);
2299 }
2300
2301 #[test]
2302 #[cfg_attr(miri, ignore)] fn test_scientific_notation_with_inference() {
2304 let mut file = File::open("test/data/scientific_notation_test.csv").unwrap();
2305 let format = Format::default().with_header(false).with_delimiter(b',');
2306
2307 let (schema, _) = format.infer_schema(&mut file, None).unwrap();
2308 file.rewind().unwrap();
2309
2310 let builder = ReaderBuilder::new(Arc::new(schema))
2311 .with_format(format)
2312 .with_batch_size(512)
2313 .with_projection(vec![0, 1]);
2314
2315 let mut csv = builder.build(file).unwrap();
2316 let batch = csv.next().unwrap().unwrap();
2317
2318 let schema = batch.schema();
2319
2320 assert_eq!(&DataType::Float64, schema.field(0).data_type());
2321 }
2322
2323 fn invalid_csv_helper(file_name: &str) -> String {
2324 let file = File::open(file_name).unwrap();
2325 let schema = Schema::new(vec![
2326 Field::new("c_int", DataType::UInt64, false),
2327 Field::new("c_float", DataType::Float32, false),
2328 Field::new("c_string", DataType::Utf8, false),
2329 Field::new("c_bool", DataType::Boolean, false),
2330 ]);
2331
2332 let builder = ReaderBuilder::new(Arc::new(schema))
2333 .with_header(true)
2334 .with_delimiter(b'|')
2335 .with_batch_size(512)
2336 .with_projection(vec![0, 1, 2, 3]);
2337
2338 let mut csv = builder.build(file).unwrap();
2339
2340 csv.next().unwrap().unwrap_err().to_string()
2341 }
2342
2343 #[test]
2344 fn test_parse_invalid_csv_float() {
2345 let file_name = "test/data/various_invalid_types/invalid_float.csv";
2346
2347 let error = invalid_csv_helper(file_name);
2348 assert_eq!(
2349 "Parser error: Error while parsing value '4.x4' as type 'Float32' for column 1 at line 4. Row data: '[4,4.x4,,false]'",
2350 error
2351 );
2352 }
2353
2354 #[test]
2355 fn test_parse_invalid_csv_int() {
2356 let file_name = "test/data/various_invalid_types/invalid_int.csv";
2357
2358 let error = invalid_csv_helper(file_name);
2359 assert_eq!(
2360 "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]'",
2361 error
2362 );
2363 }
2364
2365 #[test]
2366 fn test_parse_invalid_csv_bool() {
2367 let file_name = "test/data/various_invalid_types/invalid_bool.csv";
2368
2369 let error = invalid_csv_helper(file_name);
2370 assert_eq!(
2371 "Parser error: Error while parsing value 'none' as type 'Boolean' for column 3 at line 2. Row data: '[2,2.2,2.22,none]'",
2372 error
2373 );
2374 }
2375
2376 fn infer_field_schema(string: &str) -> DataType {
2378 let mut v = InferredDataType::default();
2379 v.update(string);
2380 v.get()
2381 }
2382
2383 #[test]
2384 #[cfg_attr(miri, ignore)] fn test_infer_field_schema() {
2386 assert_eq!(infer_field_schema("A"), DataType::Utf8);
2387 assert_eq!(infer_field_schema("\"123\""), DataType::Utf8);
2388 assert_eq!(infer_field_schema("10"), DataType::Int64);
2389 assert_eq!(infer_field_schema("10.2"), DataType::Float64);
2390 assert_eq!(infer_field_schema(".2"), DataType::Float64);
2391 assert_eq!(infer_field_schema("2."), DataType::Float64);
2392 assert_eq!(infer_field_schema("NaN"), DataType::Float64);
2393 assert_eq!(infer_field_schema("nan"), DataType::Float64);
2394 assert_eq!(infer_field_schema("inf"), DataType::Float64);
2395 assert_eq!(infer_field_schema("-inf"), DataType::Float64);
2396 assert_eq!(infer_field_schema("true"), DataType::Boolean);
2397 assert_eq!(infer_field_schema("trUe"), DataType::Boolean);
2398 assert_eq!(infer_field_schema("false"), DataType::Boolean);
2399 assert_eq!(infer_field_schema("2020-11-08"), DataType::Date32);
2400 assert_eq!(
2401 infer_field_schema("2020-11-08T14:20:01"),
2402 DataType::Timestamp(TimeUnit::Second, None)
2403 );
2404 assert_eq!(
2405 infer_field_schema("2020-11-08 14:20:01"),
2406 DataType::Timestamp(TimeUnit::Second, None)
2407 );
2408 assert_eq!(
2409 infer_field_schema("2020-11-08 14:20:01"),
2410 DataType::Timestamp(TimeUnit::Second, None)
2411 );
2412 assert_eq!(infer_field_schema("-5.13"), DataType::Float64);
2413 assert_eq!(infer_field_schema("0.1300"), DataType::Float64);
2414 assert_eq!(
2415 infer_field_schema("2021-12-19 13:12:30.921"),
2416 DataType::Timestamp(TimeUnit::Millisecond, None)
2417 );
2418 assert_eq!(
2419 infer_field_schema("2021-12-19T13:12:30.123456789"),
2420 DataType::Timestamp(TimeUnit::Nanosecond, None)
2421 );
2422 assert_eq!(infer_field_schema("–9223372036854775809"), DataType::Utf8);
2423 assert_eq!(infer_field_schema("9223372036854775808"), DataType::Utf8);
2424 }
2425
2426 #[test]
2427 fn parse_date32() {
2428 assert_eq!(Date32Type::parse("1970-01-01").unwrap(), 0);
2429 assert_eq!(Date32Type::parse("2020-03-15").unwrap(), 18336);
2430 assert_eq!(Date32Type::parse("1945-05-08").unwrap(), -9004);
2431 }
2432
2433 #[test]
2434 fn parse_time() {
2435 assert_eq!(
2436 Time64NanosecondType::parse("12:10:01.123456789 AM"),
2437 Some(601_123_456_789)
2438 );
2439 assert_eq!(
2440 Time64MicrosecondType::parse("12:10:01.123456 am"),
2441 Some(601_123_456)
2442 );
2443 assert_eq!(
2444 Time32MillisecondType::parse("2:10:01.12 PM"),
2445 Some(51_001_120)
2446 );
2447 assert_eq!(Time32SecondType::parse("2:10:01 pm"), Some(51_001));
2448 }
2449
2450 #[test]
2451 fn parse_date64() {
2452 assert_eq!(Date64Type::parse("1970-01-01T00:00:00").unwrap(), 0);
2453 assert_eq!(
2454 Date64Type::parse("2018-11-13T17:11:10").unwrap(),
2455 1542129070000
2456 );
2457 assert_eq!(
2458 Date64Type::parse("2018-11-13T17:11:10.011").unwrap(),
2459 1542129070011
2460 );
2461 assert_eq!(
2462 Date64Type::parse("1900-02-28T12:34:56").unwrap(),
2463 -2203932304000
2464 );
2465 assert_eq!(
2466 Date64Type::parse_formatted("1900-02-28 12:34:56", "%Y-%m-%d %H:%M:%S").unwrap(),
2467 -2203932304000
2468 );
2469 assert_eq!(
2470 Date64Type::parse_formatted("1900-02-28 12:34:56+0030", "%Y-%m-%d %H:%M:%S%z").unwrap(),
2471 -2203932304000 - (30 * 60 * 1000)
2472 );
2473 }
2474
2475 fn test_parse_timestamp_impl<T: ArrowTimestampType>(
2476 timezone: Option<Arc<str>>,
2477 expected: &[i64],
2478 ) {
2479 let csv = [
2480 "1970-01-01T00:00:00",
2481 "1970-01-01T00:00:00Z",
2482 "1970-01-01T00:00:00+02:00",
2483 ]
2484 .join("\n");
2485 let schema = Arc::new(Schema::new(vec![Field::new(
2486 "field",
2487 DataType::Timestamp(T::UNIT, timezone.clone()),
2488 true,
2489 )]));
2490
2491 let mut decoder = ReaderBuilder::new(schema).build_decoder();
2492
2493 let decoded = decoder.decode(csv.as_bytes()).unwrap();
2494 assert_eq!(decoded, csv.len());
2495 decoder.decode(&[]).unwrap();
2496
2497 let batch = decoder.flush().unwrap().unwrap();
2498 assert_eq!(batch.num_columns(), 1);
2499 assert_eq!(batch.num_rows(), 3);
2500 let col = batch.column(0).as_primitive::<T>();
2501 assert_eq!(col.values(), expected);
2502 assert_eq!(col.data_type(), &DataType::Timestamp(T::UNIT, timezone));
2503 }
2504
2505 #[test]
2506 fn test_parse_timestamp() {
2507 test_parse_timestamp_impl::<TimestampNanosecondType>(None, &[0, 0, -7_200_000_000_000]);
2508 test_parse_timestamp_impl::<TimestampNanosecondType>(
2509 Some("+00:00".into()),
2510 &[0, 0, -7_200_000_000_000],
2511 );
2512 test_parse_timestamp_impl::<TimestampNanosecondType>(
2513 Some("-05:00".into()),
2514 &[18_000_000_000_000, 0, -7_200_000_000_000],
2515 );
2516 test_parse_timestamp_impl::<TimestampMicrosecondType>(
2517 Some("-03".into()),
2518 &[10_800_000_000, 0, -7_200_000_000],
2519 );
2520 test_parse_timestamp_impl::<TimestampMillisecondType>(
2521 Some("-03".into()),
2522 &[10_800_000, 0, -7_200_000],
2523 );
2524 test_parse_timestamp_impl::<TimestampSecondType>(Some("-03".into()), &[10_800, 0, -7_200]);
2525 }
2526
2527 #[test]
2528 #[cfg_attr(miri, ignore)] fn test_infer_schema_from_multiple_files() {
2530 let mut csv1 = NamedTempFile::new().unwrap();
2531 let mut csv2 = NamedTempFile::new().unwrap();
2532 let csv3 = NamedTempFile::new().unwrap(); let mut csv4 = NamedTempFile::new().unwrap();
2534 writeln!(csv1, "c1,c2,c3").unwrap();
2535 writeln!(csv1, "1,\"foo\",0.5").unwrap();
2536 writeln!(csv1, "3,\"bar\",1").unwrap();
2537 writeln!(csv1, "3,\"bar\",2e-06").unwrap();
2538 writeln!(csv2, "c1,c2,c3,c4").unwrap();
2540 writeln!(csv2, "10,,3.14,true").unwrap();
2541 writeln!(csv4, "c1,c2,c3").unwrap();
2543 writeln!(csv4, "10,\"foo\",").unwrap();
2544
2545 let schema = infer_schema_from_files(
2546 &[
2547 csv3.path().to_str().unwrap().to_string(),
2548 csv1.path().to_str().unwrap().to_string(),
2549 csv2.path().to_str().unwrap().to_string(),
2550 csv4.path().to_str().unwrap().to_string(),
2551 ],
2552 b',',
2553 Some(4), true,
2555 )
2556 .unwrap();
2557
2558 assert_eq!(schema.fields().len(), 4);
2559 assert!(schema.field(0).is_nullable());
2560 assert!(schema.field(1).is_nullable());
2561 assert!(schema.field(2).is_nullable());
2562 assert!(schema.field(3).is_nullable());
2563
2564 assert_eq!(&DataType::Int64, schema.field(0).data_type());
2565 assert_eq!(&DataType::Utf8, schema.field(1).data_type());
2566 assert_eq!(&DataType::Float64, schema.field(2).data_type());
2567 assert_eq!(&DataType::Boolean, schema.field(3).data_type());
2568 }
2569
2570 #[test]
2571 fn test_bounded() {
2572 let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2573 let data = [
2574 vec!["0"],
2575 vec!["1"],
2576 vec!["2"],
2577 vec!["3"],
2578 vec!["4"],
2579 vec!["5"],
2580 vec!["6"],
2581 ];
2582
2583 let data = data
2584 .iter()
2585 .map(|x| x.join(","))
2586 .collect::<Vec<_>>()
2587 .join("\n");
2588 let data = data.as_bytes();
2589
2590 let reader = std::io::Cursor::new(data);
2591
2592 let mut csv = ReaderBuilder::new(Arc::new(schema))
2593 .with_batch_size(2)
2594 .with_projection(vec![0])
2595 .with_bounds(2, 6)
2596 .build_buffered(reader)
2597 .unwrap();
2598
2599 let batch = csv.next().unwrap().unwrap();
2600 let a = batch.column(0);
2601 let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2602 assert_eq!(a, &UInt32Array::from(vec![2, 3]));
2603
2604 let batch = csv.next().unwrap().unwrap();
2605 let a = batch.column(0);
2606 let a = a.as_any().downcast_ref::<UInt32Array>().unwrap();
2607 assert_eq!(a, &UInt32Array::from(vec![4, 5]));
2608
2609 assert!(csv.next().is_none());
2610 }
2611
2612 #[test]
2613 fn test_empty_projection() {
2614 let schema = Schema::new(vec![Field::new("int", DataType::UInt32, false)]);
2615 let data = [vec!["0"], vec!["1"]];
2616
2617 let data = data
2618 .iter()
2619 .map(|x| x.join(","))
2620 .collect::<Vec<_>>()
2621 .join("\n");
2622
2623 let mut csv = ReaderBuilder::new(Arc::new(schema))
2624 .with_batch_size(2)
2625 .with_projection(vec![])
2626 .build_buffered(Cursor::new(data.as_bytes()))
2627 .unwrap();
2628
2629 let batch = csv.next().unwrap().unwrap();
2630 assert_eq!(batch.columns().len(), 0);
2631 assert_eq!(batch.num_rows(), 2);
2632
2633 assert!(csv.next().is_none());
2634 }
2635
2636 #[test]
2637 fn test_parsing_bool() {
2638 assert_eq!(Some(true), parse_bool("true"));
2640 assert_eq!(Some(true), parse_bool("tRUe"));
2641 assert_eq!(Some(true), parse_bool("True"));
2642 assert_eq!(Some(true), parse_bool("TRUE"));
2643 assert_eq!(None, parse_bool("t"));
2644 assert_eq!(None, parse_bool("T"));
2645 assert_eq!(None, parse_bool(""));
2646
2647 assert_eq!(Some(false), parse_bool("false"));
2648 assert_eq!(Some(false), parse_bool("fALse"));
2649 assert_eq!(Some(false), parse_bool("False"));
2650 assert_eq!(Some(false), parse_bool("FALSE"));
2651 assert_eq!(None, parse_bool("f"));
2652 assert_eq!(None, parse_bool("F"));
2653 assert_eq!(None, parse_bool(""));
2654 }
2655
2656 #[test]
2657 fn test_parsing_float() {
2658 assert_eq!(Some(12.34), Float64Type::parse("12.34"));
2659 assert_eq!(Some(-12.34), Float64Type::parse("-12.34"));
2660 assert_eq!(Some(12.0), Float64Type::parse("12"));
2661 assert_eq!(Some(0.0), Float64Type::parse("0"));
2662 assert_eq!(Some(2.0), Float64Type::parse("2."));
2663 assert_eq!(Some(0.2), Float64Type::parse(".2"));
2664 assert!(Float64Type::parse("nan").unwrap().is_nan());
2665 assert!(Float64Type::parse("NaN").unwrap().is_nan());
2666 assert!(Float64Type::parse("inf").unwrap().is_infinite());
2667 assert!(Float64Type::parse("inf").unwrap().is_sign_positive());
2668 assert!(Float64Type::parse("-inf").unwrap().is_infinite());
2669 assert!(Float64Type::parse("-inf").unwrap().is_sign_negative());
2670 assert_eq!(None, Float64Type::parse(""));
2671 assert_eq!(None, Float64Type::parse("dd"));
2672 assert_eq!(None, Float64Type::parse("12.34.56"));
2673 }
2674
2675 #[test]
2676 fn test_non_std_quote() {
2677 let schema = Schema::new(vec![
2678 Field::new("text1", DataType::Utf8, false),
2679 Field::new("text2", DataType::Utf8, false),
2680 ]);
2681 let builder = ReaderBuilder::new(Arc::new(schema))
2682 .with_header(false)
2683 .with_quote(b'~'); let mut csv_text = Vec::new();
2686 let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2687 for index in 0..10 {
2688 let text1 = format!("id{index:}");
2689 let text2 = format!("value{index:}");
2690 csv_writer
2691 .write_fmt(format_args!("~{text1}~,~{text2}~\r\n"))
2692 .unwrap();
2693 }
2694 let mut csv_reader = std::io::Cursor::new(&csv_text);
2695 let mut reader = builder.build(&mut csv_reader).unwrap();
2696 let batch = reader.next().unwrap().unwrap();
2697 let col0 = batch.column(0);
2698 assert_eq!(col0.len(), 10);
2699 let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2700 assert_eq!(col0_arr.value(0), "id0");
2701 let col1 = batch.column(1);
2702 assert_eq!(col1.len(), 10);
2703 let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2704 assert_eq!(col1_arr.value(5), "value5");
2705 }
2706
2707 #[test]
2708 fn test_non_std_escape() {
2709 let schema = Schema::new(vec![
2710 Field::new("text1", DataType::Utf8, false),
2711 Field::new("text2", DataType::Utf8, false),
2712 ]);
2713 let builder = ReaderBuilder::new(Arc::new(schema))
2714 .with_header(false)
2715 .with_escape(b'\\'); let mut csv_text = Vec::new();
2718 let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2719 for index in 0..10 {
2720 let text1 = format!("id{index:}");
2721 let text2 = format!("value\\\"{index:}");
2722 csv_writer
2723 .write_fmt(format_args!("\"{text1}\",\"{text2}\"\r\n"))
2724 .unwrap();
2725 }
2726 let mut csv_reader = std::io::Cursor::new(&csv_text);
2727 let mut reader = builder.build(&mut csv_reader).unwrap();
2728 let batch = reader.next().unwrap().unwrap();
2729 let col0 = batch.column(0);
2730 assert_eq!(col0.len(), 10);
2731 let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2732 assert_eq!(col0_arr.value(0), "id0");
2733 let col1 = batch.column(1);
2734 assert_eq!(col1.len(), 10);
2735 let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2736 assert_eq!(col1_arr.value(5), "value\"5");
2737 }
2738
2739 #[test]
2740 fn test_non_std_terminator() {
2741 let schema = Schema::new(vec![
2742 Field::new("text1", DataType::Utf8, false),
2743 Field::new("text2", DataType::Utf8, false),
2744 ]);
2745 let builder = ReaderBuilder::new(Arc::new(schema))
2746 .with_header(false)
2747 .with_terminator(b'\n'); let mut csv_text = Vec::new();
2750 let mut csv_writer = std::io::Cursor::new(&mut csv_text);
2751 for index in 0..10 {
2752 let text1 = format!("id{index:}");
2753 let text2 = format!("value{index:}");
2754 csv_writer
2755 .write_fmt(format_args!("\"{text1}\",\"{text2}\"\n"))
2756 .unwrap();
2757 }
2758 let mut csv_reader = std::io::Cursor::new(&csv_text);
2759 let mut reader = builder.build(&mut csv_reader).unwrap();
2760 let batch = reader.next().unwrap().unwrap();
2761 let col0 = batch.column(0);
2762 assert_eq!(col0.len(), 10);
2763 let col0_arr = col0.as_any().downcast_ref::<StringArray>().unwrap();
2764 assert_eq!(col0_arr.value(0), "id0");
2765 let col1 = batch.column(1);
2766 assert_eq!(col1.len(), 10);
2767 let col1_arr = col1.as_any().downcast_ref::<StringArray>().unwrap();
2768 assert_eq!(col1_arr.value(5), "value5");
2769 }
2770
2771 #[test]
2772 fn test_header_bounds() {
2773 let csv = "a,b\na,b\na,b\na,b\na,b\n";
2774 let tests = [
2775 (None, false, 5),
2776 (None, true, 4),
2777 (Some((0, 4)), false, 4),
2778 (Some((1, 4)), false, 3),
2779 (Some((0, 4)), true, 4),
2780 (Some((1, 4)), true, 3),
2781 ];
2782 let schema = Arc::new(Schema::new(vec![
2783 Field::new("a", DataType::Utf8, false),
2784 Field::new("a", DataType::Utf8, false),
2785 ]));
2786
2787 for (idx, (bounds, has_header, expected)) in tests.into_iter().enumerate() {
2788 let mut reader = ReaderBuilder::new(schema.clone()).with_header(has_header);
2789 if let Some((start, end)) = bounds {
2790 reader = reader.with_bounds(start, end);
2791 }
2792 let b = reader
2793 .build_buffered(Cursor::new(csv.as_bytes()))
2794 .unwrap()
2795 .next()
2796 .unwrap()
2797 .unwrap();
2798 assert_eq!(b.num_rows(), expected, "{idx}");
2799 }
2800 }
2801
2802 #[test]
2803 fn test_header_validation() {
2804 let schema = Arc::new(Schema::new(vec![
2805 Field::new("a", DataType::Int32, false),
2806 Field::new("b", DataType::Int32, false),
2807 ]));
2808
2809 let csv = "a,c\n1,2\n";
2810 let err = ReaderBuilder::new(schema.clone())
2811 .with_header(true)
2812 .with_header_validation(true)
2813 .build_buffered(Cursor::new(csv.as_bytes()))
2814 .unwrap()
2815 .next()
2816 .unwrap()
2817 .unwrap_err()
2818 .to_string();
2819 assert_eq!(
2820 err,
2821 "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"c\""
2822 );
2823
2824 let batch = ReaderBuilder::new(schema)
2825 .with_header(true)
2826 .with_header_validation(false)
2827 .build_buffered(Cursor::new(csv.as_bytes()))
2828 .unwrap()
2829 .next()
2830 .unwrap()
2831 .unwrap();
2832 assert_eq!(batch.num_rows(), 1);
2833 }
2834
2835 #[test]
2836 fn test_header_validation_with_buffered_reader() {
2837 let schema = Arc::new(Schema::new(vec![
2838 Field::new("a", DataType::Int32, false),
2839 Field::new("b", DataType::Int32, false),
2840 ]));
2841
2842 let csv = "a,b\n1,2\n";
2843 let buffered = std::io::BufReader::with_capacity(1, Cursor::new(csv.as_bytes()));
2844 let batch = ReaderBuilder::new(schema)
2845 .with_header(true)
2846 .with_header_validation(true)
2847 .build_buffered(buffered)
2848 .unwrap()
2849 .next()
2850 .unwrap()
2851 .unwrap();
2852
2853 assert_eq!(batch.num_rows(), 1);
2854 let a = batch.column(0).as_primitive::<Int32Type>();
2855 assert_eq!(a.value(0), 1);
2856 }
2857
2858 #[test]
2859 fn test_header_validation_with_truncated_rows() {
2860 let schema = Arc::new(Schema::new(vec![
2861 Field::new("a", DataType::Int32, true),
2862 Field::new("b", DataType::Int32, true),
2863 ]));
2864
2865 let csv = "a\n1\n";
2866 let err = ReaderBuilder::new(schema.clone())
2867 .with_header(true)
2868 .with_header_validation(true)
2869 .with_truncated_rows(true)
2870 .build_buffered(Cursor::new(csv.as_bytes()))
2871 .unwrap()
2872 .next()
2873 .unwrap()
2874 .unwrap_err()
2875 .to_string();
2876 assert_eq!(
2877 err,
2878 "Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"\"",
2879 )
2880 }
2881
2882 #[test]
2883 fn test_null_boolean() {
2884 let csv = "true,false\nFalse,True\n,True\nFalse,";
2885 let schema = Arc::new(Schema::new(vec![
2886 Field::new("a", DataType::Boolean, true),
2887 Field::new("a", DataType::Boolean, true),
2888 ]));
2889
2890 let b = ReaderBuilder::new(schema)
2891 .build_buffered(Cursor::new(csv.as_bytes()))
2892 .unwrap()
2893 .next()
2894 .unwrap()
2895 .unwrap();
2896
2897 assert_eq!(b.num_rows(), 4);
2898 assert_eq!(b.num_columns(), 2);
2899
2900 let c = b.column(0).as_boolean();
2901 assert_eq!(c.null_count(), 1);
2902 assert!(c.value(0));
2903 assert!(!c.value(1));
2904 assert!(c.is_null(2));
2905 assert!(!c.value(3));
2906
2907 let c = b.column(1).as_boolean();
2908 assert_eq!(c.null_count(), 1);
2909 assert!(!c.value(0));
2910 assert!(c.value(1));
2911 assert!(c.value(2));
2912 assert!(c.is_null(3));
2913 }
2914
2915 #[test]
2916 fn test_truncated_rows() {
2917 let data = "a,b,c\n1,2,3\n4,5\n\n6,7,8";
2918 let schema = Arc::new(Schema::new(vec![
2919 Field::new("a", DataType::Int32, true),
2920 Field::new("b", DataType::Int32, true),
2921 Field::new("c", DataType::Int32, true),
2922 ]));
2923
2924 let reader = ReaderBuilder::new(schema.clone())
2925 .with_header(true)
2926 .with_truncated_rows(true)
2927 .build(Cursor::new(data))
2928 .unwrap();
2929
2930 let batches = reader.collect::<Result<Vec<_>, _>>();
2931 assert!(batches.is_ok());
2932 let batch = batches.unwrap().into_iter().next().unwrap();
2933 assert_eq!(batch.num_rows(), 3);
2935
2936 let reader = ReaderBuilder::new(schema.clone())
2937 .with_header(true)
2938 .with_truncated_rows(false)
2939 .build(Cursor::new(data))
2940 .unwrap();
2941
2942 let batches = reader.collect::<Result<Vec<_>, _>>();
2943 assert!(match batches {
2944 Err(ArrowError::CsvError(e)) => e.contains("incorrect number of fields"),
2945 _ => false,
2946 });
2947 }
2948
2949 #[test]
2950 fn test_truncated_rows_csv() {
2951 let file = File::open("test/data/truncated_rows.csv").unwrap();
2952 let schema = Arc::new(Schema::new(vec![
2953 Field::new("Name", DataType::Utf8, true),
2954 Field::new("Age", DataType::UInt32, true),
2955 Field::new("Occupation", DataType::Utf8, true),
2956 Field::new("DOB", DataType::Date32, true),
2957 ]));
2958 let reader = ReaderBuilder::new(schema.clone())
2959 .with_header(true)
2960 .with_batch_size(24)
2961 .with_truncated_rows(true);
2962 let csv = reader.build(file).unwrap();
2963 let batches = csv.collect::<Result<Vec<_>, _>>().unwrap();
2964
2965 assert_eq!(batches.len(), 1);
2966 let batch = &batches[0];
2967 assert_eq!(batch.num_rows(), 6);
2968 assert_eq!(batch.num_columns(), 4);
2969 let name = batch
2970 .column(0)
2971 .as_any()
2972 .downcast_ref::<StringArray>()
2973 .unwrap();
2974 let age = batch
2975 .column(1)
2976 .as_any()
2977 .downcast_ref::<UInt32Array>()
2978 .unwrap();
2979 let occupation = batch
2980 .column(2)
2981 .as_any()
2982 .downcast_ref::<StringArray>()
2983 .unwrap();
2984 let dob = batch
2985 .column(3)
2986 .as_any()
2987 .downcast_ref::<Date32Array>()
2988 .unwrap();
2989
2990 assert_eq!(name.value(0), "A1");
2991 assert_eq!(name.value(1), "B2");
2992 assert!(name.is_null(2));
2993 assert_eq!(name.value(3), "C3");
2994 assert_eq!(name.value(4), "D4");
2995 assert_eq!(name.value(5), "E5");
2996
2997 assert_eq!(age.value(0), 34);
2998 assert_eq!(age.value(1), 29);
2999 assert!(age.is_null(2));
3000 assert_eq!(age.value(3), 45);
3001 assert!(age.is_null(4));
3002 assert_eq!(age.value(5), 31);
3003
3004 assert_eq!(occupation.value(0), "Engineer");
3005 assert_eq!(occupation.value(1), "Doctor");
3006 assert!(occupation.is_null(2));
3007 assert_eq!(occupation.value(3), "Artist");
3008 assert!(occupation.is_null(4));
3009 assert!(occupation.is_null(5));
3010
3011 assert_eq!(dob.value(0), 5675);
3012 assert!(dob.is_null(1));
3013 assert!(dob.is_null(2));
3014 assert_eq!(dob.value(3), -1858);
3015 assert!(dob.is_null(4));
3016 assert!(dob.is_null(5));
3017 }
3018
3019 fn truncated_row_count_schema() -> SchemaRef {
3021 Arc::new(Schema::new(vec![
3022 Field::new("name", DataType::Utf8, true),
3023 Field::new("age", DataType::Int32, true),
3024 Field::new("city", DataType::Utf8, true),
3025 ]))
3026 }
3027
3028 #[test]
3029 fn test_truncated_row_count_counts_padded_rows() {
3030 let data = "name,age,city\nAlice,25,Rome\nBob,30\n";
3031
3032 let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3033 .with_header(true)
3034 .with_truncated_rows(true)
3035 .build(Cursor::new(data))
3036 .unwrap();
3037
3038 let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3039 assert_eq!(batches[0].num_rows(), 2);
3040 assert_eq!(reader.truncated_row_count(), 1);
3041 }
3042
3043 #[test]
3044 fn test_truncated_row_count_ignores_empty_trailing_field() {
3045 let data = "name,age,city\nAlice,25,Rome\nCarol,35,\n";
3049
3050 let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3051 .with_header(true)
3052 .with_truncated_rows(true)
3053 .build(Cursor::new(data))
3054 .unwrap();
3055
3056 let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3057 let batch = &batches[0];
3058 assert_eq!(batch.num_rows(), 2);
3059 assert!(batch.column(2).is_null(1));
3060 assert_eq!(reader.truncated_row_count(), 0);
3061 }
3062
3063 #[test]
3064 fn test_truncated_row_count_clean_file() {
3065 let data = "name,age,city\nAlice,25,Rome\nBob,30,Milan\n";
3066
3067 let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3068 .with_header(true)
3069 .with_truncated_rows(true)
3070 .build(Cursor::new(data))
3071 .unwrap();
3072
3073 let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3074 assert_eq!(batches[0].num_rows(), 2);
3075 assert_eq!(reader.truncated_row_count(), 0);
3076 }
3077
3078 #[test]
3079 fn test_truncated_row_count_without_truncated_rows() {
3080 let data = "name,age,city\nAlice,25,Rome\nBob,30\n";
3081
3082 let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3083 .with_header(true)
3084 .with_truncated_rows(false)
3085 .build(Cursor::new(data))
3086 .unwrap();
3087
3088 let err = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap_err();
3090 assert!(
3091 err.to_string().contains("incorrect number of fields"),
3092 "{err}"
3093 );
3094 assert_eq!(reader.truncated_row_count(), 0);
3095 }
3096
3097 #[test]
3098 fn test_truncated_row_count_accumulates_across_batches() {
3099 let data = "name,age,city\nn0,0\nn1,1\nn2,2\nn3,3\nn4,4\nn5,5\n";
3101
3102 let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3103 .with_header(true)
3104 .with_truncated_rows(true)
3105 .with_batch_size(2)
3106 .build(Cursor::new(data))
3107 .unwrap();
3108
3109 let mut running = vec![];
3110 while let Some(batch) = reader.next().transpose().unwrap() {
3111 assert_eq!(batch.num_rows(), 2);
3112 running.push(reader.truncated_row_count());
3113 }
3114
3115 assert_eq!(running, vec![2, 4, 6]);
3117 assert_eq!(reader.truncated_row_count(), 6);
3118 }
3119
3120 #[test]
3121 fn test_truncated_row_count_excludes_skipped_rows() {
3122 let data = "name,age\nAlice,25,Rome\nBob,30,Milan\n";
3125
3126 let mut reader = ReaderBuilder::new(truncated_row_count_schema())
3127 .with_header(true)
3128 .with_truncated_rows(true)
3129 .build(Cursor::new(data))
3130 .unwrap();
3131
3132 let batches = reader.by_ref().collect::<Result<Vec<_>, _>>().unwrap();
3133 assert_eq!(batches[0].num_rows(), 2);
3134 assert_eq!(reader.truncated_row_count(), 0);
3135 }
3136
3137 #[test]
3138 fn test_truncated_row_count_on_decoder() {
3139 let data = "1,2\n3\n";
3140 let schema = Arc::new(Schema::new(vec![
3141 Field::new("a", DataType::Int32, true),
3142 Field::new("b", DataType::Int32, true),
3143 ]));
3144
3145 let mut decoder = ReaderBuilder::new(schema)
3146 .with_truncated_rows(true)
3147 .build_decoder();
3148
3149 assert_eq!(decoder.truncated_row_count(), 0);
3150 let decoded = decoder.decode(data.as_bytes()).unwrap();
3151 assert_eq!(decoded, data.len());
3152 decoder.flush().unwrap().unwrap();
3153 assert_eq!(decoder.truncated_row_count(), 1);
3154 }
3155
3156 #[test]
3157 fn test_truncated_rows_not_nullable_error() {
3158 let data = "a,b,c\n1,2,3\n4,5";
3159 let schema = Arc::new(Schema::new(vec![
3160 Field::new("a", DataType::Int32, false),
3161 Field::new("b", DataType::Int32, false),
3162 Field::new("c", DataType::Int32, false),
3163 ]));
3164
3165 let reader = ReaderBuilder::new(schema.clone())
3166 .with_header(true)
3167 .with_truncated_rows(true)
3168 .build(Cursor::new(data))
3169 .unwrap();
3170
3171 let batches = reader.collect::<Result<Vec<_>, _>>();
3172 assert!(match batches {
3173 Err(ArrowError::InvalidArgumentError(e)) => e.contains("contains null values"),
3174 _ => false,
3175 });
3176 }
3177
3178 #[test]
3179 #[cfg_attr(miri, ignore)] fn test_buffered() {
3181 let tests = [
3182 ("test/data/uk_cities.csv", false, 37),
3183 ("test/data/various_types.csv", true, 10),
3184 ("test/data/decimal_test.csv", false, 10),
3185 ];
3186
3187 for (path, has_header, expected_rows) in tests {
3188 let (schema, _) = Format::default()
3189 .infer_schema(File::open(path).unwrap(), None)
3190 .unwrap();
3191 let schema = Arc::new(schema);
3192
3193 for batch_size in [1, 4] {
3194 for capacity in [1, 3, 7, 100] {
3195 let reader = ReaderBuilder::new(schema.clone())
3196 .with_batch_size(batch_size)
3197 .with_header(has_header)
3198 .build(File::open(path).unwrap())
3199 .unwrap();
3200
3201 let expected = reader.collect::<Result<Vec<_>, _>>().unwrap();
3202
3203 assert_eq!(
3204 expected.iter().map(|x| x.num_rows()).sum::<usize>(),
3205 expected_rows
3206 );
3207
3208 let buffered =
3209 std::io::BufReader::with_capacity(capacity, File::open(path).unwrap());
3210
3211 let reader = ReaderBuilder::new(schema.clone())
3212 .with_batch_size(batch_size)
3213 .with_header(has_header)
3214 .build_buffered(buffered)
3215 .unwrap();
3216
3217 let actual = reader.collect::<Result<Vec<_>, _>>().unwrap();
3218 assert_eq!(expected, actual)
3219 }
3220 }
3221 }
3222 }
3223
3224 fn err_test(csv: &[u8], expected: &str) {
3225 fn err_test_with_schema(csv: &[u8], expected: &str, schema: Arc<Schema>) {
3226 let buffer = std::io::BufReader::with_capacity(2, Cursor::new(csv));
3227 let b = ReaderBuilder::new(schema)
3228 .with_batch_size(2)
3229 .build_buffered(buffer)
3230 .unwrap();
3231 let err = b.collect::<Result<Vec<_>, _>>().unwrap_err().to_string();
3232 assert_eq!(err, expected)
3233 }
3234
3235 let schema_utf8 = Arc::new(Schema::new(vec![
3236 Field::new("text1", DataType::Utf8, true),
3237 Field::new("text2", DataType::Utf8, true),
3238 ]));
3239 err_test_with_schema(csv, expected, schema_utf8);
3240
3241 let schema_utf8view = Arc::new(Schema::new(vec![
3242 Field::new("text1", DataType::Utf8View, true),
3243 Field::new("text2", DataType::Utf8View, true),
3244 ]));
3245 err_test_with_schema(csv, expected, schema_utf8view);
3246 }
3247
3248 #[test]
3249 fn test_invalid_utf8() {
3250 err_test(
3251 b"sdf,dsfg\ndfd,hgh\xFFue\n,sds\nFalhghse,",
3252 "Csv error: Encountered invalid UTF-8 data for line 2 and field 2",
3253 );
3254
3255 err_test(
3256 b"sdf,dsfg\ndksdk,jf\nd\xFFfd,hghue\n,sds\nFalhghse,",
3257 "Csv error: Encountered invalid UTF-8 data for line 3 and field 1",
3258 );
3259
3260 err_test(
3261 b"sdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
3262 "Csv error: Encountered invalid UTF-8 data for line 5 and field 2",
3263 );
3264
3265 err_test(
3266 b"\xFFsdf,dsfg\ndksdk,jf\ndsdsfd,hghue\n,sds\nFalhghse,\xFF",
3267 "Csv error: Encountered invalid UTF-8 data for line 1 and field 1",
3268 );
3269 }
3270
3271 struct InstrumentedRead<R> {
3272 r: R,
3273 fill_count: usize,
3274 fill_sizes: Vec<usize>,
3275 }
3276
3277 impl<R> InstrumentedRead<R> {
3278 fn new(r: R) -> Self {
3279 Self {
3280 r,
3281 fill_count: 0,
3282 fill_sizes: vec![],
3283 }
3284 }
3285 }
3286
3287 impl<R: Seek> Seek for InstrumentedRead<R> {
3288 fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
3289 self.r.seek(pos)
3290 }
3291 }
3292
3293 impl<R: BufRead> Read for InstrumentedRead<R> {
3294 fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
3295 self.r.read(buf)
3296 }
3297 }
3298
3299 impl<R: BufRead> BufRead for InstrumentedRead<R> {
3300 fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
3301 self.fill_count += 1;
3302 let buf = self.r.fill_buf()?;
3303 self.fill_sizes.push(buf.len());
3304 Ok(buf)
3305 }
3306
3307 fn consume(&mut self, amt: usize) {
3308 self.r.consume(amt)
3309 }
3310 }
3311
3312 #[test]
3313 fn test_io() {
3314 let schema = Arc::new(Schema::new(vec![
3315 Field::new("a", DataType::Utf8, false),
3316 Field::new("b", DataType::Utf8, false),
3317 ]));
3318 let csv = "foo,bar\nbaz,foo\na,b\nc,d";
3319 let mut read = InstrumentedRead::new(Cursor::new(csv.as_bytes()));
3320 let reader = ReaderBuilder::new(schema)
3321 .with_batch_size(3)
3322 .build_buffered(&mut read)
3323 .unwrap();
3324
3325 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3326 assert_eq!(batches.len(), 2);
3327 assert_eq!(batches[0].num_rows(), 3);
3328 assert_eq!(batches[1].num_rows(), 1);
3329
3330 assert_eq!(&read.fill_sizes, &[23, 3, 0, 0]);
3336 assert_eq!(read.fill_count, 4);
3337 }
3338
3339 #[test]
3340 #[cfg_attr(miri, ignore)] fn test_inference() {
3342 let cases: &[(&[&str], DataType)] = &[
3343 (&[], DataType::Null),
3344 (&["false", "12"], DataType::Utf8),
3345 (&["12", "cupcakes"], DataType::Utf8),
3346 (&["12", "12.4"], DataType::Float64),
3347 (&["14050", "24332"], DataType::Int64),
3348 (&["14050.0", "true"], DataType::Utf8),
3349 (&["14050", "2020-03-19 00:00:00"], DataType::Utf8),
3350 (&["14050", "2340.0", "2020-03-19 00:00:00"], DataType::Utf8),
3351 (
3352 &["2020-03-19 02:00:00", "2020-03-19 00:00:00"],
3353 DataType::Timestamp(TimeUnit::Second, None),
3354 ),
3355 (&["2020-03-19", "2020-03-20"], DataType::Date32),
3356 (
3357 &["2020-03-19", "2020-03-19 02:00:00", "2020-03-19 00:00:00"],
3358 DataType::Timestamp(TimeUnit::Second, None),
3359 ),
3360 (
3361 &[
3362 "2020-03-19",
3363 "2020-03-19 02:00:00",
3364 "2020-03-19 00:00:00.000",
3365 ],
3366 DataType::Timestamp(TimeUnit::Millisecond, None),
3367 ),
3368 (
3369 &[
3370 "2020-03-19",
3371 "2020-03-19 02:00:00",
3372 "2020-03-19 00:00:00.000000",
3373 ],
3374 DataType::Timestamp(TimeUnit::Microsecond, None),
3375 ),
3376 (
3377 &["2020-03-19 02:00:00+02:00", "2020-03-19 02:00:00Z"],
3378 DataType::Timestamp(TimeUnit::Second, None),
3379 ),
3380 (
3381 &[
3382 "2020-03-19",
3383 "2020-03-19 02:00:00+02:00",
3384 "2020-03-19 02:00:00Z",
3385 "2020-03-19 02:00:00.12Z",
3386 ],
3387 DataType::Timestamp(TimeUnit::Millisecond, None),
3388 ),
3389 (
3390 &[
3391 "2020-03-19",
3392 "2020-03-19 02:00:00.000000000",
3393 "2020-03-19 00:00:00.000000",
3394 ],
3395 DataType::Timestamp(TimeUnit::Nanosecond, None),
3396 ),
3397 ];
3398
3399 for (values, expected) in cases {
3400 let mut t = InferredDataType::default();
3401 for v in *values {
3402 t.update(v)
3403 }
3404 assert_eq!(&t.get(), expected, "{values:?}")
3405 }
3406 }
3407
3408 #[test]
3409 #[cfg_attr(miri, ignore)] fn test_record_length_mismatch() {
3411 let csv = "\
3412 a,b,c\n\
3413 1,2,3\n\
3414 4,5\n\
3415 6,7,8";
3416 let mut read = Cursor::new(csv.as_bytes());
3417 let result = Format::default()
3418 .with_header(true)
3419 .infer_schema(&mut read, None);
3420 assert!(result.is_err());
3421 assert_eq!(
3423 result.err().unwrap().to_string(),
3424 "Csv error: Encountered unequal lengths between records on CSV file. Expected 3 records, found 2 records at line 3"
3425 );
3426 }
3427
3428 #[test]
3429 fn test_comment() {
3430 let schema = Schema::new(vec![
3431 Field::new("a", DataType::Int8, false),
3432 Field::new("b", DataType::Int8, false),
3433 ]);
3434
3435 let csv = "# comment1 \n1,2\n#comment2\n11,22";
3436 let mut read = Cursor::new(csv.as_bytes());
3437 let reader = ReaderBuilder::new(Arc::new(schema))
3438 .with_comment(b'#')
3439 .build(&mut read)
3440 .unwrap();
3441
3442 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3443 assert_eq!(batches.len(), 1);
3444 let b = batches.first().unwrap();
3445 assert_eq!(b.num_columns(), 2);
3446 assert_eq!(
3447 b.column(0)
3448 .as_any()
3449 .downcast_ref::<Int8Array>()
3450 .unwrap()
3451 .values(),
3452 &vec![1, 11]
3453 );
3454 assert_eq!(
3455 b.column(1)
3456 .as_any()
3457 .downcast_ref::<Int8Array>()
3458 .unwrap()
3459 .values(),
3460 &vec![2, 22]
3461 );
3462 }
3463
3464 #[test]
3465 fn test_parse_string_view_single_column() {
3466 let csv = ["foo", "something_cannot_be_inlined", "foobar"].join("\n");
3467 let schema = Arc::new(Schema::new(vec![Field::new(
3468 "c1",
3469 DataType::Utf8View,
3470 true,
3471 )]));
3472
3473 let mut decoder = ReaderBuilder::new(schema).build_decoder();
3474
3475 let decoded = decoder.decode(csv.as_bytes()).unwrap();
3476 assert_eq!(decoded, csv.len());
3477 decoder.decode(&[]).unwrap();
3478
3479 let batch = decoder.flush().unwrap().unwrap();
3480 assert_eq!(batch.num_columns(), 1);
3481 assert_eq!(batch.num_rows(), 3);
3482 let col = batch.column(0).as_string_view();
3483 assert_eq!(col.data_type(), &DataType::Utf8View);
3484 assert_eq!(col.value(0), "foo");
3485 assert_eq!(col.value(1), "something_cannot_be_inlined");
3486 assert_eq!(col.value(2), "foobar");
3487 }
3488
3489 #[test]
3490 fn test_parse_string_view_multi_column() {
3491 let csv = ["foo,", ",something_cannot_be_inlined", "foobarfoobar,bar"].join("\n");
3492 let schema = Arc::new(Schema::new(vec![
3493 Field::new("c1", DataType::Utf8View, true),
3494 Field::new("c2", DataType::Utf8View, true),
3495 ]));
3496
3497 let mut decoder = ReaderBuilder::new(schema).build_decoder();
3498
3499 let decoded = decoder.decode(csv.as_bytes()).unwrap();
3500 assert_eq!(decoded, csv.len());
3501 decoder.decode(&[]).unwrap();
3502
3503 let batch = decoder.flush().unwrap().unwrap();
3504 assert_eq!(batch.num_columns(), 2);
3505 assert_eq!(batch.num_rows(), 3);
3506 let c1 = batch.column(0).as_string_view();
3507 let c2 = batch.column(1).as_string_view();
3508 assert_eq!(c1.data_type(), &DataType::Utf8View);
3509 assert_eq!(c2.data_type(), &DataType::Utf8View);
3510
3511 assert!(!c1.is_null(0));
3512 assert!(c1.is_null(1));
3513 assert!(!c1.is_null(2));
3514 assert_eq!(c1.value(0), "foo");
3515 assert_eq!(c1.value(2), "foobarfoobar");
3516
3517 assert!(c2.is_null(0));
3518 assert!(!c2.is_null(1));
3519 assert!(!c2.is_null(2));
3520 assert_eq!(c2.value(1), "something_cannot_be_inlined");
3521 assert_eq!(c2.value(2), "bar");
3522 }
3523
3524 #[test]
3525 #[cfg_attr(miri, ignore)] fn test_float_precision() {
3527 let data = [
3528 "f16,f32,f64",
3529 "1.5,1.5,1.5",
3530 "0.25,0.25,0.25",
3531 "1.23456789,1.23456789,1.23456789",
3532 "1.234567890123456,1.234567890123456,1.234567890123456",
3533 "-2.5,-2.5,-2.5",
3534 "0,0,0",
3535 ",,",
3536 ]
3537 .join("\n");
3538
3539 let schema = Schema::new(vec![
3540 Field::new("f16", DataType::Float16, true),
3541 Field::new("f32", DataType::Float32, true),
3542 Field::new("f64", DataType::Float64, true),
3543 ]);
3544
3545 let mut reader = ReaderBuilder::new(Arc::new(schema))
3546 .with_header(true)
3547 .build(Cursor::new(data))
3548 .unwrap();
3549
3550 let batch = reader.next().unwrap().unwrap();
3551 assert_eq!(batch.num_rows(), 7);
3552
3553 let f16_col = batch.column(0).as_primitive::<Float16Type>();
3554 let f32_col = batch.column(1).as_primitive::<Float32Type>();
3555 let f64_col = batch.column(2).as_primitive::<Float64Type>();
3556
3557 assert_eq!(f16_col.value(0), half::f16::from_f32(1.5));
3558 assert_eq!(f32_col.value(0), 1.5f32);
3559 assert_eq!(f64_col.value(0), 1.5f64);
3560
3561 assert_eq!(f16_col.value(1), half::f16::from_f32(0.25));
3562 assert_eq!(f32_col.value(1), 0.25f32);
3563 assert_eq!(f64_col.value(1), 0.25f64);
3564
3565 assert_eq!(f16_col.value(2), half::f16::from_f32(1.234_567_9));
3566 assert_eq!(f32_col.value(2), 1.234_567_9_f32);
3567 assert_eq!(f64_col.value(2), 1.23456789f64);
3568
3569 assert_eq!(f16_col.value(3), half::f16::from_f64(1.234567890123456f64));
3570 assert_eq!(f32_col.value(3), 1.234_567_9_f32);
3571 assert_eq!(f64_col.value(3), 1.234567890123456f64);
3572
3573 assert_eq!(f16_col.value(4), half::f16::from_f32(-2.5));
3574 assert_eq!(f32_col.value(4), -2.5f32);
3575 assert_eq!(f64_col.value(4), -2.5f64);
3576
3577 assert_eq!(f16_col.value(5), half::f16::from_f32(0.0));
3578 assert_eq!(f32_col.value(5), 0.0f32);
3579 assert_eq!(f64_col.value(5), 0.0f64);
3580
3581 assert!(f16_col.is_null(6));
3582 assert!(f32_col.is_null(6));
3583 assert!(f64_col.is_null(6));
3584 }
3585}