[EPIC/reference] Port ExecutionPlan serialization to try_to_proto / try_from_proto hooks (full)

datafusion/datasource-arrow/Cargo.toml

@@ -42,6 +42,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 }
 itertools = { workspace = true }
@@ -65,3 +66,10 @@ path = "src/mod.rs"
 # This feature is deprecated, as core functionality in the SpillManager requires all features
 # it enabled, and will be removed in a future version.
 compression = []
+# Enables `FileSource::try_to_proto` on `ArrowSource` and the `ArrowScan` decode
+# entry point. Mirrors the `proto` feature on `datafusion-datasource`.
+proto = [
+    "dep:datafusion-proto-models",
+    "datafusion-datasource/proto",
+    "datafusion-physical-plan/proto",
+]

datafusion/datasource-arrow/src/source.rs

@@ -392,6 +392,67 @@ impl FileSource for ArrowSource {
     fn projection(&self) -> Option<&ProjectionExprs> {
         Some(&self.projection.source)
     }
+
+    /// Emit an `ArrowScan` node wrapping the shared base config. Arrow is
+    /// base_conf-only (no format-specific fields). Kept byte-identical to the
+    /// former `try_from_data_source_exec` ArrowScan branch in `datafusion-proto`.
+    #[cfg(feature = "proto")]
+    fn try_to_proto(
+        &self,
+        base: &FileScanConfig,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
+    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_proto_models::protobuf;
+        use protobuf::physical_plan_node::PhysicalPlanType;
+
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::ArrowScan(
+                protobuf::ArrowScanExecNode {
+                    base_conf: Some(base.to_proto_conf(ctx)?),
+                },
+            )),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl ArrowSource {
+    /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over an
+    /// `ArrowSource`) from an `ArrowScan` [`PhysicalPlanNode`].
+    ///
+    /// The inverse of [`FileSource::try_to_proto`] on `ArrowSource`; kept
+    /// byte-compatible with the former `try_into_arrow_scan_physical_plan` in
+    /// `datafusion-proto`.
+    ///
+    /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
+    ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
+        use datafusion_datasource::file_scan_config::FileScanConfig;
+        use datafusion_datasource::source::DataSourceExec;
+        use datafusion_proto_models::protobuf;
+
+        let scan = match &node.physical_plan_type {
+            Some(protobuf::physical_plan_node::PhysicalPlanType::ArrowScan(scan)) => scan,
+            _ => {
+                return datafusion_common::internal_err!(
+                    "PhysicalPlanNode is not an ArrowScan"
+                );
+            }
+        };
+
+        let base_conf = scan.base_conf.as_ref().ok_or_else(|| {
+            datafusion_common::internal_datafusion_err!(
+                "base_conf in ArrowScanExecNode is missing."
+            )
+        })?;
+
+        let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;
+        let source = Arc::new(ArrowSource::new_file_source(table_schema));
+        let scan_conf = FileScanConfig::from_proto_conf(base_conf, ctx, source)?;
+        Ok(DataSourceExec::from_data_source(scan_conf))
+    }
 }
 
 /// `FileOpener` wrapper for both Arrow IPC file and stream formats

datafusion/datasource-avro/Cargo.toml

@@ -30,6 +30,15 @@ version.workspace = true
 [package.metadata.docs.rs]
 all-features = true
 
+[features]
+# Enables `FileSource::try_to_proto` on `AvroSource` and the `AvroScan` decode
+# entry point. Mirrors the `proto` feature on `datafusion-datasource`.
+proto = [
+    "dep:datafusion-proto-models",
+    "datafusion-datasource/proto",
+    "datafusion-physical-plan/proto",
+]
+
 [dependencies]
 arrow = { workspace = true }
 arrow-avro = { workspace = true }
