diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 3c66d4dcd74fb..92286c37a670a 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -423,6 +423,18 @@ impl DataSink for ParquetSink { use datafusion_proto_models::protobuf; use protobuf::physical_plan_node::PhysicalPlanType; + // Keep the active hook exhaustive while centralizing field mapping in + // the exhaustive `TryFrom<&ParquetSink>` below. + let Self { + config: _, + parquet_options: _, + // Runtime output state, not part of the plan. + written: _, + sorting_columns: _, + // Runtime metrics are recreated on decode. + metrics: _, + } = self; + let input = ctx.encode_child(exec.input())?; let sort_order = exec.encode_sort_order(ctx)?; let sink = protobuf::ParquetSink::try_from(self)?; @@ -443,9 +455,43 @@ impl TryFrom<&ParquetSink> for datafusion_proto_models::protobuf::ParquetSink { type Error = DataFusionError; fn try_from(value: &ParquetSink) -> Result { + use datafusion_proto_models::protobuf; + + let ParquetSink { + config, + parquet_options, + // Runtime output state, not part of the plan. + written: _, + sorting_columns, + // Runtime metrics are recreated on decode. + metrics: _, + } = value; + let sorting_columns = + sorting_columns + .as_ref() + .map(|columns| protobuf::ParquetSortingColumns { + columns: columns + .iter() + .map( + |&SortingColumn { + column_idx, + descending, + nulls_first, + }| { + protobuf::ParquetSortingColumn { + column_idx, + descending, + nulls_first, + } + }, + ) + .collect(), + }); + Ok(Self { - config: Some(value.config().try_into()?), - parquet_options: Some(value.parquet_options().try_into()?), + config: Some(config.try_into()?), + parquet_options: Some(parquet_options.try_into()?), + sorting_columns, }) } } @@ -455,14 +501,19 @@ impl TryFrom<&datafusion_proto_models::protobuf::ParquetSink> for ParquetSink { type Error = DataFusionError; fn try_from(value: &datafusion_proto_models::protobuf::ParquetSink) -> Result { - let config = - FileSinkConfig::try_from(value.config.as_ref().ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "ParquetSink is missing required field 'config'" - ) - })?)?; - let parquet_options = value - .parquet_options + use datafusion_proto_models::protobuf; + + let protobuf::ParquetSink { + config, + parquet_options, + sorting_columns, + } = value; + let config = FileSinkConfig::try_from(config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'config'" + ) + })?)?; + let parquet_options = parquet_options .as_ref() .ok_or_else(|| { datafusion_common::internal_datafusion_err!( @@ -470,8 +521,26 @@ impl TryFrom<&datafusion_proto_models::protobuf::ParquetSink> for ParquetSink { ) })? .try_into()?; + let sorting_columns = sorting_columns.as_ref().map( + |protobuf::ParquetSortingColumns { columns }| { + columns + .iter() + .map( + |&protobuf::ParquetSortingColumn { + column_idx, + descending, + nulls_first, + }| SortingColumn { + column_idx, + descending, + nulls_first, + }, + ) + .collect() + }, + ); - Ok(Self::new(config, parquet_options)) + Ok(Self::new(config, parquet_options).with_sorting_columns(sorting_columns)) } } @@ -489,19 +558,23 @@ impl ParquetSink { protobuf::physical_plan_node::PhysicalPlanType::ParquetSink, "ParquetSink", ); - let input = ctx.decode_required_child( - sink_node.input.as_deref(), - "ParquetSinkExecNode", - "input", - )?; - let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + let protobuf::ParquetSinkExecNode { + input, + sink, + // Recomputed by `DataSinkExec::new`. + sink_schema: _, + sort_order, + } = sink_node.as_ref(); + let input = + ctx.decode_required_child(input.as_deref(), "ParquetSinkExecNode", "input")?; + let proto_sink = sink.as_ref().ok_or_else(|| { datafusion_common::internal_datafusion_err!( "ParquetSinkExecNode is missing required field 'sink'" ) })?; let data_sink = ParquetSink::try_from(proto_sink)?; let sort_order = DataSinkExec::decode_sort_order( - sink_node.sort_order.as_ref(), + sort_order.as_ref(), ctx, input.schema().as_ref(), )?; diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 1015bdc600bdd..e1434824e4b92 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -42,6 +42,8 @@ use arrow::datatypes::TimeUnit; use datafusion_common::DataFusionError; use datafusion_common::config::TableParquetOptions; use datafusion_common::tree_node::TreeNodeRecursion; +#[cfg(feature = "proto")] +use datafusion_common::utils::{usize_from_wire, usize_to_wire}; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -1109,16 +1111,38 @@ impl FileSource for ParquetSource { use datafusion_proto_models::protobuf; use protobuf::physical_plan_node::PhysicalPlanType; - if self.schema_provider.is_some() { + let Self { + table_parquet_options, + // Runtime metrics are recreated when the source is decoded. + metrics: _, + // Carried by `base`. + table_schema: _, + predicate, + // Rebuilt from the decode context. + parquet_file_reader_factory: _, + // Requires a custom codec for serialization. + schema_provider, + // Applied by `FileScanConfig` before execution. + batch_size: _, + metadata_size_hint, + // Carried by `base` as projection expressions. + projection: _, + // Not serialized or restored by the default decoder. + #[cfg(feature = "parquet_encryption")] + encryption_factory: _, + reverse_row_groups, + sort_order_for_reorder, + } = self; + + if schema_provider.is_some() { return Ok(None); } - let predicate = self - .filter() - .map(|pred| ctx.encode_expr(&pred)) + let predicate = predicate + .as_ref() + .map(|pred| ctx.encode_expr(pred)) .transpose()?; - let sort_order_for_reorder = self - .sort_order_for_reorder + let sort_order_for_reorder = sort_order_for_reorder .as_ref() .map(|ordering| -> datafusion_common::Result<_> { Ok(protobuf::PhysicalSortExprNodeCollection { @@ -1129,13 +1153,17 @@ impl FileSource for ParquetSource { }) }) .transpose()?; + let metadata_size_hint = metadata_size_hint + .map(|hint| usize_to_wire(hint, "ParquetSource", "metadata_size_hint")) + .transpose()?; let node = protobuf::ParquetScanExecNode { base_conf: Some(base.try_to_proto(ctx)?), predicate, - parquet_options: Some(self.table_parquet_options().try_into()?), + parquet_options: Some(table_parquet_options.try_into()?), sort_order_for_reorder, - reverse_row_groups: self.reverse_row_groups, + reverse_row_groups: *reverse_row_groups, + metadata_size_hint, }; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::ParquetScan(node)), @@ -1148,6 +1176,8 @@ impl ParquetSource { /// Reconstructs a `DataSourceExec` from a protobuf `ParquetScan`. /// /// Rebuilds the reader factory from the decode context because it is not serialized. + /// Encryption factories and crypto options are not serialized or restored; + /// plans that rely on them require custom handling. pub fn try_from_proto( node: &datafusion_proto_models::protobuf::PhysicalPlanNode, ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, @@ -1168,7 +1198,16 @@ impl ParquetSource { ); }; - let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + let protobuf::ParquetScanExecNode { + base_conf, + predicate, + parquet_options, + sort_order_for_reorder, + reverse_row_groups, + metadata_size_hint, + } = scan; + + let base_conf = base_conf.as_ref().ok_or_else(|| { datafusion_common::internal_datafusion_err!( "ParquetScanExecNode is missing required field 'base_conf'" ) @@ -1200,13 +1239,11 @@ impl ParquetSource { schema }; - let predicate = scan - .predicate + let predicate = predicate .as_ref() .map(|expr| ctx.decode_expr(expr, predicate_schema.as_ref())) .transpose()?; - let sort_order_for_reorder = scan - .sort_order_for_reorder + let sort_order_for_reorder = sort_order_for_reorder .as_ref() .map(|ordering| { optional_ordering_try_from_proto( @@ -1216,9 +1253,12 @@ impl ParquetSource { }) .transpose()? .flatten(); + let metadata_size_hint = metadata_size_hint + .map(|hint| usize_from_wire(hint, "ParquetSource", "metadata_size_hint")) + .transpose()?; let mut options = TableParquetOptions::default(); - if let Some(table_options) = scan.parquet_options.as_ref() { + if let Some(table_options) = parquet_options.as_ref() { options = table_options.try_into()?; } @@ -1243,7 +1283,8 @@ impl ParquetSource { .with_parquet_file_reader_factory(reader_factory) .with_table_parquet_options(options); source.sort_order_for_reorder = sort_order_for_reorder; - source.reverse_row_groups = scan.reverse_row_groups; + source.reverse_row_groups = *reverse_row_groups; + source.metadata_size_hint = metadata_size_hint; if let Some(predicate) = predicate { source = source.with_predicate(predicate); diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index fac5ff27191cd..0f9386969ab39 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -986,9 +986,25 @@ message CsvSinkExecNode { PhysicalSortExprNodeCollection sort_order = 4; } +message ParquetSortingColumn { + // Zero-based ordinal of the leaf column in the Parquet schema. + int32 column_idx = 1; + // Whether the column is sorted in descending order. + bool descending = 2; + // Whether nulls sort before non-null values. + bool nulls_first = 3; +} + +message ParquetSortingColumns { + // The wrapper preserves the distinction between no sorting metadata and an empty list. + repeated ParquetSortingColumn columns = 1; +} + message ParquetSink { FileSinkConfig config = 1; datafusion_common.TableParquetOptions parquet_options = 2; + // Sorting-column metadata to write to each Parquet row group. + ParquetSortingColumns sorting_columns = 3; } message ParquetSinkExecNode { @@ -1308,6 +1324,9 @@ message ParquetScanExecNode { PhysicalSortExprNodeCollection sort_order_for_reorder = 5; bool reverse_row_groups = 6; + + // Source-specific footer prefetch size. Absent means no hint. + optional uint64 metadata_size_hint = 7; } message CsvScanExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index e06f5b7011504..d97912d2654d6 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -16764,6 +16764,9 @@ impl serde::Serialize for ParquetScanExecNode { if self.reverse_row_groups { len += 1; } + if self.metadata_size_hint.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ParquetScanExecNode", len)?; if let Some(v) = self.base_conf.as_ref() { struct_ser.serialize_field("baseConf", v)?; @@ -16780,6 +16783,11 @@ impl serde::Serialize for ParquetScanExecNode { if self.reverse_row_groups { struct_ser.serialize_field("reverseRowGroups", &self.reverse_row_groups)?; } + if let Some(v) = self.metadata_size_hint.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("metadataSizeHint", ToString::to_string(&v).as_str())?; + } struct_ser.end() } } @@ -16799,6 +16807,8 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { "sortOrderForReorder", "reverse_row_groups", "reverseRowGroups", + "metadata_size_hint", + "metadataSizeHint", ]; #[allow(clippy::enum_variant_names)] @@ -16808,6 +16818,7 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { ParquetOptions, SortOrderForReorder, ReverseRowGroups, + MetadataSizeHint, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -16834,6 +16845,7 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { "parquetOptions" | "parquet_options" => Ok(GeneratedField::ParquetOptions), "sortOrderForReorder" | "sort_order_for_reorder" => Ok(GeneratedField::SortOrderForReorder), "reverseRowGroups" | "reverse_row_groups" => Ok(GeneratedField::ReverseRowGroups), + "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -16858,6 +16870,7 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { let mut parquet_options__ = None; let mut sort_order_for_reorder__ = None; let mut reverse_row_groups__ = None; + let mut metadata_size_hint__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::BaseConf => { @@ -16890,6 +16903,14 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { } reverse_row_groups__ = Some(map_.next_value()?); } + GeneratedField::MetadataSizeHint => { + if metadata_size_hint__.is_some() { + return Err(serde::de::Error::duplicate_field("metadataSizeHint")); + } + metadata_size_hint__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } } } Ok(ParquetScanExecNode { @@ -16898,6 +16919,7 @@ impl<'de> serde::Deserialize<'de> for ParquetScanExecNode { parquet_options: parquet_options__, sort_order_for_reorder: sort_order_for_reorder__, reverse_row_groups: reverse_row_groups__.unwrap_or_default(), + metadata_size_hint: metadata_size_hint__, }) } } @@ -16918,6 +16940,9 @@ impl serde::Serialize for ParquetSink { if self.parquet_options.is_some() { len += 1; } + if self.sorting_columns.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ParquetSink", len)?; if let Some(v) = self.config.as_ref() { struct_ser.serialize_field("config", v)?; @@ -16925,6 +16950,9 @@ impl serde::Serialize for ParquetSink { if let Some(v) = self.parquet_options.as_ref() { struct_ser.serialize_field("parquetOptions", v)?; } + if let Some(v) = self.sorting_columns.as_ref() { + struct_ser.serialize_field("sortingColumns", v)?; + } struct_ser.end() } } @@ -16938,12 +16966,15 @@ impl<'de> serde::Deserialize<'de> for ParquetSink { "config", "parquet_options", "parquetOptions", + "sorting_columns", + "sortingColumns", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Config, ParquetOptions, + SortingColumns, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -16967,6 +16998,7 @@ impl<'de> serde::Deserialize<'de> for ParquetSink { match value { "config" => Ok(GeneratedField::Config), "parquetOptions" | "parquet_options" => Ok(GeneratedField::ParquetOptions), + "sortingColumns" | "sorting_columns" => Ok(GeneratedField::SortingColumns), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -16988,6 +17020,7 @@ impl<'de> serde::Deserialize<'de> for ParquetSink { { let mut config__ = None; let mut parquet_options__ = None; + let mut sorting_columns__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Config => { @@ -17002,11 +17035,18 @@ impl<'de> serde::Deserialize<'de> for ParquetSink { } parquet_options__ = map_.next_value()?; } + GeneratedField::SortingColumns => { + if sorting_columns__.is_some() { + return Err(serde::de::Error::duplicate_field("sortingColumns")); + } + sorting_columns__ = map_.next_value()?; + } } } Ok(ParquetSink { config: config__, parquet_options: parquet_options__, + sorting_columns: sorting_columns__, }) } } @@ -17157,6 +17197,226 @@ impl<'de> serde::Deserialize<'de> for ParquetSinkExecNode { deserializer.deserialize_struct("datafusion.ParquetSinkExecNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ParquetSortingColumn { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.column_idx != 0 { + len += 1; + } + if self.descending { + len += 1; + } + if self.nulls_first { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ParquetSortingColumn", len)?; + if self.column_idx != 0 { + struct_ser.serialize_field("columnIdx", &self.column_idx)?; + } + if self.descending { + struct_ser.serialize_field("descending", &self.descending)?; + } + if self.nulls_first { + struct_ser.serialize_field("nullsFirst", &self.nulls_first)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ParquetSortingColumn { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "column_idx", + "columnIdx", + "descending", + "nulls_first", + "nullsFirst", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + ColumnIdx, + Descending, + NullsFirst, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "columnIdx" | "column_idx" => Ok(GeneratedField::ColumnIdx), + "descending" => Ok(GeneratedField::Descending), + "nullsFirst" | "nulls_first" => Ok(GeneratedField::NullsFirst), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ParquetSortingColumn; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ParquetSortingColumn") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut column_idx__ = None; + let mut descending__ = None; + let mut nulls_first__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::ColumnIdx => { + if column_idx__.is_some() { + return Err(serde::de::Error::duplicate_field("columnIdx")); + } + column_idx__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Descending => { + if descending__.is_some() { + return Err(serde::de::Error::duplicate_field("descending")); + } + descending__ = Some(map_.next_value()?); + } + GeneratedField::NullsFirst => { + if nulls_first__.is_some() { + return Err(serde::de::Error::duplicate_field("nullsFirst")); + } + nulls_first__ = Some(map_.next_value()?); + } + } + } + Ok(ParquetSortingColumn { + column_idx: column_idx__.unwrap_or_default(), + descending: descending__.unwrap_or_default(), + nulls_first: nulls_first__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ParquetSortingColumn", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for ParquetSortingColumns { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.columns.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ParquetSortingColumns", len)?; + if !self.columns.is_empty() { + struct_ser.serialize_field("columns", &self.columns)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ParquetSortingColumns { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "columns", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Columns, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "columns" => Ok(GeneratedField::Columns), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ParquetSortingColumns; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ParquetSortingColumns") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut columns__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Columns => { + if columns__.is_some() { + return Err(serde::de::Error::duplicate_field("columns")); + } + columns__ = Some(map_.next_value()?); + } + } + } + Ok(ParquetSortingColumns { + columns: columns__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ParquetSortingColumns", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for PartialTableReference { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index 149839dacc967..f8cf7f82f425b 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -1530,6 +1530,24 @@ pub struct CsvSinkExecNode { #[prost(message, optional, tag = "4")] pub sort_order: ::core::option::Option, } +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ParquetSortingColumn { + /// Zero-based ordinal of the leaf column in the Parquet schema. + #[prost(int32, tag = "1")] + pub column_idx: i32, + /// Whether the column is sorted in descending order. + #[prost(bool, tag = "2")] + pub descending: bool, + /// Whether nulls sort before non-null values. + #[prost(bool, tag = "3")] + pub nulls_first: bool, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ParquetSortingColumns { + /// The wrapper preserves the distinction between no sorting metadata and an empty list. + #[prost(message, repeated, tag = "1")] + pub columns: ::prost::alloc::vec::Vec, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetSink { #[prost(message, optional, tag = "1")] @@ -1538,6 +1556,9 @@ pub struct ParquetSink { pub parquet_options: ::core::option::Option< super::datafusion_common::TableParquetOptions, >, + /// Sorting-column metadata to write to each Parquet row group. + #[prost(message, optional, tag = "3")] + pub sorting_columns: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetSinkExecNode { @@ -2005,6 +2026,9 @@ pub struct ParquetScanExecNode { pub sort_order_for_reorder: ::core::option::Option, #[prost(bool, tag = "6")] pub reverse_row_groups: bool, + /// Source-specific footer prefetch size. Absent means no hint. + #[prost(uint64, optional, tag = "7")] + pub metadata_size_hint: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct CsvScanExecNode { diff --git a/datafusion/proto/tests/cases/plans/mod.rs b/datafusion/proto/tests/cases/plans/mod.rs index 8dad1eff67032..40839d4b3a31a 100644 --- a/datafusion/proto/tests/cases/plans/mod.rs +++ b/datafusion/proto/tests/cases/plans/mod.rs @@ -92,6 +92,23 @@ fn roundtrip_test_and_return( Ok(result_exec_plan) } +/// Round-trip through the public JSON APIs and return the decoded plan for field assertions. +#[cfg(feature = "json")] +fn roundtrip_test_json_and_return( + exec_plan: Arc, + ctx: &SessionContext, +) -> Result> { + use datafusion_proto::bytes::{physical_plan_from_json, physical_plan_to_json}; + + let json = physical_plan_to_json(Arc::clone(&exec_plan))?; + let result_exec_plan = physical_plan_from_json(&json, ctx.task_ctx().as_ref())?; + pretty_assertions::assert_eq!( + format!("{exec_plan:?}"), + format!("{result_exec_plan:?}") + ); + Ok(result_exec_plan) +} + /// Perform a serde roundtrip and assert that the string representation of the before and after plans /// are identical. Note that this often isn't sufficient to guarantee that no information is /// lost during serde because the string representation of a plan often only shows a subset of state. diff --git a/datafusion/proto/tests/cases/plans/sinks.rs b/datafusion/proto/tests/cases/plans/sinks.rs index 8d8b736000592..1bc8c9d41e512 100644 --- a/datafusion/proto/tests/cases/plans/sinks.rs +++ b/datafusion/proto/tests/cases/plans/sinks.rs @@ -17,7 +17,7 @@ //! Data sinks and their file sink configurations. -use super::{roundtrip_test, roundtrip_test_and_return}; +use super::roundtrip_test_and_return; use arrow::csv::writer::Terminator; use arrow::csv::{QuoteStyle, WriterBuilder}; use async_trait::async_trait; @@ -33,6 +33,7 @@ use datafusion::datasource::physical_plan::{ }; use datafusion::datasource::sink::{DataSink, DataSinkExec}; use datafusion::execution::TaskContext; +use datafusion::parquet::file::metadata::SortingColumn; use datafusion::physical_expr::{LexRequirement, PhysicalSortRequirement}; use datafusion::physical_plan::expressions::Column; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; @@ -425,7 +426,7 @@ fn roundtrip_parquet_sink() -> Result<()> { let field_a = Field::new("plan_type", DataType::Utf8, false); let field_b = Field::new("plan", DataType::Utf8, false); let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let input = Arc::new(PlaceholderRowExec::new(schema.clone())); + let input: Arc = Arc::new(PlaceholderRowExec::new(schema.clone())); let file_sink_config = FileSinkConfig { original_url: String::default(), @@ -439,11 +440,7 @@ fn roundtrip_parquet_sink() -> Result<()> { file_extension: "parquet".into(), file_output_mode: FileOutputMode::Automatic, }; - let data_sink = Arc::new(ParquetSink::new( - file_sink_config, - TableParquetOptions::default(), - )); - let sort_order = [PhysicalSortRequirement::new( + let sort_order: LexRequirement = [PhysicalSortRequirement::new( Arc::new(Column::new("plan_type", 0)), Some(SortOptions { descending: true, @@ -452,9 +449,65 @@ fn roundtrip_parquet_sink() -> Result<()> { )] .into(); - roundtrip_test(Arc::new(DataSinkExec::new( - input, - data_sink, - Some(sort_order), - ))) + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let sorting_columns = vec![ + SortingColumn { + column_idx: 0, + descending: true, + nulls_first: false, + }, + SortingColumn { + column_idx: 1, + descending: false, + nulls_first: true, + }, + ]; + for sorting_columns in [ + None, + Some(vec![]), + Some(sorting_columns[..1].to_vec()), + Some(sorting_columns), + ] { + let data_sink = Arc::new( + ParquetSink::new(file_sink_config.clone(), TableParquetOptions::default()) + .with_sorting_columns(sorting_columns.clone()), + ); + let roundtripped = roundtrip_test_and_return( + Arc::new(DataSinkExec::new( + Arc::clone(&input), + data_sink, + Some(sort_order.clone()), + )), + &ctx, + &codec, + &proto_converter, + )?; + #[cfg(feature = "json")] + let roundtripped = super::roundtrip_test_json_and_return(roundtripped, &ctx)?; + let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + let Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetSink(node)) = + node.physical_plan_type + else { + panic!("expected ParquetSink node"); + }; + let actual = node + .sink + .expect("ParquetSinkExecNode should contain a sink") + .sorting_columns + .map(|columns| { + columns + .columns + .into_iter() + .map(|column| SortingColumn { + column_idx: column.column_idx, + descending: column.descending, + nulls_first: column.nulls_first, + }) + .collect::>() + }); + assert_eq!(actual, sorting_columns); + } + Ok(()) } diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs index 7f04a794f2803..c4897357ee00b 100644 --- a/datafusion/proto/tests/cases/plans/sources.rs +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -111,7 +111,7 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); let predicate = Arc::new(BinaryExpr::new( - Arc::new(Column::new("col", 1)), + Arc::new(Column::new("col", 0)), Operator::Eq, lit("1"), )); @@ -119,28 +119,51 @@ fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { let mut options = TableParquetOptions::new(); options.global.pushdown_filters = true; - let file_source = Arc::new( - ParquetSource::new(Arc::clone(&file_schema)) - .with_table_parquet_options(options) - .with_predicate(predicate), - ); - - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )])]) - .with_statistics(Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( - vec![Field::new("col", DataType::Utf8, false)], - ))), - }) - .build(); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + for metadata_size_hint in [None, Some(0), Some(8192), Some(usize::MAX)] { + let mut file_source = ParquetSource::new(Arc::clone(&file_schema)) + .with_table_parquet_options(options.clone()) + .with_predicate(predicate.clone()); + if let Some(hint) = metadata_size_hint { + file_source = file_source.with_metadata_size_hint(hint); + } + let scan_config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(file_source), + ) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&file_schema), + }) + .build(); - roundtrip_test(DataSourceExec::from_data_source(scan_config)) + let roundtripped = roundtrip_test_and_return( + DataSourceExec::from_data_source(scan_config), + &ctx, + &codec, + &proto_converter, + )?; + #[cfg(feature = "json")] + let roundtripped = super::roundtrip_test_json_and_return(roundtripped, &ctx)?; + let node = PhysicalPlanNode::try_from_physical_plan(roundtripped, &codec)?; + let Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetScan(scan)) = + node.physical_plan_type + else { + return internal_err!("Expected ParquetScan node"); + }; + assert_eq!( + scan.metadata_size_hint, + metadata_size_hint.map(|hint| hint as u64) + ); + } + Ok(()) } #[tokio::test] diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 0a9828ff7fbae..f83c8b3d5c4dc 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -227,6 +227,27 @@ The protobuf wire format remains backward compatible. See [PR #24945](https://github.com/apache/datafusion/pull/24945) for details. +### Generated protobuf Parquet structs gained state fields + +The generated public protobuf `ParquetScanExecNode` struct now carries an +optional `metadata_size_hint`, and `ParquetSink` now carries optional +`sorting_columns`. These fields preserve the corresponding Parquet source and +sink settings across physical-plan serialization. + +**Who is affected:** + +- Users constructing either generated struct with an exhaustive struct literal. + +**Migration guide:** + +Add `metadata_size_hint: None` to `ParquetScanExecNode` literals and +`sorting_columns: None` to `ParquetSink` literals to retain the previous +behavior, or use `..Default::default()` for unspecified fields. + +The protobuf wire format remains backward compatible. + +See [PR #25057](https://github.com/apache/datafusion/pull/25057) for details. + ### Generated protobuf `CsvWriterOptions` changed The generated `CsvWriterOptions` now includes `compression_level`,