From 7f41edbb371751ee897e3fdaee9065a1dc45fa7b Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 16:52:51 +0800 Subject: [PATCH 1/7] feat(proto): add DataSink serialization hook DataSinkExec needs a generic way to delegate protobuf encoding to its wrapped sink. The new hook keeps sink-specific wire logic out of the central serializer while retaining the fallback for unmigrated sinks. Add focused coverage for child and sort-order encoding and verify that existing file sinks continue to use the compatibility path. Refs #23498 --- datafusion/datasource/Cargo.toml | 10 ++- datafusion/datasource/src/sink.rs | 57 ++++++++++++ .../tests/cases/roundtrip_physical_plan.rs | 89 ++++++++++++++++++- 3 files changed, 151 insertions(+), 5 deletions(-) diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 459ca436f365d..b78ac616decee 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,10 +34,12 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] -# Enables the protobuf conversions for the file-scan leaf types owned by this -# crate (`FileRange`, `PartitionedFile`, `FileGroup`). Off by default so -# consumers that never serialize plans pay nothing. -proto = ["dep:datafusion-proto-models"] +# Enables protobuf conversions for datasource types and serialization hooks. +# Off by default so consumers that never serialize plans pay nothing. +proto = [ + "dep:datafusion-proto-models", + "datafusion-physical-plan/proto", +] [dependencies] arrow = { workspace = true } diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 18ebe80773e8a..327d80fea6f79 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -71,6 +71,25 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { data: SendableRecordBatchStream, context: &Arc, ) -> Result; + + /// Serialize this sink into a full protobuf plan node, if it knows how. + /// + /// [`DataSinkExec::try_to_proto`] encodes the shared child plan and required + /// output ordering before delegating to this hook. Implementations only need + /// to add their sink-specific fields. + /// + /// Returning `Ok(None)` preserves the legacy central serialization fallback. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _input: datafusion_proto_models::protobuf::PhysicalPlanNode, + _sort_order: Option< + datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + _sink_schema: &Schema, + ) -> Result> { + Ok(None) + } } impl dyn DataSink { @@ -268,6 +287,44 @@ impl ExecutionPlan for DataSinkExec { fn metrics(&self) -> Option { self.sink.metrics() } + + /// Encodes the shared sink plan fields before delegating the sink-specific + /// protobuf representation to [`DataSink::try_to_proto`]. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + let input = ctx.encode_child(self.input())?; + let sort_order = self + .sort_order() + .as_ref() + .map(|requirements| { + requirements + .iter() + .map(|requirement| { + let expr: PhysicalSortExpr = requirement.to_owned().into(); + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), + asc: !expr.options.descending, + nulls_first: expr.options.nulls_first, + }) + }) + .collect::>>() + .map(|physical_sort_expr_nodes| { + protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + } + }) + }) + .transpose()?; + + self.sink() + .try_to_proto(input, sort_order, self.schema().as_ref()) + } } /// Create a output record batch with a count diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 864e6d68676ee..43542ca80649c 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -23,6 +23,7 @@ use std::vec; use arrow::array::RecordBatch; use arrow::csv::WriterBuilder; use arrow::datatypes::{Fields, TimeUnit}; +use async_trait::async_trait; use datafusion::arrow::array::ArrayRef; use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema, SchemaRef}; @@ -39,7 +40,7 @@ use datafusion::datasource::physical_plan::{ FileSinkConfig, ParquetSource, wrap_partition_type_in_dict, wrap_partition_value_in_dict, }; -use datafusion::datasource::sink::DataSinkExec; +use datafusion::datasource::sink::{DataSink, DataSinkExec}; use datafusion::datasource::source::DataSourceExec; use datafusion::execution::TaskContext; use datafusion::functions_aggregate::count::count_udaf; @@ -2054,6 +2055,92 @@ fn roundtrip_explain() -> Result<()> { Ok(()) } +#[derive(Debug)] +struct ProtoHookSink { + schema: SchemaRef, +} + +impl DisplayAs for ProtoHookSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ProtoHookSink") + } +} + +#[async_trait] +impl DataSink for ProtoHookSink { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + _data: SendableRecordBatchStream, + _context: &Arc, + ) -> Result { + unreachable!("serialization test does not execute the sink") + } + + fn try_to_proto( + &self, + input: PhysicalPlanNode, + sort_order: Option, + sink_schema: &Schema, + ) -> Result> { + assert!(matches!( + input.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow(_)) + )); + assert_eq!( + sort_order + .as_ref() + .map(|ordering| ordering.physical_sort_expr_nodes.len()), + Some(1) + ); + assert_eq!(sink_schema.fields().len(), 1); + + Ok(Some(PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(sink_schema.try_into()?), + partitions: 1, + }, + ), + ), + })) + } +} + +#[test] +fn data_sink_exec_delegates_to_sink_proto_hook() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let input = Arc::new(PlaceholderRowExec::new(Arc::clone(&input_schema))); + let sink = Arc::new(ProtoHookSink { + schema: Arc::clone(&input_schema), + }); + let sort_order = [PhysicalSortRequirement::new( + Arc::new(Column::new("value", 0)), + Some(SortOptions::default()), + )] + .into(); + let plan = Arc::new(DataSinkExec::new(input, sink, Some(sort_order))); + + let node = PhysicalPlanNode::try_from_physical_plan( + plan, + &DefaultPhysicalExtensionCodec {}, + )?; + + assert!(matches!( + node.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(_)) + )); + Ok(()) +} + #[tokio::test] async fn roundtrip_json_source() -> Result<()> { let ctx = SessionContext::new(); From c74e81be24dd9705ad2c523b7a40edd0f3976d74 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 17:19:42 +0800 Subject: [PATCH 2/7] refactor(proto): move FileSinkConfig serde FileSinkConfig is owned by datafusion-datasource, but its protobuf conversion lived in the central proto crate. Co-locating it lets sink implementations reuse the wire logic without depending on datafusion-proto. Keep existing TryFromProto implementations as compatibility delegates and verify both APIs produce the same protobuf data. CLOSES #23498 --- datafusion/datasource/src/file_sink_config.rs | 6 + .../datasource/src/file_sink_config/proto.rs | 164 ++++++++++++++++++ .../proto/src/physical_plan/from_proto.rs | 51 +----- .../proto/src/physical_plan/to_proto.rs | 43 +---- .../tests/cases/roundtrip_physical_plan.rs | 31 ++++ 5 files changed, 204 insertions(+), 91 deletions(-) create mode 100644 datafusion/datasource/src/file_sink_config/proto.rs diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index 1abce86a3565f..b2e20567a15bb 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -32,6 +32,12 @@ use datafusion_expr::dml::InsertOp; use async_trait::async_trait; use object_store::ObjectStore; +#[cfg(feature = "proto")] +mod proto; + +#[cfg(feature = "proto")] +pub use proto::parse_sink_sort_order; + /// Determines how `FileSink` output paths are interpreted. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum FileOutputMode { diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs new file mode 100644 index 0000000000000..405ede03c8efe --- /dev/null +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -0,0 +1,164 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Protobuf conversion for the format-independent [`FileSinkConfig`]. + +use std::sync::Arc; + +use arrow::compute::SortOptions; +use arrow::datatypes::Schema; +use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_expr::dml::InsertOp; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr_common::sort_expr::LexRequirement; +use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; +use datafusion_proto_models::protobuf; + +use crate::file_groups::FileGroup; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; +use crate::ListingTableUrl; + +impl FileSinkConfig { + /// Serialize this shared file-sink configuration without format-specific + /// writer options. + pub fn to_proto(&self) -> Result { + let file_groups = self + .file_group + .iter() + .map(TryInto::try_into) + .collect::>>()?; + let table_paths = self + .table_paths + .iter() + .map(ToString::to_string) + .collect::>(); + let table_partition_cols = self + .table_partition_cols + .iter() + .map(|(name, data_type)| { + Ok(protobuf::PartitionColumn { + name: name.to_owned(), + arrow_type: Some(data_type.try_into()?), + }) + }) + .collect::>>()?; + let file_output_mode = match self.file_output_mode { + FileOutputMode::Automatic => protobuf::FileOutputMode::Automatic, + FileOutputMode::SingleFile => protobuf::FileOutputMode::SingleFile, + FileOutputMode::Directory => protobuf::FileOutputMode::Directory, + }; + + Ok(protobuf::FileSinkConfig { + object_store_url: self.object_store_url.to_string(), + file_groups, + table_paths, + output_schema: Some(self.output_schema.as_ref().try_into()?), + table_partition_cols, + keep_partition_by_columns: self.keep_partition_by_columns, + insert_op: self.insert_op as i32, + file_extension: self.file_extension.clone(), + file_output_mode: file_output_mode.into(), + }) + } + + /// Reconstruct a shared file-sink configuration from protobuf. + pub fn from_proto(conf: &protobuf::FileSinkConfig) -> Result { + let file_group = FileGroup::new( + conf.file_groups + .iter() + .map(TryInto::try_into) + .collect::>>()?, + ); + let table_paths = conf + .table_paths + .iter() + .map(ListingTableUrl::parse) + .collect::>>()?; + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|protobuf::PartitionColumn { name, arrow_type }| { + let data_type = arrow_type + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "PartitionColumn is missing required field 'arrow_type'" + ) + })? + .try_into()?; + Ok((name.clone(), data_type)) + }) + .collect::>>()?; + let insert_op = match conf.insert_op() { + protobuf::InsertOp::Append => InsertOp::Append, + protobuf::InsertOp::Overwrite => InsertOp::Overwrite, + protobuf::InsertOp::Replace => InsertOp::Replace, + }; + let file_output_mode = match conf.file_output_mode() { + protobuf::FileOutputMode::Automatic => FileOutputMode::Automatic, + protobuf::FileOutputMode::SingleFile => FileOutputMode::SingleFile, + protobuf::FileOutputMode::Directory => FileOutputMode::Directory, + }; + let output_schema = conf.output_schema.as_ref().ok_or_else(|| { + internal_datafusion_err!( + "FileSinkConfig is missing required field 'output_schema'" + ) + })?; + + Ok(Self { + original_url: String::default(), + object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?, + file_group, + table_paths, + output_schema: Arc::new(output_schema.try_into()?), + table_partition_cols, + insert_op, + keep_partition_by_columns: conf.keep_partition_by_columns, + file_extension: conf.file_extension.clone(), + file_output_mode, + }) + } +} + +/// Decode a sink's optional required output ordering against its input schema. +pub fn parse_sink_sort_order( + collection: Option<&protobuf::PhysicalSortExprNodeCollection>, + ctx: &ExecutionPlanDecodeCtx<'_>, + schema: &Schema, +) -> Result> { + let Some(collection) = collection else { + return Ok(None); + }; + let sort_exprs = collection + .physical_sort_expr_nodes + .iter() + .map(|node| { + let expr = node.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("Unexpected empty physical expression") + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr, schema)?, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + }) + .collect::>>()?; + Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) +} diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 34ad8c7a62fc7..5c87cc476f74b 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -30,7 +30,7 @@ use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::file_sink_config::FileSinkConfig; -use datafusion_datasource::{FileRange, ListingTableUrl, PartitionedFile, TableSchema}; +use datafusion_datasource::{FileRange, PartitionedFile, TableSchema}; use datafusion_datasource_csv::file_format::CsvSink; use datafusion_datasource_json::file_format::JsonSink; #[cfg(feature = "parquet")] @@ -38,7 +38,6 @@ use datafusion_datasource_parquet::file_format::ParquetSink; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; -use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; @@ -694,53 +693,7 @@ impl TryFromProto<&protobuf::FileSinkConfig> for FileSinkConfig { type Error = DataFusionError; fn try_from_proto(conf: &protobuf::FileSinkConfig) -> Result { - let file_group = FileGroup::new( - conf.file_groups - .iter() - .map(TryInto::try_into) - .collect::>>()?, - ); - let table_paths = conf - .table_paths - .iter() - .map(ListingTableUrl::parse) - .collect::>>()?; - let table_partition_cols = conf - .table_partition_cols - .iter() - .map(|protobuf::PartitionColumn { name, arrow_type }| { - let data_type = convert_required!(arrow_type)?; - Ok((name.clone(), data_type)) - }) - .collect::>>()?; - let insert_op = match conf.insert_op() { - protobuf::InsertOp::Append => InsertOp::Append, - protobuf::InsertOp::Overwrite => InsertOp::Overwrite, - protobuf::InsertOp::Replace => InsertOp::Replace, - }; - let file_output_mode = match conf.file_output_mode() { - protobuf::FileOutputMode::Automatic => { - datafusion_datasource::file_sink_config::FileOutputMode::Automatic - } - protobuf::FileOutputMode::SingleFile => { - datafusion_datasource::file_sink_config::FileOutputMode::SingleFile - } - protobuf::FileOutputMode::Directory => { - datafusion_datasource::file_sink_config::FileOutputMode::Directory - } - }; - Ok(Self { - original_url: String::default(), - object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?, - file_group, - table_paths, - output_schema: Arc::new(convert_required!(conf.output_schema)?), - table_partition_cols, - insert_op, - keep_partition_by_columns: conf.keep_partition_by_columns, - file_extension: conf.file_extension.clone(), - file_output_mode, - }) + FileSinkConfig::from_proto(conf) } } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 5189972f0e200..a9637a028a415 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -611,47 +611,6 @@ impl TryFromProto<&FileSinkConfig> for protobuf::FileSinkConfig { type Error = DataFusionError; fn try_from_proto(conf: &FileSinkConfig) -> Result { - let file_groups = conf - .file_group - .iter() - .map(protobuf::PartitionedFile::try_from_proto) - .collect::>>()?; - let table_paths = conf - .table_paths - .iter() - .map(ToString::to_string) - .collect::>(); - let table_partition_cols = conf - .table_partition_cols - .iter() - .map(|(name, data_type)| { - Ok(protobuf::PartitionColumn { - name: name.to_owned(), - arrow_type: Some(data_type.try_into()?), - }) - }) - .collect::>>()?; - let file_output_mode = match conf.file_output_mode { - datafusion_datasource::file_sink_config::FileOutputMode::Automatic => { - protobuf::FileOutputMode::Automatic - } - datafusion_datasource::file_sink_config::FileOutputMode::SingleFile => { - protobuf::FileOutputMode::SingleFile - } - datafusion_datasource::file_sink_config::FileOutputMode::Directory => { - protobuf::FileOutputMode::Directory - } - }; - Ok(Self { - object_store_url: conf.object_store_url.to_string(), - file_groups, - table_paths, - output_schema: Some(conf.output_schema.as_ref().try_into()?), - table_partition_cols, - keep_partition_by_columns: conf.keep_partition_by_columns, - insert_op: conf.insert_op as i32, - file_extension: conf.file_extension.to_string(), - file_output_mode: file_output_mode.into(), - }) + conf.to_proto() } } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 43542ca80649c..5cc53883a068a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -134,6 +134,7 @@ use datafusion_proto::bytes::{ physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter, }; +use datafusion_proto::convert::TryFromProto; use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; use datafusion_proto::physical_plan::{ AsExecutionPlan, DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec, @@ -2141,6 +2142,36 @@ fn data_sink_exec_delegates_to_sink_proto_hook() -> Result<()> { Ok(()) } +#[test] +fn file_sink_config_conversion_preserves_compatibility_api() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "partition", + DataType::Utf8, + false, + )])); + let config = FileSinkConfig { + original_url: "file:///tmp/output".to_string(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp/output", 1)]), + table_paths: vec![ListingTableUrl::parse("file:///tmp/output")?], + output_schema: schema, + table_partition_cols: vec![("partition".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: true, + file_extension: "parquet".to_string(), + file_output_mode: FileOutputMode::Directory, + }; + + let direct = config.to_proto()?; + let compatibility = protobuf::FileSinkConfig::try_from_proto(&config)?; + assert_eq!(direct, compatibility); + + let direct = FileSinkConfig::from_proto(&direct)?; + let compatibility = FileSinkConfig::try_from_proto(&compatibility)?; + assert_eq!(direct.to_proto()?, compatibility.to_proto()?); + Ok(()) +} + #[tokio::test] async fn roundtrip_json_source() -> Result<()> { let ctx = SessionContext::new(); From 0cabe886c95b69e80f5d6d610f12b9da0648edde Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 27 Jul 2026 16:13:46 +0800 Subject: [PATCH 3/7] fix(proto): validate file sink config decoding Protobuf enum accessors silently replace unknown values with defaults. Rejecting malformed values prevents serialized plans from changing sink behavior and verifies compatibility decoding against the encoded config. Refs #23498 --- .../datasource/src/file_sink_config/proto.rs | 102 +++++++++++++++++- .../tests/cases/roundtrip_physical_plan.rs | 15 +-- 2 files changed, 107 insertions(+), 10 deletions(-) diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs index 405ede03c8efe..f095f2809786b 100644 --- a/datafusion/datasource/src/file_sink_config/proto.rs +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -29,9 +29,9 @@ use datafusion_physical_expr_common::sort_expr::LexRequirement; use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; use datafusion_proto_models::protobuf; +use crate::ListingTableUrl; use crate::file_groups::FileGroup; use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; -use crate::ListingTableUrl; impl FileSinkConfig { /// Serialize this shared file-sink configuration without format-specific @@ -104,12 +104,25 @@ impl FileSinkConfig { Ok((name.clone(), data_type)) }) .collect::>>()?; - let insert_op = match conf.insert_op() { + let insert_op = protobuf::InsertOp::try_from(conf.insert_op).map_err(|_| { + internal_datafusion_err!( + "Received a FileSinkConfig message with unknown InsertOp {}", + conf.insert_op + ) + })?; + let insert_op = match insert_op { protobuf::InsertOp::Append => InsertOp::Append, protobuf::InsertOp::Overwrite => InsertOp::Overwrite, protobuf::InsertOp::Replace => InsertOp::Replace, }; - let file_output_mode = match conf.file_output_mode() { + let file_output_mode = protobuf::FileOutputMode::try_from(conf.file_output_mode) + .map_err(|_| { + internal_datafusion_err!( + "Received a FileSinkConfig message with unknown FileOutputMode {}", + conf.file_output_mode + ) + })?; + let file_output_mode = match file_output_mode { protobuf::FileOutputMode::Automatic => FileOutputMode::Automatic, protobuf::FileOutputMode::SingleFile => FileOutputMode::SingleFile, protobuf::FileOutputMode::Directory => FileOutputMode::Directory, @@ -162,3 +175,86 @@ pub fn parse_sink_sort_order( .collect::>>()?; Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) } +#[cfg(test)] +mod tests { + use datafusion_common::DataFusionError; + + use super::*; + + fn valid_file_sink_config() -> protobuf::FileSinkConfig { + protobuf::FileSinkConfig { + object_store_url: ObjectStoreUrl::local_filesystem().to_string(), + output_schema: Some( + (&Schema::empty()) + .try_into() + .expect("empty schema should serialize"), + ), + insert_op: protobuf::InsertOp::Append.into(), + file_output_mode: protobuf::FileOutputMode::Automatic.into(), + ..Default::default() + } + } + + fn assert_decode_error( + mutate: impl FnOnce(&mut protobuf::FileSinkConfig), + expected: impl AsRef, + ) { + let mut conf = valid_file_sink_config(); + mutate(&mut conf); + + let error = + FileSinkConfig::from_proto(&conf).expect_err("invalid config should fail"); + match error { + DataFusionError::Internal(message) => { + let message = message + .split_once(DataFusionError::BACK_TRACE_SEP) + .map_or(message.as_str(), |(message, _)| message); + assert_eq!(message, expected.as_ref()); + } + error => panic!("expected internal error, got {error}"), + } + } + + #[test] + fn rejects_unknown_insert_op() { + assert_decode_error( + |conf| conf.insert_op = i32::MAX, + format!( + "Received a FileSinkConfig message with unknown InsertOp {}", + i32::MAX + ), + ); + } + + #[test] + fn rejects_unknown_file_output_mode() { + assert_decode_error( + |conf| conf.file_output_mode = i32::MAX, + format!( + "Received a FileSinkConfig message with unknown FileOutputMode {}", + i32::MAX + ), + ); + } + + #[test] + fn rejects_missing_output_schema() { + assert_decode_error( + |conf| conf.output_schema = None, + "FileSinkConfig is missing required field 'output_schema'", + ); + } + + #[test] + fn rejects_partition_column_without_arrow_type() { + assert_decode_error( + |conf| { + conf.table_partition_cols.push(protobuf::PartitionColumn { + name: "partition".to_string(), + arrow_type: None, + }); + }, + "PartitionColumn is missing required field 'arrow_type'", + ); + } +} diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 5cc53883a068a..8618cd2a97a50 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -2162,13 +2162,14 @@ fn file_sink_config_conversion_preserves_compatibility_api() -> Result<()> { file_output_mode: FileOutputMode::Directory, }; - let direct = config.to_proto()?; - let compatibility = protobuf::FileSinkConfig::try_from_proto(&config)?; - assert_eq!(direct, compatibility); - - let direct = FileSinkConfig::from_proto(&direct)?; - let compatibility = FileSinkConfig::try_from_proto(&compatibility)?; - assert_eq!(direct.to_proto()?, compatibility.to_proto()?); + let encoded = config.to_proto()?; + let compatibility_encoded = protobuf::FileSinkConfig::try_from_proto(&config)?; + assert_eq!(encoded, compatibility_encoded); + + let decoded = FileSinkConfig::from_proto(&encoded)?; + let compatibility_decoded = FileSinkConfig::try_from_proto(&encoded)?; + assert_eq!(decoded.to_proto()?, encoded); + assert_eq!(compatibility_decoded.to_proto()?, encoded); Ok(()) } From 99a7f622f6aff2d707b74ec01a84bbb75e54a366 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Fri, 31 Jul 2026 10:35:55 +0800 Subject: [PATCH 4/7] refactor(proto): refine data sink serde hooks Defer encoding until a sink opts into serialization so the legacy fallback does not encode child plans and sort expressions twice. Restore FileSinkConfig's standard TryFrom API and map enum values by name to preserve wire compatibility independently of discriminants. Refs #23498 --- datafusion/datasource/src/file_sink_config.rs | 3 - .../datasource/src/file_sink_config/proto.rs | 73 ++++++----------- datafusion/datasource/src/sink.rs | 81 ++++++++++--------- .../proto/src/physical_plan/from_proto.rs | 2 +- .../proto/src/physical_plan/to_proto.rs | 2 +- .../tests/cases/roundtrip_physical_plan.rs | 55 +++++++++---- 6 files changed, 110 insertions(+), 106 deletions(-) diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index b2e20567a15bb..48dce9a0cdb3e 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -35,9 +35,6 @@ use object_store::ObjectStore; #[cfg(feature = "proto")] mod proto; -#[cfg(feature = "proto")] -pub use proto::parse_sink_sort_order; - /// Determines how `FileSink` output paths are interpreted. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum FileOutputMode { diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs index f095f2809786b..ed4b5c48bd2af 100644 --- a/datafusion/datasource/src/file_sink_config/proto.rs +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -19,35 +19,32 @@ use std::sync::Arc; -use arrow::compute::SortOptions; -use arrow::datatypes::Schema; -use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::dml::InsertOp; -use datafusion_physical_expr::PhysicalSortExpr; -use datafusion_physical_expr_common::sort_expr::LexRequirement; -use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; use datafusion_proto_models::protobuf; use crate::ListingTableUrl; use crate::file_groups::FileGroup; use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; -impl FileSinkConfig { +impl TryFrom<&FileSinkConfig> for protobuf::FileSinkConfig { + type Error = DataFusionError; + /// Serialize this shared file-sink configuration without format-specific /// writer options. - pub fn to_proto(&self) -> Result { - let file_groups = self + fn try_from(config: &FileSinkConfig) -> Result { + let file_groups = config .file_group .iter() .map(TryInto::try_into) .collect::>>()?; - let table_paths = self + let table_paths = config .table_paths .iter() .map(ToString::to_string) .collect::>(); - let table_partition_cols = self + let table_partition_cols = config .table_partition_cols .iter() .map(|(name, data_type)| { @@ -57,27 +54,36 @@ impl FileSinkConfig { }) }) .collect::>>()?; - let file_output_mode = match self.file_output_mode { + let insert_op = match config.insert_op { + InsertOp::Append => protobuf::InsertOp::Append, + InsertOp::Overwrite => protobuf::InsertOp::Overwrite, + InsertOp::Replace => protobuf::InsertOp::Replace, + }; + let file_output_mode = match config.file_output_mode { FileOutputMode::Automatic => protobuf::FileOutputMode::Automatic, FileOutputMode::SingleFile => protobuf::FileOutputMode::SingleFile, FileOutputMode::Directory => protobuf::FileOutputMode::Directory, }; Ok(protobuf::FileSinkConfig { - object_store_url: self.object_store_url.to_string(), + object_store_url: config.object_store_url.to_string(), file_groups, table_paths, - output_schema: Some(self.output_schema.as_ref().try_into()?), + output_schema: Some(config.output_schema.as_ref().try_into()?), table_partition_cols, - keep_partition_by_columns: self.keep_partition_by_columns, - insert_op: self.insert_op as i32, - file_extension: self.file_extension.clone(), + keep_partition_by_columns: config.keep_partition_by_columns, + insert_op: insert_op.into(), + file_extension: config.file_extension.clone(), file_output_mode: file_output_mode.into(), }) } +} + +impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig { + type Error = DataFusionError; /// Reconstruct a shared file-sink configuration from protobuf. - pub fn from_proto(conf: &protobuf::FileSinkConfig) -> Result { + fn try_from(conf: &protobuf::FileSinkConfig) -> Result { let file_group = FileGroup::new( conf.file_groups .iter() @@ -148,36 +154,9 @@ impl FileSinkConfig { } } -/// Decode a sink's optional required output ordering against its input schema. -pub fn parse_sink_sort_order( - collection: Option<&protobuf::PhysicalSortExprNodeCollection>, - ctx: &ExecutionPlanDecodeCtx<'_>, - schema: &Schema, -) -> Result> { - let Some(collection) = collection else { - return Ok(None); - }; - let sort_exprs = collection - .physical_sort_expr_nodes - .iter() - .map(|node| { - let expr = node.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!("Unexpected empty physical expression") - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, schema)?, - options: SortOptions { - descending: !node.asc, - nulls_first: node.nulls_first, - }, - }) - }) - .collect::>>()?; - Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) -} #[cfg(test)] mod tests { - use datafusion_common::DataFusionError; + use arrow::datatypes::Schema; use super::*; @@ -203,7 +182,7 @@ mod tests { mutate(&mut conf); let error = - FileSinkConfig::from_proto(&conf).expect_err("invalid config should fail"); + FileSinkConfig::try_from(&conf).expect_err("invalid config should fail"); match error { DataFusionError::Internal(message) => { let message = message diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 327d80fea6f79..89a39c2ed4c86 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -74,19 +74,16 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { /// Serialize this sink into a full protobuf plan node, if it knows how. /// - /// [`DataSinkExec::try_to_proto`] encodes the shared child plan and required - /// output ordering before delegating to this hook. Implementations only need - /// to add their sink-specific fields. + /// Implementations can use `ctx` to encode the input plan, sink-specific + /// expressions, and [`DataSinkExec::encode_sort_order`]. /// - /// Returning `Ok(None)` preserves the legacy central serialization fallback. + /// Returning `Ok(None)` preserves the legacy central serialization fallback + /// without eagerly encoding any child plans or expressions. #[cfg(feature = "proto")] fn try_to_proto( &self, - _input: datafusion_proto_models::protobuf::PhysicalPlanNode, - _sort_order: Option< - datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, - >, - _sink_schema: &Schema, + _exec: &DataSinkExec, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { Ok(None) } @@ -164,6 +161,39 @@ impl DataSinkExec { &self.sort_order } + /// Encode the optional sink ordering for a protobuf plan node. + #[cfg(feature = "proto")] + pub fn encode_sort_order( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> + { + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + self.sort_order + .as_ref() + .map(|requirements| { + requirements + .iter() + .map(|requirement| { + let expr: PhysicalSortExpr = requirement.to_owned().into(); + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), + asc: !expr.options.descending, + nulls_first: expr.options.nulls_first, + }) + }) + .collect::>>() + .map(|physical_sort_expr_nodes| { + protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + } + }) + }) + .transpose() + } + fn create_schema( input: &Arc, schema: SchemaRef, @@ -288,42 +318,13 @@ impl ExecutionPlan for DataSinkExec { self.sink.metrics() } - /// Encodes the shared sink plan fields before delegating the sink-specific - /// protobuf representation to [`DataSink::try_to_proto`]. + /// Delegates protobuf serialization to the underlying sink. #[cfg(feature = "proto")] fn try_to_proto( &self, ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, ) -> Result> { - use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_proto_models::protobuf; - - let input = ctx.encode_child(self.input())?; - let sort_order = self - .sort_order() - .as_ref() - .map(|requirements| { - requirements - .iter() - .map(|requirement| { - let expr: PhysicalSortExpr = requirement.to_owned().into(); - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }) - }) - .collect::>>() - .map(|physical_sort_expr_nodes| { - protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes, - } - }) - }) - .transpose()?; - - self.sink() - .try_to_proto(input, sort_order, self.schema().as_ref()) + self.sink().try_to_proto(self, ctx) } } diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 5c87cc476f74b..aaf94f624b42b 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -693,7 +693,7 @@ impl TryFromProto<&protobuf::FileSinkConfig> for FileSinkConfig { type Error = DataFusionError; fn try_from_proto(conf: &protobuf::FileSinkConfig) -> Result { - FileSinkConfig::from_proto(conf) + conf.try_into() } } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index a9637a028a415..448c2e0e93a45 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -611,6 +611,6 @@ impl TryFromProto<&FileSinkConfig> for protobuf::FileSinkConfig { type Error = DataFusionError; fn try_from_proto(conf: &FileSinkConfig) -> Result { - conf.to_proto() + conf.try_into() } } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 8618cd2a97a50..345d4c0be74fc 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -82,6 +82,7 @@ use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion::physical_plan::metrics::MetricCategory; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion::physical_plan::proto::ExecutionPlanEncodeCtx; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::scalar_subquery::{ ScalarSubqueryExec, ScalarSubqueryLink, @@ -134,7 +135,6 @@ use datafusion_proto::bytes::{ physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter, }; -use datafusion_proto::convert::TryFromProto; use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; use datafusion_proto::physical_plan::{ AsExecutionPlan, DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec, @@ -2083,10 +2083,11 @@ impl DataSink for ProtoHookSink { fn try_to_proto( &self, - input: PhysicalPlanNode, - sort_order: Option, - sink_schema: &Schema, + exec: &DataSinkExec, + ctx: &ExecutionPlanEncodeCtx<'_>, ) -> Result> { + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; assert!(matches!( input.physical_plan_type, Some(protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow(_)) @@ -2097,13 +2098,13 @@ impl DataSink for ProtoHookSink { .map(|ordering| ordering.physical_sort_expr_nodes.len()), Some(1) ); - assert_eq!(sink_schema.fields().len(), 1); + assert_eq!(exec.schema().fields().len(), 1); Ok(Some(PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::Empty( protobuf::EmptyExecNode { - schema: Some(sink_schema.try_into()?), + schema: Some(exec.schema().as_ref().try_into()?), partitions: 1, }, ), @@ -2143,7 +2144,7 @@ fn data_sink_exec_delegates_to_sink_proto_hook() -> Result<()> { } #[test] -fn file_sink_config_conversion_preserves_compatibility_api() -> Result<()> { +fn file_sink_config_roundtrip_preserves_fields() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new( "partition", DataType::Utf8, @@ -2162,14 +2163,40 @@ fn file_sink_config_conversion_preserves_compatibility_api() -> Result<()> { file_output_mode: FileOutputMode::Directory, }; - let encoded = config.to_proto()?; - let compatibility_encoded = protobuf::FileSinkConfig::try_from_proto(&config)?; - assert_eq!(encoded, compatibility_encoded); + let encoded = protobuf::FileSinkConfig::try_from(&config)?; + assert_eq!(encoded.insert_op(), protobuf::InsertOp::Overwrite); + assert_eq!( + encoded.file_output_mode(), + protobuf::FileOutputMode::Directory + ); - let decoded = FileSinkConfig::from_proto(&encoded)?; - let compatibility_decoded = FileSinkConfig::try_from_proto(&encoded)?; - assert_eq!(decoded.to_proto()?, encoded); - assert_eq!(compatibility_decoded.to_proto()?, encoded); + let decoded = FileSinkConfig::try_from(&encoded)?; + assert_eq!(decoded.object_store_url, config.object_store_url); + assert_eq!(decoded.table_paths, config.table_paths); + assert_eq!( + decoded.output_schema.as_ref(), + config.output_schema.as_ref() + ); + assert_eq!(decoded.table_partition_cols, config.table_partition_cols); + assert_eq!(decoded.insert_op, config.insert_op); + assert_eq!( + decoded.keep_partition_by_columns, + config.keep_partition_by_columns + ); + assert_eq!(decoded.file_extension, config.file_extension); + assert_eq!(decoded.file_output_mode, config.file_output_mode); + + let [decoded_file] = decoded.file_group.files() else { + panic!("expected one decoded output file"); + }; + let [config_file] = config.file_group.files() else { + panic!("expected one configured output file"); + }; + assert_eq!( + decoded_file.object_meta.location, + config_file.object_meta.location + ); + assert_eq!(decoded_file.object_meta.size, config_file.object_meta.size); Ok(()) } From 62da8a553ce65fc6dce21769ee58e5ecb898e4b8 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 18:54:12 +0800 Subject: [PATCH 5/7] refactor(proto): migrate CsvSink serde CsvSink serialization still depended on central downcasts in datafusion-proto. Moving encode and decode ownership into the CSV sink exercises the DataSink hook and keeps format-specific wire logic with the concrete type. Retain the old decode helper as a deprecated compatibility shim. Refs #23519 --- Cargo.lock | 1 + datafusion/datasource-csv/Cargo.toml | 8 ++ datafusion/datasource-csv/src/file_format.rs | 85 ++++++++++++++++++++ datafusion/datasource/src/sink.rs | 36 ++++++++- datafusion/proto/Cargo.toml | 2 +- datafusion/proto/src/physical_plan/mod.rs | 59 ++++---------- 6 files changed, 145 insertions(+), 46 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9cd5d12d5bca4..7e842fae1f7a3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2017,6 +2017,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 295092512742b..4026e6e808653 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -30,6 +30,13 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -41,6 +48,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index a7f01f6ffec13..20affba89dca2 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -825,6 +825,91 @@ impl DataSink for CsvSink { ) -> Result { FileSink::write_all(self, data, context).await } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + let sink = protobuf::CsvSink { + config: Some((&self.config).try_into()?), + writer_options: Some(self.writer_options().try_into()?), + }; + let node = protobuf::CsvSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(exec.schema().as_ref().try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl CsvSink { + /// Reconstructs a [`DataSinkExec`] containing a `CsvSink` from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sink_node = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::CsvSink(sink)) => { + sink.as_ref() + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a CsvSink" + ); + } + }; + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "CsvSinkExecNode", + "input", + )?; + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSinkExecNode is missing required field 'sink'" + ) + })?; + let config = + FileSinkConfig::try_from(proto_sink.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSink is missing required field 'config'" + ) + })?)?; + let writer_options = proto_sink + .writer_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSink is missing required field 'writer_options'" + ) + })? + .try_into()?; + let data_sink = CsvSink::new(config, writer_options); + let sort_order = DataSinkExec::decode_sort_order( + sink_node.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } #[cfg(test)] diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 89a39c2ed4c86..0d6e92b70585a 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use arrow::array::{ArrayRef, RecordBatch, UInt64Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::{Distribution, EquivalenceProperties}; use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequirements}; @@ -194,6 +194,40 @@ impl DataSinkExec { .transpose() } + /// Decode the optional sink ordering from a protobuf plan node. + #[cfg(feature = "proto")] + pub fn decode_sort_order( + collection: Option< + &datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + schema: &Schema, + ) -> Result> { + use arrow::compute::SortOptions; + use datafusion_physical_expr::PhysicalSortExpr; + + let Some(collection) = collection else { + return Ok(None); + }; + let sort_exprs = collection + .physical_sort_expr_nodes + .iter() + .map(|node| { + let expr = node.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("Unexpected empty physical expression") + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr, schema)?, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + }) + .collect::>>()?; + Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) + } + fn create_schema( input: &Arc, schema: SchemaRef, diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index 037be27769f4d..e4abadba9f59d 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -57,7 +57,7 @@ datafusion-common = { workspace = true } datafusion-datasource = { workspace = true, features = ["proto"] } datafusion-datasource-arrow = { workspace = true } datafusion-datasource-avro = { workspace = true, optional = true } -datafusion-datasource-csv = { workspace = true } +datafusion-datasource-csv = { workspace = true, features = ["proto"] } datafusion-datasource-json = { workspace = true } datafusion-datasource-parquet = { workspace = true, optional = true } datafusion-execution = { workspace = true } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 79c6394933eae..03087392d2120 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -786,8 +786,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::JsonSink(sink) => { self.try_into_json_sink_physical_plan(sink, ctx, proto_converter) } - PhysicalPlanType::CsvSink(sink) => { - self.try_into_csv_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::CsvSink(_) => { + CsvSink::try_from_proto(self.node(), &decode_ctx) } #[cfg_attr(not(feature = "parquet"), allow(unused_variables))] PhysicalPlanType::ParquetSink(sink) => { @@ -1672,41 +1672,25 @@ pub trait PhysicalPlanNodeExt: Sized { ))) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `CsvSink` deserializes itself via `CsvSink::try_from_proto`" + )] fn try_into_csv_sink_physical_plan( &self, sink: &protobuf::CsvSinkExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink = CsvSink::try_from_proto( - sink.sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new(sink.clone()))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + CsvSink::try_from_proto(&node, &decode_ctx) } #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] @@ -2620,19 +2604,6 @@ pub trait PhysicalPlanNodeExt: Sized { })); } - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new( - protobuf::CsvSinkExecNode { - input: Some(Box::new(input)), - sink: Some(protobuf::CsvSink::try_from_proto(sink)?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } - #[cfg(feature = "parquet")] if let Some(sink) = exec.sink().downcast_ref::() { return Ok(Some(protobuf::PhysicalPlanNode { From f0ac7e4e89d1767862b1c937f00a8bdab68df5f3 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 19:11:58 +0800 Subject: [PATCH 6/7] refactor(proto): migrate JsonSink serde JsonSink serialization still depended on central downcasts in datafusion-proto. Moving encode and decode ownership into the JSON sink exercises the DataSink hook and keeps format-specific wire logic with the concrete type. Retain the old decode helper as a deprecated compatibility shim. Refs #23519 --- Cargo.lock | 1 + datafusion/datasource-json/Cargo.toml | 8 ++ datafusion/datasource-json/src/file_format.rs | 85 +++++++++++++++++++ datafusion/proto/Cargo.toml | 2 +- datafusion/proto/src/physical_plan/mod.rs | 59 ++++--------- 5 files changed, 110 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7e842fae1f7a3..1ff1adc39c27b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2039,6 +2039,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", diff --git a/datafusion/datasource-json/Cargo.toml b/datafusion/datasource-json/Cargo.toml index b5947ea5c4c67..7aefbb42c1a7b 100644 --- a/datafusion/datasource-json/Cargo.toml +++ b/datafusion/datasource-json/Cargo.toml @@ -30,6 +30,13 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -41,6 +48,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index 43bde2a039059..5331521b6b488 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -490,6 +490,91 @@ impl DataSink for JsonSink { ) -> Result { FileSink::write_all(self, data, context).await } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + let sink = protobuf::JsonSink { + config: Some((&self.config).try_into()?), + writer_options: Some(self.writer_options().try_into()?), + }; + let node = protobuf::JsonSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(exec.schema().as_ref().try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl JsonSink { + /// Reconstructs a [`DataSinkExec`] containing a `JsonSink` from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sink_node = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::JsonSink(sink)) => { + sink.as_ref() + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a JsonSink" + ); + } + }; + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "JsonSinkExecNode", + "input", + )?; + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSinkExecNode is missing required field 'sink'" + ) + })?; + let config = + FileSinkConfig::try_from(proto_sink.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSink is missing required field 'config'" + ) + })?)?; + let writer_options = proto_sink + .writer_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSink is missing required field 'writer_options'" + ) + })? + .try_into()?; + let data_sink = JsonSink::new(config, writer_options); + let sort_order = DataSinkExec::decode_sort_order( + sink_node.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } #[derive(Debug)] diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index e4abadba9f59d..7bec657542c32 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -58,7 +58,7 @@ datafusion-datasource = { workspace = true, features = ["proto"] } datafusion-datasource-arrow = { workspace = true } datafusion-datasource-avro = { workspace = true, optional = true } datafusion-datasource-csv = { workspace = true, features = ["proto"] } -datafusion-datasource-json = { workspace = true } +datafusion-datasource-json = { workspace = true, features = ["proto"] } datafusion-datasource-parquet = { workspace = true, optional = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 03087392d2120..a80c78b70fd34 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -783,8 +783,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::Analyze(_) => { AnalyzeExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::JsonSink(sink) => { - self.try_into_json_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::JsonSink(_) => { + JsonSink::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::CsvSink(_) => { CsvSink::try_from_proto(self.node(), &decode_ctx) @@ -1635,41 +1635,25 @@ pub trait PhysicalPlanNodeExt: Sized { AnalyzeExec::try_from_proto(self.node(), &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `JsonSink` deserializes itself via `JsonSink::try_from_proto`" + )] fn try_into_json_sink_physical_plan( &self, sink: &protobuf::JsonSinkExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink = JsonSink::try_from_proto( - sink.sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new(sink.clone()))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + JsonSink::try_from_proto(&node, &decode_ctx) } #[deprecated( @@ -2591,19 +2575,6 @@ pub trait PhysicalPlanNodeExt: Sized { None => None, }; - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new( - protobuf::JsonSinkExecNode { - input: Some(Box::new(input)), - sink: Some(protobuf::JsonSink::try_from_proto(sink)?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } - #[cfg(feature = "parquet")] if let Some(sink) = exec.sink().downcast_ref::() { return Ok(Some(protobuf::PhysicalPlanNode { From d3a99c328f903c27d021c8551c3ad42108d47887 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 19:48:09 +0800 Subject: [PATCH 7/7] refactor(proto): migrate ParquetSink serde ParquetSink serialization still depended on the final concrete sink downcast in datafusion-proto. Moving encode and decode ownership into the Parquet sink completes the DataSink migration and keeps format-specific wire logic with the concrete type. Remove the now-unused central dispatch path while retaining its public helpers as deprecated compatibility shims. Refs #23519 Signed-off-by: Jiawei Zhao --- Cargo.lock | 1 + datafusion/datasource-parquet/Cargo.toml | 6 + datafusion/datasource-parquet/src/sink.rs | 89 ++++++++++++++ datafusion/datasource/src/sink.rs | 3 +- datafusion/proto/Cargo.toml | 2 +- datafusion/proto/src/physical_plan/mod.rs | 136 ++++++---------------- 6 files changed, 132 insertions(+), 105 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1ff1adc39c27b..b72742135b830 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2070,6 +2070,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-pruning", "datafusion-session", "futures", diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index 32424069c17a0..a2589af19a6ee 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -46,6 +46,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-pruning = { workspace = true } datafusion-session = { workspace = true } futures = { workspace = true } @@ -74,6 +75,11 @@ name = "datafusion_datasource_parquet" path = "src/mod.rs" [features] +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] parquet_encryption = [ "parquet/encryption", "datafusion-common/parquet_encryption", diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index df2f17c6be22d..edfe6c40d3c20 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -33,6 +33,8 @@ use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; use datafusion_datasource::sink::DataSink; +#[cfg(feature = "proto")] +use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::{ ObjectWriterBuilder, SharedBuffer, get_writer_schema, @@ -40,6 +42,8 @@ use datafusion_datasource::write::{ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +#[cfg(feature = "proto")] +use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::metrics::{ ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet, Time, @@ -409,6 +413,91 @@ impl DataSink for ParquetSink { ) -> Result { FileSink::write_all(self, data, context).await } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + let sink = protobuf::ParquetSink { + config: Some((&self.config).try_into()?), + parquet_options: Some(self.parquet_options().try_into()?), + }; + let node = protobuf::ParquetSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(exec.schema().as_ref().try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl ParquetSink { + /// Reconstructs a [`DataSinkExec`] containing a `ParquetSink` from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sink_node = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetSink(sink)) => { + sink.as_ref() + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a 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(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSinkExecNode is missing required field 'sink'" + ) + })?; + let config = + FileSinkConfig::try_from(proto_sink.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'config'" + ) + })?)?; + let parquet_options = proto_sink + .parquet_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'parquet_options'" + ) + })? + .try_into()?; + let data_sink = ParquetSink::new(config, parquet_options); + let sort_order = DataSinkExec::decode_sort_order( + sink_node.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } /// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 0d6e92b70585a..d7128700da1ed 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -77,8 +77,7 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { /// Implementations can use `ctx` to encode the input plan, sink-specific /// expressions, and [`DataSinkExec::encode_sort_order`]. /// - /// Returning `Ok(None)` preserves the legacy central serialization fallback - /// without eagerly encoding any child plans or expressions. + /// Returning `Ok(None)` lets the caller try its extension codec instead. #[cfg(feature = "proto")] fn try_to_proto( &self, diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index 7bec657542c32..34258db2abf00 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -59,7 +59,7 @@ datafusion-datasource-arrow = { workspace = true } datafusion-datasource-avro = { workspace = true, optional = true } datafusion-datasource-csv = { workspace = true, features = ["proto"] } datafusion-datasource-json = { workspace = true, features = ["proto"] } -datafusion-datasource-parquet = { workspace = true, optional = true } +datafusion-datasource-parquet = { workspace = true, optional = true, features = ["proto"] } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-functions-table = { workspace = true } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index a80c78b70fd34..29e05ea68bb6d 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -54,7 +54,7 @@ use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; -use datafusion_physical_expr::{LexOrdering, LexRequirement}; +use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; use datafusion_physical_plan::aggregates::AggregateExec; @@ -70,7 +70,6 @@ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::coop::CooperativeExec; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::explain::ExplainExec; -use datafusion_physical_plan::expressions::PhysicalSortExpr; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::{ CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, @@ -96,7 +95,6 @@ use prost::Message; use prost::bytes::BufMut; use crate::common::{byte_to_string, str_to_byte}; -use crate::convert::TryFromProto; use crate::convert_required; use crate::physical_plan::from_proto::{ parse_physical_expr_with_converter, parse_physical_sort_exprs, @@ -789,9 +787,15 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::CsvSink(_) => { CsvSink::try_from_proto(self.node(), &decode_ctx) } - #[cfg_attr(not(feature = "parquet"), allow(unused_variables))] - PhysicalPlanType::ParquetSink(sink) => { - self.try_into_parquet_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::ParquetSink(_) => { + #[cfg(feature = "parquet")] + { + ParquetSink::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "parquet"))] + panic!( + "Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled" + ) } PhysicalPlanType::Unnest(_) => { UnnestExec::try_from_proto(self.node(), &decode_ctx) @@ -855,16 +859,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(exec) = plan.downcast_ref::() - && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_sink_exec( - exec, - codec, - proto_converter, - )? - { - return Ok(node); - } - if let Some(exec) = plan.downcast_ref::() && let Some(node) = protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)? @@ -1678,6 +1672,10 @@ pub trait PhysicalPlanNodeExt: Sized { } #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `ParquetSink` deserializes itself via `ParquetSink::try_from_proto`" + )] fn try_into_parquet_sink_physical_plan( &self, sink: &protobuf::ParquetSinkExecNode, @@ -1686,35 +1684,17 @@ pub trait PhysicalPlanNodeExt: Sized { ) -> Result> { #[cfg(feature = "parquet")] { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink = ParquetSink::try_from_proto( - sink.sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))?, - )?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new( + sink.clone(), + ))), + }; + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + ParquetSink::try_from_proto(&node, &decode_ctx) } #[cfg(not(feature = "parquet"))] panic!("Trying to use ParquetSink without `parquet` feature enabled"); @@ -2540,57 +2520,21 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `DataSinkExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_data_sink_exec( exec: &DataSinkExec, codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let input: protobuf::PhysicalPlanNode = - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let sort_order = match exec.sort_order() { - Some(requirements) => { - let expr = requirements - .iter() - .map(|requirement| { - let expr: PhysicalSortExpr = requirement.to_owned().into(); - let sort_expr = protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }; - Ok(sort_expr) - }) - .collect::>>()?; - Some(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: expr, - }) - } - None => None, + let encoder = ConverterPlanEncoder { + codec, + proto_converter, }; - - #[cfg(feature = "parquet")] - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new( - protobuf::ParquetSinkExecNode { - input: Some(Box::new(input)), - sink: Some(protobuf::ParquetSink::try_from_proto(sink)?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } - - // If unknown DataSink then let extension handle it - Ok(None) + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + exec.try_to_proto(&encode_ctx) } #[deprecated( @@ -3320,18 +3264,6 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { } } -fn into_physical_plan( - node: &Option>, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result> { - if let Some(field) = node { - proto_converter.proto_to_execution_plan(field, ctx) - } else { - Err(proto_error("Missing required field in protobuf")) - } -} - /// Adapter backing [`ExecutionPlanEncodeCtx`] for plans migrated to the /// `try_to_proto` hook (#22419). Routes child-plan and child-expr encoding back /// through the central converter so nested plans honor their own hooks.