diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index d5fc6288eaaa3..c867d9fcb438d 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -362,6 +362,10 @@ impl FileSource for CsvSource { .transpose()?, newlines_in_values: self.newlines_in_values(), truncate_rows: self.truncate_rows(), + terminator: self + .terminator() + .map(|terminator| proto_byte_to_string(terminator, "terminator")) + .transpose()?, }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), @@ -586,16 +590,13 @@ fn proto_str_to_byte(s: &str, description: &str) -> Result { impl CsvSource { /// Reconstructs a `DataSourceExec` from a protobuf `CsvScan`. /// - /// Custom line terminators are not represented in the wire format. + /// Payloads without a terminator use the default newline terminator. 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::file_scan_config::FileScanConfig; use datafusion_datasource::source::DataSourceExec; use datafusion_proto_models::protobuf; @@ -626,6 +627,11 @@ impl CsvSource { } None => None, }; + let terminator = scan + .terminator + .as_deref() + .map(|terminator| proto_str_to_byte(terminator, "terminator")) + .transpose()?; let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; @@ -641,16 +647,11 @@ impl CsvSource { CsvSource::new(table_schema) .with_csv_options(csv_options) .with_escape(escape) - .with_comment(comment), + .with_comment(comment) + .with_terminator(terminator), ); - // 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(); + let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?; Ok(DataSourceExec::from_data_source(conf)) } } diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index c6d420bafb2f7..3c9e4db49e285 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -159,6 +159,11 @@ impl JsonSource { self.newline_delimited = newline_delimited; self } + + /// Returns whether this source reads newline-delimited JSON. + pub fn is_newline_delimited(&self) -> bool { + self.newline_delimited + } } impl From for Arc { @@ -260,6 +265,11 @@ impl FileSource for JsonSource { let node = protobuf::JsonScanExecNode { base_conf: Some(base.try_to_proto(ctx)?), + newline_delimited: if self.newline_delimited { + None + } else { + Some(false) + }, }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::JsonScan(node)), @@ -271,7 +281,7 @@ impl FileSource for JsonSource { impl JsonSource { /// Reconstructs a `DataSourceExec` from a protobuf `JsonScan`. /// - /// Defaults to newline-delimited JSON because protobuf does not encode the mode. + /// Payloads without a mode default to newline-delimited JSON. pub fn try_from_proto( node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, @@ -296,7 +306,10 @@ impl JsonSource { })?; let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; - let source = Arc::new(JsonSource::new(table_schema)); + let source = Arc::new( + JsonSource::new(table_schema) + .with_newline_delimited(scan.newline_delimited.unwrap_or(true)), + ); let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?; Ok(DataSourceExec::from_data_source(conf)) diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index d7135173c8934..efd17e4f57091 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -27,8 +27,8 @@ //! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its //! `*ScanExecNode` around [`FileScanConfig::try_to_proto`] and decodes with //! [`FileScanConfig::try_from_proto`], keeping a single copy of the shared -//! wire logic. The wire format is byte-for-byte identical to the old central -//! serializer. +//! wire logic. Existing fields remain wire-compatible with the old central +//! serializer; new options use optional fields with legacy defaults. //! //! Child physical expressions (sort orderings, hash/range partitioning, and //! projection expressions) are (de)serialized through `ctx.encode_expr` / @@ -38,6 +38,7 @@ use std::sync::Arc; use arrow::datatypes::Schema; +use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; @@ -46,9 +47,11 @@ use datafusion_physical_expr_common::sort_expr::{ sort_exprs_try_from_proto, sort_exprs_try_to_proto, }; use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; +use datafusion_proto_models::datafusion_common::CompressionTypeVariant as ProtoCompressionTypeVariant; use datafusion_proto_models::protobuf; use crate::file::FileSource; +use crate::file_compression_type::FileCompressionType; use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use crate::table_schema::TableSchema; @@ -57,8 +60,9 @@ impl FileScanConfig { /// [`protobuf::FileScanExecConf`]. /// /// Each concrete [`FileSource::try_to_proto`] - /// wraps the returned value in its own `*ScanExecNode`. Byte-compatible with - /// the former `serialize_file_scan_config` in `datafusion-proto`. + /// wraps the returned value in its own `*ScanExecNode`. Existing fields are + /// byte-compatible with the former `serialize_file_scan_config` in + /// `datafusion-proto`. pub fn try_to_proto( &self, ctx: &ExecutionPlanEncodeCtx<'_>, @@ -114,6 +118,13 @@ impl FileScanConfig { }) .transpose()?; + let file_compression_type = + self.file_compression_type.is_compressed().then(|| { + let compression: ProtoCompressionTypeVariant = + (*self.file_compression_type.get_variant()).into(); + compression as i32 + }); + Ok(protobuf::FileScanExecConf { file_groups, statistics: Some((&self.statistics()).into()), @@ -131,6 +142,7 @@ impl FileScanConfig { batch_size: self.batch_size.map(|s| s as u64), projection_exprs, output_partitioning, + file_compression_type, }) } @@ -138,7 +150,8 @@ impl FileScanConfig { /// and a `file_source` the caller has already rebuilt (typically from the /// table schema via [`FileScanConfig::parse_table_schema_from_proto`]). /// - /// Byte-compatible with the former `parse_protobuf_file_scan_config`. + /// Existing fields are byte-compatible with the former + /// `parse_protobuf_file_scan_config`. pub fn try_from_proto( conf: &protobuf::FileScanExecConf, ctx: &ExecutionPlanDecodeCtx<'_>, @@ -194,6 +207,19 @@ impl FileScanConfig { .transpose()? .flatten(); + let file_compression_type = conf + .file_compression_type + .map(|value| { + let compression = + ProtoCompressionTypeVariant::try_from(value).map_err(|_| { + internal_datafusion_err!("Unknown file compression type: {value}") + })?; + let compression: CompressionTypeVariant = compression.into(); + Ok::<_, DataFusionError>(FileCompressionType::from(compression)) + }) + .transpose()? + .unwrap_or(FileCompressionType::UNCOMPRESSED); + // Parse projection expressions if present and apply to the file source. let file_source = if let Some(proto_projection_exprs) = &conf.projection_exprs { let projection_exprs: Vec = proto_projection_exprs @@ -226,7 +252,8 @@ impl FileScanConfig { .with_limit(conf.limit.as_ref().map(|sl| sl.limit as usize)) .with_output_ordering(output_ordering) .with_output_partitioning(output_partitioning) - .with_batch_size(conf.batch_size.map(|s| s as usize)); + .with_batch_size(conf.batch_size.map(|s| s as usize)) + .with_file_compression_type(file_compression_type); Ok(config_builder.build()) } diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index 99b4ef6272b2f..e84f893319f3c 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1254,6 +1254,9 @@ message FileScanExecConf { reserved 14; reserved "partitioned_by_file_group"; optional Partitioning output_partitioning = 15; + // Compression used by formats such as CSV and JSON. Absent means uncompressed + // for compatibility with payloads written before this field existed. + optional datafusion_common.CompressionTypeVariant file_compression_type = 16; } message ParquetScanExecNode { @@ -1280,10 +1283,14 @@ message CsvScanExecNode { } bool newlines_in_values = 7; bool truncate_rows = 8; + // Custom line terminator. Absent means the default newline terminator. + optional string terminator = 9; } message JsonScanExecNode { FileScanExecConf base_conf = 1; + // Absent means newline-delimited JSON for compatibility with older payloads. + optional bool newline_delimited = 2; } message AvroScanExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 61b1ea3ff1043..e63614f00f9ff 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -4381,6 +4381,9 @@ impl serde::Serialize for CsvScanExecNode { if self.truncate_rows { len += 1; } + if self.terminator.is_some() { + len += 1; + } if self.optional_escape.is_some() { len += 1; } @@ -4406,6 +4409,9 @@ impl serde::Serialize for CsvScanExecNode { if self.truncate_rows { struct_ser.serialize_field("truncateRows", &self.truncate_rows)?; } + if let Some(v) = self.terminator.as_ref() { + struct_ser.serialize_field("terminator", v)?; + } if let Some(v) = self.optional_escape.as_ref() { match v { csv_scan_exec_node::OptionalEscape::Escape(v) => { @@ -4440,6 +4446,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "newlinesInValues", "truncate_rows", "truncateRows", + "terminator", "escape", "comment", ]; @@ -4452,6 +4459,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { Quote, NewlinesInValues, TruncateRows, + Terminator, Escape, Comment, } @@ -4481,6 +4489,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "quote" => Ok(GeneratedField::Quote), "newlinesInValues" | "newlines_in_values" => Ok(GeneratedField::NewlinesInValues), "truncateRows" | "truncate_rows" => Ok(GeneratedField::TruncateRows), + "terminator" => Ok(GeneratedField::Terminator), "escape" => Ok(GeneratedField::Escape), "comment" => Ok(GeneratedField::Comment), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), @@ -4508,6 +4517,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { let mut quote__ = None; let mut newlines_in_values__ = None; let mut truncate_rows__ = None; + let mut terminator__ = None; let mut optional_escape__ = None; let mut optional_comment__ = None; while let Some(k) = map_.next_key()? { @@ -4548,6 +4558,12 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { } truncate_rows__ = Some(map_.next_value()?); } + GeneratedField::Terminator => { + if terminator__.is_some() { + return Err(serde::de::Error::duplicate_field("terminator")); + } + terminator__ = map_.next_value()?; + } GeneratedField::Escape => { if optional_escape__.is_some() { return Err(serde::de::Error::duplicate_field("escape")); @@ -4569,6 +4585,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { quote: quote__.unwrap_or_default(), newlines_in_values: newlines_in_values__.unwrap_or_default(), truncate_rows: truncate_rows__.unwrap_or_default(), + terminator: terminator__, optional_escape: optional_escape__, optional_comment: optional_comment__, }) @@ -7019,6 +7036,9 @@ impl serde::Serialize for FileScanExecConf { if self.output_partitioning.is_some() { len += 1; } + if self.file_compression_type.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.FileScanExecConf", len)?; if !self.file_groups.is_empty() { struct_ser.serialize_field("fileGroups", &self.file_groups)?; @@ -7058,6 +7078,11 @@ impl serde::Serialize for FileScanExecConf { if let Some(v) = self.output_partitioning.as_ref() { struct_ser.serialize_field("outputPartitioning", v)?; } + if let Some(v) = self.file_compression_type.as_ref() { + let v = super::datafusion_common::CompressionTypeVariant::try_from(*v) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", *v)))?; + struct_ser.serialize_field("fileCompressionType", &v)?; + } struct_ser.end() } } @@ -7087,6 +7112,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "projectionExprs", "output_partitioning", "outputPartitioning", + "file_compression_type", + "fileCompressionType", ]; #[allow(clippy::enum_variant_names)] @@ -7103,6 +7130,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { BatchSize, ProjectionExprs, OutputPartitioning, + FileCompressionType, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -7136,6 +7164,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "batchSize" | "batch_size" => Ok(GeneratedField::BatchSize), "projectionExprs" | "projection_exprs" => Ok(GeneratedField::ProjectionExprs), "outputPartitioning" | "output_partitioning" => Ok(GeneratedField::OutputPartitioning), + "fileCompressionType" | "file_compression_type" => Ok(GeneratedField::FileCompressionType), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -7167,6 +7196,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { let mut batch_size__ = None; let mut projection_exprs__ = None; let mut output_partitioning__ = None; + let mut file_compression_type__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::FileGroups => { @@ -7246,6 +7276,12 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { } output_partitioning__ = map_.next_value()?; } + GeneratedField::FileCompressionType => { + if file_compression_type__.is_some() { + return Err(serde::de::Error::duplicate_field("fileCompressionType")); + } + file_compression_type__ = map_.next_value::<::std::option::Option>()?.map(|x| x as i32); + } } } Ok(FileScanExecConf { @@ -7261,6 +7297,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { batch_size: batch_size__, projection_exprs: projection_exprs__, output_partitioning: output_partitioning__, + file_compression_type: file_compression_type__, }) } } @@ -11192,10 +11229,16 @@ impl serde::Serialize for JsonScanExecNode { if self.base_conf.is_some() { len += 1; } + if self.newline_delimited.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.JsonScanExecNode", len)?; if let Some(v) = self.base_conf.as_ref() { struct_ser.serialize_field("baseConf", v)?; } + if let Some(v) = self.newline_delimited.as_ref() { + struct_ser.serialize_field("newlineDelimited", v)?; + } struct_ser.end() } } @@ -11208,11 +11251,14 @@ impl<'de> serde::Deserialize<'de> for JsonScanExecNode { const FIELDS: &[&str] = &[ "base_conf", "baseConf", + "newline_delimited", + "newlineDelimited", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { BaseConf, + NewlineDelimited, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -11235,6 +11281,7 @@ impl<'de> serde::Deserialize<'de> for JsonScanExecNode { { match value { "baseConf" | "base_conf" => Ok(GeneratedField::BaseConf), + "newlineDelimited" | "newline_delimited" => Ok(GeneratedField::NewlineDelimited), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -11255,6 +11302,7 @@ impl<'de> serde::Deserialize<'de> for JsonScanExecNode { V: serde::de::MapAccess<'de>, { let mut base_conf__ = None; + let mut newline_delimited__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::BaseConf => { @@ -11263,10 +11311,17 @@ impl<'de> serde::Deserialize<'de> for JsonScanExecNode { } base_conf__ = map_.next_value()?; } + GeneratedField::NewlineDelimited => { + if newline_delimited__.is_some() { + return Err(serde::de::Error::duplicate_field("newlineDelimited")); + } + newline_delimited__ = map_.next_value()?; + } } } Ok(JsonScanExecNode { base_conf: base_conf__, + newline_delimited: newline_delimited__, }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 233b5fee1b29e..826bbaf2819c1 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1929,6 +1929,14 @@ pub struct FileScanExecConf { pub projection_exprs: ::core::option::Option, #[prost(message, optional, tag = "15")] pub output_partitioning: ::core::option::Option, + /// Compression used by formats such as CSV and JSON. Absent means uncompressed + /// for compatibility with payloads written before this field existed. + #[prost( + enumeration = "super::datafusion_common::CompressionTypeVariant", + optional, + tag = "16" + )] + pub file_compression_type: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetScanExecNode { @@ -1955,6 +1963,9 @@ pub struct CsvScanExecNode { pub newlines_in_values: bool, #[prost(bool, tag = "8")] pub truncate_rows: bool, + /// Custom line terminator. Absent means the default newline terminator. + #[prost(string, optional, tag = "9")] + pub terminator: ::core::option::Option<::prost::alloc::string::String>, #[prost(oneof = "csv_scan_exec_node::OptionalEscape", tags = "5")] pub optional_escape: ::core::option::Option, #[prost(oneof = "csv_scan_exec_node::OptionalComment", tags = "6")] @@ -1977,6 +1988,9 @@ pub mod csv_scan_exec_node { pub struct JsonScanExecNode { #[prost(message, optional, tag = "1")] pub base_conf: ::core::option::Option, + /// Absent means newline-delimited JSON for compatibility with older payloads. + #[prost(bool, optional, tag = "2")] + pub newline_delimited: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct AvroScanExecNode { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index de684857f4446..cbcb4ea199b2f 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -108,6 +108,7 @@ mod file_scan_config_serde { use arrow::datatypes::{DataType, Field}; use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; use datafusion_datasource::file::FileSource; + use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::{ FileScanConfig, FileScanConfigBuilder, @@ -262,6 +263,7 @@ mod file_scan_config_serde { .with_statistics(table_statistics) .with_limit(Some(17)) .with_batch_size(Some(256)) + .with_file_compression_type(FileCompressionType::GZIP) .with_output_ordering(vec![ordering]) .with_output_partitioning(output_partitioning) .build() @@ -370,10 +372,27 @@ mod file_scan_config_serde { assert_eq!(decoded.file_groups[1].len(), 1); assert!(decoded.file_groups[0].files()[0].arrow_schema.is_some()); assert!(decoded.file_groups[0].files()[1].arrow_schema.is_none()); + assert_eq!(decoded.file_compression_type, FileCompressionType::GZIP); Ok(()) } + #[test] + fn new_file_scan_config_decode_without_compression_uses_legacy_default() -> Result<()> + { + let serde = FileScanSerdeHarness::new(); + let mut encoded = serde.encode(&test_config(None))?; + assert!(encoded.file_compression_type.is_some()); + + encoded.file_compression_type = None; + let decoded = serde.decode(&encoded)?; + assert_eq!( + decoded.file_compression_type, + FileCompressionType::UNCOMPRESSED + ); + Ok(()) + } + #[test] fn new_file_scan_config_serde_preserves_projection_presence() -> Result<()> { let serde = FileScanSerdeHarness::new(); @@ -454,6 +473,16 @@ mod file_scan_config_serde { "unexpected error: {err}" ); + let mut unknown_compression = valid; + unknown_compression.file_compression_type = Some(i32::MAX); + let err = serde + .decode(&unknown_compression) + .expect_err("unknown compression type must fail"); + assert!( + err.to_string().contains("Unknown file compression type"), + "unexpected error: {err}" + ); + Ok(()) } diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index 04708dec6439c..eacae87d8afb7 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -53,6 +53,7 @@ use datafusion::scalar::ScalarValue; use datafusion_common::config::TableParquetOptions; use datafusion_common::stats::Precision; use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::{TableSchema, TableSchemaBuilder}; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; @@ -61,6 +62,7 @@ use datafusion_proto::physical_plan::{ AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalProtoConverterExtension, }; +use datafusion_proto::protobuf; use datafusion_proto::protobuf::PhysicalPlanNode; use prost::Message; use std::collections::HashMap; @@ -181,18 +183,69 @@ fn roundtrip_arrow_scan() -> Result<()> { } #[test] -fn roundtrip_json_scan() -> Result<()> { +fn roundtrip_json_scan_preserves_format_options() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - let file_source = Arc::new(JsonSource::new(TableSchema::from(&file_schema))); + let file_source = Arc::new( + JsonSource::new(TableSchema::from(&file_schema)).with_newline_delimited(false), + ); let scan_config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.json".to_string(), + "/path/to/file.json.gz".to_string(), 1024, )])]) + .with_file_compression_type(FileCompressionType::GZIP) .build(); - roundtrip_test(DataSourceExec::from_data_source(scan_config)) + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let plan: Arc = DataSourceExec::from_data_source(scan_config); + let roundtripped = roundtrip_test_and_return( + Arc::clone(&plan), + &ctx, + &codec, + &DefaultPhysicalProtoConverter {}, + )?; + let file_scan = roundtripped + .downcast_ref::() + .and_then(|exec| exec.data_source().downcast_ref::()) + .ok_or_else(|| internal_datafusion_err!("Expected FileScanConfig"))?; + let json_source = file_scan + .file_source() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected JsonSource"))?; + assert!(!json_source.is_newline_delimited()); + assert_eq!(file_scan.file_compression_type, FileCompressionType::GZIP); + + // Payloads written before these fields existed must keep their historical + // defaults: newline-delimited JSON without compression. + let mut node = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; + match node.physical_plan_type.as_mut() { + Some(protobuf::physical_plan_node::PhysicalPlanType::JsonScan(scan)) => { + scan.newline_delimited = None; + scan.base_conf + .as_mut() + .expect("JSON scan has a base config") + .file_compression_type = None; + } + other => return internal_err!("Expected JsonScan node, got {other:?}"), + } + let decoded = node.try_into_physical_plan(ctx.task_ctx().as_ref(), &codec)?; + let file_scan = decoded + .downcast_ref::() + .and_then(|exec| exec.data_source().downcast_ref::()) + .ok_or_else(|| internal_datafusion_err!("Expected FileScanConfig"))?; + let json_source = file_scan + .file_source() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected JsonSource"))?; + assert!(json_source.is_newline_delimited()); + assert_eq!( + file_scan.file_compression_type, + FileCompressionType::UNCOMPRESSED + ); + Ok(()) } #[cfg(feature = "avro")] @@ -227,6 +280,7 @@ fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { quote: b'\'', escape: Some(b'\\'), comment: Some(b'#'), + terminator: Some(b'\r'), newlines_in_values: Some(true), truncated_rows: Some(true), ..Default::default() @@ -235,9 +289,10 @@ fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { 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(), + "/path/to/file.csv.gz".to_string(), 1024, )])]) + .with_file_compression_type(FileCompressionType::GZIP) .build(); let ctx = SessionContext::new(); @@ -264,8 +319,10 @@ fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { assert_eq!(csv_source.quote(), b'\''); assert_eq!(csv_source.escape(), Some(b'\\')); assert_eq!(csv_source.comment(), Some(b'#')); + assert_eq!(csv_source.terminator(), Some(b'\r')); assert!(csv_source.newlines_in_values()); assert!(csv_source.truncate_rows()); + assert_eq!(file_scan.file_compression_type, FileCompressionType::GZIP); Ok(()) }