diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 4026e6e808653..7e7195dfda9d5 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -31,6 +31,7 @@ version.workspace = true all-features = true [features] +# Enables protobuf serialization hooks for CSV sources and sinks. proto = [ "dep:datafusion-proto-models", "datafusion-datasource/proto", diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 25ec311880405..f498a3c5b7fbe 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -308,6 +308,49 @@ impl FileSource for CsvSource { DisplayFormatType::TreeRender => Ok(()), } } + + /// Emit a `CsvScan` node wrapping the shared base config and CSV options. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::CsvScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + has_header: self.has_header(), + delimiter: proto_byte_to_string(self.delimiter(), "delimiter")?, + quote: proto_byte_to_string(self.quote(), "quote")?, + optional_escape: self + .escape() + .map(|escape| { + Ok::<_, DataFusionError>( + protobuf::csv_scan_exec_node::OptionalEscape::Escape( + proto_byte_to_string(escape, "escape")?, + ), + ) + }) + .transpose()?, + optional_comment: self + .comment() + .map(|comment| { + Ok::<_, DataFusionError>( + protobuf::csv_scan_exec_node::OptionalComment::Comment( + proto_byte_to_string(comment, "comment")?, + ), + ) + }) + .transpose()?, + newlines_in_values: self.newlines_in_values(), + truncate_rows: self.truncate_rows(), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), + })) + } } impl FileOpener for CsvOpener { @@ -501,3 +544,97 @@ pub async fn plan_to_csv( Ok(()) } + +#[cfg(feature = "proto")] +fn proto_byte_to_string(b: u8, description: &str) -> Result { + let bytes = &[b]; + let s = std::str::from_utf8(bytes).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Invalid CSV {description}: can not represent {bytes:0x?} as utf8" + ) + })?; + Ok(s.to_owned()) +} + +#[cfg(feature = "proto")] +fn proto_str_to_byte(s: &str, description: &str) -> Result { + datafusion_common::assert_eq_or_internal_err!( + s.len(), + 1, + "Invalid CSV {description}: expected single character, got {s}" + ); + Ok(s.as_bytes()[0]) +} + +#[cfg(feature = "proto")] +impl CsvSource { + /// Reconstructs a `DataSourceExec` from a protobuf `CsvScan`. + /// + /// Custom line terminators are not represented in the wire format. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::config::CsvOptions; + use datafusion_datasource::file_compression_type::FileCompressionType; + use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, + }; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::CsvScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a CsvScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvScanExecNode is missing required field 'base_conf'" + ) + })?; + + let escape = match &scan.optional_escape { + Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) => { + Some(proto_str_to_byte(escape, "escape")?) + } + None => None, + }; + let comment = match &scan.optional_comment { + Some(protobuf::csv_scan_exec_node::OptionalComment::Comment(comment)) => { + Some(proto_str_to_byte(comment, "comment")?) + } + None => None, + }; + + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + + let csv_options = CsvOptions { + has_header: Some(scan.has_header), + delimiter: proto_str_to_byte(&scan.delimiter, "delimiter")?, + quote: proto_str_to_byte(&scan.quote, "quote")?, + newlines_in_values: Some(scan.newlines_in_values), + truncated_rows: Some(scan.truncate_rows), + ..Default::default() + }; + let source = Arc::new( + CsvSource::new(table_schema) + .with_csv_options(csv_options) + .with_escape(escape) + .with_comment(comment), + ); + + // The compression type is not on the wire; CSV scans always + // deserialize as uncompressed. + let conf = FileScanConfigBuilder::from(FileScanConfig::try_from_proto( + base_conf, ctx, source, + )?) + .with_file_compression_type(FileCompressionType::UNCOMPRESSED) + .build(); + Ok(DataSourceExec::from_data_source(conf)) + } +} diff --git a/datafusion/proto/src/common.rs b/datafusion/proto/src/common.rs index bff017edbc998..dd9af97781114 100644 --- a/datafusion/proto/src/common.rs +++ b/datafusion/proto/src/common.rs @@ -15,27 +15,6 @@ // specific language governing permissions and limitations // under the License. -use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err}; - -pub(crate) fn str_to_byte(s: &String, description: &str) -> Result { - assert_eq_or_internal_err!( - s.len(), - 1, - "Invalid CSV {description}: expected single character, got {s}" - ); - Ok(s.as_bytes()[0]) -} - -pub(crate) fn byte_to_string(b: u8, description: &str) -> Result { - let b = &[b]; - let b = std::str::from_utf8(b).map_err(|_| { - internal_datafusion_err!( - "Invalid CSV {description}: can not represent {b:0x?} as utf8" - ) - })?; - Ok(b.to_owned()) -} - #[macro_export] macro_rules! convert_required { ($PB:expr) => {{ diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 7e162bf95454a..22f623c8aa076 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -23,14 +23,12 @@ use std::sync::Arc; use arrow::datatypes::{IntervalMonthDayNanoType, Schema, SchemaRef}; use datafusion_catalog::memory::MemorySourceConfig; -use datafusion_common::config::CsvOptions; use datafusion_common::{ DataFusionError, Result, internal_datafusion_err, internal_err, not_impl_err, }; #[cfg(feature = "parquet")] use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_compression_type::FileCompressionType; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::{DataSource, DataSourceExec}; use datafusion_datasource_arrow::source::ArrowSource; @@ -94,7 +92,6 @@ use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use prost::Message; use prost::bytes::BufMut; -use crate::common::{byte_to_string, str_to_byte}; use crate::convert_required; use crate::physical_plan::from_proto::{ parse_physical_expr_with_converter, parse_physical_sort_exprs, @@ -126,6 +123,7 @@ mod file_scan_config_serde { use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; + use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::file_stream::FileOpener; use datafusion_datasource::{PartitionedFile, TableSchema}; use datafusion_execution::object_store::ObjectStoreUrl; @@ -1084,8 +1082,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::Filter(_) => { FilterExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CsvScan(scan) => { - self.try_into_csv_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::CsvScan(_) => { + CsvSource::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::JsonScan(scan) => { self.try_into_json_scan_physical_plan(scan, ctx, proto_converter) @@ -1340,57 +1338,25 @@ pub trait PhysicalPlanNodeExt: Sized { FilterExec::try_from_proto(&node, &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CsvSource` deserializes itself via `CsvSource::try_from_proto`" + )] fn try_into_csv_scan_physical_plan( &self, scan: &protobuf::CsvScanExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let escape = - if let Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) = - &scan.optional_escape - { - Some(str_to_byte(escape, "escape")?) - } else { - None - }; - - let comment = if let Some( - protobuf::csv_scan_exec_node::OptionalComment::Comment(comment), - ) = &scan.optional_comment - { - Some(str_to_byte(comment, "comment")?) - } else { - None - }; - - // Parse table schema with partition columns - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; - - let csv_options = CsvOptions { - has_header: Some(scan.has_header), - delimiter: str_to_byte(&scan.delimiter, "delimiter")?, - quote: str_to_byte(&scan.quote, "quote")?, - newlines_in_values: Some(scan.newlines_in_values), - ..Default::default() - }; - let source = Arc::new( - CsvSource::new(table_schema) - .with_csv_options(csv_options) - .with_escape(escape) - .with_comment(comment), - ); - - let conf = FileScanConfigBuilder::from(parse_protobuf_file_scan_config( - scan.base_conf.as_ref().unwrap(), + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvScan(scan.clone())), + }; + let decoder = ConverterPlanDecoder { ctx, proto_converter, - source, - )?) - .with_file_compression_type(FileCompressionType::UNCOMPRESSED) - .build(); - Ok(DataSourceExec::from_data_source(conf)) + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + CsvSource::try_from_proto(&node, &decode_ctx) } fn try_into_json_scan_physical_plan( @@ -2561,47 +2527,6 @@ pub trait PhysicalPlanNodeExt: Sized { proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { let data_source = data_source_exec.data_source(); - if let Some(maybe_csv) = data_source.downcast_ref::() { - let source = maybe_csv.file_source(); - if let Some(csv_config) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvScan( - protobuf::CsvScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_csv, - codec, - proto_converter, - )?), - has_header: csv_config.has_header(), - delimiter: byte_to_string( - csv_config.delimiter(), - "delimiter", - )?, - quote: byte_to_string(csv_config.quote(), "quote")?, - optional_escape: if let Some(escape) = csv_config.escape() { - Some( - protobuf::csv_scan_exec_node::OptionalEscape::Escape( - byte_to_string(escape, "escape")?, - ), - ) - } else { - None - }, - optional_comment: if let Some(comment) = csv_config.comment() - { - Some(protobuf::csv_scan_exec_node::OptionalComment::Comment( - byte_to_string(comment, "comment")?, - )) - } else { - None - }, - newlines_in_values: csv_config.newlines_in_values(), - truncate_rows: csv_config.truncate_rows(), - }, - )), - })); - } - } if let Some(scan_conf) = data_source.downcast_ref::() { let source = scan_conf.file_source(); diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index afd4057d0457f..65a817b4cb428 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -36,8 +36,8 @@ use datafusion::datasource::listing::{ }; use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{ - ArrowSource, FileGroup, FileOutputMode, FileScanConfig, FileScanConfigBuilder, - FileSinkConfig, ParquetSource, wrap_partition_type_in_dict, + ArrowSource, CsvSource, FileGroup, FileOutputMode, FileScanConfig, + FileScanConfigBuilder, FileSinkConfig, ParquetSource, wrap_partition_type_in_dict, wrap_partition_value_in_dict, }; use datafusion::datasource::sink::{DataSink, DataSinkExec}; @@ -1363,6 +1363,62 @@ fn roundtrip_arrow_scan() -> Result<()> { roundtrip_test(DataSourceExec::from_data_source(scan_config)) } +#[test] +fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { + use datafusion::common::config::CsvOptions; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let table_schema = TableSchema::from(&file_schema); + let file_source = + Arc::new(CsvSource::new(table_schema).with_csv_options(CsvOptions { + has_header: Some(false), + delimiter: b'|', + quote: b'\'', + escape: Some(b'\\'), + comment: Some(b'#'), + newlines_in_values: Some(true), + truncated_rows: Some(true), + ..Default::default() + })); + + 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 data_source = roundtripped + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected DataSourceExec"))?; + let file_scan = data_source + .data_source() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected FileScanConfig"))?; + let csv_source = file_scan + .file_source() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected CsvSource"))?; + + assert!(!csv_source.has_header()); + assert_eq!(csv_source.delimiter(), b'|'); + assert_eq!(csv_source.quote(), b'\''); + assert_eq!(csv_source.escape(), Some(b'\\')); + assert_eq!(csv_source.comment(), Some(b'#')); + assert!(csv_source.newlines_in_values()); + assert!(csv_source.truncate_rows()); + Ok(()) +} + #[tokio::test] async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { let mut file_group =