diff --git a/Cargo.lock b/Cargo.lock index 4b4c40791e517..6d13619d52e87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1997,6 +1997,7 @@ dependencies = [ "datafusion-datasource", "datafusion-physical-expr-adapter", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml index adc2be1cb8f24..70b675d63f427 100644 --- a/datafusion/datasource-avro/Cargo.toml +++ b/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 } diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index e3be9d8a401d0..cef85c58dfa39 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/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> { + 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> { + 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 { diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index fb3308ae52a20..cac64cf588a14 100644 --- a/datafusion/proto/Cargo.toml +++ b/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"] } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 3cfd7fc3188c0..ee3fc32aac20a 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/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> { #[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::() { - let source = maybe_avro.file_source(); - if source.downcast_ref::().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) } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index c3dcaa8814c4d..99ba5cf0abc8a 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/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;