Proto: add DataSink serialization hook

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 }

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 {

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<Self> {
+        let file_groups = config
+            .file_group
+            .iter()
+            .map(TryInto::try_into)
+            .collect::<Result<Vec<_>>>()?;
+        let table_paths = config
+            .table_paths
+            .iter()
+            .map(ToString::to_string)
+            .collect::<Vec<_>>();
+        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::<Result<Vec<_>>>()?;
+        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<Self> {
+        let file_group = FileGroup::new(
+            conf.file_groups
+                .iter()
+                .map(TryInto::try_into)
+                .collect::<Result<Vec<_>>>()?,
+        );
+        let table_paths = conf
+            .table_paths
+            .iter()
+            .map(ListingTableUrl::parse)
+            .collect::<Result<Vec<_>>>()?;
+        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::<Result<Vec<_>>>()?;
+        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<str>,
+    ) {
+        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'",
+        );
+    }
+}

datafusion/datasource/src/sink.rs

@@ -71,6 +71,22 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync {
         data: SendableRecordBatchStream,
         context: &Arc<TaskContext>,
     ) -> Result<u64>;
+
+    /// 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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        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<Option<datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection>>
+    {
+        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::<Result<Vec<_>>>()
+                    .map(|physical_sort_expr_nodes| {
+                        protobuf::PhysicalSortExprNodeCollection {
+                            physical_sort_expr_nodes,
+                        }
+                    })
+            })
+            .transpose()
+    }
+
     fn create_schema(
         input: &Arc<dyn ExecutionPlan>,
         schema: SchemaRef,
@@ -268,6 +317,15 @@ impl ExecutionPlan for DataSinkExec {
     fn metrics(&self) -> Option<MetricsSet> {
         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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        self.sink().try_to_proto(self, ctx)
+    }
 }
 
 /// Create a output record batch with a count