kosiew commented on code in PR #27021:
URL: https://github.com/apache/datafusion/pull/27021#discussion_r4547422329
##########
datafusion/datasource-csv/src/source.rs:
##########
@@ -211,6 +247,301 @@ impl CsvSource {
Ok(builder)
}
+
+ /// A parser configured like the one arrow's CSV decoder builds, so the
+ /// header is read with the same quoting, escaping, BOM, blank line and
+ /// comment handling as the rows that follow it.
+ fn header_parser(&self) -> csv_core::Reader {
+ let mut builder = csv_core::ReaderBuilder::new();
+ builder
+ .delimiter(self.delimiter())
+ .quote(self.quote())
+ .escape(self.escape())
+ .comment(self.comment());
+ if let Some(terminator) = self.terminator() {
+ builder.terminator(csv_core::Terminator::Any(terminator));
+ }
+ builder.build()
+ }
+
+ /// Parses the header row at the start of `buf`. Returns `None` when the
+ /// header is not yet complete and more input may follow.
+ fn find_header(&self, buf: &[u8], eof: bool) -> Result<Option<HeaderRow>> {
+ // csv-core treats empty input as the end of the file, including what
+ // is left after it strips a BOM, so it must only see empty input at
+ // EOF.
+ if !eof && buf.strip_prefix(UTF8_BOM).unwrap_or(buf).is_empty() {
+ return Ok(None);
+ }
+ let mut parser = self.header_parser();
+ let mut input = buf;
+ let mut output = vec![0; 1024];
+ let mut ends = vec![0; 32];
+ let (mut output_len, mut ends_len) = (0, 0);
+ loop {
+ if !eof && input.is_empty() {
+ return Ok(None);
+ }
+ let (result, read, written, ends_written) = parser.read_record(
+ input,
+ &mut output[output_len..],
+ &mut ends[ends_len..],
+ );
+ input = &input[read..];
+ output_len += written;
+ ends_len += ends_written;
+ match result {
+ ReadRecordResult::InputEmpty if !eof => return Ok(None),
+ ReadRecordResult::InputEmpty => {}
+ ReadRecordResult::OutputFull => output.resize(output.len() *
2, 0),
+ ReadRecordResult::OutputEndsFull => ends.resize(ends.len() *
2, 0),
+ ReadRecordResult::Record => break,
+ ReadRecordResult::End => {
+ return Ok(Some(HeaderRow {
+ names: vec![],
+ start: buf.len(),
+ }));
+ }
+ }
+ }
+
+ let mut field_start = 0;
+ let names = ends[..ends_len]
+ .iter()
+ .map(|&field_end| {
+ let name =
std::str::from_utf8(&output[field_start..field_end]).map_err(
+ |e| exec_datafusion_err!("Invalid UTF-8 in CSV header:
{e}"),
+ )?;
+ field_start = field_end;
+ Ok(name.to_string())
+ })
+ .collect::<Result<_>>()?;
+ Ok(Some(HeaderRow {
+ names,
+ start: self.first_record_start(buf),
+ }))
+ }
+
+ /// Byte offset of the first record in `buf`, after a leading UTF-8 BOM,
+ /// blank records and comment lines. None of these can contain quotes, so
+ /// a plain byte scan finds the same position the parser starts reading
+ /// the header at. With no explicit terminator both `\r` and `\n` end a
+ /// record.
+ fn first_record_start(&self, buf: &[u8]) -> usize {
+ let terminator = self.terminator();
+ let comment = self.comment();
+ let is_terminator = |b: u8| match terminator {
+ Some(t) => b == t,
+ None => b == b'\n' || b == b'\r',
+ };
+ // csv-core's DFA only ends a comment at `\n` when `\n` is also a
+ // terminator; with any other terminator the comment runs to the end
+ // of the input.
+ let comment_ends_at_newline = matches!(terminator, None | Some(b'\n'));
+
+ let mut start = if buf.starts_with(UTF8_BOM) {
+ UTF8_BOM.len()
+ } else {
+ 0
+ };
+ while let Some(&first) = buf.get(start) {
+ if is_terminator(first) {
+ start += 1;
+ } else if comment == Some(first) {
+ let end = comment_ends_at_newline
+ .then(|| buf[start..].iter().position(|b| *b == b'\n'))
+ .flatten();
+ match end {
+ Some(end) => start += end + 1,
+ None => return buf.len(),
+ }
+ } else {
+ break;
+ }
+ }
+ start
+ }
+
+ /// Builds the mapping for a file with the given header, or `None` if the
+ /// file can be read positionally.
+ fn header_mapping(
+ &self,
+ header: &[String],
+ expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
+ ) -> Result<Option<Arc<HeaderMapping>>> {
+ // csv-core reads an unterminated comment at the end of the input as a
+ // record with one empty field
+ if header.is_empty() || header == [""] {
Review Comment:
One narrow edge case: `header == [""]` also matches a real quoted-empty
header, so `""
42
` with nullable `a INT` falls back to positional decoding and would emit
`42` instead of `NULL`. Could we distinguish comment-only input in
`find_header` using `self.first_record_start(buf) == buf.len()`, remove this
shortcut, and add a public read regression while preserving comment-only
behavior?
##########
datafusion/proto/tests/cases/plans/sources.rs:
##########
@@ -609,6 +609,44 @@ fn roundtrip_csv_scan_preserves_format_options() ->
Result<()> {
Ok(())
}
+#[test]
+fn roundtrip_csv_scan_preserves_match_columns_by_header() -> Result<()> {
Review Comment:
This test covers the scan-node boolean, but not the separate
`CsvOptions::match_columns_by_header` conversions in the logical file-format
codec and proto-common. As optional coverage, could we extend a logical CSV
roundtrip or add conversion tests for `None`, `Some(false)`, and `Some(true)`
to protect unset-versus-false behavior?
##########
datafusion/datasource-csv/src/source.rs:
##########
@@ -211,6 +247,301 @@ impl CsvSource {
Ok(builder)
}
+
+ /// A parser configured like the one arrow's CSV decoder builds, so the
+ /// header is read with the same quoting, escaping, BOM, blank line and
+ /// comment handling as the rows that follow it.
+ fn header_parser(&self) -> csv_core::Reader {
+ let mut builder = csv_core::ReaderBuilder::new();
+ builder
+ .delimiter(self.delimiter())
+ .quote(self.quote())
+ .escape(self.escape())
+ .comment(self.comment());
+ if let Some(terminator) = self.terminator() {
+ builder.terminator(csv_core::Terminator::Any(terminator));
+ }
+ builder.build()
+ }
+
+ /// Parses the header row at the start of `buf`. Returns `None` when the
+ /// header is not yet complete and more input may follow.
+ fn find_header(&self, buf: &[u8], eof: bool) -> Result<Option<HeaderRow>> {
+ // csv-core treats empty input as the end of the file, including what
+ // is left after it strips a BOM, so it must only see empty input at
+ // EOF.
+ if !eof && buf.strip_prefix(UTF8_BOM).unwrap_or(buf).is_empty() {
+ return Ok(None);
+ }
+ let mut parser = self.header_parser();
+ let mut input = buf;
+ let mut output = vec![0; 1024];
+ let mut ends = vec![0; 32];
+ let (mut output_len, mut ends_len) = (0, 0);
+ loop {
+ if !eof && input.is_empty() {
+ return Ok(None);
+ }
+ let (result, read, written, ends_written) = parser.read_record(
+ input,
+ &mut output[output_len..],
+ &mut ends[ends_len..],
+ );
+ input = &input[read..];
+ output_len += written;
+ ends_len += ends_written;
+ match result {
+ ReadRecordResult::InputEmpty if !eof => return Ok(None),
+ ReadRecordResult::InputEmpty => {}
+ ReadRecordResult::OutputFull => output.resize(output.len() *
2, 0),
+ ReadRecordResult::OutputEndsFull => ends.resize(ends.len() *
2, 0),
+ ReadRecordResult::Record => break,
+ ReadRecordResult::End => {
+ return Ok(Some(HeaderRow {
+ names: vec![],
+ start: buf.len(),
+ }));
+ }
+ }
+ }
+
+ let mut field_start = 0;
+ let names = ends[..ends_len]
+ .iter()
+ .map(|&field_end| {
+ let name =
std::str::from_utf8(&output[field_start..field_end]).map_err(
+ |e| exec_datafusion_err!("Invalid UTF-8 in CSV header:
{e}"),
+ )?;
+ field_start = field_end;
+ Ok(name.to_string())
+ })
+ .collect::<Result<_>>()?;
+ Ok(Some(HeaderRow {
+ names,
+ start: self.first_record_start(buf),
+ }))
+ }
+
+ /// Byte offset of the first record in `buf`, after a leading UTF-8 BOM,
+ /// blank records and comment lines. None of these can contain quotes, so
+ /// a plain byte scan finds the same position the parser starts reading
+ /// the header at. With no explicit terminator both `\r` and `\n` end a
+ /// record.
+ fn first_record_start(&self, buf: &[u8]) -> usize {
+ let terminator = self.terminator();
+ let comment = self.comment();
+ let is_terminator = |b: u8| match terminator {
+ Some(t) => b == t,
+ None => b == b'\n' || b == b'\r',
+ };
+ // csv-core's DFA only ends a comment at `\n` when `\n` is also a
+ // terminator; with any other terminator the comment runs to the end
+ // of the input.
+ let comment_ends_at_newline = matches!(terminator, None | Some(b'\n'));
+
+ let mut start = if buf.starts_with(UTF8_BOM) {
+ UTF8_BOM.len()
+ } else {
+ 0
+ };
+ while let Some(&first) = buf.get(start) {
+ if is_terminator(first) {
+ start += 1;
+ } else if comment == Some(first) {
+ let end = comment_ends_at_newline
+ .then(|| buf[start..].iter().position(|b| *b == b'\n'))
+ .flatten();
+ match end {
+ Some(end) => start += end + 1,
+ None => return buf.len(),
+ }
+ } else {
+ break;
+ }
+ }
+ start
+ }
+
+ /// Builds the mapping for a file with the given header, or `None` if the
+ /// file can be read positionally.
+ fn header_mapping(
+ &self,
+ header: &[String],
+ expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
+ ) -> Result<Option<Arc<HeaderMapping>>> {
+ // csv-core reads an unterminated comment at the end of the input as a
+ // record with one empty field
+ if header.is_empty() || header == [""] {
+ return Ok(None);
+ }
+ HeaderMapping::try_new(
+ header,
+ self.table_schema.file_schema(),
+ &self.projection.file_indices,
+ expr_adapter_factory,
+ )
+ .map(|mapping| mapping.map(Arc::new))
+ }
+}
+
+/// Maps the columns of a single CSV file onto the table's file schema by
+/// header name.
+///
+/// Only the header columns that are declared in the table are decoded; the
+/// decoded batch is then adapted to the projected file schema with the scan's
+/// [`PhysicalExprAdapterFactory`], which is the same mechanism Parquet uses
+/// to fill missing columns and reorder by name.
+#[derive(Debug)]
+struct HeaderMapping {
+ /// One field per header column, typed from the table schema where the
+ /// name is declared there
+ physical_schema: SchemaRef,
+ /// Header positions to decode, in the order of the table projection
+ physical_projection: Vec<usize>,
+ adapter: BatchAdapter,
+}
+
+impl HeaderMapping {
+ fn try_new(
+ header: &[String],
+ file_schema: &Schema,
+ file_indices: &[usize],
+ expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,
+ ) -> Result<Option<Self>> {
+ let matches_positionally = header.len() == file_schema.fields().len()
+ && header
+ .iter()
+ .zip(file_schema.fields())
+ .all(|(name, field)| name == field.name());
+ if matches_positionally {
+ return Ok(None);
+ }
+
+ let physical_fields: Vec<Field> = header
+ .iter()
+ .map(|name| match file_schema.field_with_name(name) {
+ Ok(field) => field.clone(),
+ Err(_) => Field::new(name, DataType::Utf8, true),
+ })
+ .collect();
+
+ let physical_projection: Vec<usize> = file_indices
Review Comment:
Could we add an optional public read regression for duplicate headers? With
`a,a
1,not_an_int
` and nullable integer `a`, the first occurrence should produce `1` without
parsing the second; local and chunked variants would also protect both opener
paths.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
