From c992af1276ac5aabac2df59b89d6993525b8dbfc Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 18 Sep 2026 01:05:31 +0300 Subject: [PATCH] Add recoverable CSV record error handler --- arrow-csv/src/lib.rs | 6 +- arrow-csv/src/reader/mod.rs | 36 ++++- arrow-csv/src/reader/records.rs | 250 ++++++++++++++++++++++++++++++++ 3 files changed, 288 insertions(+), 4 deletions(-) diff --git a/arrow-csv/src/lib.rs b/arrow-csv/src/lib.rs index 4c4b04098175..ea70f205f7e5 100644 --- a/arrow-csv/src/lib.rs +++ b/arrow-csv/src/lib.rs @@ -29,9 +29,9 @@ pub mod reader; pub mod writer; -pub use self::reader::Reader; -pub use self::reader::ReaderBuilder; -pub use self::reader::infer_schema_from_files; +pub use self::reader::{ + CsvRecordError, CsvRecordErrorHandler, Reader, ReaderBuilder, infer_schema_from_files, +}; pub use self::writer::QuoteStyle; pub use self::writer::Writer; pub use self::writer::WriterBuilder; diff --git a/arrow-csv/src/reader/mod.rs b/arrow-csv/src/reader/mod.rs index 51e855f7dfc3..56f916512936 100644 --- a/arrow-csv/src/reader/mod.rs +++ b/arrow-csv/src/reader/mod.rs @@ -182,6 +182,27 @@ use crate::map_csv_error; use crate::reader::records::{RecordDecoder, StringRecords}; use arrow_array::timezone::Tz; +/// Metadata for a CSV record whose field count does not match the reader schema. +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +pub struct CsvRecordError<'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. + pub byte_offset: usize, + /// Number of fields required by the reader schema. + pub expected_fields: usize, + /// Number of fields found in the record. + pub actual_fields: usize, + /// Original record bytes, including its record terminator when present. + pub record: &'a [u8], +} + +/// Receives malformed CSV records that should be skipped instead of aborting the scan. +pub trait CsvRecordErrorHandler: Debug + Send + Sync { + /// Handle one malformed record. Returning an error aborts the scan. + fn handle(&self, error: &CsvRecordError<'_>) -> Result<(), ArrowError>; +} + /// Order should match [`InferredDataType`] static REGEX_SET: LazyLock = LazyLock::new(|| { RegexSet::new([ @@ -1162,6 +1183,8 @@ pub struct ReaderBuilder { bounds: Bounds, /// Optional projection for which columns to load (zero-based column indices) projection: Option>, + /// Optional handler for records whose field count differs from the schema. + record_error_handler: Option>, } impl ReaderBuilder { @@ -1194,6 +1217,7 @@ impl ReaderBuilder { batch_size: 1024, bounds: None, projection: None, + record_error_handler: None, } } @@ -1283,6 +1307,15 @@ impl ReaderBuilder { self } + /// Skip records whose field count differs from the schema and report them to `handler`. + /// + /// The default strict path does not retain raw record bytes and is unchanged when no + /// handler is configured. + pub fn with_record_error_handler(mut self, handler: Arc) -> Self { + self.record_error_handler = Some(handler); + self + } + /// Create a new `Reader` from a non-buffered reader /// /// If `R: BufRead` consider using [`Self::build_buffered`] to avoid unnecessary additional @@ -1306,7 +1339,8 @@ impl ReaderBuilder { delimiter, self.schema.fields().len(), self.format.truncated_rows, - ); + ) + .with_record_error_handler(self.record_error_handler); let header = self.format.header as usize; diff --git a/arrow-csv/src/reader/records.rs b/arrow-csv/src/reader/records.rs index a6d54f867c1d..a55aaf341829 100644 --- a/arrow-csv/src/reader/records.rs +++ b/arrow-csv/src/reader/records.rs @@ -17,6 +17,9 @@ use arrow_schema::ArrowError; use csv_core::{ReadRecordResult, Reader}; +use std::sync::Arc; + +use super::{CsvRecordError, CsvRecordErrorHandler}; /// The estimated length of a field in bytes const AVERAGE_FIELD_SIZE: usize = 8; @@ -24,6 +27,12 @@ const AVERAGE_FIELD_SIZE: usize = 8; /// The minimum amount of data in a single read const MIN_CAPACITY: usize = 1024; +/// Prevent malformed records from growing the field-offset buffer without bound. +/// +/// This is larger than typical schema limits and only applies after a record has already +/// exceeded the configured schema width. +const MAX_RECOVERABLE_EXCESS_FIELDS: usize = 16_384; + /// [`RecordDecoder`] provides a push-based interface to decoder [`StringRecords`] #[derive(Debug)] pub struct RecordDecoder { @@ -62,6 +71,20 @@ pub struct RecordDecoder { /// Default value is false /// When enabled fills in missing columns with null truncated_rows: bool, + + /// Optional recovery hook. Keeping this `None` selects the allocation-free strict path. + record_error_handler: Option>, + + /// Raw bytes for the current record, retained only by the recovery path. + record_bytes: Vec, + + /// Total number of input bytes consumed by the recovery path. + stream_offset: usize, + + /// Start offsets for rolling back a malformed record from decoded output. + record_byte_offset: usize, + record_data_start: usize, + record_offsets_start: usize, } impl RecordDecoder { @@ -77,13 +100,40 @@ impl RecordDecoder { data: vec![], num_rows: 0, truncated_rows, + record_error_handler: None, + record_bytes: vec![], + stream_offset: 0, + record_byte_offset: 0, + record_data_start: 0, + record_offsets_start: 1, } } + pub fn with_record_error_handler( + mut self, + handler: Option>, + ) -> Self { + self.record_error_handler = handler; + self + } + /// Decodes records from `input` returning the number of records and bytes read /// /// 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> { + match self.record_error_handler.clone() { + Some(handler) => { + self.decode_with_record_error_handler(input, to_read, handler.as_ref()) + } + None => self.decode_strict(input, to_read), + } + } + + fn decode_strict( + &mut self, + input: &[u8], + to_read: usize, + ) -> Result<(usize, usize), ArrowError> { if to_read == 0 { return Ok((0, 0)); } @@ -170,6 +220,122 @@ impl RecordDecoder { } } + fn decode_with_record_error_handler( + &mut self, + input: &[u8], + to_read: usize, + handler: &dyn CsvRecordErrorHandler, + ) -> Result<(usize, usize), ArrowError> { + if to_read == 0 { + return Ok((0, 0)); + } + + self.offsets + .resize(self.offsets_len + to_read * self.num_columns, 0); + let max_offsets_len = self + .offsets + .len() + .checked_add(MAX_RECOVERABLE_EXCESS_FIELDS) + .ok_or_else(|| ArrowError::CsvError("CSV field offset capacity overflowed".into()))?; + let mut input_offset = 0; + let mut read = 0; + + loop { + let remaining_rows = to_read - read; + let capacity = remaining_rows * self.num_columns * AVERAGE_FIELD_SIZE; + let estimated_data = capacity.max(MIN_CAPACITY); + self.data.resize(self.data_len + estimated_data, 0); + + loop { + let record_input_start = input_offset; + let (result, bytes_read, bytes_written, end_positions) = + self.delimiter.read_record( + &input[input_offset..], + &mut self.data[self.data_len..], + &mut self.offsets[self.offsets_len..], + ); + + self.current_field += end_positions; + self.offsets_len += end_positions; + input_offset += bytes_read; + self.data_len += bytes_written; + self.stream_offset = + self.stream_offset.checked_add(bytes_read).ok_or_else(|| { + ArrowError::CsvError("CSV input byte offset overflowed".into()) + })?; + self.record_bytes + .extend_from_slice(&input[record_input_start..input_offset]); + + match result { + ReadRecordResult::End | ReadRecordResult::InputEmpty => { + return Ok((read, input_offset)); + } + ReadRecordResult::OutputFull => break, + ReadRecordResult::OutputEndsFull => { + let new_len = self + .offsets + .len() + .checked_add(self.num_columns.max(1)) + .ok_or_else(|| { + ArrowError::CsvError("CSV field offset capacity overflowed".into()) + })? + .min(max_offsets_len); + if new_len == self.offsets.len() { + return Err(ArrowError::CsvError(format!( + "malformed CSV record on line {} exceeds the recovery limit of {} excess fields", + self.line_number, MAX_RECOVERABLE_EXCESS_FIELDS + ))); + } + self.offsets.resize(new_len, 0); + } + ReadRecordResult::Record => { + if self.current_field != self.num_columns + && !(self.truncated_rows && self.current_field < self.num_columns) + { + handler.handle(&CsvRecordError { + line_number: self.line_number, + byte_offset: self.record_byte_offset, + expected_fields: self.num_columns, + actual_fields: self.current_field, + record: &self.record_bytes, + })?; + self.data_len = self.record_data_start; + self.offsets_len = self.record_offsets_start; + self.current_field = 0; + self.line_number += 1; + self.record_bytes.clear(); + self.record_byte_offset = self.stream_offset; + continue; + } + + if self.current_field < self.num_columns { + let fill_count = self.num_columns - self.current_field; + let fill_value = self.offsets[self.offsets_len - 1]; + self.offsets[self.offsets_len..self.offsets_len + fill_count] + .fill(fill_value); + self.offsets_len += fill_count; + } + read += 1; + self.current_field = 0; + self.line_number += 1; + self.num_rows += 1; + self.record_bytes.clear(); + self.record_byte_offset = self.stream_offset; + self.record_data_start = self.data_len; + self.record_offsets_start = self.offsets_len; + + if read == to_read { + return Ok((read, input_offset)); + } + if input.len() == input_offset { + return Ok((read, input_offset)); + } + } + } + } + } + } + /// Returns the current number of buffered records pub fn len(&self) -> usize { self.num_rows @@ -186,6 +352,8 @@ impl RecordDecoder { self.offsets_len = 1; self.data_len = 0; self.num_rows = 0; + self.record_data_start = 0; + self.record_offsets_start = 1; } /// Flushes the current contents of the reader @@ -240,6 +408,8 @@ impl RecordDecoder { self.offsets_len = 1; self.data_len = 0; self.num_rows = 0; + self.record_data_start = 0; + self.record_offsets_start = 1; Ok(StringRecords { num_rows, @@ -313,8 +483,36 @@ impl std::fmt::Display for StringRecord<'_> { #[cfg(test)] mod tests { use crate::reader::records::RecordDecoder; + use crate::reader::{CsvRecordError, CsvRecordErrorHandler}; + use arrow_schema::ArrowError; use csv_core::Reader; use std::io::{BufRead, BufReader, Cursor}; + use std::sync::{Arc, Mutex}; + + #[derive(Debug, Clone, Eq, PartialEq)] + struct OwnedRecordError { + line_number: usize, + byte_offset: usize, + expected_fields: usize, + actual_fields: usize, + record: Vec, + } + + #[derive(Debug, Default)] + struct CollectRecordErrors(Mutex>); + + impl CsvRecordErrorHandler for CollectRecordErrors { + fn handle(&self, error: &CsvRecordError<'_>) -> Result<(), ArrowError> { + self.0.lock().unwrap().push(OwnedRecordError { + line_number: error.line_number, + byte_offset: error.byte_offset, + expected_fields: error.expected_fields, + actual_fields: error.actual_fields, + record: error.record.to_vec(), + }); + Ok(()) + } + } #[test] fn test_basic() { @@ -387,6 +585,58 @@ mod tests { assert_eq!(err, expected); } + #[test] + fn test_invalid_fields_handler_skips_records_across_input_chunks() { + let csv = b"1,ok\n2,extra,value\n3\n4,after\n"; + let handler = Arc::new(CollectRecordErrors::default()); + let mut decoder = RecordDecoder::new(Reader::new(), 2, false) + .with_record_error_handler(Some(handler.clone())); + let mut reader = BufReader::with_capacity(3, Cursor::new(csv)); + + loop { + let buf = reader.fill_buf().unwrap(); + let (_, bytes) = decoder.decode(buf, 4 - decoder.len()).unwrap(); + reader.consume(bytes); + if bytes == 0 || decoder.len() == 4 { + break; + } + } + + let records = decoder.flush().unwrap(); + let actual = records + .iter() + .map(|record| [record.get(0).to_owned(), record.get(1).to_owned()]) + .collect::>(); + assert_eq!( + actual, + [ + [String::from("1"), String::from("ok")], + [String::from("4"), String::from("after")], + ] + ); + + let errors = handler.0.lock().unwrap(); + assert_eq!( + *errors, + [ + OwnedRecordError { + line_number: 2, + byte_offset: 5, + expected_fields: 2, + actual_fields: 3, + record: b"2,extra,value\n".to_vec(), + }, + OwnedRecordError { + line_number: 3, + byte_offset: 19, + expected_fields: 2, + actual_fields: 1, + record: b"3\n".to_vec(), + }, + ] + ); + } + #[test] fn test_skip_insufficient_rows() { let csv = "a\nv\n";