diff --git a/Cargo.lock b/Cargo.lock index c013544b4e1d4..8f68f51237eab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2552,11 +2552,13 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "doc-comment", + "flate2", "object_store", "pretty_assertions", "prost", "recursive", "serde_json", + "tempfile", "tokio", ] diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 04533d8c7882d..0bfa499d7a111 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -356,6 +356,7 @@ impl FileSource for CsvSource { .transpose()?, newlines_in_values: self.newlines_in_values(), truncate_rows: self.truncate_rows(), + terminator: self.terminator().map(|terminator| vec![terminator]), }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), @@ -576,20 +577,28 @@ fn proto_str_to_byte(s: &str, description: &str) -> Result { Ok(s.as_bytes()[0]) } +#[cfg(feature = "proto")] +fn proto_bytes_to_byte(bytes: &[u8], description: &str) -> Result { + let [byte] = bytes else { + return datafusion_common::internal_err!( + "Invalid CSV {description}: expected exactly one byte, got {}", + bytes.len() + ); + }; + Ok(*byte) +} + #[cfg(feature = "proto")] 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; @@ -617,6 +626,11 @@ impl CsvSource { } None => None, }; + let terminator = scan + .terminator + .as_deref() + .map(|terminator| proto_bytes_to_byte(terminator, "terminator")) + .transpose()?; let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; @@ -632,16 +646,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 15840223ccb20..5b576c1166dc1 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 { @@ -254,6 +259,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)), @@ -265,7 +275,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<'_>, @@ -289,7 +299,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 d98b67a66e0a9..e6d600e7d0641 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1265,6 +1265,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 { @@ -1291,10 +1294,14 @@ message CsvScanExecNode { } bool newlines_in_values = 7; bool truncate_rows = 8; + // Custom one-byte line terminator. Absent means the default newline terminator. + optional bytes 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 21309bb2d0941..e6c9117705f1e 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -4471,6 +4471,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; } @@ -4496,6 +4499,11 @@ 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() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("terminator", pbjson::private::base64::encode(&v).as_str())?; + } if let Some(v) = self.optional_escape.as_ref() { match v { csv_scan_exec_node::OptionalEscape::Escape(v) => { @@ -4530,6 +4538,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { "newlinesInValues", "truncate_rows", "truncateRows", + "terminator", "escape", "comment", ]; @@ -4542,6 +4551,7 @@ impl<'de> serde::Deserialize<'de> for CsvScanExecNode { Quote, NewlinesInValues, TruncateRows, + Terminator, Escape, Comment, } @@ -4571,6 +4581,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)), @@ -4598,6 +4609,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()? { @@ -4638,6 +4650,14 @@ 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::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) + ; + } GeneratedField::Escape => { if optional_escape__.is_some() { return Err(serde::de::Error::duplicate_field("escape")); @@ -4659,6 +4679,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__, }) @@ -7109,6 +7130,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)?; @@ -7148,6 +7172,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() } } @@ -7177,6 +7206,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "projectionExprs", "output_partitioning", "outputPartitioning", + "file_compression_type", + "fileCompressionType", ]; #[allow(clippy::enum_variant_names)] @@ -7193,6 +7224,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 @@ -7226,6 +7258,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)), } } @@ -7257,6 +7290,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 => { @@ -7336,6 +7370,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 { @@ -7351,6 +7391,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__, }) } } @@ -11282,10 +11323,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() } } @@ -11298,11 +11345,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 @@ -11325,6 +11375,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)), } } @@ -11345,6 +11396,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 => { @@ -11353,10 +11405,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 d830624322e14..694d32d7616a1 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1947,6 +1947,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 { @@ -1973,6 +1981,9 @@ pub struct CsvScanExecNode { pub newlines_in_values: bool, #[prost(bool, tag = "8")] pub truncate_rows: bool, + /// 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(oneof = "csv_scan_exec_node::OptionalEscape", tags = "5")] pub optional_escape: ::core::option::Option, #[prost(oneof = "csv_scan_exec_node::OptionalComment", tags = "6")] @@ -1995,6 +2006,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/Cargo.toml b/datafusion/proto/Cargo.toml index b6e9d258681e8..4b11bea01103f 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -83,6 +83,7 @@ serde_json = { workspace = true, optional = true } async-trait = { workspace = true } datafusion = { workspace = true, default-features = false, features = [ "sql", + "compression", "datetime_expressions", "nested_expressions", "unicode_expressions", @@ -91,5 +92,7 @@ datafusion-functions = { workspace = true, default-features = true } datafusion-functions-aggregate = { workspace = true } datafusion-functions-window-common = { workspace = true } doc-comment = { workspace = true } +flate2 = { workspace = true } pretty_assertions = "1.4" +tempfile = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread"] } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 222901aff5211..5f534ef63336d 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -106,6 +106,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, @@ -257,6 +258,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() @@ -365,10 +367,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(); @@ -449,6 +468,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 ca8dca74d7ed0..4392ef52381ce 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; @@ -280,18 +281,119 @@ fn arrow_scan_without_format_field_decodes_as_file_format() -> 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(()) +} + +#[tokio::test] +async fn roundtrip_compressed_json_array_scan_executes() -> Result<()> { + use datafusion::prelude::JsonReadOptions; + use flate2::Compression; + use flate2::write::GzEncoder; + use std::io::Write; + + let tmp_dir = tempfile::TempDir::new()?; + let path = tmp_dir.path().join("array.json.gz"); + let file = std::fs::File::create(&path)?; + let mut encoder = GzEncoder::new(file, Compression::default()); + encoder.write_all(br#"[{"a": 1, "b": "hello"}, {"a": 2, "b": "world"}]"#)?; + encoder.finish()?; + + let ctx = SessionContext::new(); + let options = JsonReadOptions::default() + .newline_delimited(false) + .file_compression_type(FileCompressionType::GZIP) + .file_extension(".json.gz"); + ctx.register_json("test_table", path.to_string_lossy(), options) + .await?; + + let initial_plan = ctx + .sql("SELECT a, b FROM test_table ORDER BY a") + .await? + .create_physical_plan() + .await?; + let roundtripped = roundtrip_test_and_return( + initial_plan, + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let batches = + datafusion::physical_plan::collect(roundtripped, ctx.task_ctx()).await?; + + datafusion::assert_batches_eq!( + &[ + "+---+-------+", + "| a | b |", + "+---+-------+", + "| 1 | hello |", + "| 2 | world |", + "+---+-------+", + ], + &batches + ); + Ok(()) } #[cfg(feature = "avro")] @@ -326,6 +428,7 @@ fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { quote: b'\'', escape: Some(b'\\'), comment: Some(b'#'), + terminator: Some(0xff), newlines_in_values: Some(true), truncated_rows: Some(true), ..Default::default() @@ -334,16 +437,19 @@ 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(); + let codec = DefaultPhysicalExtensionCodec {}; + let plan: Arc = DataSourceExec::from_data_source(scan_config); let roundtripped = roundtrip_test_and_return( - DataSourceExec::from_data_source(scan_config), + Arc::clone(&plan), &ctx, - &DefaultPhysicalExtensionCodec {}, + &codec, &DefaultPhysicalProtoConverter {}, )?; let data_source = roundtripped @@ -363,8 +469,28 @@ 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(0xff)); assert!(csv_source.newlines_in_values()); assert!(csv_source.truncate_rows()); + assert_eq!(file_scan.file_compression_type, FileCompressionType::GZIP); + + for invalid_terminator in [vec![], vec![b'\r', b'\n']] { + let mut node = + PhysicalPlanNode::try_from_physical_plan(Arc::clone(&plan), &codec)?; + match node.physical_plan_type.as_mut() { + Some(protobuf::physical_plan_node::PhysicalPlanType::CsvScan(scan)) => { + scan.terminator = Some(invalid_terminator); + } + other => return internal_err!("Expected CsvScan node, got {other:?}"), + } + let err = node + .try_into_physical_plan(ctx.task_ctx().as_ref(), &codec) + .expect_err("invalid terminator length must fail"); + assert!( + err.to_string().contains("expected exactly one byte"), + "unexpected error: {err}" + ); + } Ok(()) }