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/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index 1abce86a3565f..48dce9a0cdb3e 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -32,6 +32,9 @@ use datafusion_expr::dml::InsertOp; use async_trait::async_trait; use object_store::ObjectStore; +#[cfg(feature = "proto")] +mod proto; + /// 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..ed4b5c48bd2af --- /dev/null +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -0,0 +1,239 @@ +// 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 datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_expr::dml::InsertOp; +use datafusion_proto_models::protobuf; + +use crate::ListingTableUrl; +use crate::file_groups::FileGroup; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; + +impl TryFrom<&FileSinkConfig> for protobuf::FileSinkConfig { + type Error = DataFusionError; + + /// Serialize this shared file-sink configuration without format-specific + /// writer options. + fn try_from(config: &FileSinkConfig) -> Result { + let file_groups = config + .file_group + .iter() + .map(TryInto::try_into) + .collect::>>()?; + let table_paths = config + .table_paths + .iter() + .map(ToString::to_string) + .collect::>(); + let table_partition_cols = config + .table_partition_cols + .iter() + .map(|(name, data_type)| { + Ok(protobuf::PartitionColumn { + name: name.to_owned(), + arrow_type: Some(data_type.try_into()?), + }) + }) + .collect::>>()?; + 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: config.object_store_url.to_string(), + file_groups, + table_paths, + output_schema: Some(config.output_schema.as_ref().try_into()?), + table_partition_cols, + 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. + fn try_from(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 = 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 = 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, + }; + 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, + }) + } +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::Schema; + + 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::try_from(&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/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 18ebe80773e8a..89a39c2ed4c86 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -71,6 +71,22 @@ 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. + /// + /// 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. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _exec: &DataSinkExec, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn DataSink { @@ -145,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, @@ -268,6 +317,15 @@ impl ExecutionPlan for DataSinkExec { fn metrics(&self) -> Option { self.sink.metrics() } + + /// Delegates protobuf serialization to the underlying sink. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + self.sink().try_to_proto(self, ctx) + } } /// Create a output record batch with a count diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 34ad8c7a62fc7..aaf94f624b42b 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, - }) + conf.try_into() } } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 5189972f0e200..448c2e0e93a45 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.try_into() } } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 864e6d68676ee..345d4c0be74fc 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; @@ -81,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, @@ -2054,6 +2056,150 @@ 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, + 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(_)) + )); + assert_eq!( + sort_order + .as_ref() + .map(|ordering| ordering.physical_sort_expr_nodes.len()), + Some(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(exec.schema().as_ref().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(()) +} + +#[test] +fn file_sink_config_roundtrip_preserves_fields() -> 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 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::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(()) +} + #[tokio::test] async fn roundtrip_json_source() -> Result<()> { let ctx = SessionContext::new();