Jefffrey commented on code in PR #10903:
URL: https://github.com/apache/arrow-rs/pull/10903#discussion_r3912018934
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -351,6 +362,16 @@ impl Format {
self
}
+ /// Whether to ignore extra fields when parsing.
Review Comment:
```suggestion
/// How to deal with rows that have more columns than the expected
schema when parsing.
```
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -2681,6 +2859,128 @@ mod tests {
});
}
+ #[test]
+ fn test_truncated_rows_with_extra_fields() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ Field::new("c", DataType::Int32, true),
+ ]));
+
+ let data_short = "a,b,c\n1,2";
+ let data_extra = "a,b,c\n4,5,6,7";
+ let data_mixed = "a,b,c\n1,2\n4,5,6,7";
+
+ // 1. with_truncated_rows(true) + ExtraFields::Ignore
+ // - a row with fewer fields is padded
+ // - a row with extra fields has the extra fields ignored
+ let reader = ReaderBuilder::new(schema.clone())
+ .with_header(true)
+ .with_truncated_rows(true)
+ .with_extra_fields(ExtraFields::Ignore)
+ .build(Cursor::new(data_mixed))
+ .unwrap();
+ let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
+ let batch = &batches[0];
+ assert_eq!(batch.num_rows(), 2);
+ let col_a = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<arrow_array::Int32Array>()
+ .unwrap();
+ let col_c = batch
+ .column(2)
+ .as_any()
+ .downcast_ref::<arrow_array::Int32Array>()
+ .unwrap();
+ assert_eq!(col_a.value(0), 1);
+ assert!(col_c.is_null(0));
+ assert_eq!(col_a.value(1), 4);
+ assert_eq!(col_c.value(1), 6);
+
+ // 2. with_truncated_rows(true) + ExtraFields::Error
+ // - a short row is still padded
+ let reader = ReaderBuilder::new(schema.clone())
+ .with_header(true)
+ .with_truncated_rows(true)
+ .with_extra_fields(ExtraFields::Error)
+ .build(Cursor::new(data_short))
+ .unwrap();
+ let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
+ assert_eq!(batches[0].num_rows(), 1);
+
+ // - a row with extra fields still returns the exact field-count error
+ let reader = ReaderBuilder::new(schema.clone())
+ .with_header(true)
+ .with_truncated_rows(true)
+ .with_extra_fields(ExtraFields::Error)
+ .build(Cursor::new(data_extra))
+ .unwrap();
+ let err = reader.collect::<Result<Vec<_>, _>>().unwrap_err();
+ assert!(
+ err.to_string().contains("incorrect number of fields")
+ || err.to_string().contains("Expected 3"),
+ "{}",
Review Comment:
we should not have a `||` in the assert; we should clearly assert what we
expect, not two things we might expect
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -2647,6 +2675,156 @@ 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 3"),
Review Comment:
could we assert more of the error message to be a little more informative
here
##########
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:
i had codex do a review and it identified a similar issue that still
persists:
```rust
#[test]
fn test_infer_schema_extra_fields_ignore_rejects_truncated_rows() {
let csv = "a,b\n1\n";
let result = Format::default()
.with_header(true)
.with_extra_fields(ExtraFields::Ignore)
.infer_schema(Cursor::new(csv), None);
dbg!(&result);
}
```
ideally this schema inference would fail, since we have a row too short, but
when i tried it out it succeeds and infers columns `a` & `b`
so i think this still needs to be resolved?
##########
arrow-csv/src/reader/mod.rs:
##########
@@ -2681,6 +2859,128 @@ mod tests {
});
}
+ #[test]
+ fn test_truncated_rows_with_extra_fields() {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ Field::new("c", DataType::Int32, true),
+ ]));
+
+ let data_short = "a,b,c\n1,2";
+ let data_extra = "a,b,c\n4,5,6,7";
+ let data_mixed = "a,b,c\n1,2\n4,5,6,7";
+
+ // 1. with_truncated_rows(true) + ExtraFields::Ignore
+ // - a row with fewer fields is padded
+ // - a row with extra fields has the extra fields ignored
+ let reader = ReaderBuilder::new(schema.clone())
+ .with_header(true)
+ .with_truncated_rows(true)
+ .with_extra_fields(ExtraFields::Ignore)
+ .build(Cursor::new(data_mixed))
+ .unwrap();
+ let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
+ let batch = &batches[0];
+ assert_eq!(batch.num_rows(), 2);
+ let col_a = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<arrow_array::Int32Array>()
+ .unwrap();
Review Comment:
```suggestion
.as_primitive::<Int32Type>();
```
above we were already downcasting like this, so lets keep it consistent
(same for below)
--
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]
