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
4 changes: 3 additions & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

48 changes: 48 additions & 0 deletions datafusion/datasource-csv/src/file_format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -926,6 +926,54 @@ impl CsvSink {
}
}

/// Encode a [`CsvFormatFactory`]'s options as their protobuf form.
///
/// The reverse direction is `From<&protobuf::CsvOptions> for CsvOptions` in
/// `datafusion-proto-models`: `CsvOptions` is a `datafusion-common` type, so
/// that half cannot live here.
#[cfg(feature = "proto")]
impl From<&CsvFormatFactory> for datafusion_proto_models::protobuf::CsvOptions {
fn from(factory: &CsvFormatFactory) -> Self {
if let Some(options) = &factory.options {
datafusion_proto_models::protobuf::CsvOptions {
has_header: options.has_header.map_or(vec![], |v| vec![v as u8]),
delimiter: vec![options.delimiter],
quote: vec![options.quote],
terminator: options.terminator.map_or(vec![], |v| vec![v]),
escape: options.escape.map_or(vec![], |v| vec![v]),
double_quote: options.double_quote.map_or(vec![], |v| vec![v as u8]),
compression: options.compression as i32,
schema_infer_max_rec: options.schema_infer_max_rec.map(|v| v as u64),
date_format: options.date_format.clone().unwrap_or_default(),
datetime_format: options.datetime_format.clone().unwrap_or_default(),
timestamp_format: options.timestamp_format.clone().unwrap_or_default(),
timestamp_tz_format: options
.timestamp_tz_format
.clone()
.unwrap_or_default(),
time_format: options.time_format.clone().unwrap_or_default(),
null_value: options.null_value.clone().unwrap_or_default(),
null_regex: options.null_regex.clone().unwrap_or_default(),
comment: options.comment.map_or(vec![], |v| vec![v]),
newlines_in_values: options
.newlines_in_values
.map_or(vec![], |v| vec![v as u8]),
truncated_rows: options.truncated_rows.map_or(vec![], |v| vec![v as u8]),
compression_level: options.compression_level,
quote_style: options.quote_style as i32,
ignore_leading_whitespace: options
.ignore_leading_whitespace
.map_or(vec![], |v| vec![v as u8]),
ignore_trailing_whitespace: options
.ignore_trailing_whitespace
.map_or(vec![], |v| vec![v as u8]),
}
} else {
datafusion_proto_models::protobuf::CsvOptions::default()
}
}
}

#[cfg(test)]
mod tests {
use super::build_schema_helper;
Expand Down
21 changes: 21 additions & 0 deletions datafusion/datasource-json/src/file_format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -615,3 +615,24 @@ impl Decoder for JsonDecoder {
false
}
}

/// Encode a [`JsonFormatFactory`]'s options as their protobuf form.
///
/// The reverse direction is `From<&protobuf::JsonOptions> for JsonOptions` in
/// `datafusion-proto-models`: `JsonOptions` is a `datafusion-common` type, so
/// that half cannot live here.
#[cfg(feature = "proto")]
impl From<&JsonFormatFactory> for datafusion_proto_models::protobuf::JsonOptions {
fn from(factory: &JsonFormatFactory) -> Self {
if let Some(options) = &factory.options {
datafusion_proto_models::protobuf::JsonOptions {
compression: options.compression as i32,
schema_infer_max_rec: options.schema_infer_max_rec.map(|v| v as u64),
compression_level: options.compression_level,
newline_delimited: Some(options.newline_delimited),
}
} else {
datafusion_proto_models::protobuf::JsonOptions::default()
}
}
}
126 changes: 126 additions & 0 deletions datafusion/datasource-parquet/src/file_format.rs
Original file line number Diff line number Diff line change
Expand Up @@ -680,3 +680,129 @@ pub fn statistics_from_parquet_meta_calc(
) -> Result<Statistics> {
DFParquetMetadata::statistics_from_parquet_metadata(metadata, &table_schema)
}

#[cfg(feature = "proto")]
use datafusion_proto_models::protobuf::{self, parquet_column_options, parquet_options};