@@ -39,6 +48,7 @@ datafusion-common = { workspace = true, features = ["object_store"] }
 datafusion-datasource = { workspace = true }
 datafusion-physical-expr-adapter = { 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-avro/src/source.rs

@@ -168,6 +168,67 @@ impl FileSource for AvroSource {
         // Avro OCF does not support safe byte-range splitting in this reader path.
         false
     }
+
+    /// Emit an `AvroScan` node wrapping the shared base config. Avro carries no
+    /// format-specific fields, so this is just the base conf. Kept
+    /// byte-identical to the former `try_from_data_source_exec` AvroScan branch
+    /// in `datafusion-proto`.
+    #[cfg(feature = "proto")]
+    fn try_to_proto(
+        &self,
+        base: &FileScanConfig,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
+    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
+        use datafusion_proto_models::protobuf;
+        use protobuf::physical_plan_node::PhysicalPlanType;
+
+        let node = protobuf::AvroScanExecNode {
+            base_conf: Some(base.to_proto_conf(ctx)?),
+        };
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::AvroScan(node)),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl AvroSource {
+    /// Reconstruct a `DataSourceExec` (wrapping a `FileScanConfig` over an
+    /// `AvroSource`) from an `AvroScan`
+    /// [`PhysicalPlanNode`](datafusion_proto_models::protobuf::PhysicalPlanNode).
+    ///
+    /// The inverse of [`FileSource::try_to_proto`] on `AvroSource`; kept
+    /// byte-compatible with the former `try_into_avro_scan_physical_plan` in
+    /// `datafusion-proto`.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
+        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
+    ) -> Result<Arc<dyn datafusion_physical_plan::ExecutionPlan>> {
+        use datafusion_datasource::source::DataSourceExec;
+        use datafusion_proto_models::protobuf;
+
+        let scan = match &node.physical_plan_type {
+            Some(protobuf::physical_plan_node::PhysicalPlanType::AvroScan(scan)) => scan,
+            _ => {
+                return datafusion_common::internal_err!(
+                    "PhysicalPlanNode is not an AvroScan"
+                );
+            }
+        };
+
+        let base_conf = scan.base_conf.as_ref().ok_or_else(|| {
+            datafusion_common::internal_datafusion_err!(
+                "AvroScanExecNode is missing required field 'base_conf'"
+            )
+        })?;
+
+        // Parse table schema with partition columns, then rebuild the source.
+        let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;
+        let source = Arc::new(AvroSource::new(table_schema));
+
+        let conf = FileScanConfig::from_proto_conf(base_conf, ctx, source)?;
+        Ok(DataSourceExec::from_data_source(conf))
+    }
 }
 
 mod private {

datafusion/datasource-csv/Cargo.toml

@@ -30,6 +30,15 @@ version.workspace = true
 [package.metadata.docs.rs]
 all-features = true
 
+[features]
+# Enables `FileSource::try_to_proto` on `CsvSource` and the `CsvScan` decode
+# entry point. Mirrors the `proto` feature on `datafusion-datasource`.
+proto = [
+    "dep:datafusion-proto-models",
+    "datafusion-datasource/proto",
+    "datafusion-physical-plan/proto",
+]
+
 [dependencies]
 arrow = { workspace = true }
 async-trait = { workspace = true }
@@ -41,6 +50,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,104 @@ impl DataSink for CsvSink {
     ) -> Result<u64> {
         FileSink::write_all(self, data, context).await
     }
+
+    /// Emit a `CsvSink` node wrapping the shared sink config. Kept
+    /// byte-identical to the former `try_from_data_sink_exec` CsvSink branch in
+    /// `datafusion-proto`.
+    #[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 {
+    /// Reconstruct a `DataSinkExec` (wrapping a `CsvSink`) from a `CsvSink`
+    /// [`PhysicalPlanNode`].
+    ///
+    /// The inverse of [`DataSink::try_to_proto`] on `CsvSink`; kept
+    /// byte-compatible with the former `try_into_csv_sink_physical_plan` in
+    /// `datafusion-proto`.
+    ///
+    /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode
+    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 sink_schema = input.schema();
+        let sort_order =
+            parse_sink_sort_order(sink_node.sort_order.as_ref(), ctx, &sink_schema)?;
+
+        Ok(Arc::new(DataSinkExec::new(
+            input,
+            Arc::new(data_sink),
+            sort_order,
+        )))
+    }
 }
 
 #[cfg(test)]