diff --git a/Cargo.lock b/Cargo.lock index 053c7dfbd6819..ec850f7e35cea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2042,11 +2042,13 @@ dependencies = [ "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-proto-models", "datafusion-session", "futures", + "log", "object_store", "regex", "tokio", diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 5106ea51ea6ba..d2c33c47f89db 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -4087,6 +4087,17 @@ config_namespace! { /// of columns, it creates a union schema containing all columns found across the files, and will /// pad any files missing columns with null values for their rows. pub truncated_rows: Option, default = None + /// Whether to map file columns to the table schema by header name instead of by position. + /// + /// Only applies when `has_header` is true. When enabled, each file's header row is read + /// and columns are matched by name: table columns missing from a file are filled with + /// nulls, and file columns not declared in the table are ignored. Names are matched exactly + /// (case-sensitive) and duplicate header names resolve to the first occurrence. Header + /// names containing quoted newlines are not supported. The header is located the same + /// way the CSV parser does: a leading UTF-8 BOM, blank lines, and comment lines are skipped. + /// + /// Defaults to `false` when `None`. + pub match_columns_by_header: Option, default = None } } @@ -4212,6 +4223,13 @@ impl CsvOptions { self } + /// Whether to map file columns to the table schema by header name instead of by position. + /// Only applies when the file has a header. + pub fn with_match_columns_by_header(mut self, match_by_header: bool) -> Self { + self.match_columns_by_header = Some(match_by_header); + self + } + /// Set the compression level for the output file. /// The valid range depends on the compression algorithm. /// If not specified, the default level for the algorithm is used. diff --git a/datafusion/core/src/datasource/file_format/options.rs b/datafusion/core/src/datasource/file_format/options.rs index 297f33cda4a3c..72e737ba78172 100644 --- a/datafusion/core/src/datasource/file_format/options.rs +++ b/datafusion/core/src/datasource/file_format/options.rs @@ -86,6 +86,12 @@ pub struct CsvReadOptions<'a> { /// When set to true then it will allow records with less than the expected number of columns and fill the missing columns with nulls. /// If the record’s schema is not nullable, then it will still return an error. pub truncated_rows: bool, + /// Whether to map file columns to the table schema by header name instead of by position. + /// + /// Only applies when `has_header` is true. Table columns missing from a file are + /// filled with nulls and file columns not declared in the table are ignored. + /// Defaults to `false`. + pub match_columns_by_header: bool, } impl Default for CsvReadOptions<'_> { @@ -113,6 +119,7 @@ impl<'a> CsvReadOptions<'a> { comment: None, null_regex: None, truncated_rows: false, + match_columns_by_header: false, } } @@ -228,6 +235,13 @@ impl<'a> CsvReadOptions<'a> { self.truncated_rows = truncated_rows; self } + + /// Configure whether file columns are matched to the table schema by header name + /// instead of by position. Only applies when `has_header` is true. + pub fn match_columns_by_header(mut self, match_columns_by_header: bool) -> Self { + self.match_columns_by_header = match_columns_by_header; + self + } } /// Options that control the reading of Parquet files. @@ -643,7 +657,8 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { .with_schema_infer_max_rec(self.schema_infer_max_records) .with_file_compression_type(self.file_compression_type.to_owned()) .with_null_regex(self.null_regex.clone()) - .with_truncated_rows(self.truncated_rows); + .with_truncated_rows(self.truncated_rows) + .with_match_columns_by_header(self.match_columns_by_header); ListingOptions::new(Arc::new(file_format)) .with_file_extension(self.file_extension) diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs index 004e1d60cb22d..9dd66bab46bf8 100644 --- a/datafusion/core/src/datasource/physical_plan/csv.rs +++ b/datafusion/core/src/datasource/physical_plan/csv.rs @@ -802,4 +802,339 @@ mod tests { Arc::new(schema) } + + /// Files sharing the table `(a INT, b TEXT, c INT, d TEXT)` but written + /// with different header layouts + const HEADER_LAYOUT_FILES: [(&str, &str); 3] = [ + ("full.csv", "a,b,c,d\n1,x,10,p\n4,y,40,s\n"), + ("missing_middle.csv", "a,d\n2,q\n5,t\n"), + ("reordered.csv", "d,extra,c,a\nr,ignored,30,3\n"), + ]; + + #[derive(Debug, Clone, Copy)] + enum HeaderLayoutRead { + Local, + /// Tiny object store chunks so the header spans several chunks + Chunked, + /// Byte-range partitions that start mid-file + Repartitioned, + } + + fn write_header_layout_files(dir: &std::path::Path) -> Result<()> { + for (name, contents) in HEADER_LAYOUT_FILES { + fs::write(dir.join(name), contents)?; + } + Ok(()) + } + + async fn header_layout_table( + ctx: &SessionContext, + location: &str, + options: &str, + ) -> Result<()> { + ctx.sql(&format!( + "CREATE EXTERNAL TABLE t (a INT, b TEXT, c INT, d TEXT) \ + STORED AS CSV LOCATION '{location}' \ + OPTIONS ('format.has_header' 'true'{options})" + )) + .await? + .collect() + .await?; + Ok(()) + } + + #[rstest] + #[case::local(HeaderLayoutRead::Local)] + #[case::chunked(HeaderLayoutRead::Chunked)] + #[case::repartitioned(HeaderLayoutRead::Repartitioned)] + #[tokio::test] + async fn csv_match_columns_by_header(#[case] read: HeaderLayoutRead) -> Result<()> { + let tmp_dir = TempDir::new()?; + write_header_layout_files(tmp_dir.path())?; + + let ctx = match read { + HeaderLayoutRead::Local | HeaderLayoutRead::Chunked => SessionContext::new(), + HeaderLayoutRead::Repartitioned => SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_min_size(1), + ), + }; + if matches!(read, HeaderLayoutRead::Chunked) { + ctx.register_object_store( + &Url::parse("file://").unwrap(), + Arc::new(ChunkedStore::new(Arc::new(LocalFileSystem::new()), 3)), + ); + } + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true'", + ) + .await?; + + let plan = ctx + .sql("SELECT * FROM t") + .await? + .create_physical_plan() + .await?; + let plan = datafusion_physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string(); + assert!(plan.contains("match_columns_by_header=true"), "{plan}"); + if matches!(read, HeaderLayoutRead::Repartitioned) { + assert!(plan.contains("missing_middle.csv:"), "{plan}"); + } + + let all = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + let projected = ctx + .sql("SELECT d, c FROM t ORDER BY d") + .await? + .collect() + .await?; + let missing_only = ctx + .sql("SELECT count(b), count(*) FROM t") + .await? + .collect() + .await?; + + insta::allow_duplicates! { + assert_snapshot!(batches_to_string(&all), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 1 | x | 10 | p | + | 2 | | | q | + | 3 | | 30 | r | + | 4 | y | 40 | s | + | 5 | | | t | + +---+---+----+---+ + "); + assert_snapshot!(batches_to_string(&projected), @r" + +---+----+ + | d | c | + +---+----+ + | p | 10 | + | q | | + | r | 30 | + | s | 40 | + | t | | + +---+----+ + "); + assert_snapshot!(batches_to_string(&missing_only), @r" + +------------+----------+ + | count(t.b) | count(*) | + +------------+----------+ + | 2 | 5 | + +------------+----------+ + "); + } + Ok(()) + } + + #[rstest] + #[case::local(HeaderLayoutRead::Local)] + #[case::chunked(HeaderLayoutRead::Chunked)] + #[case::repartitioned(HeaderLayoutRead::Repartitioned)] + #[tokio::test] + async fn csv_match_columns_by_header_with_comments( + #[case] read: HeaderLayoutRead, + ) -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write( + tmp_dir.path().join("a.csv"), + "# generated\n# by test\nd,c,a\nr,30,3\n# trailing\ns,40,4\n", + )?; + + let ctx = match read { + HeaderLayoutRead::Local | HeaderLayoutRead::Chunked => SessionContext::new(), + HeaderLayoutRead::Repartitioned => SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_min_size(1), + ), + }; + if matches!(read, HeaderLayoutRead::Chunked) { + ctx.register_object_store( + &Url::parse("file://").unwrap(), + Arc::new(ChunkedStore::new(Arc::new(LocalFileSystem::new()), 3)), + ); + } + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.comment' '#'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + insta::allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 3 | | 30 | r | + | 4 | | 40 | s | + +---+---+----+---+ + "); + } + Ok(()) + } + + /// The header is found the way the CSV parser finds it: after a UTF-8 + /// BOM, comment lines and blank lines + #[rstest] + #[case::local(HeaderLayoutRead::Local)] + #[case::chunked(HeaderLayoutRead::Chunked)] + #[case::repartitioned(HeaderLayoutRead::Repartitioned)] + #[tokio::test] + async fn csv_match_columns_by_header_with_bom_and_blank_lines( + #[case] read: HeaderLayoutRead, + ) -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write( + tmp_dir.path().join("a.csv"), + "\u{feff}# generated\n\n\r\nd,c,a\nr,30,3\n\ns,40,4\n", + )?; + + let ctx = match read { + HeaderLayoutRead::Local | HeaderLayoutRead::Chunked => SessionContext::new(), + HeaderLayoutRead::Repartitioned => SessionContext::new_with_config( + SessionConfig::new() + .with_target_partitions(4) + .with_repartition_file_min_size(1), + ), + }; + if matches!(read, HeaderLayoutRead::Chunked) { + ctx.register_object_store( + &Url::parse("file://").unwrap(), + Arc::new(ChunkedStore::new(Arc::new(LocalFileSystem::new()), 3)), + ); + } + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.comment' '#'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + insta::allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 3 | | 30 | r | + | 4 | | 40 | s | + +---+---+----+---+ + "); + } + Ok(()) + } + + #[tokio::test] + async fn csv_match_columns_by_header_with_truncated_rows() -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write(tmp_dir.path().join("a.csv"), "d,a,c\nr,1,10\ns,2\n")?; + + let ctx = SessionContext::new(); + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.truncated_rows' 'true'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 1 | | 10 | r | + | 2 | | | s | + +---+---+----+---+ + "); + Ok(()) + } + + #[tokio::test] + async fn csv_columns_mapped_by_position_by_default() -> Result<()> { + let tmp_dir = TempDir::new()?; + fs::write(tmp_dir.path().join("a.csv"), "a,d\n2,q\n")?; + + let ctx = SessionContext::new(); + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.truncated_rows' 'true'", + ) + .await?; + + let batches = ctx.sql("SELECT * FROM t").await?.collect().await?; + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+---+---+ + | a | b | c | d | + +---+---+---+---+ + | 2 | q | | | + +---+---+---+---+ + "); + Ok(()) + } + + #[cfg(feature = "compression")] + #[tokio::test] + async fn csv_match_columns_by_header_gzip() -> Result<()> { + use flate2::Compression; + use flate2::write::GzEncoder; + + let tmp_dir = TempDir::new()?; + for (name, contents) in HEADER_LAYOUT_FILES { + let file = File::create(tmp_dir.path().join(format!("{name}.gz")))?; + let mut encoder = GzEncoder::new(file, Compression::default()); + encoder.write_all(contents.as_bytes())?; + encoder.finish()?; + } + + let ctx = SessionContext::new(); + header_layout_table( + &ctx, + &format!("{}/", tmp_dir.path().display()), + ", 'format.match_columns_by_header' 'true', 'format.compression' 'gzip'", + ) + .await?; + + let batches = ctx + .sql("SELECT * FROM t ORDER BY a") + .await? + .collect() + .await?; + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+----+---+ + | a | b | c | d | + +---+---+----+---+ + | 1 | x | 10 | p | + | 2 | | | q | + | 3 | | 30 | r | + | 4 | y | 40 | s | + | 5 | | | t | + +---+---+----+---+ + "); + Ok(()) + } } diff --git a/datafusion/core/tests/datasource/csv.rs b/datafusion/core/tests/datasource/csv.rs index 93cf4ae6a3004..46a00b924bc9b 100644 --- a/datafusion/core/tests/datasource/csv.rs +++ b/datafusion/core/tests/datasource/csv.rs @@ -4,6 +4,7 @@ //! Test for CSV schema inference with different column counts (GitHub issue #17516) +use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::error::Result; use datafusion::prelude::*; use datafusion_common::test_util::batches_to_sort_string; @@ -107,3 +108,40 @@ async fn test_csv_schema_inference_different_column_counts() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn test_csv_read_options_match_columns_by_header() -> Result<()> { + let temp_dir = TempDir::new()?; + fs::write(temp_dir.path().join("full.csv"), "a,b,c\n1,x,10\n")?; + fs::write( + temp_dir.path().join("reordered.csv"), + "c,extra,a\n20,ignored,2\n", + )?; + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Int32, true), + ]); + let ctx = SessionContext::new(); + let batches = ctx + .read_csv( + temp_dir.path().to_str().unwrap(), + CsvReadOptions::new() + .schema(&schema) + .match_columns_by_header(true), + ) + .await? + .collect() + .await?; + + assert_snapshot!(batches_to_sort_string(&batches), @r" + +---+---+----+ + | a | b | c | + +---+---+----+ + | 1 | x | 10 | + | 2 | | 20 | + +---+---+----+ + "); + Ok(()) +} diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index be178436fb2bc..4a6abc0d46a20 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -34,11 +34,13 @@ datafusion-common-runtime = { workspace = true } datafusion-datasource = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } +datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } +log = { workspace = true } object_store = { workspace = true } regex = { workspace = true } tokio = { workspace = true } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 578694c18a283..07a75591a3a16 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -293,6 +293,14 @@ impl CsvFormat { self } + /// Set whether file columns are matched to the table schema by header name + /// instead of by position + /// - defaults to false + pub fn with_match_columns_by_header(mut self, match_by_header: bool) -> Self { + self.options.match_columns_by_header = Some(match_by_header); + self + } + /// The delimiter character. pub fn delimiter(&self) -> u8 { self.options.delimiter @@ -998,6 +1006,9 @@ impl From<&CsvFormatFactory> for datafusion_proto_models::protobuf::CsvOptions { ignore_trailing_whitespace: options .ignore_trailing_whitespace .map_or(vec![], |v| vec![v as u8]), + match_columns_by_header: options + .match_columns_by_header + .map_or(vec![], |v| vec![v as u8]), } } else { datafusion_proto_models::protobuf::CsvOptions::default() diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index cbc8df820f222..de0799e55420c 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -9,6 +9,7 @@ use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use datafusion_physical_plan::projection::ProjectionExprs; use std::fmt; use std::io::Read; +use std::ops::Range; use std::sync::Arc; use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream}; @@ -18,7 +19,10 @@ use datafusion_datasource::{ FileRange, ListingTableUrl, PartitionedFile, TableSchema, as_file_source, }; +use arrow::array::RecordBatch; use arrow::csv; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use bytes::Bytes; use datafusion_common::config::CsvOptions; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{DataFusionError, Result, exec_datafusion_err}; @@ -26,15 +30,21 @@ use datafusion_common_runtime::JoinSet; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::TaskContext; +use datafusion_physical_expr_adapter::{ + BatchAdapter, BatchAdapterFactory, DefaultPhysicalExprAdapterFactory, + PhysicalExprAdapterFactory, +}; use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet}; use datafusion_physical_plan::{ DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, }; use crate::file_format::CsvDecoder; +use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; use object_store::buffered::BufWriter; -use object_store::{GetOptions, GetResultPayload, ObjectStore}; +use object_store::path::Path; +use object_store::{GetOptions, GetResultPayload, ObjectStore, ObjectStoreExt}; use regex::Regex; use tokio::io::AsyncWriteExt; @@ -174,28 +184,53 @@ impl CsvSource { pub fn newlines_in_values(&self) -> bool { self.options.newlines_in_values.unwrap_or(false) } + + /// true if file columns are matched to the table schema by header name + pub fn match_columns_by_header(&self) -> bool { + self.options.match_columns_by_header.unwrap_or(false) && self.has_header() + } + + /// Whether to match file columns to the table schema by header name + pub fn with_match_columns_by_header(&self, match_by_header: bool) -> Self { + let mut conf = self.clone(); + conf.options.match_columns_by_header = Some(match_by_header); + conf + } } impl CsvSource { - fn open(&self, reader: R) -> Result> { - Ok(self.builder()?.build(reader)?) + fn open( + &self, + reader: R, + mapping: Option<&HeaderMapping>, + ) -> Result> { + Ok(self.builder(mapping)?.build(reader)?) } - fn builder(&self) -> Result { - let mut builder = - csv::ReaderBuilder::new(Arc::clone(self.table_schema.file_schema())) - .with_delimiter(self.delimiter()) - .with_batch_size( - self.batch_size - .expect("Batch size must be set before initializing builder"), - ) - .with_header(self.has_header()) - .with_quote(self.quote()) - .with_truncated_rows(self.truncate_rows()); + fn builder(&self, mapping: Option<&HeaderMapping>) -> Result { + let (schema, projection) = match mapping { + Some(mapping) => ( + Arc::clone(&mapping.physical_schema), + mapping.physical_projection.clone(), + ), + None => ( + Arc::clone(self.table_schema.file_schema()), + self.projection.file_indices.clone(), + ), + }; + let mut builder = csv::ReaderBuilder::new(schema) + .with_delimiter(self.delimiter()) + .with_batch_size( + self.batch_size + .expect("Batch size must be set before initializing builder"), + ) + .with_header(self.has_header()) + .with_quote(self.quote()) + .with_truncated_rows(self.truncate_rows()); if let Some(terminator) = self.terminator() { builder = builder.with_terminator(terminator); } - builder = builder.with_projection(self.projection.file_indices.clone()); + builder = builder.with_projection(projection); if let Some(escape) = self.escape() { builder = builder.with_escape(escape) } @@ -211,6 +246,217 @@ impl CsvSource { Ok(builder) } + + fn parse_header(&self, line: &[u8]) -> Result> { + let mut format = csv::reader::Format::default() + .with_header(true) + .with_delimiter(self.delimiter()) + .with_quote(self.quote()); + if let Some(terminator) = self.terminator() { + format = format.with_terminator(terminator); + } + if let Some(escape) = self.escape() { + format = format.with_escape(escape); + } + if let Some(comment) = self.comment() { + format = format.with_comment(comment); + } + let (schema, _) = format.infer_schema(line, Some(0))?; + Ok(schema.fields().iter().map(|f| f.name().clone()).collect()) + } + + /// Builds the mapping for a file whose header row is `line`, or `None` if + /// the file can be read positionally. + fn header_mapping( + &self, + line: &[u8], + expr_adapter_factory: Arc, + ) -> Result>> { + if line.is_empty() { + return Ok(None); + } + let header = self.parse_header(line)?; + 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, + adapter: BatchAdapter, +} + +impl HeaderMapping { + fn try_new( + header: &[String], + file_schema: &Schema, + file_indices: &[usize], + expr_adapter_factory: Arc, + ) -> Result> { + 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 = 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 = file_indices + .iter() + .filter_map(|&index| { + let name = file_schema.field(index).name(); + header.iter().position(|column| column == name) + }) + .collect(); + + let physical_schema = Arc::new(Schema::new(physical_fields)); + let decoded_schema = Arc::new(physical_schema.project(&physical_projection)?); + let target_schema = Arc::new(file_schema.project(file_indices)?); + let adapter = BatchAdapterFactory::new(target_schema) + .with_adapter_factory(expr_adapter_factory) + .make_adapter(&decoded_schema)?; + + Ok(Some(Self { + physical_schema, + physical_projection, + adapter, + })) + } + + fn remap(&self, batch: &RecordBatch) -> Result { + self.adapter.adapt_batch(batch) + } +} + +/// Returns the byte range of the header row in `buf`, following the same +/// rules as the CSV parser: a leading UTF-8 BOM is ignored, empty records are +/// skipped, and comment lines (which always end at `\n`) are skipped. With no +/// explicit terminator both `\r` and `\n` end a record. Returns `None` when +/// the header is not yet complete and more input may follow. +fn find_header_line( + buf: &[u8], + terminator: Option, + comment: Option, + eof: bool, +) -> Option> { + const BOM: &[u8] = b"\xef\xbb\xbf"; + let is_terminator = |b: u8| match terminator { + Some(t) => b == t, + None => b == b'\n' || b == b'\r', + }; + + let mut start = 0; + if buf.starts_with(BOM) { + start = BOM.len(); + } else if !eof && buf.len() < BOM.len() && BOM.starts_with(buf) { + return None; + } + + loop { + let Some(&first) = buf.get(start) else { + return eof.then_some(buf.len()..buf.len()); + }; + if is_terminator(first) { + start += 1; + } else if comment == Some(first) { + match buf[start..].iter().position(|b| *b == b'\n') { + Some(end) => start += end + 1, + None if eof => return Some(buf.len()..buf.len()), + None => return None, + } + } else { + return match buf[start..].iter().position(|b| is_terminator(*b)) { + Some(end) => Some(start..start + end), + None if eof => Some(start..buf.len()), + None => None, + }; + } + } +} + +/// Reads the header row of `input`, returning it together with a stream that +/// still yields every byte of `input`. +async fn peek_header_line( + mut input: BoxStream<'static, Result>, + terminator: Option, + comment: Option, +) -> Result<(Vec, BoxStream<'static, Result>)> { + let mut prefix = Vec::new(); + let mut buffered = Vec::new(); + loop { + let chunk = input.next().await.transpose()?; + let eof = chunk.is_none(); + if let Some(chunk) = chunk { + prefix.extend_from_slice(&chunk); + buffered.push(Ok(chunk)); + } + if let Some(header) = find_header_line(&prefix, terminator, comment, eof) { + let header = prefix[header].to_vec(); + let buffered = futures::stream::iter(buffered); + let stream = if eof { + buffered.boxed() + } else { + buffered.chain(input).boxed() + }; + return Ok((header, stream)); + } + } +} + +/// Reads the header row of an uncompressed file, growing the requested range +/// until it is complete. Also returns the byte offset where the header starts. +async fn fetch_header_line( + store: &dyn ObjectStore, + location: &Path, + file_size: u64, + terminator: Option, + comment: Option, +) -> Result<(Vec, u64)> { + let mut len = file_size.min(64 * 1024); + loop { + let bytes = store.get_range(location, 0..len).await?; + let eof = len >= file_size; + if let Some(header) = find_header_line(&bytes, terminator, comment, eof) { + return Ok((bytes[header.clone()].to_vec(), header.start as u64)); + } + len = (len * 2).min(file_size); + } +} + +fn remap_stream( + stream: BoxStream<'static, Result>, + mapping: Option>, +) -> BoxStream<'static, Result> { + match mapping { + Some(mapping) => stream.map(move |batch| mapping.remap(&batch?)).boxed(), + None => stream, + } } /// A [`FileOpener`] that opens a CSV file and yields a [`FileOpenFuture`] @@ -219,6 +465,7 @@ pub struct CsvOpener { file_compression_type: FileCompressionType, object_store: Arc, partition_index: usize, + expr_adapter_factory: Arc, } impl CsvOpener { @@ -233,8 +480,19 @@ impl CsvOpener { file_compression_type, object_store, partition_index: 0, + expr_adapter_factory: Arc::new(DefaultPhysicalExprAdapterFactory), } } + + /// Set the [`PhysicalExprAdapterFactory`] used to adapt files whose + /// columns are matched to the table schema by header name + pub fn with_expr_adapter_factory( + mut self, + expr_adapter_factory: Arc, + ) -> Self { + self.expr_adapter_factory = expr_adapter_factory; + self + } } impl From for Arc { @@ -250,11 +508,21 @@ impl FileSource for CsvSource { base_config: &FileScanConfig, partition_index: usize, ) -> Result> { + if self.options.match_columns_by_header == Some(true) && !self.has_header() { + log::warn!( + "CSV option match_columns_by_header is ignored because has_header is false; columns are matched by position" + ); + } + let expr_adapter_factory = base_config + .expr_adapter_factory + .clone() + .unwrap_or_else(|| Arc::new(DefaultPhysicalExprAdapterFactory) as _); let mut opener = Arc::new(CsvOpener { config: Arc::new(self.clone()), file_compression_type: base_config.file_compression_type, object_store, partition_index, + expr_adapter_factory, }) as Arc; opener = ProjectionOpener::try_new( self.projection.clone(), @@ -307,7 +575,11 @@ impl FileSource for CsvSource { fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, ", has_header={}", self.has_header()) + write!(f, ", has_header={}", self.has_header())?; + if self.match_columns_by_header() { + write!(f, ", match_columns_by_header=true")?; + } + Ok(()) } DisplayFormatType::TreeRender => Ok(()), } @@ -360,6 +632,7 @@ impl FileSource for CsvSource { newlines_in_values: self.options.newlines_in_values.unwrap_or(false), truncate_rows: self.truncate_rows(), terminator: self.terminator().map(|terminator| vec![terminator]), + match_columns_by_header: self.match_columns_by_header(), }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), @@ -395,6 +668,8 @@ impl FileOpener for CsvOpener { // `self.config.has_header` controls whether to skip reading the 1st line header // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle // partition, then don't skip first line + let match_by_header = self.config.match_columns_by_header(); + let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory); let mut csv_has_header = self.config.has_header(); if let Some(FileRange { start, .. }) = partitioned_file.range && start != 0 @@ -417,6 +692,8 @@ impl FileOpener for CsvOpener { let store = Arc::clone(&self.object_store); let terminator = self.config.terminator(); + let line_end = terminator.unwrap_or(b'\n'); + let comment = self.config.comment(); let baseline_metrics = BaselineMetrics::new(&self.config.metrics, self.partition_index); @@ -440,18 +717,40 @@ impl FileOpener for CsvOpener { ) })?; + // A partition may start mid-file, so the header is fetched + // separately. Every byte-range partition of a file repeats this + // small request; cache per file if that becomes costly. + let mapping = if match_by_header { + let (header, header_start) = fetch_header_line( + store.as_ref(), + &location, + file_size, + terminator, + comment, + ) + .await?; + // A partition starting within leading comment lines owns + // the header row, so it has to skip it. + if raw_start <= header_start { + config.options.has_header = Some(true); + } + config.header_mapping(&header, Arc::clone(&expr_adapter_factory))? + } else { + None + }; + let aligned_stream = AlignedBoundaryStream::new( Arc::clone(&store), location.clone(), raw_start, raw_end, file_size, - terminator.unwrap_or(b'\n'), + line_end, ) .await? .map_err(DataFusionError::from); - let decoder = config.builder()?.build_decoder(); + let decoder = config.builder(mapping.as_deref())?.build_decoder(); let input = file_compression_type .convert_stream(aligned_stream.boxed())? .fuse(); @@ -459,7 +758,7 @@ impl FileOpener for CsvOpener { input, DecoderDeserializer::new(CsvDecoder::new(decoder)), ); - return Ok(stream.map_err(Into::into).boxed()); + return Ok(remap_stream(stream.map_err(Into::into).boxed(), mapping)); } // No range specified — read the entire file @@ -469,8 +768,28 @@ impl FileOpener for CsvOpener { match result.payload { #[cfg(not(target_arch = "wasm32"))] GetResultPayload::File(file, _) => { + use std::io::{BufRead, BufReader, Cursor}; + let decoder = file_compression_type.convert_read(file)?; - let mut reader = config.open(decoder)?; + let (decoder, mapping): (Box, _) = if match_by_header + { + let mut decoder = BufReader::new(decoder); + let mut prefix = Vec::new(); + let header = loop { + let eof = decoder.read_until(line_end, &mut prefix)? == 0; + if let Some(header) = + find_header_line(&prefix, terminator, comment, eof) + { + break prefix[header].to_vec(); + } + }; + let mapping = config + .header_mapping(&header, Arc::clone(&expr_adapter_factory))?; + (Box::new(Cursor::new(prefix).chain(decoder)), mapping) + } else { + (decoder, None) + }; + let mut reader = config.open(decoder, mapping.as_deref())?; // Use std::iter::from_fn to wrap execution of iterator's next() method. let iterator = std::iter::from_fn(move || { @@ -480,20 +799,29 @@ impl FileOpener for CsvOpener { result }); - Ok(futures::stream::iter(iterator) + let stream = futures::stream::iter(iterator) .map(|r| r.map_err(Into::into)) - .boxed()) + .boxed(); + Ok(remap_stream(stream, mapping)) } GetResultPayload::Stream(s) => { - let decoder = config.builder()?.build_decoder(); let s = s.map_err(DataFusionError::from); - let input = file_compression_type.convert_stream(s.boxed())?.fuse(); - + let mut input = file_compression_type.convert_stream(s.boxed())?; + let mut mapping = None; + if match_by_header { + let (header, rest) = + peek_header_line(input, terminator, comment).await?; + mapping = config + .header_mapping(&header, Arc::clone(&expr_adapter_factory))?; + input = rest; + } + + let decoder = config.builder(mapping.as_deref())?.build_decoder(); let stream = deserialize_stream( - input, + input.fuse(), DecoderDeserializer::new(CsvDecoder::new(decoder)), ); - Ok(stream.map_err(Into::into).boxed()) + Ok(remap_stream(stream.map_err(Into::into).boxed(), mapping)) } } })) @@ -519,7 +847,7 @@ pub async fn plan_to_csv( let storeref = Arc::clone(&store); let plan: Arc = Arc::clone(&plan); let filename = format!("{}/part-{i}.csv", parsed.prefix()); - let file = object_store::path::Path::parse(filename)?; + let file = Path::parse(filename)?; let mut stream = plan.execute(i, Arc::clone(&task_ctx))?; join_set.spawn(async move { @@ -643,6 +971,7 @@ impl CsvSource { quote: proto_str_to_byte(&scan.quote, "quote")?, newlines_in_values: Some(scan.newlines_in_values), truncated_rows: Some(scan.truncate_rows), + match_columns_by_header: Some(scan.match_columns_by_header), ..Default::default() }; let source = Arc::new( @@ -685,7 +1014,7 @@ mod tests { let data = "id,name\n1,alice\n2,N/A\n3,carol\n"; let batch = source - .open(data.as_bytes()) + .open(data.as_bytes(), None) .unwrap() .next() .unwrap() @@ -713,7 +1042,7 @@ mod tests { let data = "id,value\n1,10\n2,N/A\n3,30\n"; let batch = source - .open(data.as_bytes()) + .open(data.as_bytes(), None) .unwrap() .next() .unwrap() @@ -739,7 +1068,7 @@ mod tests { let data = "id,name\n1,alice\n2,N/A\n3,carol\n"; let batch = source - .open(data.as_bytes()) + .open(data.as_bytes(), None) .unwrap() .next() .unwrap() @@ -758,10 +1087,231 @@ mod tests { let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, true)])); let source = csv_source(schema, Some("(")); - let err = source.open(b"id\n1\n".as_slice()).unwrap_err(); + let err = source.open(b"id\n1\n".as_slice(), None).unwrap_err(); assert!( err.to_string().contains("Unable to parse CSV null regex"), "unexpected error: {err}" ); } + + use arrow::array::{Int32Array, RecordBatchOptions}; + + fn table_schema() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Int32, true), + ]) + } + + fn header(names: &[&str]) -> Vec { + names.iter().map(|name| name.to_string()).collect() + } + + fn adapter_factory() -> Arc { + Arc::new(DefaultPhysicalExprAdapterFactory) + } + + #[test] + fn header_mapping_skipped_when_header_matches() -> Result<()> { + let mapping = HeaderMapping::try_new( + &header(&["a", "b", "c"]), + &table_schema(), + &[0, 2], + adapter_factory(), + )?; + assert!(mapping.is_none()); + Ok(()) + } + + #[test] + fn header_mapping_by_name() -> Result<()> { + let mapping = HeaderMapping::try_new( + &header(&["c", "extra", "a"]), + &table_schema(), + &[0, 1, 2], + adapter_factory(), + )? + .unwrap(); + + let physical_types: Vec<_> = mapping + .physical_schema + .fields() + .iter() + .map(|f| f.data_type().clone()) + .collect(); + assert_eq!( + physical_types, + vec![DataType::Int32, DataType::Utf8, DataType::Int32] + ); + assert_eq!(mapping.physical_projection, vec![2, 0]); + + let decoded_schema = mapping.physical_schema.project(&[2, 0])?; + let decoded = RecordBatch::try_new( + Arc::new(decoded_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ], + )?; + let remapped = mapping.remap(&decoded)?; + assert_eq!(remapped.schema().as_ref(), &table_schema()); + assert_eq!(remapped.column(1).null_count(), 2); + assert_eq!( + remapped.column(0).as_ref(), + &Int32Array::from(vec![1, 2]) as &dyn Array + ); + assert_eq!( + remapped.column(2).as_ref(), + &Int32Array::from(vec![10, 20]) as &dyn Array + ); + Ok(()) + } + + #[test] + fn header_mapping_only_missing_columns() -> Result<()> { + let mapping = HeaderMapping::try_new( + &header(&["a", "c"]), + &table_schema(), + &[1], + adapter_factory(), + )? + .unwrap(); + assert!(mapping.physical_projection.is_empty()); + + let decoded = RecordBatch::try_new_with_options( + Arc::new(Schema::empty()), + vec![], + &RecordBatchOptions::new().with_row_count(Some(3)), + )?; + let remapped = mapping.remap(&decoded)?; + assert_eq!(remapped.num_rows(), 3); + assert_eq!(remapped.column(0).null_count(), 3); + Ok(()) + } + + #[test] + fn header_mapping_missing_non_nullable_column() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, false), + ]); + let err = + HeaderMapping::try_new(&header(&["a"]), &schema, &[0, 1], adapter_factory()) + .unwrap_err(); + assert!( + err.to_string() + .contains("Non-nullable column 'b' is missing from the physical schema"), + "{err}" + ); + } + + #[test] + fn parse_header_handles_quotes() -> Result<()> { + let source = CsvSource::new(Arc::new(table_schema())); + assert_eq!( + source.parse_header(b"\"a,1\",b\r")?, + vec!["a,1".to_string(), "b".to_string()] + ); + Ok(()) + } + + #[test] + fn find_header_line_skips_comments() { + let comment = Some(b'#'); + let newline = Some(b'\n'); + assert_eq!( + find_header_line(b"# c1\n#c2\na,b\n1,2\n", newline, comment, false), + Some(9..12) + ); + assert_eq!( + find_header_line(b"#c\na,b\n", newline, None, false), + Some(0..2) + ); + assert_eq!(find_header_line(b"# c1\na,", newline, comment, false), None); + assert_eq!( + find_header_line(b"# c1\na,b", newline, comment, true), + Some(5..8) + ); + assert_eq!( + find_header_line(b"# c1\n# c2", newline, comment, true), + Some(9..9) + ); + assert_eq!( + find_header_line(b"# c1\n# c2", newline, comment, false), + None + ); + assert_eq!(find_header_line(b"", newline, comment, true), Some(0..0)); + } + + #[test] + fn find_header_line_skips_leading_blank_records() { + assert_eq!( + find_header_line(b"\n\na,b\n1,2\n", Some(b'\n'), None, false), + Some(2..5) + ); + // Without an explicit terminator both `\r` and `\n` end a record + assert_eq!( + find_header_line(b"\r\n\n\ra,b\r\n1,2\r\n", None, None, false), + Some(4..7) + ); + // A custom terminator keeps `\n` and `\r` as ordinary bytes + assert_eq!( + find_header_line(b";;a,b\n;1,2;", Some(b';'), None, false), + Some(2..6) + ); + // Whitespace-only lines are records, not blank lines + assert_eq!( + find_header_line(b" \na,b\n", Some(b'\n'), None, false), + Some(0..2) + ); + assert_eq!(find_header_line(b"\n\n", Some(b'\n'), None, false), None); + assert_eq!( + find_header_line(b"\n\n", Some(b'\n'), None, true), + Some(2..2) + ); + } + + #[test] + fn find_header_line_strips_bom() { + let comment = Some(b'#'); + assert_eq!( + find_header_line(b"\xef\xbb\xbfa,b\n1,2\n", Some(b'\n'), comment, false), + Some(3..6) + ); + assert_eq!( + find_header_line(b"\xef\xbb\xbf# c1\n\na,b\n", Some(b'\n'), comment, false), + Some(9..12) + ); + // An incomplete BOM needs more input before it can be classified + assert_eq!( + find_header_line(b"\xef\xbb", Some(b'\n'), comment, false), + None + ); + assert_eq!( + find_header_line(b"\xef\xbb", Some(b'\n'), comment, true), + Some(0..2) + ); + // Only a BOM at the very start of the file is ignored + assert_eq!( + find_header_line(b"\n\xef\xbb\xbfa,b\n", Some(b'\n'), comment, false), + Some(1..7) + ); + } + + #[test] + fn find_header_line_comment_ends_at_newline_with_custom_terminator() { + assert_eq!( + find_header_line(b"# c1;ignored\na,b;1,2;", Some(b';'), Some(b'#'), false), + Some(13..16) + ); + } + + #[test] + fn match_columns_by_header_requires_header() { + let source = CsvSource::new(Arc::new(table_schema())) + .with_csv_options(CsvOptions::default().with_has_header(false)) + .with_match_columns_by_header(true); + assert!(!source.match_columns_by_header()); + } } diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 49070439c3d8f..7c10dae1693e0 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -493,6 +493,8 @@ message CsvOptions { bytes ignore_leading_whitespace = 21; // Whether to ignore trailing whitespace in string values bytes ignore_trailing_whitespace = 22; + // Whether to map file columns to the table schema by header name + bytes match_columns_by_header = 23; } // Options controlling CSV format diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index fd94fc38d5901..2c7598273535f 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -1163,6 +1163,10 @@ impl TryFrom<&protobuf::CsvOptions> for CsvOptions { .ignore_trailing_whitespace .first() .map(|h| *h != 0), + match_columns_by_header: proto_opts + .match_columns_by_header + .first() + .map(|h| *h != 0), }) } } diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 09386a01eb556..0b06993145d0e 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -1755,6 +1755,9 @@ impl serde::Serialize for CsvOptions { if !self.ignore_trailing_whitespace.is_empty() { len += 1; } + if !self.match_columns_by_header.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion_common.CsvOptions", len)?; if !self.has_header.is_empty() { #[allow(clippy::needless_borrow)] @@ -1850,6 +1853,11 @@ impl serde::Serialize for CsvOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("ignoreTrailingWhitespace", pbjson::private::base64::encode(&self.ignore_trailing_whitespace).as_str())?; } + if !self.match_columns_by_header.is_empty() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("matchColumnsByHeader", pbjson::private::base64::encode(&self.match_columns_by_header).as_str())?; + } struct_ser.end() } } @@ -1898,6 +1906,8 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { "ignoreLeadingWhitespace", "ignore_trailing_whitespace", "ignoreTrailingWhitespace", + "match_columns_by_header", + "matchColumnsByHeader", ]; #[allow(clippy::enum_variant_names)] @@ -1924,6 +1934,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { QuoteStyle, IgnoreLeadingWhitespace, IgnoreTrailingWhitespace, + MatchColumnsByHeader, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1967,6 +1978,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { "quoteStyle" | "quote_style" => Ok(GeneratedField::QuoteStyle), "ignoreLeadingWhitespace" | "ignore_leading_whitespace" => Ok(GeneratedField::IgnoreLeadingWhitespace), "ignoreTrailingWhitespace" | "ignore_trailing_whitespace" => Ok(GeneratedField::IgnoreTrailingWhitespace), + "matchColumnsByHeader" | "match_columns_by_header" => Ok(GeneratedField::MatchColumnsByHeader), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -2008,6 +2020,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { let mut quote_style__ = None; let mut ignore_leading_whitespace__ = None; let mut ignore_trailing_whitespace__ = None; + let mut match_columns_by_header__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::HasHeader => { @@ -2168,6 +2181,14 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) ; } + GeneratedField::MatchColumnsByHeader => { + if match_columns_by_header__.is_some() { + return Err(serde::de::Error::duplicate_field("matchColumnsByHeader")); + } + match_columns_by_header__ = + Some(map_.next_value::<::pbjson::private::BytesDeserialize<_>>()?.0) + ; + } } } Ok(CsvOptions { @@ -2193,6 +2214,7 @@ impl<'de> serde::Deserialize<'de> for CsvOptions { quote_style: quote_style__.unwrap_or_default(), ignore_leading_whitespace: ignore_leading_whitespace__.unwrap_or_default(), ignore_trailing_whitespace: ignore_trailing_whitespace__.unwrap_or_default(), + match_columns_by_header: match_columns_by_header__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index ee586fa75ee37..8cfb07c02afbb 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -711,6 +711,9 @@ pub struct CsvOptions { /// Whether to ignore trailing whitespace in string values #[prost(bytes = "vec", tag = "22")] pub ignore_trailing_whitespace: ::prost::alloc::vec::Vec, + /// Whether to map file columns to the table schema by header name + #[prost(bytes = "vec", tag = "23")] + pub match_columns_by_header: ::prost::alloc::vec::Vec, } /// Options controlling CSV format #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 80a9dc662db5e..2404bba41fff9 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -1076,6 +1076,9 @@ impl TryFrom<&CsvOptions> for protobuf::CsvOptions { ignore_trailing_whitespace: opts .ignore_trailing_whitespace .map_or_else(Vec::new, |h| vec![h as u8]), + match_columns_by_header: opts + .match_columns_by_header + .map_or_else(Vec::new, |h| vec![h as u8]), }) } } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 341b2dbd3d1c3..57461796a22a0 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1399,6 +1399,7 @@ message CsvScanExecNode { bool truncate_rows = 8; // Custom one-byte line terminator. Absent means the default newline terminator. optional bytes terminator = 9; + bool match_columns_by_header = 10; } message JsonScanExecNode { diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs index 5ec8939f8b530..4f34657984b7c 100644 --- a/datafusion/proto-models/src/from_proto.rs +++ b/datafusion/proto-models/src/from_proto.rs @@ -283,6 +283,11 @@ impl TryFrom<&CsvOptionsProto> for CsvOptions { } else { Some(proto.ignore_trailing_whitespace[0] != 0) }, + match_columns_by_header: if proto.match_columns_by_header.is_empty() { + None + } else { + Some(proto.match_columns_by_header[0] != 0) + }, }) } } diff --git a/datafusion/proto-models/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs index ee586fa75ee37..8cfb07c02afbb 100644 --- a/datafusion/proto-models/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -711,6 +711,9 @@ pub struct CsvOptions { /// Whether to ignore trailing whitespace in string values #[prost(bytes = "vec", tag = "22")] pub ignore_trailing_whitespace: ::prost::alloc::vec::Vec, + /// Whether to map file columns to the table schema by header name + #[prost(bytes = "vec", tag = "23")] + pub match_columns_by_header: ::prost::alloc::vec::Vec, } /// Options controlling CSV format #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 34f2a293ae816..ff385073f638d 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -5015,6 +5015,9 @@ impl serde::Serialize for CsvScanExecNode { if self.terminator.is_some() { len += 1; } + if self.match_columns_by_header { + len += 1; + } if self.optional_escape.is_some() { len += 1; } @@ -5045,6 +5048,9 @@ impl serde::Serialize for CsvScanExecNode { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("terminator", pbjson::private::base64::encode(&v).as_str())?; } + if self.match_columns_by_header { + struct_ser.serialize_field("matchColumnsByHeader", &self.match_columns_by_header)?; + } if let Some(v) = self.optional_escape.as_ref() { match v { csv_scan_exec_node::OptionalEscape::Escape(v) => { @@ -5080,6 +5086,8 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "truncate_rows", "truncateRows", "terminator", + "match_columns_by_header", + "matchColumnsByHeader", "escape", "comment", ]; @@ -5093,6 +5101,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { NewlinesInValues, TruncateRows, Terminator, + MatchColumnsByHeader, Escape, Comment, } @@ -5123,6 +5132,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "newlinesInValues" | "newlines_in_values" => Ok(GeneratedField::NewlinesInValues), "truncateRows" | "truncate_rows" => Ok(GeneratedField::TruncateRows), "terminator" => Ok(GeneratedField::Terminator), + "matchColumnsByHeader" | "match_columns_by_header" => Ok(GeneratedField::MatchColumnsByHeader), "escape" => Ok(GeneratedField::Escape), "comment" => Ok(GeneratedField::Comment), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), @@ -5151,6 +5161,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { let mut newlines_in_values__ = None; let mut truncate_rows__ = None; let mut terminator__ = None; + let mut match_columns_by_header__ = None; let mut optional_escape__ = None; let mut optional_comment__ = None; while let Some(k) = map_.next_key()? { @@ -5199,6 +5210,12 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) ; } + GeneratedField::MatchColumnsByHeader => { + if match_columns_by_header__.is_some() { + return Err(serde::de::Error::duplicate_field("matchColumnsByHeader")); + } + match_columns_by_header__ = Some(map_.next_value()?); + } GeneratedField::Escape => { if optional_escape__.is_some() { return Err(serde::de::Error::duplicate_field("escape")); @@ -5221,6 +5238,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { newlines_in_values: newlines_in_values__.unwrap_or_default(), truncate_rows: truncate_rows__.unwrap_or_default(), terminator: terminator__, + match_columns_by_header: match_columns_by_header__.unwrap_or_default(), optional_escape: optional_escape__, optional_comment: optional_comment__, }) diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 142e560baea77..c27417d0b497f 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2144,6 +2144,8 @@ pub struct CsvScanExecNode { /// Custom one-byte line terminator. Absent means the default newline terminator. #[prost(bytes = "vec", optional, tag = "9")] pub terminator: ::core::option::Option<::prost::alloc::vec::Vec>, + #[prost(bool, tag = "10")] + pub match_columns_by_header: bool, #[prost(oneof = "csv_scan_exec_node::OptionalEscape", tags = "5")] pub optional_escape: ::core::option::Option, #[prost(oneof = "csv_scan_exec_node::OptionalComment", tags = "6")] diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index 32909afddf83b..0bafea72fa31e 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/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<()> { + use datafusion::common::config::CsvOptions; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, true)])); + let file_source = Arc::new( + CsvSource::new(TableSchema::from(&file_schema)).with_csv_options( + CsvOptions::default() + .with_has_header(true) + .with_match_columns_by_header(true), + ), + ); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.csv".to_string(), + 1024, + )])]) + .build(); + + let ctx = SessionContext::new(); + let roundtripped = roundtrip_test_and_return( + DataSourceExec::from_data_source(scan_config), + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let csv_source = roundtripped + .downcast_ref::() + .and_then(|exec| exec.data_source().downcast_ref::()) + .and_then(|scan| scan.file_source().downcast_ref::()) + .ok_or_else(|| internal_datafusion_err!("Expected CsvSource"))?; + + assert!(csv_source.match_columns_by_header()); + Ok(()) +} + #[tokio::test] async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { let mut file_group = diff --git a/datafusion/sqllogictest/test_files/csv_files.slt b/datafusion/sqllogictest/test_files/csv_files.slt index 549a74c4eb54a..d8ce7e91567dd 100644 --- a/datafusion/sqllogictest/test_files/csv_files.slt +++ b/datafusion/sqllogictest/test_files/csv_files.slt @@ -597,3 +597,43 @@ DROP TABLE stored_trim_both; statement ok DROP TABLE whitespace_source; + +# Match file columns to the table schema by header name +query I +COPY (SELECT 1 AS a, 'x' AS b, 10 AS c) +TO 'test_files/scratch/csv_files/header_layouts/full.csv' +STORED AS csv +OPTIONS ('format.has_header' 'true'); +---- +1 + +query I +COPY (SELECT 20 AS c, 'ignored' AS undeclared, 'y' AS b) +TO 'test_files/scratch/csv_files/header_layouts/reordered.csv' +STORED AS csv +OPTIONS ('format.has_header' 'true'); +---- +1 + +statement ok +CREATE EXTERNAL TABLE header_layouts ( + a INT, + b TEXT, + c INT +) STORED AS CSV +LOCATION 'test_files/scratch/csv_files/header_layouts/' +OPTIONS ('format.has_header' 'true', 'format.match_columns_by_header' 'true'); + +query ITI +SELECT a, b, c FROM header_layouts ORDER BY c; +---- +1 x 10 +NULL y 20 + +query I +SELECT c FROM header_layouts WHERE a IS NULL; +---- +20 + +statement ok +DROP TABLE header_layouts; diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index b567e505e27e4..555e7145b7f88 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1546,3 +1546,26 @@ fn is_nullable(&self) -> bool { ``` See [issue #18882](https://github.com/apache/datafusion/issues/18882) for details. + +### New `match_columns_by_header` field on CSV options + +`datafusion_common::config::CsvOptions` and `datafusion::prelude::CsvReadOptions` +have a new public field, `match_columns_by_header`, which maps CSV file columns +to the table schema by header name instead of by position. The protobuf +messages `CsvOptions` and `CsvScanExecNode` gained the same field. + +Code that constructs these structs with an exhaustive struct literal must add +the field or use `..Default::default()`. Builder methods are unaffected: + +```rust +// 55.0.0 +let options = CsvReadOptions { + has_header: true, + ..Default::default() +}; + +// 56.0.0 +let options = CsvReadOptions::new() + .has_header(true) + .match_columns_by_header(true); +``` diff --git a/docs/source/user-guide/sql/format_options.md b/docs/source/user-guide/sql/format_options.md index a4fc32fe8e277..50079d8b62d89 100644 --- a/docs/source/user-guide/sql/format_options.md +++ b/docs/source/user-guide/sql/format_options.md @@ -105,25 +105,26 @@ OPTIONS('COMPRESSION' 'gzip'); The following options are available when reading or writing CSV files. Note: If any unsupported option is specified, an error will be raised and the query will fail. -| Option | Description | Default Value | -| -------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------ | -| COMPRESSION | Sets the compression that should be applied to the entire CSV file. Supported values are GZIP, BZIP2, XZ, ZSTD, and UNCOMPRESSED. | UNCOMPRESSED | -| HAS_HEADER | Sets if the CSV file should include column headers. If not set, uses session or system default. | None | -| DELIMITER | Sets the character which should be used as the column delimiter within the CSV file. | `,` (comma) | -| QUOTE | Sets the character which should be used for quoting values within the CSV file. | `"` (double quote) | -| TERMINATOR | Sets the character which should be used as the line terminator within the CSV file. | None | -| ESCAPE | Sets the character which should be used for escaping special characters within the CSV file. | None | -| DOUBLE_QUOTE | Sets if quotes within quoted fields should be escaped by doubling them (e.g., `"aaa""bbb"`). | None | -| NEWLINES_IN_VALUES | Sets if newlines in quoted values are supported. If not set, uses session or system default. | None | -| DATE_FORMAT | Sets the format that dates should be encoded in within the CSV file. | None | -| DATETIME_FORMAT | Sets the format that datetimes should be encoded in within the CSV file. | None | -| TIMESTAMP_FORMAT | Sets the format that timestamps should be encoded in within the CSV file. | None | -| TIMESTAMP_TZ_FORMAT | Sets the format that timestamps with timezone should be encoded in within the CSV file. | None | -| TIME_FORMAT | Sets the format that times should be encoded in within the CSV file. | None | -| NULL_VALUE | Sets the string which should be used to indicate null values within the CSV file. | None | -| NULL_REGEX | Sets the regex pattern to match null values when loading CSVs. | None | -| SCHEMA_INFER_MAX_REC | Sets the maximum number of records to scan to infer the schema. If set to 0, schema inference is disabled and all fields will be inferred as Utf8 (string) type. | None | -| COMMENT | Sets the character which should be used to indicate comment lines in the CSV file. | None | +| Option | Description | Default Value | +| ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ------------------ | +| COMPRESSION | Sets the compression that should be applied to the entire CSV file. Supported values are GZIP, BZIP2, XZ, ZSTD, and UNCOMPRESSED. | UNCOMPRESSED | +| HAS_HEADER | Sets if the CSV file should include column headers. If not set, uses session or system default. | None | +| DELIMITER | Sets the character which should be used as the column delimiter within the CSV file. | `,` (comma) | +| QUOTE | Sets the character which should be used for quoting values within the CSV file. | `"` (double quote) | +| TERMINATOR | Sets the character which should be used as the line terminator within the CSV file. | None | +| ESCAPE | Sets the character which should be used for escaping special characters within the CSV file. | None | +| DOUBLE_QUOTE | Sets if quotes within quoted fields should be escaped by doubling them (e.g., `"aaa""bbb"`). | None | +| NEWLINES_IN_VALUES | Sets if newlines in quoted values are supported. If not set, uses session or system default. | None | +| DATE_FORMAT | Sets the format that dates should be encoded in within the CSV file. | None | +| DATETIME_FORMAT | Sets the format that datetimes should be encoded in within the CSV file. | None | +| TIMESTAMP_FORMAT | Sets the format that timestamps should be encoded in within the CSV file. | None | +| TIMESTAMP_TZ_FORMAT | Sets the format that timestamps with timezone should be encoded in within the CSV file. | None | +| TIME_FORMAT | Sets the format that times should be encoded in within the CSV file. | None | +| NULL_VALUE | Sets the string which should be used to indicate null values within the CSV file. | None | +| NULL_REGEX | Sets the regex pattern to match null values when loading CSVs. | None | +| SCHEMA_INFER_MAX_REC | Sets the maximum number of records to scan to infer the schema. If set to 0, schema inference is disabled and all fields will be inferred as Utf8 (string) type. | None | +| COMMENT | Sets the character which should be used to indicate comment lines in the CSV file. | None | +| MATCH_COLUMNS_BY_HEADER | Only applies when `HAS_HEADER` is true. Matches columns to the table schema by header name instead of by position. Missing columns are read as null and undeclared columns are ignored. Names are matched exactly (case-sensitive); duplicate header names resolve to the first occurrence. A leading UTF-8 BOM, blank lines, and comment lines before the header are skipped, as the CSV parser does. | None | **Example:**