/// Encode a [`ParquetFormatFactory`]'s options as their protobuf form.
///
/// The reverse direction is `TryFrom<&protobuf::TableParquetOptions> for
/// TableParquetOptions` in `datafusion-proto-models`: `TableParquetOptions` is
/// a `datafusion-common` type, so that half cannot live here.
#[cfg(feature = "proto")]
impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions {
fn from(factory: &ParquetFormatFactory) -> Self {
let global_options = if let Some(ref options) = factory.options {
options.clone()
} else {
return protobuf::TableParquetOptions::default();
};

let column_specific_options = global_options.column_specific_options;
protobuf::TableParquetOptions {
global: Some(protobuf::ParquetOptions {
enable_page_index: global_options.global.enable_page_index,
pruning: global_options.global.pruning,
skip_metadata: global_options.global.skip_metadata,
metadata_size_hint_opt: global_options.global.metadata_size_hint.map(|size| {
parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size as u64)
}),
pushdown_filters: global_options.global.pushdown_filters,
reorder_filters: global_options.global.reorder_filters,
force_filter_selections: global_options.global.force_filter_selections,
data_pagesize_limit: global_options.global.data_pagesize_limit as u64,
write_batch_size: global_options.global.write_batch_size as u64,
writer_version: global_options.global.writer_version.to_string(),
compression_opt: global_options.global.compression.map(|compression| {
parquet_options::CompressionOpt::Compression(compression)
}),
dictionary_enabled_opt: global_options.global.dictionary_enabled.map(|enabled| {
parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled)
}),
dictionary_page_size_limit: global_options.global.dictionary_page_size_limit as u64,
statistics_enabled_opt: global_options.global.statistics_enabled.map(|enabled| {
parquet_options::StatisticsEnabledOpt::StatisticsEnabled(enabled)
}),
max_row_group_size: global_options.global.max_row_group_size as u64,
max_in_list_size: global_options.global.max_in_list_size as u64,
created_by: global_options.global.created_by.clone(),
column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| {
parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64)
}),
statistics_truncate_length_opt: global_options.global.statistics_truncate_length.map(|length| {
parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length as u64)
}),
data_page_row_count_limit: global_options.global.data_page_row_count_limit as u64,
encoding_opt: global_options.global.encoding.map(|encoding| {
parquet_options::EncodingOpt::Encoding(encoding)
}),
bloom_filter_on_read: global_options.global.bloom_filter_on_read,
bloom_filter_on_write: global_options.global.bloom_filter_on_write,
bloom_filter_fpp_opt: global_options.global.bloom_filter_fpp.map(|fpp| {
parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp)
}),
bloom_filter_ndv_opt: global_options.global.bloom_filter_ndv.map(|ndv| {
parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv)
}),
allow_single_file_parallelism: global_options.global.allow_single_file_parallelism,
maximum_parallel_row_group_writers: global_options.global.maximum_parallel_row_group_writers as u64,
maximum_buffered_record_batches_per_stream: global_options.global.maximum_buffered_record_batches_per_stream as u64,
schema_force_view_types: global_options.global.schema_force_view_types,
binary_as_string: global_options.global.binary_as_string,
skip_arrow_metadata: global_options.global.skip_arrow_metadata,
coerce_int96_opt: global_options.global.coerce_int96.map(|compression| {
parquet_options::CoerceInt96Opt::CoerceInt96(compression)
}),
coerce_int96_tz_opt: global_options.global.coerce_int96_tz.map(|tz| {
parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz)
}),
max_predicate_cache_size_opt: global_options.global.max_predicate_cache_size.map(|size| {
parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(size as u64)
}),
max_row_group_bytes_opt: global_options.global.max_row_group_bytes.map(|size| {
parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(size.get() as u64)
}),
content_defined_chunking: Some(protobuf::ParquetCdcOptions {
enabled: global_options.global.content_defined_chunking.enabled,
min_chunk_size: global_options.global.content_defined_chunking.min_chunk_size as u64,
max_chunk_size: global_options.global.content_defined_chunking.max_chunk_size as u64,
norm_level: global_options.global.content_defined_chunking.norm_level,
}),
}),
column_specific_options: column_specific_options.into_iter().map(|(column_name, options)| {
protobuf::ParquetColumnSpecificOptions {
column_name,
options: Some(protobuf::ParquetColumnOptions {
bloom_filter_enabled_opt: options.bloom_filter_enabled.map(|enabled| {
parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(enabled)
}),
encoding_opt: options.encoding.map(|encoding| {
parquet_column_options::EncodingOpt::Encoding(encoding)
}),
dictionary_enabled_opt: options.dictionary_enabled.map(|enabled| {
parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(enabled)
}),
compression_opt: options.compression.map(|compression| {
parquet_column_options::CompressionOpt::Compression(compression)
}),
statistics_enabled_opt: options.statistics_enabled.map(|enabled| {
parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(enabled)
}),
bloom_filter_fpp_opt: options.bloom_filter_fpp.map(|fpp| {
parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(fpp)
}),
bloom_filter_ndv_opt: options.bloom_filter_ndv.map(|ndv| {
parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(ndv)
}),
})
}
}).collect(),
key_value_metadata: global_options.key_value_metadata
.iter()
.filter_map(|(key, value)| {
value.as_ref().map(|v| (key.clone(), v.clone()))
})
.collect(),
}
}
}
73 changes: 69 additions & 4 deletions datafusion/datasource/src/proto.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,9 @@
//! Protobuf conversions for the file-scan leaf types owned by this crate:
//! [`FileRange`], [`PartitionedFile`] and [`FileGroup`].
//!
//! These are the single copy of that wire logic. `datafusion-proto`'s
//! `TryFromProto` implementations for the same types are thin shims that
//! delegate here, so the format cannot drift between the central serializer and
//! the per-source `try_to_proto` hooks.
//! These are the single copy of that wire logic, used both by the central
//! serializer in `datafusion-proto` and by the per-source `try_to_proto` hooks,
//! so the format cannot drift between them.
//!
//! None of these conversions need a codec or an encode/decode context: every
//! field is plain data or goes through `datafusion-proto-common`. That is why
Expand Down Expand Up @@ -204,6 +203,53 @@ mod tests {
Ok(())
}

