diff --git a/arrow-csv/src/lib.rs b/arrow-csv/src/lib.rs index bee6c56670a4..0a090ae0c62e 100644 --- a/arrow-csv/src/lib.rs +++ b/arrow-csv/src/lib.rs @@ -30,8 +30,8 @@ pub mod reader; pub mod writer; pub use self::reader::{ - CsvFilterRecord, CsvInvalidUtf8Error, CsvRecord, CsvRecordError, CsvRecordErrorHandler, Reader, - ReaderBuilder, infer_schema_from_files, + CsvBlankRecordPolicy, CsvFilterRecord, CsvInvalidUtf8Error, CsvRecord, CsvRecordError, + CsvRecordErrorHandler, Reader, ReaderBuilder, infer_schema_from_files, }; pub use self::writer::QuoteStyle; pub use self::writer::Writer; diff --git a/arrow-csv/src/reader/mod.rs b/arrow-csv/src/reader/mod.rs index a7cd71388c23..302e4036800d 100644 --- a/arrow-csv/src/reader/mod.rs +++ b/arrow-csv/src/reader/mod.rs @@ -177,7 +177,7 @@ use std::collections::HashMap; use std::fmt::{self, Debug}; use std::fs::File; use std::io::{self, BufRead, BufReader as StdBufReader, Cursor, Read}; -use std::sync::{Arc, LazyLock}; +use std::sync::{Arc, LazyLock, atomic::AtomicU64}; use crate::map_csv_error; use crate::reader::records::{RecordDecoder, StringRecords}; @@ -190,7 +190,8 @@ pub struct CsvRecordError<'a> { pub line_number: usize, /// One-based physical line containing the start of the malformed record. pub source_line_number: usize, - /// Zero-based byte offset of the start of the record in the input stream. + /// Zero-based byte offset in the supplied input, or in the original source + /// when [`ReaderBuilder::with_initial_source_position`] is set. pub byte_offset: usize, /// Number of fields required by the reader schema. pub expected_fields: usize, @@ -205,7 +206,8 @@ pub struct CsvRecordError<'a> { pub struct CsvRecord<'a> { /// One-based record number, including any header row. pub line_number: usize, - /// Zero-based byte offset of the start of the record in the input stream. + /// Zero-based byte offset in the supplied input, or in the original source + /// when [`ReaderBuilder::with_initial_source_position`] is set. pub byte_offset: usize, /// Original record bytes, including its record terminator when present. pub record: &'a [u8], @@ -218,7 +220,8 @@ pub struct CsvFilterRecord<'a> { pub line_number: usize, /// One-based physical line containing the start of the record. pub source_line_number: usize, - /// Zero-based byte offset of the start of the record in the input stream. + /// Zero-based byte offset in the supplied input, or in the original source + /// when [`ReaderBuilder::with_initial_source_position`] is set. pub byte_offset: usize, /// Original record bytes, including its record terminator when present. pub record: &'a [u8], @@ -231,7 +234,8 @@ pub struct CsvInvalidUtf8Error<'a> { pub line_number: usize, /// One-based physical line containing the start of the malformed record. pub source_line_number: usize, - /// Zero-based byte offset of the start of the record in the input stream. + /// Zero-based byte offset in the supplied input, or in the original source + /// when [`ReaderBuilder::with_initial_source_position`] is set. pub byte_offset: usize, /// Zero-based index of the decoded field containing invalid UTF-8. pub field_index: usize, @@ -263,6 +267,22 @@ pub trait CsvRecordErrorHandler: Debug + Send + Sync { Ok(()) } + /// Whether successful-record callbacks need the physical source line. + fn needs_source_line_number(&self) -> bool { + false + } + + /// Observe a decoded record with its one-based physical source line. + /// + /// The default keeps existing handlers compatible with this callback. + fn handle_record_with_source_line( + &self, + record: &CsvRecord<'_>, + _source_line_number: usize, + ) -> Result<(), ArrowError> { + self.handle_record(record) + } + /// Opt in to filtering decoded records before they enter a batch. fn filters_records(&self) -> bool { false @@ -283,9 +303,8 @@ pub trait CsvRecordErrorHandler: Debug + Send + Sync { false } - /// Opt in to treating physical blank lines as empty one-column records. - /// This is supported only for a single-column schema. The default reader - /// still skips blank lines. + /// Opt in to treating physical blank lines as records with empty fields. + /// This requires a nonempty schema. The default reader still skips blank lines. fn blank_records_are_rows(&self) -> bool { false } @@ -301,6 +320,22 @@ pub trait CsvRecordErrorHandler: Debug + Send + Sync { } } +/// Opt-in handling of physical blank CSV records in the strict decoder. +/// +/// The default is [`Self::Ignore`], preserving the ordinary CSV scan path. +#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)] +pub enum CsvBlankRecordPolicy { + /// Preserve the CSV parser's default behavior of ignoring blank records. + #[default] + Ignore, + /// Fail when a physical blank record is encountered. + Reject, + /// Emit one SQL NULL row for a blank record when all source columns are nullable. + EmitNullRow, + /// Omit blank data records, counting each one in the optional counter. + Skip, +} + /// Order should match [`InferredDataType`] static REGEX_SET: LazyLock = LazyLock::new(|| { RegexSet::new([ @@ -1219,6 +1254,9 @@ impl Decoder { validate_header(&rows, self.schema.fields())?; self.header_validation = false; self.to_skip -= 1; + if bytes == 0 && !buf.is_empty() { + return self.decode(buf); + } return Ok(bytes); } @@ -1227,6 +1265,9 @@ impl Decoder { let (skipped, bytes) = self.record_decoder.decode_skipping(buf, to_skip)?; self.to_skip -= skipped; self.record_decoder.clear(); + if bytes == 0 && skipped > 0 && !buf.is_empty() { + return self.decode(buf); + } return Ok(bytes); } @@ -1754,8 +1795,12 @@ pub struct ReaderBuilder { projection: Option>, /// Optional handler for records whose field count differs from the schema. record_error_handler: Option>, + blank_record_policy: CsvBlankRecordPolicy, + skipped_blank_records_counter: Option>, ignore_extra_fields: bool, numeric_boolean_values: bool, + source_line_offset: usize, + source_byte_offset: usize, } impl ReaderBuilder { @@ -1789,8 +1834,12 @@ impl ReaderBuilder { bounds: None, projection: None, record_error_handler: None, + blank_record_policy: CsvBlankRecordPolicy::Ignore, + skipped_blank_records_counter: None, ignore_extra_fields: false, numeric_boolean_values: false, + source_line_offset: 0, + source_byte_offset: 0, } } @@ -1800,6 +1849,15 @@ impl ReaderBuilder { self } + /// Set the physical lines and bytes consumed before the input passed to + /// this reader. Record diagnostics retain offsets into the original source; + /// record ordinals still begin at one for the supplied input. + pub fn with_initial_source_position(mut self, lines: usize, bytes: usize) -> Self { + self.source_line_offset = lines; + self.source_byte_offset = bytes; + self + } + /// Set whether to validate the CSV header against the schema /// /// This option only applies when [`Self::with_header`] is set to `true`, and defaults to `false` @@ -1916,6 +1974,28 @@ impl ReaderBuilder { self } + /// Configure strict, no-handler handling of physical blank CSV records. + /// + /// This option does not change the recovery path selected by + /// [`Self::with_record_error_handler`]. `EmitNullRow` requires a nonempty schema + /// with all columns nullable and emits missing fields, independent of configured + /// null values. + pub fn with_blank_record_policy(mut self, policy: CsvBlankRecordPolicy) -> Self { + self.blank_record_policy = policy; + self + } + + /// Count records omitted by [`CsvBlankRecordPolicy::Skip`]. + /// + /// The counter is incremented only for blank records in the data region, + /// not for blanks encountered while skipping a header or bounded prefix. + /// Input removed before this reader (for example, a physical header skip) + /// is likewise not counted. + pub fn with_skipped_blank_records_counter(mut self, counter: Arc) -> Self { + self.skipped_blank_records_counter = Some(counter); + self + } + /// Replace malformed UTF-8 in decoded CSV fields with U+FFFD. /// Disabled by default; malformed batches may allocate replacement data. pub fn with_replace_invalid_utf8(mut self, replace: bool) -> Self { @@ -1949,8 +2029,15 @@ impl ReaderBuilder { ) .with_ignore_extra_fields(self.ignore_extra_fields) .with_record_error_handler(self.record_error_handler) + .with_blank_record_policy( + self.blank_record_policy, + !self.schema.fields().is_empty() + && self.schema.fields().iter().all(|field| field.is_nullable()), + self.skipped_blank_records_counter, + ) .with_replace_invalid_utf8(self.format.replace_invalid_utf8) .with_ignored_prefix_options(self.format.comment, self.format.terminator) + .with_initial_source_position(self.source_line_offset, self.source_byte_offset) .with_utf8_preserving_syntax( self.format.delimiter.is_none_or(|byte| byte.is_ascii()) && self.format.quote.is_none_or(|byte| byte.is_ascii()) @@ -1994,6 +2081,85 @@ impl ReaderBuilder { #[cfg(test)] mod tests { use super::*; + use std::sync::Mutex; + + #[derive(Debug, Default)] + struct CaptureSourcePositions(Mutex>); + + impl CsvRecordErrorHandler for CaptureSourcePositions { + fn handle(&self, error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + self.0.lock().unwrap().push(( + error.line_number, + error.source_line_number, + error.byte_offset, + )); + Ok(()) + } + } + + #[test] + fn test_reader_initial_source_position_preserves_original_diagnostics() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("value", DataType::Utf8, false), + ])); + let positions = Arc::new(CaptureSourcePositions::default()); + let reader = ReaderBuilder::new(schema) + .with_header(false) + .with_initial_source_position(1, 8) + .with_record_error_handler(positions.clone()) + .build("1,bad,extra\n2,ok\n".as_bytes()) + .unwrap(); + let batches = reader.collect::, _>>().unwrap(); + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 1 + ); + assert_eq!(*positions.0.lock().unwrap(), [(1, 2, 8)]); + + let positions = Arc::new(CaptureSourcePositions::default()); + let reader = ReaderBuilder::new(Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("value", DataType::Utf8, false), + ]))) + .with_initial_source_position(2, 10) + .with_record_error_handler(positions.clone()) + .build("1,bad,extra\n".as_bytes()) + .unwrap(); + reader.collect::, _>>().unwrap(); + assert_eq!(*positions.0.lock().unwrap(), [(1, 3, 10)]); + } + + #[test] + fn test_positioned_reader_handles_bom_and_large_line_offset() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("value", DataType::Utf8, false), + ])); + let prefix = b"\xef\xbb\xbf# comment\n1,ok\n"; + let positions = Arc::new(CaptureSourcePositions::default()); + let reader = ReaderBuilder::new(schema.clone()) + .with_comment(b'#') + .with_initial_source_position(1, 8) + .with_record_error_handler(positions.clone()) + .build(b"\xef\xbb\xbf# comment\n1,ok\n2,extra,value\n".as_slice()) + .unwrap(); + reader.collect::, _>>().unwrap(); + let positions = positions.0.lock().unwrap(); + assert_eq!(positions.len(), 1); + assert_eq!(positions[0].1, 4); + assert_eq!(positions[0].2, 8 + prefix.len()); + drop(positions); + + let positions = Arc::new(CaptureSourcePositions::default()); + let reader = ReaderBuilder::new(schema) + .with_initial_source_position(usize::MAX, 0) + .with_record_error_handler(positions.clone()) + .build("1,bad,extra\n".as_bytes()) + .unwrap(); + reader.collect::, _>>().unwrap(); + assert_eq!(*positions.0.lock().unwrap(), [(1, usize::MAX, 0)]); + } #[derive(Debug)] struct SkipMalformed; @@ -2008,6 +2174,641 @@ mod tests { } } + #[test] + fn test_multicolumn_blank_record_materializes_null_fields() { + #[derive(Debug)] + struct BlankRows; + + impl CsvRecordErrorHandler for BlankRows { + fn handle(&self, _error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + Ok(()) + } + + fn blank_records_are_rows(&self) -> bool { + true + } + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + let batch = ReaderBuilder::new(schema) + .with_record_error_handler(Arc::new(BlankRows)) + .build(Cursor::new(b"1,a\n\n2,b\n")) + .unwrap() + .next() + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 3); + assert!(batch.column(0).is_null(1)); + assert!(batch.column(1).is_null(1)); + } + + #[test] + fn test_strict_blank_rejects_cr_lf_crlf_across_chunks() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + for input in [ + b"1,a\n\n2,b\n".as_slice(), + b"1,a\r\r2,b\r", + b"1,a\r\n\r\n2,b\r\n", + ] { + for capacity in 1..=4 { + let buffered = StdBufReader::with_capacity(capacity, Cursor::new(input)); + let error = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(buffered) + .unwrap() + .collect::, _>>() + .unwrap_err(); + assert!( + error.to_string().contains("blank CSV record at line 2"), + "{input:?}, capacity={capacity}: {error}" + ); + } + } + } + + #[test] + fn test_strict_blank_ignores_quoted_newlines() { + let schema = Arc::new(Schema::new(vec![ + Field::new("first", DataType::Utf8, true), + Field::new("second", DataType::Utf8, true), + ])); + for capacity in 1..=4 { + let buffered = + StdBufReader::with_capacity(capacity, Cursor::new(b"\"a\n\nb\",x\r\nnext,y\r\n")); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(buffered) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_eq!(batches[0].column(0).as_string::().value(0), "a\n\nb"); + } + } + + #[test] + fn test_strict_blank_emits_missing_field_independent_of_null_values() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for separator in ["\n", "\r", "\r\n"] { + let input = format!("alpha{separator}{separator}\"\"{separator}"); + for capacity in 1..=4 { + let buffered = StdBufReader::with_capacity(capacity, Cursor::new(input.as_bytes())); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::EmitNullRow) + .with_null_values(vec!["NULL".to_owned()]) + .with_preserve_quoted_empty(true) + .build_buffered(buffered) + .unwrap() + .collect::, _>>() + .unwrap(); + let rows: usize = batches.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(rows, 3, "{separator:?}, capacity={capacity}"); + let values = batches[0].column(0).as_string::(); + assert_eq!(values.value(0), "alpha"); + assert!(values.is_null(1)); + assert!(!values.is_null(2)); + assert_eq!(values.value(2), ""); + } + } + } + + #[test] + fn test_strict_blank_skip_counts_only_physical_blanks() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + for capacity in 1..=4 { + let counter = Arc::new(AtomicU64::new(0)); + let buffered = + StdBufReader::with_capacity(capacity, Cursor::new(b"1,a\n\n\r\n2,b\r\n")); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .with_skipped_blank_records_counter(Arc::clone(&counter)) + .build_buffered(buffered) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_eq!(counter.load(std::sync::atomic::Ordering::Relaxed), 2); + } + } + + #[test] + fn test_strict_blank_emit_null_requires_nullable_columns() { + for schema in [ + Schema::new(vec![Field::new("value", DataType::Utf8, false)]), + Schema::new(vec![ + Field::new("first", DataType::Utf8, true), + Field::new("second", DataType::Utf8, false), + ]), + ] { + let error = ReaderBuilder::new(Arc::new(schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::EmitNullRow) + .build(Cursor::new(b"\n")) + .unwrap() + .collect::, _>>() + .unwrap_err(); + assert!( + error + .to_string() + .contains("require all columns to be nullable") + ); + } + } + + #[test] + fn test_strict_blank_emits_all_null_source_fields() { + let schema = Arc::new(Schema::new(vec![ + Field::new("first", DataType::Int64, true), + Field::new("second", DataType::Utf8, true), + ])); + for separator in ["\n", "\r", "\r\n"] { + let input = format!("{separator}1,alpha{separator}2,\"\"{separator}"); + for capacity in 1..=4 { + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::EmitNullRow) + .with_null_values(vec!["sentinel".to_owned()]) + .with_preserve_quoted_empty(true) + .with_batch_size(1) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(input.as_bytes()), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 3, "{separator:?}, capacity={capacity}"); + assert!(batches[0].column(0).is_null(0)); + assert!(batches[0].column(1).is_null(0)); + assert_eq!(batches[1].column(0).as_primitive::().value(0), 1); + assert_eq!(batches[1].column(1).as_string::().value(0), "alpha"); + assert!(!batches[2].column(1).is_null(0)); + assert_eq!(batches[2].column(1).as_string::().value(0), ""); + } + } + } + + #[test] + fn test_strict_blank_skip_with_bom_comments_and_header() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + let counter = Arc::new(AtomicU64::new(0)); + let batches = ReaderBuilder::new(schema) + .with_header(true) + .with_comment(b'#') + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .with_skipped_blank_records_counter(Arc::clone(&counter)) + .build(Cursor::new( + b"\xef\xbb\xbfid,value\n#ignored\n\n1,a\n\n2,b\n", + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + assert_eq!(counter.load(std::sync::atomic::Ordering::Relaxed), 2); + } + + #[test] + fn test_strict_blank_validated_crlf_header_across_chunks() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + for capacity in 1..=4 { + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_header(true) + .with_header_validation(true) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"id,value\r\n1,a\r\n"), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + } + } + + #[test] + fn test_strict_blank_split_bom_validated_header_and_blank() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + for capacity in 1..=4 { + let error = ReaderBuilder::new(Arc::clone(&schema)) + .with_header(true) + .with_header_validation(true) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"\xef\xbb\xbfid,value\r\n\r\n1,a\r\n"), + )) + .unwrap() + .collect::, _>>() + .unwrap_err(); + assert!( + error.to_string().contains("blank CSV record at line 2"), + "capacity={capacity}: {error}" + ); + } + } + + #[test] + fn test_strict_blank_split_bom_preserves_second_bom() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for capacity in 1..=4 { + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"\xef\xbb\xbf\xef\xbb\xbfvalue\n"), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!( + batches[0].column(0).as_string::().value(0), + "\u{feff}value" + ); + } + } + + #[test] + fn test_strict_blank_split_bom_before_blank_and_at_eof() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for capacity in 1..=4 { + for ending in ["\n", "\r", "\r\n"] { + let input = format!("\u{feff}{ending}value{ending}"); + let counter = Arc::new(AtomicU64::new(0)); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .with_skipped_blank_records_counter(Arc::clone(&counter)) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(input.as_bytes()), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(batches[0].column(0).as_string::().value(0), "value"); + assert_eq!(counter.load(std::sync::atomic::Ordering::Relaxed), 1); + } + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"\xef\xbb\xbf"), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert!(batches.is_empty()); + } + } + + #[test] + fn test_strict_blank_split_bom_before_comment() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for capacity in 1..=4 { + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_comment(b'#') + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"\xef\xbb\xbf#comment\nvalue\n"), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(batches[0].column(0).as_string::().value(0), "value"); + } + } + + #[test] + fn test_strict_blank_fast_path_split_crlf_and_quoted_newline() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for capacity in 1..=4 { + let counter = Arc::new(AtomicU64::new(0)); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_batch_size(1) + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .with_skipped_blank_records_counter(Arc::clone(&counter)) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"\"one\r\ntwo\"\r\nthree\r\n\r\nfour\r\n"), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + let values: Vec<_> = batches + .iter() + .map(|batch| batch.column(0).as_string::().value(0)) + .collect(); + assert_eq!(values, ["one\r\ntwo", "three", "four"]); + assert_eq!(counter.load(std::sync::atomic::Ordering::Relaxed), 1); + } + } + + #[test] + fn test_strict_blank_fast_path_preserves_source_line() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for input in [ + b"\"a\r\nb\"\r\nvalue\r\n\r\n".as_slice(), + b"\"a\rb\"\r\nvalue\r\n\r\n", + ] { + for capacity in [1, 2, 4, 64] { + let error = ReaderBuilder::new(Arc::clone(&schema)) + .with_batch_size(1) + .with_blank_record_policy(CsvBlankRecordPolicy::Reject) + .build_buffered(StdBufReader::with_capacity(capacity, Cursor::new(input))) + .unwrap() + .collect::, _>>() + .unwrap_err(); + assert!( + error.to_string().contains("source line 4"), + "input={input:?}, capacity={capacity}: {error}" + ); + } + } + } + + #[test] + fn test_strict_blank_skip_counter_excludes_skipped_header_and_bounds() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for capacity in 1..=4 { + let counter = Arc::new(AtomicU64::new(0)); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_header(true) + .with_bounds(1, 2) + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .with_skipped_blank_records_counter(Arc::clone(&counter)) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(b"\nheader\n\nignored\n\nvalue\n"), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(batches[0].column(0).as_string::().value(0), "value"); + assert_eq!(counter.load(std::sync::atomic::Ordering::Relaxed), 1); + } + } + + #[test] + fn test_strict_blank_skip_counter_excludes_pre_header_blanks_with_validation() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for validate_header in [false, true] { + for separator in ["\n", "\r", "\r\n"] { + let input = + format!("{separator}{separator}value{separator}{separator}alpha{separator}"); + for capacity in 1..=4 { + let counter = Arc::new(AtomicU64::new(0)); + let batches = ReaderBuilder::new(Arc::clone(&schema)) + .with_header(true) + .with_header_validation(validate_header) + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .with_skipped_blank_records_counter(Arc::clone(&counter)) + .build_buffered(StdBufReader::with_capacity( + capacity, + Cursor::new(input.as_bytes()), + )) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + assert_eq!(batches[0].column(0).as_string::().value(0), "alpha"); + assert_eq!( + counter.load(std::sync::atomic::Ordering::Relaxed), + 1, + "validate_header={validate_header}, separator={separator:?}, capacity={capacity}" + ); + } + } + } + } + + #[test] + fn test_strict_blank_batch_boundary_and_trailing_cr() { + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + let batches = ReaderBuilder::new(schema) + .with_batch_size(1) + .with_blank_record_policy(CsvBlankRecordPolicy::EmitNullRow) + .build_buffered(StdBufReader::with_capacity(1, Cursor::new(b"a\r\n\r\nb\r"))) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 3); + assert!(!batches[0].column(0).is_null(0)); + assert!(batches[1].column(0).is_null(0)); + assert!(!batches[2].column(0).is_null(0)); + } + + #[test] + fn test_strict_blank_skip_preserves_following_field_count_line() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + let error = ReaderBuilder::new(schema) + .with_blank_record_policy(CsvBlankRecordPolicy::Skip) + .build(Cursor::new(b"1,a\n\nbad\n")) + .unwrap() + .collect::, _>>() + .unwrap_err(); + assert!(error.to_string().contains("line 3"), "{error}"); + } + + #[test] + fn test_default_strict_reader_still_ignores_blank_lines() { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("value", DataType::Utf8, true), + ])); + let batches = ReaderBuilder::new(schema) + .build(Cursor::new(b"1,a\n\n2,b\n")) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); + } + + #[test] + #[ignore = "local strict CSV blank-policy throughput benchmark"] + fn bench_strict_blank_record_policies() { + const ROWS: usize = 200_000; + const BLANK_INTERVAL: usize = 10_000; + const REPETITIONS: usize = 5; + + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + for ending in ["\n", "\r\n"] { + let row = format!("some-long-value{ending}"); + for sparse in [false, true] { + let mut input = Vec::with_capacity( + ROWS * row.len() + usize::from(sparse) * (ROWS / BLANK_INTERVAL) * ending.len(), + ); + for i in 0..ROWS { + input.extend_from_slice(row.as_bytes()); + if sparse && (i + 1) % BLANK_INTERVAL == 0 { + input.extend_from_slice(ending.as_bytes()); + } + } + let blanks = if sparse { ROWS / BLANK_INTERVAL } else { 0 }; + for policy in [ + CsvBlankRecordPolicy::Ignore, + CsvBlankRecordPolicy::Reject, + CsvBlankRecordPolicy::EmitNullRow, + CsvBlankRecordPolicy::Skip, + ] { + let mut times = Vec::with_capacity(REPETITIONS); + for _ in 0..REPETITIONS { + let counter = Arc::new(AtomicU64::new(0)); + let mut builder = ReaderBuilder::new(Arc::clone(&schema)) + .with_batch_size(8192) + .with_null_values(vec!["NULL".to_owned()]); + if policy != CsvBlankRecordPolicy::Ignore { + builder = builder.with_blank_record_policy(policy); + } + if policy == CsvBlankRecordPolicy::Skip { + builder = + builder.with_skipped_blank_records_counter(Arc::clone(&counter)); + } + let started = std::time::Instant::now(); + let result = builder + .build(Cursor::new(input.as_slice())) + .unwrap() + .collect::, _>>(); + times.push(started.elapsed()); + if policy == CsvBlankRecordPolicy::Reject && sparse { + assert!( + result + .unwrap_err() + .to_string() + .contains("blank CSV record at line") + ); + } else { + let loaded: usize = + result.unwrap().iter().map(RecordBatch::num_rows).sum(); + assert_eq!( + loaded, + ROWS + usize::from(policy == CsvBlankRecordPolicy::EmitNullRow) + * blanks + ); + } + if policy == CsvBlankRecordPolicy::Skip { + assert_eq!( + counter.load(std::sync::atomic::Ordering::Relaxed), + blanks as u64 + ); + } + } + times.sort_unstable(); + let median = times[REPETITIONS / 2].as_secs_f64(); + let (bytes, rows) = if policy == CsvBlankRecordPolicy::Reject && sparse { + (BLANK_INTERVAL * row.len() + ending.len(), BLANK_INTERVAL) + } else { + (input.len(), ROWS) + }; + println!( + "strict blank policy={policy:?} ending={ending:?} sparse={sparse}: {:.1} MiB/s, {:.2} Mrows/s (median {median:.3}s)", + bytes as f64 / median / (1024.0 * 1024.0), + rows as f64 / median / 1_000_000.0, + ); + } + } + } + } + + #[test] + fn test_split_cr_blank_header_does_not_hide_first_data_row() { + #[derive(Debug)] + struct BlankRows; + + impl CsvRecordErrorHandler for BlankRows { + fn handle(&self, _error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + Ok(()) + } + + fn blank_records_are_rows(&self) -> bool { + true + } + } + + for validate in [false, true] { + let field_name = if validate { "" } else { "value" }; + let schema = Arc::new(Schema::new(vec![Field::new( + field_name, + DataType::Utf8, + true, + )])); + let input = std::io::BufReader::with_capacity(1, Cursor::new(b"\rX\n")); + let batches = ReaderBuilder::new(schema) + .with_header(true) + .with_header_validation(validate) + .with_record_error_handler(Arc::new(BlankRows)) + .build_buffered(input) + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.len(), 1, "validate={validate}"); + assert_eq!(batches[0].num_rows(), 1); + assert_eq!(batches[0].column(0).as_string::().value(0), "X"); + } + } + + #[test] + fn test_multicolumn_blank_record_preserves_quoted_empty_across_crlf_chunks() { + #[derive(Debug)] + struct BlankRows; + + impl CsvRecordErrorHandler for BlankRows { + fn handle(&self, _error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + Ok(()) + } + + fn blank_records_are_rows(&self) -> bool { + true + } + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("first", DataType::Utf8, true), + Field::new("second", DataType::Utf8, true), + ])); + let input = std::io::BufReader::with_capacity(1, Cursor::new(b"\"\",x\r\n\r\nz,\"\"\r\n")); + let batch = ReaderBuilder::new(schema) + .with_preserve_quoted_empty(true) + .with_record_error_handler(Arc::new(BlankRows)) + .build_buffered(input) + .unwrap() + .next() + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 3); + let first = batch.column(0).as_string::(); + let second = batch.column(1).as_string::(); + assert!(!first.is_null(0)); + assert_eq!(first.value(0), ""); + assert!(first.is_null(1)); + assert!(second.is_null(1)); + assert!(!second.is_null(2)); + assert_eq!(second.value(2), ""); + } + #[test] fn test_ignore_extra_fields_keeps_strict_reader_default() { let schema = Arc::new(Schema::new(vec![Field::new("first", DataType::Utf8, true)])); diff --git a/arrow-csv/src/reader/records.rs b/arrow-csv/src/reader/records.rs index 4520e37a3f85..610f9c225934 100644 --- a/arrow-csv/src/reader/records.rs +++ b/arrow-csv/src/reader/records.rs @@ -17,11 +17,14 @@ use arrow_schema::ArrowError; use csv_core::{ReadRecordResult, Reader}; -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicU64, Ordering}, +}; use super::{ - CsvFilterRecord, CsvInvalidUtf8Error, CsvRecord, CsvRecordError, CsvRecordErrorHandler, - NullRegex, + CsvBlankRecordPolicy, CsvFilterRecord, CsvInvalidUtf8Error, CsvRecord, CsvRecordError, + CsvRecordErrorHandler, NullRegex, }; #[derive(Debug)] @@ -43,6 +46,63 @@ fn is_comment_start(byte: u8, comment: Option, terminator: Option) -> bo Some(byte) == comment && !is_record_terminator(byte, terminator) } +#[cfg(target_arch = "x86_64")] +fn scan_record_terminators(input: &[u8]) -> (bool, bool) { + use std::arch::x86_64::{ + __m128i, _mm_cmpeq_epi8, _mm_loadu_si128, _mm_movemask_epi8, _mm_set1_epi8, + }; + + // SSE2 is part of the x86-64 baseline. Compare 16 bytes at once and + // exclude CRLF, which is one record delimiter rather than two. + let (cr_byte, lf_byte) = unsafe { (_mm_set1_epi8(b'\r' as i8), _mm_set1_epi8(b'\n' as i8)) }; + let mut offset = 0; + let mut previous = None; + let mut bare_cr = false; + while input.len() - offset >= 16 { + let (cr, lf) = unsafe { + let word = _mm_loadu_si128(input.as_ptr().add(offset).cast::<__m128i>()); + ( + _mm_movemask_epi8(_mm_cmpeq_epi8(word, cr_byte)) as u32, + _mm_movemask_epi8(_mm_cmpeq_epi8(word, lf_byte)) as u32, + ) + }; + let terminators = cr | lf; + bare_cr |= cr & 0x7fff & !(lf >> 1) != 0; + bare_cr |= previous == Some(b'\r') && lf & 1 == 0; + if (terminators & (terminators << 1) & !((cr << 1) & lf)) != 0 { + return (true, bare_cr); + } + if previous.is_some_and(|last| terminators & 1 != 0 && !(last == b'\r' && lf & 1 != 0)) { + return (true, bare_cr); + } + previous = is_record_terminator(input[offset + 15], None).then_some(input[offset + 15]); + offset += 16; + } + for &byte in &input[offset..] { + bare_cr |= previous == Some(b'\r') && byte != b'\n'; + if previous.is_some_and(|last| { + is_record_terminator(byte, None) && !(last == b'\r' && byte == b'\n') + }) { + return (true, bare_cr); + } + previous = is_record_terminator(byte, None).then_some(byte); + } + (false, bare_cr) +} + +#[cfg(not(target_arch = "x86_64"))] +fn scan_record_terminators(input: &[u8]) -> (bool, bool) { + let mut blank = false; + let mut bare_cr = false; + for pair in input.windows(2) { + blank |= is_record_terminator(pair[0], None) + && is_record_terminator(pair[1], None) + && pair != b"\r\n"; + bare_cr |= pair[0] == b'\r' && pair[1] != b'\n'; + } + (blank, bare_cr) +} + fn ignored_record_prefix(raw: &[u8], comment: Option, terminator: Option) -> usize { let mut offset = 0; while let Some(&byte) = raw.get(offset) { @@ -73,6 +133,7 @@ fn count_source_line_breaks(raw: &[u8]) -> usize { #[derive(Clone, Copy, Debug)] struct SourceLineCounter { + initial_line_offset: usize, lf_count: usize, count_lf_independently: bool, bare_cr_count: usize, @@ -87,6 +148,11 @@ enum PendingCr { source_line_number: usize, byte_offset: usize, }, + MaterializedBlank { + line_number: usize, + source_line_number: usize, + byte_offset: usize, + }, } impl SourceLineCounter { @@ -122,7 +188,9 @@ impl SourceLineCounter { .saturating_add(usize::from( self.pending_cr && record.last() == Some(&b'\r'), )); - end_line.saturating_sub(count_source_line_breaks(record)) + end_line + .saturating_sub(count_source_line_breaks(record)) + .saturating_add(self.initial_line_offset) } } @@ -285,8 +353,18 @@ pub struct RecordDecoder { /// Optional recovery hook. Keeping this `None` selects the allocation-free strict path. record_error_handler: Option>, + /// Strict blank handling is entirely opt-in; ordinary scans keep the original path. + blank_record_policy: CsvBlankRecordPolicy, + blank_null_row_supported: bool, + skipped_blank_records_counter: Option>, + strict_at_record_start: bool, + strict_pending_cr: bool, + strict_comment: bool, + strict_bom_len: usize, + /// Cache this capability once so the common recovery path avoids another virtual call. filter_records: bool, + record_source_lines: bool, /// Keep the ordinary recovery path free of blank-record callbacks. report_blank_records: bool, @@ -320,6 +398,7 @@ pub struct RecordDecoder { /// Total number of input bytes consumed by the recovery path. stream_offset: usize, + initial_stream_offset: usize, /// Start offsets for rolling back a malformed record from decoded output. record_byte_offset: usize, @@ -345,13 +424,22 @@ impl RecordDecoder { ignore_extra_fields: false, skipping_extra_fields: false, record_error_handler: None, + blank_record_policy: CsvBlankRecordPolicy::Ignore, + blank_null_row_supported: false, + skipped_blank_records_counter: None, + strict_at_record_start: true, + strict_pending_cr: false, + strict_comment: false, + strict_bom_len: 0, filter_records: false, + record_source_lines: false, report_blank_records: false, blank_records_are_rows: false, pending_cr: None, max_record_bytes: None, replace_invalid_utf8: false, source_lines: SourceLineCounter { + initial_line_offset: 0, lf_count: 0, count_lf_independently: true, bare_cr_count: 0, @@ -366,12 +454,25 @@ impl RecordDecoder { record_needs_utf8_validation: false, record_bytes: vec![], stream_offset: 0, + initial_stream_offset: 0, record_byte_offset: 0, record_data_start: 0, record_offsets_start: 1, } } + pub(crate) fn with_initial_source_position( + mut self, + skipped_lines: usize, + skipped_bytes: usize, + ) -> Self { + self.source_lines.initial_line_offset = skipped_lines; + self.stream_offset = skipped_bytes; + self.initial_stream_offset = skipped_bytes; + self.record_byte_offset = skipped_bytes; + self + } + pub fn with_ignore_extra_fields(mut self, ignore: bool) -> Self { self.ignore_extra_fields = ignore; self @@ -384,6 +485,9 @@ impl RecordDecoder { self.filter_records = handler .as_ref() .is_some_and(|handler| handler.filters_records()); + self.record_source_lines = handler + .as_ref() + .is_some_and(|handler| handler.needs_source_line_number()); self.blank_records_are_rows = handler .as_ref() .is_some_and(|handler| handler.blank_records_are_rows()); @@ -398,6 +502,21 @@ impl RecordDecoder { self } + pub(super) fn with_blank_record_policy( + mut self, + policy: CsvBlankRecordPolicy, + all_columns_nullable: bool, + counter: Option>, + ) -> Self { + self.blank_record_policy = policy; + self.blank_null_row_supported = all_columns_nullable; + self.skipped_blank_records_counter = counter; + if policy == CsvBlankRecordPolicy::EmitNullRow { + self.field_counts.get_or_insert_with(Vec::new); + } + self + } + pub fn with_replace_invalid_utf8(mut self, replace: bool) -> Self { self.replace_invalid_utf8 = replace; self @@ -442,15 +561,14 @@ impl RecordDecoder { /// /// Note: this expects to be called with an empty `input` to signal EOF pub fn decode(&mut self, input: &[u8], to_read: usize) -> Result<(usize, usize), ArrowError> { - if self.blank_records_are_rows && self.num_columns != 1 { + if self.blank_records_are_rows && self.num_columns == 0 { return Err(ArrowError::CsvError( - "CSV blank records can be rows only with a one-column schema".to_owned(), + "CSV blank records as rows require at least one column".to_owned(), )); } match self.record_error_handler.clone() { Some(handler) - if self.blank_records_are_rows - && self.input_may_have_one_column_blank(input, to_read) => + if self.blank_records_are_rows && self.input_may_have_blank(input, to_read) => { self.decode_with_record_error_handler::( input, @@ -467,8 +585,17 @@ impl RecordDecoder { to_read, handler.as_ref(), ), - None if self.ignore_extra_fields => self.decode_strict::(input, to_read), - None => self.decode_strict::(input, to_read), + None if self.blank_record_policy != CsvBlankRecordPolicy::Ignore => { + if self.ignore_extra_fields { + self.decode_strict_blank::(input, to_read) + } else { + self.decode_strict_blank::(input, to_read) + } + } + None if self.ignore_extra_fields => { + self.decode_strict::(input, to_read) + } + None => self.decode_strict::(input, to_read), } } @@ -487,6 +614,11 @@ impl RecordDecoder { &mut self, input: &[u8], ) -> Result<(usize, usize), ArrowError> { + if self.record_error_handler.is_none() + && self.blank_record_policy != CsvBlankRecordPolicy::Ignore + { + return self.decode_skipping(input, 1); + } if self.blank_records_are_rows && self.terminator.is_some() && self @@ -507,11 +639,14 @@ impl RecordDecoder { return Ok((0, input_offset)); } } - if self.blank_records_are_rows && self.header_starts_with_blank(input) { + if self.blank_records_are_rows + && (matches!(self.pending_cr, Some(PendingCr::MaterializedBlank { .. })) + || self.header_starts_with_blank(input)) + { return self.decode(input, 1); } let scan_input = if self.blank_records_are_rows - && self.stream_offset > 0 + && self.stream_offset > self.initial_stream_offset && input.starts_with(b"\xef\xbb\xbf") { &input[..1] @@ -519,9 +654,9 @@ impl RecordDecoder { input }; let (rows, bytes) = if self.ignore_extra_fields { - self.decode_strict::(scan_input, 1)? + self.decode_strict::(scan_input, 1)? } else { - self.decode_strict::(scan_input, 1)? + self.decode_strict::(scan_input, 1)? }; self.source_lines.advance(&input[..bytes]); self.stream_offset = self @@ -548,7 +683,7 @@ impl RecordDecoder { if self.current_field != 0 || (!self.record_bytes.is_empty() && !pending_comment) { return false; } - let mut remaining = if self.stream_offset == 0 { + let mut remaining = if self.stream_offset == self.initial_stream_offset { input.strip_prefix(b"\xef\xbb\xbf").unwrap_or(input) } else { input @@ -592,11 +727,18 @@ impl RecordDecoder { handler: &dyn CsvRecordErrorHandler, with_lf: bool, ) -> Result<(), ArrowError> { - if let Some(PendingCr::Blank { + let Some(pending) = *pending_cr else { + return Ok(()); + }; + if matches!(pending, PendingCr::MaterializedBlank { .. }) { + return Ok(()); + } + *pending_cr = None; + if let PendingCr::Blank { line_number, source_line_number, byte_offset, - }) = pending_cr.take() + } = pending { handler.handle_blank_record(&CsvFilterRecord { line_number, @@ -686,8 +828,62 @@ impl RecordDecoder { Ok(()) } + fn materialize_blank_record( + &mut self, + handler: &dyn CsvRecordErrorHandler, + raw: &[u8], + line_number: usize, + source_line_number: usize, + byte_offset: usize, + read: &mut usize, + ) -> Result<(), ArrowError> { + let skipped = + if !self.skipping_records && self.filter_records && handler.has_record_error(raw) { + handler.skip_record(&CsvFilterRecord { + line_number, + source_line_number, + byte_offset, + record: raw, + })? + } else { + false + }; + if !self.skipping_records && !skipped { + let record = CsvRecord { + line_number, + byte_offset, + record: raw, + }; + if self.record_source_lines { + handler.handle_record_with_source_line(&record, source_line_number)?; + } else { + handler.handle_record(&record)?; + } + self.offsets[self.offsets_len..self.offsets_len + self.num_columns].fill(0); + self.offsets_len += self.num_columns; + if let Some(counts) = &mut self.field_counts { + if counts.len() == self.num_rows { + counts.push(self.num_columns); + } else { + counts[self.num_rows] = self.num_columns; + } + } + if let Some(quoted) = &mut self.quoted_empty { + quoted.record(self.num_rows, self.num_columns, 1, false, raw)?; + } + self.num_rows += 1; + } + if !skipped { + *read += 1; + } + self.line_number = line_number.saturating_add(1); + self.record_data_start = self.data_len; + self.record_offsets_start = self.offsets_len; + Ok(()) + } + #[inline(always)] - fn consume_one_column_blank_prefix( + fn consume_blank_prefix( &mut self, input: &[u8], input_offset: &mut usize, @@ -699,6 +895,58 @@ impl RecordDecoder { if self.current_field != 0 || !self.record_bytes.is_empty() { return Ok(false); } + if let Some(PendingCr::MaterializedBlank { + line_number, + source_line_number, + byte_offset, + }) = self.pending_cr + { + if *input_offset == input.len() && !input.is_empty() { + return Ok(false); + } + let with_lf = self.terminator.is_none() && input.get(*input_offset) == Some(&b'\n'); + if with_lf && self.max_record_bytes.is_some_and(|limit| limit < 2) { + return Err(ArrowError::CsvError( + "CSV record exceeds the configured 1 byte limit".to_owned(), + )); + } + let raw = if with_lf { + b"\r\n".as_slice() + } else { + b"\r".as_slice() + }; + self.pending_cr = None; + self.materialize_blank_record( + handler, + raw, + line_number, + source_line_number, + byte_offset, + read, + )?; + if with_lf { + let mut discarded_data = [0_u8; 1]; + let mut discarded_offsets = [0_usize; 1]; + let (_, consumed, _, _) = self.delimiter.read_record( + &input[*input_offset..*input_offset + 1], + &mut discarded_data, + &mut discarded_offsets, + ); + if consumed != 1 { + return Err(ArrowError::CsvError( + "CSV blank record CRLF continuation was not consumed".to_owned(), + )); + } + *input_offset += 1; + self.stream_offset = self.stream_offset.checked_add(1).ok_or_else(|| { + ArrowError::CsvError("CSV input byte offset overflowed".into()) + })?; + self.record_byte_offset = self.stream_offset; + source_lines.advance(&input[*counted_until..*input_offset]); + *counted_until = *input_offset; + } + return Ok(true); + } let Some(&byte) = input.get(*input_offset) else { return Ok(false); }; @@ -709,7 +957,9 @@ impl RecordDecoder { } let continuation = matches!(self.pending_cr, Some(PendingCr::Data)) && byte == b'\n'; self.pending_cr = None; - if self.stream_offset == 0 && input[*input_offset..].starts_with(b"\xef\xbb\xbf") { + if self.stream_offset == self.initial_stream_offset + && input[*input_offset..].starts_with(b"\xef\xbb\xbf") + { return Ok(false); } @@ -743,54 +993,7 @@ impl RecordDecoder { "CSV blank record terminator was not consumed".to_owned(), )); } - - if blank { - let skipped = - if !self.skipping_records && self.filter_records && handler.has_record_error(raw) { - source_lines.advance(&input[*counted_until..*input_offset + 1]); - *counted_until = *input_offset + 1; - handler.skip_record(&CsvFilterRecord { - line_number: self.line_number, - source_line_number: source_lines.record_start_line( - usize::try_from(self.delimiter.line()).unwrap_or(usize::MAX), - raw, - ), - byte_offset: self.stream_offset, - record: raw, - })? - } else { - false - }; - if !self.skipping_records && !skipped { - handler.handle_record(&CsvRecord { - line_number: self.line_number, - byte_offset: self.stream_offset, - record: raw, - })?; - self.offsets[self.offsets_len] = 0; - self.offsets_len += 1; - if let Some(counts) = &mut self.field_counts { - if counts.len() == self.num_rows { - counts.push(1); - } else { - counts[self.num_rows] = 1; - } - } - if let Some(quoted) = &mut self.quoted_empty { - quoted.record(self.num_rows, 1, 1, false, raw)?; - } - self.num_rows += 1; - } - if !skipped { - *read += 1; - } - self.line_number += 1; - self.record_data_start = self.data_len; - self.record_offsets_start = self.offsets_len; - self.pending_cr = - (self.terminator.is_none() && byte == b'\r').then_some(PendingCr::Data); - } - + let byte_offset = self.stream_offset; *input_offset += 1; self.stream_offset = self .stream_offset @@ -799,6 +1002,28 @@ impl RecordDecoder { self.record_byte_offset = self.stream_offset; source_lines.advance(&input[*counted_until..*input_offset]); *counted_until = *input_offset; + if blank { + let source_line_number = source_lines.record_start_line( + usize::try_from(self.delimiter.line()).unwrap_or(usize::MAX), + raw, + ); + if self.terminator.is_none() && byte == b'\r' { + self.pending_cr = Some(PendingCr::MaterializedBlank { + line_number: self.line_number, + source_line_number, + byte_offset, + }); + } else { + self.materialize_blank_record( + handler, + raw, + self.line_number, + source_line_number, + byte_offset, + read, + )?; + } + } Ok(true) } @@ -819,7 +1044,8 @@ impl RecordDecoder { if self.current_field != 0 || !self.record_bytes.is_empty() || input.get(*input_offset) != Some(&comment) - || (self.stream_offset == 0 && input[*input_offset..].starts_with(b"\xef\xbb\xbf")) + || (self.stream_offset == self.initial_stream_offset + && input[*input_offset..].starts_with(b"\xef\xbb\xbf")) { return Ok(false); } @@ -857,8 +1083,11 @@ impl RecordDecoder { Ok(true) } - fn input_may_have_one_column_blank(&self, input: &[u8], to_read: usize) -> bool { + fn input_may_have_blank(&self, input: &[u8], to_read: usize) -> bool { const MAX_PREFLIGHT_BYTES_PER_ROW: usize = 64; + if matches!(self.pending_cr, Some(PendingCr::MaterializedBlank { .. })) { + return true; + } if self.skipping_custom_comment { return true; } @@ -888,7 +1117,180 @@ impl RecordDecoder { } } - fn decode_strict( + fn advance_strict_position(&mut self, input: &[u8]) -> Result<(), ArrowError> { + self.source_lines.advance(input); + self.stream_offset = self + .stream_offset + .checked_add(input.len()) + .ok_or_else(|| ArrowError::CsvError("CSV input byte offset overflowed".into()))?; + Ok(()) + } + + fn advance_strict_position_fast( + &mut self, + input: &[u8], + bare_cr: bool, + ) -> Result<(), ArrowError> { + if self.blank_record_policy == CsvBlankRecordPolicy::Reject { + if bare_cr || self.terminator.is_some() { + return self.advance_strict_position(input); + } + if let Some(&first) = input.first() { + if self.source_lines.pending_cr && first != b'\n' { + self.source_lines.bare_cr_count = + self.source_lines.bare_cr_count.saturating_add(1); + } + self.source_lines.pending_cr = input.last() == Some(&b'\r'); + } + } + self.stream_offset = self + .stream_offset + .checked_add(input.len()) + .ok_or_else(|| ArrowError::CsvError("CSV input byte offset overflowed".into()))?; + Ok(()) + } + + fn finish_strict( + &mut self, + input: &[u8], + read: usize, + consumed: usize, + ) -> Result<(usize, usize), ArrowError> { + if TRACK_BLANK { + self.advance_strict_position(&input[..consumed])?; + } + Ok((read, consumed)) + } + + fn consume_strict_ignored(&mut self, input: &[u8]) -> Result<(), ArrowError> { + let mut discarded_data = [0_u8; 1]; + let mut discarded_offsets = [0_usize; 1]; + let (result, bytes, _, ends) = + self.delimiter + .read_record(input, &mut discarded_data, &mut discarded_offsets); + if bytes != input.len() || ends != 0 || matches!(result, ReadRecordResult::Record) { + return Err(ArrowError::CsvError( + "CSV blank record terminator was not consumed".to_owned(), + )); + } + Ok(()) + } + + fn handle_strict_blank( + &mut self, + raw: &[u8], + input_prefix: &[u8], + read: &mut usize, + ) -> Result<(), ArrowError> { + match self.blank_record_policy { + CsvBlankRecordPolicy::Ignore => unreachable!(), + CsvBlankRecordPolicy::Reject => { + let mut source_lines = self.source_lines; + source_lines.advance(input_prefix); + let source_line = source_lines.record_start_line( + usize::try_from(self.delimiter.line()).unwrap_or(usize::MAX), + raw, + ); + return Err(ArrowError::CsvError(format!( + "blank CSV record at line {}, source line {source_line}", + self.line_number, + ))); + } + CsvBlankRecordPolicy::EmitNullRow => { + if !self.blank_null_row_supported { + return Err(ArrowError::CsvError( + "CSV blank records as NULL rows require all columns to be nullable" + .to_owned(), + )); + } + self.offsets[self.offsets_len..self.offsets_len + self.num_columns].fill(0); + self.offsets_len += self.num_columns; + let counts = self + .field_counts + .as_mut() + .expect("initialized for NULL rows"); + if counts.len() == self.num_rows { + counts.push(0); + } else { + counts[self.num_rows] = 0; + } + if let Some(quoted) = &mut self.quoted_empty { + let end = self.offsets_len - 1; + quoted.fields.resize(end, 0); + quoted.fields[end - self.num_columns..end].fill(0); + } + self.num_rows += 1; + *read += 1; + } + CsvBlankRecordPolicy::Skip => { + if !self.skipping_records { + if let Some(counter) = &self.skipped_blank_records_counter { + counter.fetch_add(1, Ordering::Relaxed); + } + } + } + } + self.line_number = self.line_number.saturating_add(1); + self.strict_at_record_start = true; + self.strict_pending_cr = self.terminator.is_none() && raw == b"\r"; + self.record_data_start = self.data_len; + self.record_offsets_start = self.offsets_len; + Ok(()) + } + + fn strict_fast_input_bare_cr(&self, input: &[u8], to_read: usize) -> Option { + const MAX_PREFLIGHT_BYTES_PER_ROW: usize = 64; + if self.comment.is_some() + || self.strict_comment + || self.strict_bom_len != 0 + || self.quoted_empty.is_some() + || self.stream_offset == self.initial_stream_offset + || input.len() > to_read.saturating_mul(MAX_PREFLIGHT_BYTES_PER_ROW) + { + return None; + } + + let prefix = usize::from(self.strict_pending_cr && input.first() == Some(&b'\n')); + if self.strict_at_record_start + && input + .get(prefix) + .is_some_and(|byte| is_record_terminator(*byte, self.terminator)) + { + return None; + } + + if let Some(term) = self.terminator { + memchr::memmem::find(input, &[term, term]) + .is_none() + .then_some(true) + } else { + let (blank, bare_cr) = scan_record_terminators(input); + (!blank).then_some(bare_cr) + } + } + + fn decode_strict_blank( + &mut self, + input: &[u8], + to_read: usize, + ) -> Result<(usize, usize), ArrowError> { + let fast_bare_cr = self.strict_fast_input_bare_cr(input, to_read); + if let Some(bare_cr) = fast_bare_cr { + let (rows, bytes) = + self.decode_strict::(input, to_read)?; + self.advance_strict_position_fast(&input[..bytes], bare_cr)?; + Ok((rows, bytes)) + } else { + self.decode_strict::(input, to_read) + } + } + + fn decode_strict< + const IGNORE_EXTRA: bool, + const TRACK_BLANK: bool, + const CHECK_BLANK: bool, + const CAPTURE_STATE: bool, + >( &mut self, input: &[u8], to_read: usize, @@ -901,10 +1303,26 @@ impl RecordDecoder { "CSV cannot ignore extra fields with an empty schema".to_owned(), )); } + if TRACK_BLANK && input.is_empty() && self.strict_bom_len != 0 { + let prefix_len = self.strict_bom_len; + self.strict_bom_len = 0; + let (rows, bytes) = self.decode_strict::( + &b"\xef\xbb\xbf"[..prefix_len], + to_read, + )?; + debug_assert_eq!((rows, bytes), (0, prefix_len)); + self.strict_at_record_start = false; + } - // Reserve sufficient capacity in offsets + // Opt-in input cannot contain more records than bytes. The default + // reservation remains unchanged for ordinary scans. + let reserve_rows = if TRACK_BLANK { + to_read.min(input.len().saturating_add(1)) + } else { + to_read + }; self.offsets - .resize(self.offsets_len + to_read * self.num_columns, 0); + .resize(self.offsets_len + reserve_rows * self.num_columns, 0); // The current offset into `input` let mut input_offset = 0; @@ -912,25 +1330,138 @@ impl RecordDecoder { // The number of rows decoded in this pass let mut read = 0; + let mut fast_boundary_offset = 0; + let mut fast_at_boundary = self.strict_at_record_start; + let mut fast_boundary_cr = self.strict_pending_cr; loop { // Reserve necessary space in output data based on best estimate let remaining_rows = to_read - read; let capacity = remaining_rows * self.num_columns * AVERAGE_FIELD_SIZE; - let estimated_data = capacity.max(MIN_CAPACITY); + let estimated_data = if TRACK_BLANK { + capacity.min(input.len() - input_offset).max(64) + } else { + capacity.max(MIN_CAPACITY) + }; self.data.resize(self.data_len + estimated_data, 0); // Try to read a record loop { + let mut continuing_crlf = false; + if CHECK_BLANK && !input.is_empty() { + if input_offset == input.len() { + return self.finish_strict::(input, read, input_offset); + } + let mut remaining = &input[input_offset..]; + if self.strict_bom_len != 0 + || (self.stream_offset == self.initial_stream_offset && input_offset == 0) + { + const BOM: &[u8; 3] = b"\xef\xbb\xbf"; + while self.strict_bom_len < 3 + && input.get(input_offset) == Some(&BOM[self.strict_bom_len]) + { + self.strict_bom_len += 1; + input_offset += 1; + } + if self.strict_bom_len == 3 { + record_start = input_offset; + if input_offset == input.len() { + return self.finish_strict::( + input, + read, + input_offset, + ); + } + remaining = &input[input_offset..]; + } + if self.strict_bom_len != 0 && self.strict_bom_len != 3 { + if input_offset == input.len() { + return self.finish_strict::( + input, + read, + input_offset, + ); + } + let prefix_len = self.strict_bom_len; + self.strict_bom_len = 0; + let (rows, bytes) = self + .decode_strict::( + &BOM[..prefix_len], + to_read - read, + )?; + debug_assert_eq!((rows, bytes), (0, prefix_len)); + self.strict_at_record_start = false; + record_start = input_offset; + continue; + } + } + if self.strict_bom_len != 3 && self.strict_pending_cr { + self.strict_pending_cr = false; + if remaining[0] == b'\n' { + if remaining.len() == 1 + || is_record_terminator(remaining[1], self.terminator) + || is_comment_start(remaining[1], self.comment, self.terminator) + { + self.consume_strict_ignored(&remaining[..1])?; + input_offset += 1; + record_start = input_offset; + continue; + } + continuing_crlf = true; + record_start = input_offset + 1; + } + } + if self.strict_bom_len != 3 + && self.strict_at_record_start + && !self.strict_comment + && !continuing_crlf + { + if is_comment_start(remaining[0], self.comment, self.terminator) { + self.strict_comment = true; + } else if is_record_terminator(remaining[0], self.terminator) { + let blank_len = + if self.terminator.is_none() && remaining.starts_with(b"\r\n") { + 2 + } else { + 1 + }; + let raw = &remaining[..blank_len]; + self.consume_strict_ignored(raw)?; + input_offset += blank_len; + self.handle_strict_blank(raw, &input[..input_offset], &mut read)?; + record_start = input_offset; + if read == to_read { + return self.finish_strict::( + input, + read, + input_offset, + ); + } + continue; + } + } + } let skipping_extra_fields = IGNORE_EXTRA && self.skipping_extra_fields; + let mut bom_prefix = [0xef, 0xbb, 0xbf, 0]; + let feeding_bom = CHECK_BLANK && self.strict_bom_len == 3 && !input.is_empty(); + let scan_input = if feeding_bom { + // csv_core must consume the first BOM so it retains a later BOM as data. + bom_prefix[3] = input[input_offset]; + &bom_prefix[..] + } else if TRACK_BLANK && self.strict_comment { + let remaining = &input[input_offset..]; + let end = + memchr::memchr(b'\n', remaining).map_or(remaining.len(), |index| index + 1); + &remaining[..end] + } else { + &input[input_offset..] + }; + let input_before_read = input_offset; let (result, bytes_read, bytes_written, end_positions) = if skipping_extra_fields { let mut skipped_data = [0_u8; 1024]; let mut skipped_offsets = [0_usize; 16]; - self.delimiter.read_record( - &input[input_offset..], - &mut skipped_data, - &mut skipped_offsets, - ) + self.delimiter + .read_record(scan_input, &mut skipped_data, &mut skipped_offsets) } else { let offsets_end = if IGNORE_EXTRA { self.record_offsets_start + self.num_columns @@ -938,13 +1469,60 @@ impl RecordDecoder { self.offsets.len() }; self.delimiter.read_record( - &input[input_offset..], + scan_input, &mut self.data[self.data_len..], &mut self.offsets[self.offsets_len..offsets_end], ) }; self.current_field += end_positions; - input_offset += bytes_read; + let input_bytes_read = if feeding_bom { + self.strict_bom_len = 0; + debug_assert!(bytes_read >= 3); + bytes_read - 3 + } else { + bytes_read + }; + input_offset += input_bytes_read; + if continuing_crlf { + if input_bytes_read == 0 { + record_start = input_before_read; + self.strict_pending_cr = true; + } else if input_bytes_read == 1 { + record_start = input_offset; + self.strict_at_record_start = true; + } + } + if TRACK_BLANK && input_bytes_read != 0 { + if feeding_bom + && self.strict_at_record_start + && !self.strict_comment + && is_comment_start(input[input_before_read], self.comment, self.terminator) + { + self.strict_comment = true; + } + if self.strict_comment { + if memchr::memchr(b'\n', &input[input_before_read..input_offset]).is_some() + { + self.strict_comment = false; + self.strict_at_record_start = true; + } + } else { + if !continuing_crlf || input_bytes_read > 1 { + self.strict_at_record_start = false; + } + } + } + if feeding_bom && input_bytes_read != 0 { + let first = input[input_before_read]; + if is_record_terminator(first, self.terminator) { + let raw = &input[input_before_read..input_offset]; + self.handle_strict_blank(raw, &input[..input_offset], &mut read)?; + record_start = input_offset; + if read == to_read { + return self.finish_strict::(input, read, input_offset); + } + } + } if skipping_extra_fields { if self.quoted_empty.is_some() { self.record_bytes @@ -958,12 +1536,29 @@ impl RecordDecoder { match result { ReadRecordResult::End | ReadRecordResult::InputEmpty => { + if TRACK_BLANK && input_offset < input.len() { + if self.quoted_empty.is_some() { + self.record_bytes + .extend_from_slice(&input[record_start..input_offset]); + record_start = input_offset; + } + continue; + } + if CAPTURE_STATE { + let untouched = input_offset == fast_boundary_offset; + let crlf_continuation = fast_boundary_cr + && input.get(fast_boundary_offset) == Some(&b'\n') + && input_offset == fast_boundary_offset + 1; + self.strict_at_record_start = + (untouched && fast_at_boundary) || crlf_continuation; + self.strict_pending_cr = untouched && fast_boundary_cr; + } // Reached end of input if self.quoted_empty.is_some() { self.record_bytes .extend_from_slice(&input[record_start..input_offset]); } - return Ok((read, input_offset)); + return self.finish_strict::(input, read, input_offset); } // Need to allocate more capacity ReadRecordResult::OutputFull => { @@ -994,6 +1589,17 @@ impl RecordDecoder { ))); } ReadRecordResult::Record => { + if CAPTURE_STATE { + fast_boundary_offset = input_offset; + fast_at_boundary = true; + fast_boundary_cr = self.terminator.is_none() + && input.get(input_offset.saturating_sub(1)) == Some(&b'\r'); + } + if TRACK_BLANK { + self.strict_at_record_start = true; + self.strict_pending_cr = self.terminator.is_none() + && input.get(input_offset.saturating_sub(1)) == Some(&b'\r'); + } self.skipping_extra_fields = false; if self.current_field != self.num_columns && !(IGNORE_EXTRA && self.current_field > self.num_columns) @@ -1058,14 +1664,22 @@ impl RecordDecoder { if read == to_read { // Read sufficient rows - return Ok((read, input_offset)); + if CAPTURE_STATE { + self.strict_at_record_start = true; + self.strict_pending_cr = fast_boundary_cr; + } + return self.finish_strict::(input, read, input_offset); } if input.len() == input_offset { // Input exhausted, need to read more // Without this read_record will interpret the empty input // byte array as indicating the end of the file - return Ok((read, input_offset)); + if CAPTURE_STATE { + self.strict_at_record_start = true; + self.strict_pending_cr = fast_boundary_cr; + } + return self.finish_strict::(input, read, input_offset); } } } @@ -1087,9 +1701,9 @@ impl RecordDecoder { "CSV cannot ignore extra fields with an empty schema".to_owned(), )); } - if BLANK_AS_ROW && self.num_columns != 1 { + if BLANK_AS_ROW && self.num_columns == 0 { return Err(ArrowError::CsvError( - "CSV blank records can be rows only with a one-column schema".to_owned(), + "CSV blank records as rows require at least one column".to_owned(), )); } @@ -1129,7 +1743,7 @@ impl RecordDecoder { continue; } if BLANK_AS_ROW - && self.consume_one_column_blank_prefix( + && self.consume_blank_prefix( input, &mut input_offset, &mut counted_until, @@ -1164,12 +1778,12 @@ impl RecordDecoder { }; let record_input_start = input_offset; let remaining_input = &input[input_offset..]; - let initial_bom = - self.stream_offset == 0 && remaining_input.starts_with(b"\xef\xbb\xbf"); + let initial_bom = self.stream_offset == self.initial_stream_offset + && remaining_input.starts_with(b"\xef\xbb\xbf"); let scan_input = if BLANK_AS_ROW && initial_bom { &remaining_input[..3] } else if self.blank_records_are_rows - && self.stream_offset > 0 + && self.stream_offset > self.initial_stream_offset && remaining_input.starts_with(b"\xef\xbb\xbf") { &remaining_input[..1] @@ -1485,11 +2099,24 @@ impl RecordDecoder { } } - handler.handle_record(&CsvRecord { + let record = CsvRecord { line_number: self.line_number, byte_offset: self.record_byte_offset, record: &self.record_bytes, - })?; + }; + if self.record_source_lines { + source_lines.advance(&input[counted_until..input_offset]); + counted_until = input_offset; + handler.handle_record_with_source_line( + &record, + source_lines.record_start_line( + usize::try_from(self.delimiter.line()).unwrap_or(usize::MAX), + &self.record_bytes, + ), + )?; + } else { + handler.handle_record(&record)?; + } if self.current_field < self.num_columns { let fill_count = self.num_columns - self.current_field; @@ -1805,7 +2432,76 @@ mod tests { use std::sync::{Arc, Mutex}; use std::time::Instant; - use super::{MAX_RECOVERABLE_EXCESS_FIELDS, RecordDecoder, recovery_limit_error}; + use super::{ + MAX_RECOVERABLE_EXCESS_FIELDS, RecordDecoder, recovery_limit_error, scan_record_terminators, + }; + + #[test] + fn test_strict_blank_vector_preflight_across_boundaries() { + for len in 0usize..=34 { + for index in 0..len.saturating_sub(1) { + for &first in b"x\r\n" { + for &second in b"x\r\n" { + let mut input = vec![b'x'; len]; + input[index] = first; + input[index + 1] = second; + let expected = input.windows(2).any(|pair| { + (pair[0] == b'\r' || pair[0] == b'\n') + && (pair[1] == b'\r' || pair[1] == b'\n') + && pair != b"\r\n" + }); + let bare_cr = input + .windows(2) + .any(|pair| pair[0] == b'\r' && pair[1] != b'\n'); + let actual = scan_record_terminators(&input); + assert_eq!( + actual.0, expected, + "len={len}, index={index}, bytes={input:?}" + ); + if !expected { + assert_eq!( + actual.1, bare_cr, + "len={len}, index={index}, bytes={input:?}" + ); + } + } + } + } + } + } + + #[test] + fn test_strict_blank_preflight_crlf_boundaries() { + let mut decoder = RecordDecoder::new(Reader::new(), 1, false); + decoder.stream_offset = 1; + assert_eq!( + decoder.strict_fast_input_bare_cr(b"one\r\ntwo\r\n", 100), + Some(false) + ); + assert!( + decoder + .strict_fast_input_bare_cr(b"one\r\n\r\ntwo", 100) + .is_none() + ); + assert!( + decoder + .strict_fast_input_bare_cr(b"one\n\ntwo", 100) + .is_none() + ); + + decoder.strict_pending_cr = true; + assert_eq!(decoder.strict_fast_input_bare_cr(b"\n", 100), Some(false)); + assert_eq!( + decoder.strict_fast_input_bare_cr(b"\ntwo\r\n", 100), + Some(false) + ); + assert!( + decoder + .strict_fast_input_bare_cr(b"\n\r\ntwo", 100) + .is_none() + ); + assert!(decoder.strict_fast_input_bare_cr(b"\n\ntwo", 100).is_none()); + } #[derive(Debug, Clone, Eq, PartialEq)] struct OwnedRecordError { @@ -1841,6 +2537,14 @@ mod tests { record: Vec, } + #[derive(Debug, Clone, Eq, PartialEq)] + struct OwnedSourceRecord { + line_number: usize, + source_line_number: usize, + byte_offset: usize, + record: Vec, + } + #[derive(Debug, Default)] struct CollectRecords { errors: Mutex>, @@ -1852,6 +2556,32 @@ mod tests { max_record_bytes: Option, } + #[derive(Debug, Default)] + struct SourceLines(Mutex>); + + impl CsvRecordErrorHandler for SourceLines { + fn handle(&self, _error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + Ok(()) + } + + fn needs_source_line_number(&self) -> bool { + true + } + + fn handle_record_with_source_line( + &self, + record: &CsvRecord<'_>, + source_line_number: usize, + ) -> Result<(), ArrowError> { + self.0.lock().unwrap().push(( + record.line_number, + source_line_number, + record.byte_offset, + )); + Ok(()) + } + } + impl CsvRecordErrorHandler for CollectRecords { fn handle(&self, error: &CsvRecordError<'_>) -> Result<(), ArrowError> { self.errors.lock().unwrap().push(OwnedRecordError { @@ -2038,12 +2768,17 @@ mod tests { #[test] fn test_one_column_blank_records_are_rows_across_batches_and_chunks() { - for (csv, expected, blank_offset) in [ - (b"a\n\nb\n".as_slice(), ["a", "", "b"], 2), - (b"a\n\nb".as_slice(), ["a", "", "b"], 2), - (b"a\r\n\r\nb\r\n".as_slice(), ["a", "", "b"], 3), - (b"\n\n\n".as_slice(), ["", "", ""], 1), - (b"a\r\rb\r".as_slice(), ["a", "", "b"], 2), + for (csv, expected, blank_offset, blank_raw) in [ + (b"a\n\nb\n".as_slice(), ["a", "", "b"], 2, b"\n".as_slice()), + (b"a\n\nb".as_slice(), ["a", "", "b"], 2, b"\n".as_slice()), + ( + b"a\r\n\r\nb\r\n".as_slice(), + ["a", "", "b"], + 3, + b"\r\n".as_slice(), + ), + (b"\n\n\n".as_slice(), ["", "", ""], 1, b"\n".as_slice()), + (b"a\r\rb\r".as_slice(), ["a", "", "b"], 2, b"\r".as_slice()), ] { for capacity in [1, 2, 5, 64] { for batch_size in [1, 2, 8] { @@ -2057,7 +2792,9 @@ mod tests { let mut actual = Vec::new(); loop { let buf = reader.fill_buf().unwrap(); - let (_, bytes) = decoder.decode(buf, batch_size - decoder.len()).unwrap(); + let input_empty = buf.is_empty(); + let (decoded, bytes) = + decoder.decode(buf, batch_size - decoder.len()).unwrap(); reader.consume(bytes); if decoder.len() == batch_size || bytes == 0 { actual.extend( @@ -2068,9 +2805,10 @@ mod tests { .map(|record| record.get(0).to_owned()), ); } - if bytes == 0 { + if bytes == 0 && input_empty { break; } + assert!(bytes > 0 || decoded > 0); } assert_eq!(actual, expected, "capacity={capacity}, batch={batch_size}"); assert!(handler.blanks.lock().unwrap().is_empty()); @@ -2079,21 +2817,64 @@ mod tests { assert_eq!(records.len(), 3); assert_eq!(records[1].line_number, 2); assert_eq!(records[1].byte_offset, blank_offset); + assert_eq!(records[1].record, blank_raw); + } + } + } + } + + #[test] + fn test_blank_records_as_rows_with_multiple_columns() { + for csv in [b"1,2\n\n3,4\n".as_slice(), b"1,2\r\n\r\n3,4\r\n".as_slice()] { + for capacity in [1, 2, 5, 64] { + for batch_size in [1, 2, 8] { + let handler = Arc::new(CollectRecords { + materialize_blanks: true, + ..Default::default() + }); + let mut decoder = RecordDecoder::new(Reader::new(), 2, false) + .with_record_error_handler(Some(handler.clone())); + let mut reader = BufReader::with_capacity(capacity, Cursor::new(csv)); + let mut actual = Vec::new(); + loop { + let buf = reader.fill_buf().unwrap(); + let (_, bytes) = decoder.decode(buf, batch_size - decoder.len()).unwrap(); + reader.consume(bytes); + if decoder.len() == batch_size || bytes == 0 { + actual.extend(decoder.flush().unwrap().iter().map(|record| { + (record.get(0).to_owned(), record.get(1).to_owned()) + })); + } + if bytes == 0 { + break; + } + } + assert_eq!( + actual, + [ + ("1".into(), "2".into()), + ("".into(), "".into()), + ("3".into(), "4".into()) + ], + "capacity={capacity}, batch={batch_size}" + ); + assert!(handler.blanks.lock().unwrap().is_empty()); + assert!(handler.errors.lock().unwrap().is_empty()); } } } } #[test] - fn test_blank_records_as_rows_requires_one_column() { + fn test_blank_records_as_rows_requires_nonempty_schema() { let handler = Arc::new(CollectRecords { materialize_blanks: true, ..Default::default() }); let mut decoder = - RecordDecoder::new(Reader::new(), 2, false).with_record_error_handler(Some(handler)); + RecordDecoder::new(Reader::new(), 0, false).with_record_error_handler(Some(handler)); let error = decoder.decode(b"\n", 1).unwrap_err(); - assert!(error.to_string().contains("one-column schema")); + assert!(error.to_string().contains("at least one column")); } #[test] @@ -2449,6 +3230,34 @@ mod tests { assert_eq!(decoder.flush().unwrap().get(0).get(0), ""); } + #[test] + fn test_positioned_one_column_bom_with_and_without_blank_header() { + let handler = Arc::new(CollectRecords { + materialize_blanks: true, + ..Default::default() + }); + let mut decoder = RecordDecoder::new(Reader::new(), 1, false) + .with_initial_source_position(2, 11) + .with_record_error_handler(Some(handler)); + let csv = b"\xef\xbb\xbfvalue\n"; + assert_eq!(decoder.decode(csv, 1).unwrap(), (1, csv.len())); + assert_eq!(decoder.flush().unwrap().get(0).get(0), "value"); + + let handler = Arc::new(CollectRecords { + materialize_blanks: true, + ..Default::default() + }); + let mut decoder = RecordDecoder::new(Reader::new(), 1, false) + .with_initial_source_position(2, 11) + .with_record_error_handler(Some(handler)); + let csv = b"\xef\xbb\xbf\nvalue\n"; + assert_eq!(decoder.decode_validating_header(csv).unwrap(), (0, 3)); + assert_eq!(decoder.decode_validating_header(&csv[3..]).unwrap(), (1, 1)); + assert_eq!(decoder.flush().unwrap().get(0).get(0), ""); + assert_eq!(decoder.decode(&csv[4..], 1).unwrap(), (1, 6)); + assert_eq!(decoder.flush().unwrap().get(0).get(0), "value"); + } + #[test] fn test_one_column_blank_record_respects_filter() { #[derive(Debug)] @@ -2528,6 +3337,137 @@ mod tests { } } + #[test] + fn test_recovery_reports_lf_blank_after_split_crlf_data_row() { + let csv = b"1,ok\r\n\n2,bad,extra\n"; + for capacity in [1, 2, 64] { + let handler = Arc::new(CollectRecords { + report_blanks: true, + ..Default::default() + }); + let mut decoder = RecordDecoder::new(Reader::new(), 2, false) + .with_record_error_handler(Some(handler.clone())); + let mut reader = BufReader::with_capacity(capacity, Cursor::new(csv)); + loop { + let buf = reader.fill_buf().unwrap(); + let (_, bytes) = decoder.decode(buf, 10 - decoder.len()).unwrap(); + reader.consume(bytes); + if bytes == 0 { + break; + } + } + assert_eq!( + *handler.blanks.lock().unwrap(), + [OwnedBlankRecord { + line_number: 2, + source_line_number: 2, + byte_offset: 6, + record: b"\n".to_vec(), + }], + "capacity={capacity}" + ); + let errors = handler.errors.lock().unwrap(); + assert_eq!(errors[0].line_number, 3); + assert_eq!(errors[0].source_line_number, 3); + assert_eq!(errors[0].byte_offset, 7); + } + } + + #[test] + fn test_blank_record_source_position_after_skipped_prefix() { + let handler = Arc::new(CollectRecords { + report_blanks: true, + ..Default::default() + }); + let mut decoder = RecordDecoder::new(Reader::new(), 2, false) + .with_initial_source_position(2, 11) + .with_record_error_handler(Some(handler.clone())); + let csv = b"\r\n1,ok\r\n"; + assert_eq!(decoder.decode(csv, 2).unwrap(), (1, csv.len())); + assert_eq!( + *handler.blanks.lock().unwrap(), + [OwnedBlankRecord { + line_number: 1, + source_line_number: 3, + byte_offset: 11, + record: b"\r\n".to_vec(), + }] + ); + } + + #[test] + fn test_materialized_blank_record_has_source_line_after_skipped_prefix() { + #[derive(Debug, Default)] + struct MaterializedSourceLines(Mutex>); + + impl CsvRecordErrorHandler for MaterializedSourceLines { + fn handle(&self, _error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + Ok(()) + } + + fn blank_records_are_rows(&self) -> bool { + true + } + + fn needs_source_line_number(&self) -> bool { + true + } + + fn handle_record_with_source_line( + &self, + record: &CsvRecord<'_>, + source_line_number: usize, + ) -> Result<(), ArrowError> { + self.0.lock().unwrap().push(OwnedSourceRecord { + line_number: record.line_number, + source_line_number, + byte_offset: record.byte_offset, + record: record.record.to_vec(), + }); + Ok(()) + } + } + + for (columns, csv) in [(1, b"\r\nx\n".as_slice()), (2, b"\r\nx,y\n".as_slice())] { + for capacity in [1, 2, 64] { + let handler = Arc::new(MaterializedSourceLines::default()); + let mut decoder = RecordDecoder::new(Reader::new(), columns, false) + .with_initial_source_position(2, 11) + .with_record_error_handler(Some(handler.clone())); + let mut reader = BufReader::with_capacity(capacity, Cursor::new(csv)); + while decoder.len() < 2 { + let buf = reader.fill_buf().unwrap(); + let (decoded, bytes) = decoder.decode(buf, 2 - decoder.len()).unwrap(); + assert!(bytes > 0 || decoded > 0); + reader.consume(bytes); + } + let second_raw = if columns == 1 { + b"x\n".as_slice() + } else { + b"x,y\n".as_slice() + }; + assert_eq!( + *handler.0.lock().unwrap(), + [ + OwnedSourceRecord { + line_number: 1, + source_line_number: 3, + byte_offset: 11, + record: b"\r\n".to_vec(), + }, + OwnedSourceRecord { + line_number: 2, + source_line_number: 4, + byte_offset: 13, + record: second_raw.to_vec(), + }, + ], + "columns={columns}, capacity={capacity}" + ); + } + } + } + #[test] fn test_blank_records_do_not_exhaust_per_record_byte_limit() { let csv = b"\n\n\n\n\n\n\n\n\n1,ok\n"; @@ -2611,8 +3551,7 @@ mod tests { }); let mut decoder = RecordDecoder::new(Reader::new(), 1, false).with_record_error_handler(Some(handler)); - assert_eq!(decoder.decode(b"\r", 1).unwrap(), (1, 1)); - decoder.flush().unwrap(); + assert_eq!(decoder.decode(b"\r", 1).unwrap(), (0, 1)); assert!( decoder .decode(b"\n", 1) @@ -3046,6 +3985,42 @@ mod tests { } } + #[test] + fn test_successful_record_source_lines_after_skipped_prefix() { + let csv = b"\"value\ncontinued\",ok\n2,good\n"; + let handler = Arc::new(SourceLines::default()); + let mut decoder = RecordDecoder::new(Reader::new(), 2, false) + .with_initial_source_position(1, 7) + .with_record_error_handler(Some(handler.clone())); + assert_eq!(decoder.decode(csv, 2).unwrap(), (2, csv.len())); + assert_eq!( + *handler.0.lock().unwrap(), + [(1, 2, 7), (2, 4, 7 + b"\"value\ncontinued\",ok\n".len())] + ); + } + + #[test] + fn test_successful_record_source_lines_split_crlf_custom_terminator() { + let csv = b"1,a\r\nb|2,ok|"; + let mut builder = csv_core::ReaderBuilder::new(); + builder.terminator(csv_core::Terminator::Any(b'|')); + let handler = Arc::new(SourceLines::default()); + let mut decoder = RecordDecoder::new(builder.build(), 2, false) + .with_ignored_prefix_options(None, Some(b'|')) + .with_record_error_handler(Some(handler.clone())); + let mut reader = BufReader::with_capacity(1, Cursor::new(csv)); + while decoder.len() < 2 { + let buf = reader.fill_buf().unwrap(); + let (_, consumed) = decoder.decode(buf, 2 - decoder.len()).unwrap(); + reader.consume(consumed); + if consumed == 0 { + break; + } + } + assert_eq!(decoder.len(), 2); + assert_eq!(*handler.0.lock().unwrap(), [(1, 1, 0), (2, 2, 7)]); + } + #[test] fn test_invalid_utf8_handler_skips_only_malformed_record_across_chunks() { let csv = b"1,ok\n2,\xff\n3,after\n"; @@ -3587,6 +4562,45 @@ mod tests { } } + #[test] + #[ignore = "focused performance check for opt-in multi-column CSV recovery"] + fn bench_multicolumn_blank_as_rows_fast_path() { + #[derive(Debug)] + struct Noop { + materialize: bool, + } + + impl CsvRecordErrorHandler for Noop { + fn handle(&self, _error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + Ok(()) + } + + fn blank_records_are_rows(&self) -> bool { + self.materialize + } + } + + const ROWS: usize = 100_000; + let csv = b"1,value\n".repeat(ROWS); + for materialize in [false, true] { + let handler = Arc::new(Noop { materialize }); + let mut elapsed = Vec::new(); + for _ in 0..8 { + let mut decoder = RecordDecoder::new(Reader::new(), 2, false) + .with_utf8_preserving_syntax(true) + .with_record_error_handler(Some(handler.clone())); + let started = Instant::now(); + let (rows, bytes) = decoder.decode(&csv, ROWS).unwrap(); + let records = decoder.flush().unwrap(); + assert_eq!(rows, ROWS); + assert_eq!(bytes, csv.len()); + assert_eq!(records.len(), ROWS); + elapsed.push(started.elapsed()); + } + println!("CSV multi-column materialize={materialize}: {elapsed:?}"); + } + } + #[test] #[ignore = "focused performance check for bounded CSV comment preflight"] fn bench_one_column_blank_as_rows_small_batches_with_comments() {