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