Proto: migrate file sink serialization

datafusion/datasource-csv/Cargo.toml

@@ -30,6 +30,13 @@ version.workspace = true
 [package.metadata.docs.rs]
 all-features = true
 
+[features]
+proto = [
+    "dep:datafusion-proto-models",
+    "datafusion-datasource/proto",
+    "datafusion-physical-plan/proto",
+]
+
 [dependencies]
 arrow = { workspace = true }
 async-trait = { workspace = true }
@@ -41,6 +48,7 @@ datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+datafusion-proto-models = { workspace = true, optional = true }
 datafusion-session = { workspace = true }
 futures = { workspace = true }
 object_store = { workspace = true }

datafusion/datasource-csv/src/file_format.rs

@@ -826,6 +826,93 @@ impl DataSink for CsvSink {
     ) -> Result<u64> {
         FileSink::write_all(self, data, context).await
     }
+
+    #[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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_proto_models::protobuf;
+        use protobuf::physical_plan_node::PhysicalPlanType;
+
+        let sink = protobuf::CsvSink {
+            config: Some(self.config.to_proto()?),
+            writer_options: Some(self.writer_options().try_into()?),
+        };
+        let node = protobuf::CsvSinkExecNode {
+            input: Some(Box::new(input)),
+            sink: Some(sink),
+            sink_schema: Some(sink_schema.try_into()?),
+            sort_order,
+        };
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new(node))),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl CsvSink {
+    /// Reconstructs a [`DataSinkExec`] containing a `CsvSink` from protobuf.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        use datafusion_datasource::file_sink_config::parse_sink_sort_order;
+        use datafusion_proto_models::protobuf;
+
+        let sink_node = match &node.physical_plan_type {
+            Some(protobuf::physical_plan_node::PhysicalPlanType::CsvSink(sink)) => {
+                sink.as_ref()
+            }
+            _ => {
+                return datafusion_common::internal_err!(
+                    "PhysicalPlanNode is not a CsvSink"
+                );
+            }
+        };
+        let input = ctx.decode_required_child(
+            sink_node.input.as_deref(),
+            "CsvSinkExecNode",
+            "input",
+        )?;
+        let proto_sink = sink_node.sink.as_ref().ok_or_else(|| {
+            datafusion_common::internal_datafusion_err!(
+                "CsvSinkExecNode is missing required field 'sink'"
+            )
+        })?;
+        let config =
+            FileSinkConfig::from_proto(proto_sink.config.as_ref().ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "CsvSink is missing required field 'config'"
+                )
+            })?)?;
+        let writer_options = proto_sink
+            .writer_options
+            .as_ref()
+            .ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "CsvSink is missing required field 'writer_options'"
+                )
+            })?
+            .try_into()?;
+        let data_sink = CsvSink::new(config, writer_options);
+        let sort_order = parse_sink_sort_order(
+            sink_node.sort_order.as_ref(),
+            ctx,
+            input.schema().as_ref(),
+        )?;
+
+        Ok(Arc::new(DataSinkExec::new(
+            input,
+            Arc::new(data_sink),
+            sort_order,
+        )))
+    }
 }
 
 #[cfg(test)]

datafusion/datasource-json/Cargo.toml

@@ -30,6 +30,13 @@ version.workspace = true
 [package.metadata.docs.rs]
 all-features = true
 
+[features]
+proto = [
+    "dep:datafusion-proto-models",
+    "datafusion-datasource/proto",
+    "datafusion-physical-plan/proto",
+]
+
 [dependencies]
 arrow = { workspace = true }
 async-trait = { workspace = true }
@@ -41,6 +48,7 @@ datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+datafusion-proto-models = { workspace = true, optional = true }
 datafusion-session = { workspace = true }
 futures = { workspace = true }
 object_store = { workspace = true }

datafusion/datasource-json/src/file_format.rs

@@ -490,6 +490,93 @@ impl DataSink for JsonSink {
     ) -> Result<u64> {
         FileSink::write_all(self, data, context).await
     }