#[test]
fn partitioned_file_path_roundtrip_percent_encoded() -> Result<()> {
// The wire format carries the *encoded* path, so a location that already
// contains percent escapes must survive without a second round of
// encoding or decoding.
let path_str = "foo/foo%2Fbar/baz%252Fqux";
let pf = PartitionedFile::new_from_meta(ObjectMeta {
location: Path::parse(path_str)?,
last_modified: Utc.timestamp_nanos(1_000),
size: 42,
e_tag: None,
version: None,
});

let encoded = protobuf::PartitionedFile::try_from(&pf)?;
assert_eq!(encoded.path, path_str);

let decoded = PartitionedFile::try_from(&encoded)?;
assert_eq!(decoded.object_meta.location.as_ref(), path_str);
assert_eq!(decoded.object_meta.location, pf.object_meta.location);
Ok(())
}

#[test]
fn partitioned_file_arrow_schema_roundtrip_preserves_metadata() -> Result<()> {
use std::collections::HashMap;

let arrow_schema = Arc::new(Schema::new_with_metadata(
vec![
Field::new("id", DataType::Int64, false),
Field::new("value", DataType::Utf8, true).with_metadata(HashMap::from([
("field_meta".to_string(), "field_value".to_string()),
])),
],
HashMap::from([("schema_meta".to_string(), "schema_value".to_string())]),
));
let pf = PartitionedFile::new("foo/bar.parquet", 10)
.with_arrow_schema(Arc::clone(&arrow_schema));

let encoded = protobuf::PartitionedFile::try_from(&pf)?;
assert!(encoded.arrow_schema.is_some());

let decoded = PartitionedFile::try_from(&encoded)?;
assert_eq!(decoded.arrow_schema.as_deref(), Some(arrow_schema.as_ref()));
Ok(())
}

#[test]
fn partitioned_file_from_proto_rejects_invalid_path() {
let proto = protobuf::PartitionedFile {
Expand All @@ -218,6 +264,25 @@ mod tests {
);
}

#[test]
fn file_group_from_slice_matches_file_group() -> Result<()> {
// `protobuf::FileGroup: TryFrom<&[T]>` lives in `datafusion-proto-models`,
// generic over the element so that crate never names `PartitionedFile`.
// This is the caller-visible half: the bound resolves via
// `TryFrom<&PartitionedFile> for protobuf::PartitionedFile` above.
let files = vec![
PartitionedFile::new("a.parquet", 1),
PartitionedFile::new("b.parquet", 2),
];

let from_slice = protobuf::FileGroup::try_from(&files[..])?;
let from_group = protobuf::FileGroup::try_from(&FileGroup::new(files))?;

assert_eq!(from_slice, from_group);
assert_eq!(from_slice.files.len(), 2);
Ok(())
}

#[test]
fn file_group_roundtrip() -> Result<()> {
let group = FileGroup::new(vec![
Expand Down
6 changes: 6 additions & 0 deletions datafusion/expr/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,10 @@ name = "datafusion_expr"

[features]
default = ["sql"]
# Enables protobuf conversions for the expression types owned by this crate.
# Off by default so consumers that never serialize plans pay nothing. Mirrors
# the `proto` feature on `datafusion-datasource` and friends.
proto = ["dep:datafusion-proto-common", "dep:datafusion-proto-models"]
recursive_protection = ["dep:recursive"]
sql = ["sqlparser"]

Expand All @@ -56,6 +60,8 @@ datafusion-expr-common = { workspace = true }
datafusion-functions-aggregate-common = { workspace = true }
datafusion-functions-window-common = { workspace = true }
datafusion-physical-expr-common = { workspace = true }
datafusion-proto-common = { workspace = true, optional = true }
datafusion-proto-models = { workspace = true, optional = true }
indexmap = { workspace = true }
itertools = { workspace = true }
recursive = { workspace = true, optional = true }
Expand Down
5 changes: 5 additions & 0 deletions datafusion/expr/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,11 @@ pub mod dml {
pub use crate::logical_plan::dml::*;
}
pub mod planner;
/// Protobuf conversions for [`WindowFrame`], [`WindowFrameBound`],
/// [`WindowFrameUnits`], [`MergeIntoClauseKind`](dml::MergeIntoClauseKind) and
/// [`NullTreatment`](expr::NullTreatment), gated on the `proto` feature.
#[cfg(feature = "proto")]
mod proto;
pub mod registry;
pub mod simplify;
pub mod sort_properties {
Expand Down
Loading
Loading