From 7f41edbb371751ee897e3fdaee9065a1dc45fa7b Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 16:52:51 +0800 Subject: [PATCH 1/4] 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/4] 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/4] 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/4] 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(()) }