AndreaBozzo commented on code in PR #10903:
URL: https://github.com/apache/arrow-rs/pull/10903#discussion_r3882697907
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -351,6 +362,16 @@ impl Format {
self
}
+ /// Whether to ignore extra fields when parsing.
+ ///
+ /// By default this is set to `ExtraFields::Error` and will error if the
CSV rows have more
+ /// columns than expected. When set to `ExtraFields::Ignore` then it will
allow records with
+ /// more than the expected number of columns and ignore the extra fields.
+ pub fn with_extra_fields(mut self, extra_fields: ExtraFields) -> Self {
+ self.extra_fields = extra_fields;
+ self
Review Comment:
`ExtraFields` is exposed here on `Format`, but `Format::infer_schema()`
builds its `csv::Reader` through `build_reader()`, which still configures
`flexible(self.truncated_rows)` only (`mod.rs`, unchanged in this PR). So
`with_extra_fields(ExtraFields::Ignore)` still rejects an over-long row during
inference, before the configured `RecordDecoder` is ever used:
```
Format::default().with_header(true).with_extra_fields(ExtraFields::Ignore)
.infer_schema(Cursor::new("a,b
1,2,3
4,5
"), None)
Err(CsvError("Encountered unequal lengths between records on CSV file.
Expected 2 records, found 3 records at line 2"))
```
Inference already ignores values past the header width
(`take(header_length)` with `record.get(i)`), so I think one line covers it:
```rust
builder.flexible(self.truncated_rows || self.extra_fields ==
ExtraFields::Ignore);
```
I tried this locally and the full `arrow-csv` suite still passes, with
inference returning the expected 2-field schema.
Alternatively, the docs here could clarify that `with_extra_fields` applies
only when decoding against an explicit schema.
##########
arrow-csv/src/reader/records.rs:
##########
@@ -115,11 +160,12 @@ impl RecordDecoder {
// Try to read a record
loop {
+ let ends_bound = self.offsets_len + (self.num_columns -
self.current_field);
Review Comment:
Clamping the ends slice to the schema width also changes the default
`ExtraFields::Error` path. `OutputEndsFull` now fires on the first extra field,
so the error is produced by the new branch below rather than by the `Record`
arm's exact-count check:
| row | before | after |
| --- | --- | --- |
| `1,2,3` | `expected 2 got 3` | `expected 2 got more than 2` |
| `1,2,3,4,5,6` | `expected 2 got 6` | `expected 2 got more than 2` |
The actual field count is lost, and rows with different numbers of extra
fields now report identically.
If preserving that diagnostic matters, the bound could be conditional:
```rust
let ends_bound = match self.extra_fields {
ExtraFields::Ignore => self.offsets_len + (self.num_columns -
self.current_field),
ExtraFields::Error => self.offsets.len(),
};
```
I tried this locally: it restores both messages above and leaves the
`Ignore` path unchanged (the extra-field tests and the rest of the `arrow-csv`
suite still pass, apart from the two assertions noted below).
##########
arrow-csv/src/reader/records.rs:
##########
@@ -96,16 +110,47 @@ impl RecordDecoder {
return Ok((0, 0));
}
- // Reserve sufficient capacity in offsets
- self.offsets
- .resize(self.offsets_len + to_read * self.num_columns, 0);
-
// The current offset into `input`
let mut input_offset = 0;
// The number of rows decoded in this pass
let mut read = 0;
+ // Resume skipping extra fields if we were in the middle of it from a
previous chunk
+ if self.skipping_extra_fields {
+ let mut dummy_data = [0u8; 128];
+ let mut dummy_ends = [0usize; 1];
+ loop {
+ let (res, b_read, _, _) = self.delimiter.read_record(
+ &input[input_offset..],
+ &mut dummy_data,
+ &mut dummy_ends,
+ );
+ input_offset += b_read;
+ match res {
+ ReadRecordResult::Record => {
+ self.skipping_extra_fields = false;
+ read += 1;
+ self.current_field = 0;
+ self.line_number += 1;
+ self.num_rows += 1;
+ break;
+ }
+ ReadRecordResult::OutputFull |
ReadRecordResult::OutputEndsFull => {}
+ ReadRecordResult::End | ReadRecordResult::InputEmpty => {
+ return Ok((read, input_offset));
+ }
+ }
+ }
+ if read == to_read || input.len() == input_offset {
+ return Ok((read, input_offset));
Review Comment:
Minor: this resume loop is duplicated almost verbatim in the
`ExtraFields::Ignore` arm below. A small helper returning the same `(read,
input_offset)` decision would keep the two copies of the state machine from
drifting.
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -2647,6 +2675,132 @@ mod tests {
assert!(c.is_null(3));
}
+ #[test]
+ fn test_extra_fields_ignore() {
+ let data = "a,b\n1,2,3,4\n5,6\n7,8,9";
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ ]));
+
+ let mut reader = ReaderBuilder::new(schema.clone())
+ .with_header(true)
+ .with_extra_fields(ExtraFields::Ignore)
+ .build(Cursor::new(data))
+ .unwrap();
+
+ let batch = reader.next().unwrap().unwrap();
+ assert_eq!(batch.num_rows(), 3);
+ assert_eq!(batch.num_columns(), 2);
+
+ let col_a = batch.column(0).as_primitive::<Int32Type>();
+ assert_eq!(col_a.value(0), 1);
+ assert_eq!(col_a.value(1), 5);
+ assert_eq!(col_a.value(2), 7);
+
+ let col_b = batch.column(1).as_primitive::<Int32Type>();
+ assert_eq!(col_b.value(0), 2);
+ assert_eq!(col_b.value(1), 6);
+ assert_eq!(col_b.value(2), 8);
+ }
+
+ #[test]
+ fn test_extra_fields_default_errors() {
+ let data = "a,b\n1,2,3\n4,5";
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ ]));
+
+ // No with_extra_fields called (should default to Error)
+ let mut reader = ReaderBuilder::new(schema.clone())
+ .with_header(true)
+ .build(Cursor::new(data))
+ .unwrap();
+
+ let result = reader.next();
+ assert!(match result {
+ Some(Err(ArrowError::CsvError(e))) => e.contains("got more than"),
Review Comment:
These assertions encode the new wording, so the change to the default
`Error` path described in my comment on `records.rs` is locked in rather than
caught. With a mode-conditional `ends_bound` these would go back to `expected 2
got 3`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
