From 0c4f9f6e886f9c7301a679af1b9b198821cc1e5e Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 16:52:51 +0800 Subject: [PATCH 1/3] 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 --- Cargo.lock | 1 + datafusion/datasource/Cargo.toml | 7 ++ datafusion/datasource/src/sink.rs | 57 ++++++++++++ datafusion/proto/Cargo.toml | 2 +- .../tests/cases/roundtrip_physical_plan.rs | 89 ++++++++++++++++++- 5 files changed, 154 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ed63f60b41519..2ed6c2d6e902a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1936,6 +1936,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 2ac42ed900095..dd73966398115 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,6 +34,12 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] +# Enables `DataSink::try_to_proto` 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 } @@ -56,6 +62,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-session = { workspace = true } flate2 = { workspace = true, optional = true } futures = { 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/Cargo.toml b/datafusion/proto/Cargo.toml index cfff8a949418a..037be27769f4d 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -54,7 +54,7 @@ chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true } -datafusion-datasource = { 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 } 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 457c647d874f377e82d57308203372883124780e Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Tue, 21 Jul 2026 17:19:42 +0800 Subject: [PATCH 2/3] 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 | 232 ++++++++++++++++++ .../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, 272 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..0481a4c66a452 --- /dev/null +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -0,0 +1,232 @@ +// 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 chrono::{TimeZone, Utc}; +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 object_store::ObjectMeta; +use object_store::path::Path; + +use crate::file_groups::FileGroup; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; +use crate::{ListingTableUrl, PartitionedFile}; + +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(partitioned_file_to_proto) + .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(partitioned_file_from_proto) + .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))) +} + +fn partitioned_file_to_proto( + file: &PartitionedFile, +) -> Result { + let last_modified = file.object_meta.last_modified; + let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { + DataFusionError::Plan(format!( + "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" + )) + })? as u64; + + Ok(protobuf::PartitionedFile { + arrow_schema: file + .arrow_schema + .as_ref() + .map(|schema| schema.as_ref().try_into()) + .transpose()?, + path: file.object_meta.location.as_ref().to_owned(), + size: file.object_meta.size, + last_modified_ns, + partition_values: file + .partition_values + .iter() + .map(TryInto::try_into) + .collect::, _>>()?, + range: file.range.as_ref().map(|range| protobuf::FileRange { + start: range.start, + end: range.end, + }), + statistics: file.statistics.as_ref().map(|stats| stats.as_ref().into()), + }) +} + +fn partitioned_file_from_proto( + file: &protobuf::PartitionedFile, +) -> Result { + let mut partitioned_file = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse(file.path.as_str()).map_err(|error| { + internal_datafusion_err!("Invalid object_store path: {error}") + })?, + last_modified: Utc.timestamp_nanos(file.last_modified_ns as i64), + size: file.size, + e_tag: None, + version: None, + }) + .with_partition_values( + file.partition_values + .iter() + .map(TryInto::try_into) + .collect::, _>>()?, + ); + if let Some(schema) = file.arrow_schema.as_ref() { + partitioned_file = partitioned_file.with_arrow_schema(Arc::new( + schema.try_into().map_err(DataFusionError::from)?, + )); + } + if let Some(range) = file.range.as_ref() { + partitioned_file = partitioned_file.with_range(range.start, range.end); + } + if let Some(statistics) = file.statistics.as_ref() { + partitioned_file = + partitioned_file.with_statistics(Arc::new(statistics.try_into()?)); + } + Ok(partitioned_file) +} diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index d5a1e0efac6b6..c1e6b6be8780c 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -31,7 +31,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")] @@ -39,7 +39,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; @@ -728,53 +727,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(PartitionedFile::try_from_proto) - .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 e13923dbb9519..179bd36e1b399 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -632,47 +632,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 91da434decb886b89743264f7208a802658d0c11 Mon Sep 17 00:00:00 2001 From: Jiawei Zhao Date: Mon, 27 Jul 2026 16:13:46 +0800 Subject: [PATCH 3/3] 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 | 99 ++++++++++++++++++- .../tests/cases/roundtrip_physical_plan.rs | 15 +-- 2 files changed, 105 insertions(+), 9 deletions(-) diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs index 0481a4c66a452..9410b83b2b0cf 100644 --- a/datafusion/datasource/src/file_sink_config/proto.rs +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -107,12 +107,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, @@ -230,3 +243,85 @@ fn partitioned_file_from_proto( } Ok(partitioned_file) } + +#[cfg(test)] +mod tests { + 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(()) }