From 4ae3ceb58400ecb99680a07a82b06c85a9c392f2 Mon Sep 17 00:00:00 2001 From: buraksenn Date: Sun, 9 Aug 2026 11:19:46 +0300 Subject: [PATCH] refactor(proto): migrate ArrowSource serde --- Cargo.lock | 1 + datafusion/datasource-arrow/Cargo.toml | 8 +++ datafusion/datasource-arrow/src/source.rs | 58 +++++++++++++++++ datafusion/proto/Cargo.toml | 2 +- datafusion/proto/src/physical_plan/mod.rs | 78 ++++++++--------------- 5 files changed, 96 insertions(+), 51 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6d13619d52e87..c4128d30f4e0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1978,6 +1978,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "itertools 0.15.0", diff --git a/datafusion/datasource-arrow/Cargo.toml b/datafusion/datasource-arrow/Cargo.toml index 2718e424c6386..6f50135403d69 100644 --- a/datafusion/datasource-arrow/Cargo.toml +++ b/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", +] diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 27533052ce03f..dbdc1b1cf0f11 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -392,6 +392,64 @@ impl FileSource for ArrowSource { fn projection(&self) -> Option<&ProjectionExprs> { Some(&self.projection.source) } + + /// Emit an `ArrowScan` node wrapping the shared base config. + /// + /// Decoding defaults to the IPC file format because protobuf does not + /// distinguish it from the IPC stream format. + #[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; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ArrowScan( + protobuf::ArrowScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl ArrowSource { + /// Reconstructs a `DataSourceExec` from a protobuf `ArrowScan`. + /// + /// Defaults to the IPC file format because protobuf does not distinguish it + /// from the IPC stream format. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + 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!( + "ArrowScanExecNode is missing required field 'base_conf'" + ) + })?; + + 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::try_from_proto(base_conf, ctx, source)?; + Ok(DataSourceExec::from_data_source(scan_conf)) + } } /// `FileOpener` wrapper for both Arrow IPC file and stream formats diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index cac64cf588a14..008f75422a69c 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -62,7 +62,7 @@ datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true } datafusion-datasource = { workspace = true, features = ["proto"] } -datafusion-datasource-arrow = { workspace = true } +datafusion-datasource-arrow = { workspace = true, features = ["proto"] } datafusion-datasource-avro = { workspace = true, optional = true, features = ["proto"] } datafusion-datasource-csv = { workspace = true, features = ["proto"] } datafusion-datasource-json = { workspace = true, features = ["proto"] } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index ee3fc32aac20a..258c775846f60 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -26,7 +26,6 @@ use datafusion_catalog::memory::MemorySourceConfig; use datafusion_common::{ DataFusionError, Result, internal_datafusion_err, internal_err, not_impl_err, }; -use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource_arrow::source::ArrowSource; @@ -86,13 +85,8 @@ use prost::Message; use prost::bytes::BufMut; use crate::convert_required; -use crate::physical_plan::from_proto::{ - parse_physical_expr_with_converter, parse_protobuf_file_scan_config, - parse_table_schema_from_proto, -}; -use crate::physical_plan::to_proto::{ - serialize_file_scan_config, serialize_physical_expr_with_converter, -}; +use crate::physical_plan::from_proto::parse_physical_expr_with_converter; +use crate::physical_plan::to_proto::serialize_physical_expr_with_converter; use crate::protobuf::physical_plan_node::PhysicalPlanType; use crate::protobuf::{self, SortMergeJoinExecNode, proto_error}; @@ -115,7 +109,9 @@ mod file_scan_config_serde { use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; - use datafusion_datasource::file_scan_config::FileScanConfigBuilder; + use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, + }; use datafusion_datasource::file_stream::FileOpener; use datafusion_datasource::{PartitionedFile, TableSchema}; use datafusion_execution::object_store::ObjectStoreUrl; @@ -1103,8 +1099,8 @@ pub trait PhysicalPlanNodeExt: Sized { PhysicalPlanType::MemoryScan(_) => { MemorySourceConfig::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::ArrowScan(scan) => { - self.try_into_arrow_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::ArrowScan(_) => { + ArrowSource::try_from_proto(self.node(), &decode_ctx) } #[expect( deprecated, @@ -1233,16 +1229,6 @@ pub trait PhysicalPlanNodeExt: Sized { return Ok(node); } - if let Some(data_source_exec) = plan.downcast_ref::() - && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_source_exec( - data_source_exec, - codec, - proto_converter, - )? - { - return Ok(node); - } - if let Some(exec) = plan.downcast_ref::() && let Some(node) = protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)? @@ -1386,23 +1372,25 @@ pub trait PhysicalPlanNodeExt: Sized { JsonSource::try_from_proto(&node, &decode_ctx) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `ArrowSource` deserializes itself via `ArrowSource::try_from_proto`" + )] fn try_into_arrow_scan_physical_plan( &self, scan: &protobuf::ArrowScanExecNode, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let base_conf = scan.base_conf.as_ref().ok_or_else(|| { - internal_datafusion_err!("base_conf in ArrowScanExecNode is missing.") - })?; - let table_schema = parse_table_schema_from_proto(base_conf)?; - let scan_conf = parse_protobuf_file_scan_config( - base_conf, + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ArrowScan(scan.clone())), + }; + let decoder = ConverterPlanDecoder { ctx, proto_converter, - Arc::new(ArrowSource::new_file_source(table_schema)), - )?; - Ok(DataSourceExec::from_data_source(scan_conf)) + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&decoder); + ArrowSource::try_from_proto(&node, &decode_ctx) } #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] @@ -2459,31 +2447,21 @@ pub trait PhysicalPlanNodeExt: Sized { }) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `DataSourceExec` serializes itself via `ExecutionPlan::try_to_proto`" + )] fn try_from_data_source_exec( data_source_exec: &DataSourceExec, codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let data_source = data_source_exec.data_source(); - - if let Some(scan_conf) = data_source.downcast_ref::() { - let source = scan_conf.file_source(); - if let Some(_arrow_source) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ArrowScan( - protobuf::ArrowScanExecNode { - base_conf: Some(serialize_file_scan_config( - scan_conf, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - Ok(None) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + data_source_exec.try_to_proto(&encode_ctx) } #[deprecated(