Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 14 additions & 13 deletions datafusion/datasource-csv/src/source.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)),
Expand Down Expand Up @@ -586,16 +590,13 @@ fn proto_str_to_byte(s: &str, description: &str) -> Result<u8> {
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<Arc<dyn ExecutionPlan>> {
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;

Expand Down Expand Up @@ -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)?;

Expand All @@ -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))
}
}
17 changes: 15 additions & 2 deletions datafusion/datasource-json/src/source.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<JsonSource> for Arc<dyn FileSource> {
Expand Down Expand Up @@ -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)),
Expand All @@ -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<'_>,
Expand All @@ -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))
Expand Down
39 changes: 33 additions & 6 deletions datafusion/datasource/src/file_scan_config/proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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` /
Expand All @@ -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};
Expand All @@ -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;

Expand All @@ -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<'_>,
Expand Down Expand Up @@ -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()),
Expand All @@ -131,14 +142,16 @@ impl FileScanConfig {
batch_size: self.batch_size.map(|s| s as u64),
projection_exprs,
output_partitioning,
file_compression_type,
})
}

/// Reconstruct a [`FileScanConfig`] from a [`protobuf::FileScanExecConf`]
/// 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<'_>,
Expand Down Expand Up @@ -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<ProjectionExpr> = proto_projection_exprs
Expand Down Expand Up @@ -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())
}

Expand Down
7 changes: 7 additions & 0 deletions datafusion/proto-models/proto/datafusion.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down
Loading
Loading