refactor(proto): migrate AvroSource serde

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,57 @@ 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.
+    #[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.try_to_proto(ctx)?),
+        };
+        Ok(Some(protobuf::PhysicalPlanNode {
+            physical_plan_type: Some(PhysicalPlanType::AvroScan(node)),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl AvroSource {
+    /// Reconstructs a `DataSourceExec` from a protobuf `AvroScan`.
+    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'"
+            )
+        })?;
+
+        let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;
+        let source = Arc::new(AvroSource::new(table_schema));
+
+        let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?;
+        Ok(DataSourceExec::from_data_source(conf))
+    }
 }
 
 mod private {

datafusion/proto/Cargo.toml

@@ -63,7 +63,7 @@ datafusion-catalog-listing = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-datasource = { workspace = true, features = ["proto"] }
 datafusion-datasource-arrow = { workspace = true }
-datafusion-datasource-avro = { workspace = true, optional = true }
+datafusion-datasource-avro = { workspace = true, optional = true, features = ["proto"] }
 datafusion-datasource-csv = { workspace = true, features = ["proto"] }
 datafusion-datasource-json = { workspace = true, features = ["proto"] }
 datafusion-datasource-parquet = { workspace = true, optional = true, features = ["proto"] }

datafusion/proto/src/physical_plan/mod.rs

@@ -1090,8 +1090,15 @@ pub trait PhysicalPlanNodeExt: Sized {
                     "Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled"
                 )
             }
-            PhysicalPlanType::AvroScan(scan) => {
-                self.try_into_avro_scan_physical_plan(scan, ctx, proto_converter)
+            PhysicalPlanType::AvroScan(_) => {
+                #[cfg(feature = "avro")]
+                {
+                    AvroSource::try_from_proto(self.node(), &decode_ctx)
+                }
+                #[cfg(not(feature = "avro"))]
+                panic!(
+                    "Unable to process a Avro PhysicalPlan when `avro` feature is not enabled"
+                )
             }
             PhysicalPlanType::MemoryScan(_) => {
                 MemorySourceConfig::try_from_proto(self.node(), &decode_ctx)
@@ -1429,6 +1436,10 @@ pub trait PhysicalPlanNodeExt: Sized {
     }
 
     #[cfg_attr(not(feature = "avro"), expect(unused_variables))]
+    #[deprecated(
+        since = "55.0.0",
+        note = "unused by DataFusion; `AvroSource` deserializes itself via `AvroSource::try_from_proto`"
+    )]
     fn try_into_avro_scan_physical_plan(
         &self,
         scan: &protobuf::AvroScanExecNode,
@@ -1437,15 +1448,15 @@ pub trait PhysicalPlanNodeExt: Sized {
     ) -> Result<Arc<dyn ExecutionPlan>> {
         #[cfg(feature = "avro")]
         {
-            let table_schema =
-                parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?;
-            let conf = parse_protobuf_file_scan_config(
-                scan.base_conf.as_ref().unwrap(),
+            let node = protobuf::PhysicalPlanNode {
+                physical_plan_type: Some(PhysicalPlanType::AvroScan(scan.clone())),
+            };
+            let decoder = ConverterPlanDecoder {
                 ctx,
                 proto_converter,
-                Arc::new(AvroSource::new(table_schema)),
-            )?;
-            Ok(DataSourceExec::from_data_source(conf))
+            };
+            let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder);
+            AvroSource::try_from_proto(&node, &decode_ctx)
         }
 
         #[cfg(not(feature = "avro"))]
@@ -2472,24 +2483,6 @@ pub trait PhysicalPlanNodeExt: Sized {
             }
         }
 
-        #[cfg(feature = "avro")]
-        if let Some(maybe_avro) = data_source.downcast_ref::<FileScanConfig>() {
-            let source = maybe_avro.file_source();
-            if source.downcast_ref::<AvroSource>().is_some() {
-                return Ok(Some(protobuf::PhysicalPlanNode {
-                    physical_plan_type: Some(PhysicalPlanType::AvroScan(
-                        protobuf::AvroScanExecNode {
-                            base_conf: Some(serialize_file_scan_config(
-                                maybe_avro,
-                                codec,
-                                proto_converter,
-                            )?),
-                        },
-                    )),
-                }));
-            }
-        }
-
         Ok(None)
     }
 

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

@@ -1378,6 +1378,24 @@ fn roundtrip_json_scan() -> Result<()> {
     roundtrip_test(DataSourceExec::from_data_source(scan_config))
 }
 
+#[cfg(feature = "avro")]
+#[test]
+fn roundtrip_avro_scan() -> Result<()> {
+    use datafusion_datasource_avro::source::AvroSource;
+
+    let file_schema =
+        Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
+    let file_source = Arc::new(AvroSource::new(TableSchema::from(&file_schema)));
+    let scan_config =
+        FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
+            .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new(
+                "/path/to/file.avro".to_string(),
+                1024,
+            )])])
+            .build();
+    roundtrip_test(DataSourceExec::from_data_source(scan_config))
+}
+
 #[test]
 fn roundtrip_csv_scan_preserves_format_options() -> Result<()> {
     use datafusion::common::config::CsvOptions;