Skip to content

Commit 9f37683

Browse files
authored
feat(arrow_csv): add header validation option (#10144)
# Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. --> - Closes #10143. # Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> Explained in the issue. # What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> - Adds a new `ReaderBuilder` and `Format` method .with_header_validation(bool) to enable CSV header validation - Implements header row validation, which verifies each value in the first CSV row against the schema to ensure the field names match the header columns. # Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? If this PR claims a performance improvement, please include evidence such as benchmark results. --> Corresponding tests added. # Are there any user-facing changes? <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. If there are any breaking changes to public APIs, please call them out. --> The `ReaderBuilder` struct gets a new method added, `.with_header_validation(bool)`. There is no breaking change, since the validation behavior is disabled by default.
1 parent d0d3c52 commit 9f37683

1 file changed

Lines changed: 135 additions & 0 deletions

File tree

arrow-csv/src/reader/mod.rs

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,7 @@ impl InferredDataType {
273273
#[derive(Debug, Clone, Default)]
274274
pub struct Format {
275275
header: bool,
276+
header_validation: bool,
276277
delimiter: Option<u8>,
277278
escape: Option<u8>,
278279
quote: Option<u8>,
@@ -291,6 +292,16 @@ impl Format {
291292
self
292293
}
293294

295+
/// Specify whether to validate the CSV header against the schema, defaults to `false`
296+
///
297+
/// When `true`, the first row gets validated against the schema before any data is read
298+
///
299+
/// Only applies when [`Self::with_header`] is set to `true`
300+
pub fn with_header_validation(mut self, validate_header: bool) -> Self {
301+
self.header_validation = validate_header;
302+
self
303+
}
304+
294305
/// Specify a custom delimiter character, defaults to comma `','`
295306
pub fn with_delimiter(mut self, delimiter: u8) -> Self {
296307
self.delimiter = Some(delimiter);
@@ -610,6 +621,9 @@ pub struct Decoder {
610621
/// Rows to skip
611622
to_skip: usize,
612623

624+
/// Whether to validate the first skipped row against the schema
625+
header_validation: bool,
626+
613627
/// Current line number
614628
line_number: usize,
615629

@@ -635,6 +649,20 @@ impl Decoder {
635649
/// network sources such as object storage
636650
pub fn decode(&mut self, buf: &[u8]) -> Result<usize, ArrowError> {
637651
if self.to_skip != 0 {
652+
if self.header_validation {
653+
let (skipped, bytes) = self.record_decoder.decode(buf, 1)?;
654+
655+
if skipped == 0 {
656+
return Ok(bytes);
657+
}
658+
659+
let rows = self.record_decoder.flush()?;
660+
validate_header(&rows, self.schema.fields())?;
661+
self.header_validation = false;
662+
self.to_skip -= 1;
663+
return Ok(bytes);
664+
}
665+
638666
// Skip in units of `to_read` to avoid over-allocating buffers
639667
let to_skip = self.to_skip.min(self.batch_size);
640668
let (skipped, bytes) = self.record_decoder.decode(buf, to_skip)?;
@@ -678,6 +706,24 @@ impl Decoder {
678706
}
679707
}
680708

709+
fn validate_header(rows: &StringRecords<'_>, fields: &Fields) -> Result<(), ArrowError> {
710+
let header = rows.iter().next().ok_or_else(|| {
711+
ArrowError::CsvError("CSV header validation failed: no header row found".to_string())
712+
})?;
713+
714+
for (idx, field) in fields.iter().enumerate() {
715+
let actual = header.get(idx);
716+
let expected = field.name();
717+
if actual != expected {
718+
return Err(ArrowError::CsvError(format!(
719+
"CSV header does not match schema at column {idx}: expected {expected:?} but found {actual:?}"
720+
)));
721+
}
722+
}
723+
724+
Ok(())
725+
}
726+
681727
/// Parses a slice of [`StringRecords`] into a [RecordBatch]
682728
fn parse(
683729
rows: &StringRecords<'_>,
@@ -1154,6 +1200,14 @@ impl ReaderBuilder {
11541200
self
11551201
}
11561202

1203+
/// Set whether to validate the CSV header against the schema
1204+
///
1205+
/// This option only applies when [`Self::with_header`] is set to `true`, and defaults to `false`
1206+
pub fn with_header_validation(mut self, validate_header: bool) -> Self {
1207+
self.format.header_validation = validate_header;
1208+
self
1209+
}
1210+
11571211
/// Overrides the [Format] of this [ReaderBuilder]
11581212
pub fn with_format(mut self, format: Format) -> Self {
11591213
self.format = format;
@@ -1261,6 +1315,7 @@ impl ReaderBuilder {
12611315
Decoder {
12621316
schema: self.schema,
12631317
to_skip: start,
1318+
header_validation: self.format.header && self.format.header_validation,
12641319
record_decoder,
12651320
line_number: start,
12661321
end,
@@ -2351,6 +2406,86 @@ mod tests {
23512406
}
23522407
}
23532408

2409+
#[test]
2410+
fn test_header_validation() {
2411+
let schema = Arc::new(Schema::new(vec![
2412+
Field::new("a", DataType::Int32, false),
2413+
Field::new("b", DataType::Int32, false),
2414+
]));
2415+
2416+
let csv = "a,c\n1,2\n";
2417+
let err = ReaderBuilder::new(schema.clone())
2418+
.with_header(true)
2419+
.with_header_validation(true)
2420+
.build_buffered(Cursor::new(csv.as_bytes()))
2421+
.unwrap()
2422+
.next()
2423+
.unwrap()
2424+
.unwrap_err()
2425+
.to_string();
2426+
assert_eq!(
2427+
err,
2428+
"Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"c\""
2429+
);
2430+
2431+
let batch = ReaderBuilder::new(schema)
2432+
.with_header(true)
2433+
.with_header_validation(false)
2434+
.build_buffered(Cursor::new(csv.as_bytes()))
2435+
.unwrap()
2436+
.next()
2437+
.unwrap()
2438+
.unwrap();
2439+
assert_eq!(batch.num_rows(), 1);
2440+
}
2441+
2442+
#[test]
2443+
fn test_header_validation_with_buffered_reader() {
2444+
let schema = Arc::new(Schema::new(vec![
2445+
Field::new("a", DataType::Int32, false),
2446+
Field::new("b", DataType::Int32, false),
2447+
]));
2448+
2449+
let csv = "a,b\n1,2\n";
2450+
let buffered = std::io::BufReader::with_capacity(1, Cursor::new(csv.as_bytes()));
2451+
let batch = ReaderBuilder::new(schema)
2452+
.with_header(true)
2453+
.with_header_validation(true)
2454+
.build_buffered(buffered)
2455+
.unwrap()
2456+
.next()
2457+
.unwrap()
2458+
.unwrap();
2459+
2460+
assert_eq!(batch.num_rows(), 1);
2461+
let a = batch.column(0).as_primitive::<Int32Type>();
2462+
assert_eq!(a.value(0), 1);
2463+
}
2464+
2465+
#[test]
2466+
fn test_header_validation_with_truncated_rows() {
2467+
let schema = Arc::new(Schema::new(vec![
2468+
Field::new("a", DataType::Int32, true),
2469+
Field::new("b", DataType::Int32, true),
2470+
]));
2471+
2472+
let csv = "a\n1\n";
2473+
let err = ReaderBuilder::new(schema.clone())
2474+
.with_header(true)
2475+
.with_header_validation(true)
2476+
.with_truncated_rows(true)
2477+
.build_buffered(Cursor::new(csv.as_bytes()))
2478+
.unwrap()
2479+
.next()
2480+
.unwrap()
2481+
.unwrap_err()
2482+
.to_string();
2483+
assert_eq!(
2484+
err,
2485+
"Csv error: CSV header does not match schema at column 1: expected \"b\" but found \"\"",
2486+
)
2487+
}
2488+
23542489
#[test]
23552490
fn test_null_boolean() {
23562491
let csv = "true,false\nFalse,True\n,True\nFalse,";

0 commit comments

Comments
 (0)