+
+    #[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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_proto_models::protobuf;
+        use protobuf::physical_plan_node::PhysicalPlanType;
+
+        let sink = protobuf::JsonSink {
+            config: Some(self.config.to_proto()?),
+            writer_options: Some(self.writer_options().try_into()?),
+        };
+        let node = protobuf::JsonSinkExecNode {
+            input: Some(Box::new(input)),
+            sink: Some(sink),
+            sink_schema: Some(sink_schema.try_into()?),
+            sort_order,
+        };
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new(node))),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl JsonSink {
+    /// Reconstructs a [`DataSinkExec`] containing a `JsonSink` from protobuf.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        use datafusion_datasource::file_sink_config::parse_sink_sort_order;
+        use datafusion_proto_models::protobuf;
+
+        let sink_node = match &node.physical_plan_type {
+            Some(protobuf::physical_plan_node::PhysicalPlanType::JsonSink(sink)) => {
+                sink.as_ref()
+            }
+            _ => {
+                return datafusion_common::internal_err!(
+                    "PhysicalPlanNode is not a JsonSink"
+                );
+            }
+        };
+        let input = ctx.decode_required_child(
+            sink_node.input.as_deref(),
+            "JsonSinkExecNode",
+            "input",
+        )?;
+        let proto_sink = sink_node.sink.as_ref().ok_or_else(|| {
+            datafusion_common::internal_datafusion_err!(
+                "JsonSinkExecNode is missing required field 'sink'"
+            )
+        })?;
+        let config =
+            FileSinkConfig::from_proto(proto_sink.config.as_ref().ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "JsonSink is missing required field 'config'"
+                )
+            })?)?;
+        let writer_options = proto_sink
+            .writer_options
+            .as_ref()
+            .ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "JsonSink is missing required field 'writer_options'"
+                )
+            })?
+            .try_into()?;
+        let data_sink = JsonSink::new(config, writer_options);
+        let sort_order = parse_sink_sort_order(
+            sink_node.sort_order.as_ref(),
+            ctx,
+            input.schema().as_ref(),
+        )?;
+
+        Ok(Arc::new(DataSinkExec::new(
+            input,
+            Arc::new(data_sink),
+            sort_order,
+        )))
+    }
 }
 
 #[derive(Debug)]

datafusion/datasource-parquet/Cargo.toml

@@ -46,6 +46,7 @@ datafusion-physical-expr = { workspace = true }
 datafusion-physical-expr-adapter = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-plan = { workspace = true }
+datafusion-proto-models = { workspace = true, optional = true }
 datafusion-pruning = { workspace = true }
 datafusion-session = { workspace = true }
 futures = { workspace = true }
@@ -74,6 +75,11 @@ name = "datafusion_datasource_parquet"
 path = "src/mod.rs"
 
 [features]
+proto = [
+    "dep:datafusion-proto-models",
+    "datafusion-datasource/proto",
+    "datafusion-physical-plan/proto",
+]
 parquet_encryption = [
     "parquet/encryption",
     "datafusion-common/parquet_encryption",

datafusion/datasource-parquet/src/sink.rs

@@ -33,13 +33,17 @@ use datafusion_datasource::display::FileGroupDisplay;
 use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig};
 use datafusion_datasource::sink::DataSink;
+#[cfg(feature = "proto")]
+use datafusion_datasource::sink::DataSinkExec;
 use datafusion_datasource::write::demux::DemuxedStreamReceiver;
 use datafusion_datasource::write::{
     ObjectWriterBuilder, SharedBuffer, get_writer_schema,
 };
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
 use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+#[cfg(feature = "proto")]
+use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::metrics::{
     ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory,
     MetricsSet, Time,
@@ -411,6 +415,93 @@ impl DataSink for ParquetSink {
     ) -> Result<u64> {
         FileSink::write_all(self, data, context).await
     }
+
+    #[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<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_proto_models::protobuf;
+        use protobuf::physical_plan_node::PhysicalPlanType;
+
+        let sink = protobuf::ParquetSink {
+            config: Some(self.config.to_proto()?),
+            parquet_options: Some(self.parquet_options().try_into()?),
+        };
+        let node = protobuf::ParquetSinkExecNode {
+            input: Some(Box::new(input)),
+            sink: Some(sink),
+            sink_schema: Some(sink_schema.try_into()?),
+            sort_order,
+        };
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new(node))),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl ParquetSink {
+    /// Reconstructs a [`DataSinkExec`] containing a `ParquetSink` from protobuf.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        use datafusion_datasource::file_sink_config::parse_sink_sort_order;
+        use datafusion_proto_models::protobuf;
+
+        let sink_node = match &node.physical_plan_type {
+            Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetSink(sink)) => {
+                sink.as_ref()
+            }
+            _ => {
+                return datafusion_common::internal_err!(
+                    "PhysicalPlanNode is not a ParquetSink"
+                );
+            }
+        };
+        let input = ctx.decode_required_child(
+            sink_node.input.as_deref(),
+            "ParquetSinkExecNode",
+            "input",
+        )?;
+        let proto_sink = sink_node.sink.as_ref().ok_or_else(|| {
+            datafusion_common::internal_datafusion_err!(
+                "ParquetSinkExecNode is missing required field 'sink'"
+            )
+        })?;
+        let config =
+            FileSinkConfig::from_proto(proto_sink.config.as_ref().ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "ParquetSink is missing required field 'config'"
+                )
+            })?)?;
+        let parquet_options = proto_sink
+            .parquet_options
+            .as_ref()
+            .ok_or_else(|| {
+                datafusion_common::internal_datafusion_err!(
+                    "ParquetSink is missing required field 'parquet_options'"
+                )
+            })?
+            .try_into()?;
+        let data_sink = ParquetSink::new(config, parquet_options);
+        let sort_order = parse_sink_sort_order(
+            sink_node.sort_order.as_ref(),
+            ctx,
+            input.schema().as_ref(),
+        )?;
+
+        Ok(Arc::new(DataSinkExec::new(
+            input,
+            Arc::new(data_sink),
+            sort_order,
+        )))
+    }
 }
 
 /// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter]