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