This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23683-09dd8d246539ff4433d4bd62cec96bc289a5939a in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e2e8e1c88cbee9c373dab52b9ef6541d2ee6364f Author: Kumar Ujjawal <[email protected]> AuthorDate: Fri Aug 7 19:15:03 2026 +0530 Add DataSource/FileSource proto hooks and FileScanConfig serde (#23683) ## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Closes #23497. ## Rationale for this change File scan serialization currently depends on central type downcasts in `datafusion-proto`. This PR adds the foundation needed to move each data source to its own protobuf hooks while preserving the existing wire format and fallback behavior. It unblocks the concrete source migrations tracked by #23516, #23517, and #23518. <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> ## What changes are included in this PR? - Add feature-gated `try_to_proto` hooks to `DataSource` and `FileSource`. - Delegate serialization from `DataSourceExec` through `FileScanConfig` to its concrete `FileSource`. - Add shared `FileScanConfig` protobuf encoding and decoding in the datasource crate. - Keep the existing central serializer as the fallback for sources that still return `Ok(None)`. - Preserve the legacy protobuf wire format. - Add compatibility test <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? Yes <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? This adds non-breaking, feature-gated public serialization hooks for data source and file source implementations. <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. --> --------- Co-authored-by: Adrian Garcia Badaracco <[email protected]> --- datafusion/datasource/Cargo.toml | 6 +- datafusion/datasource/src/file.rs | 25 ++ datafusion/datasource/src/file_scan_config/mod.rs | 127 +++++++ .../datasource/src/file_scan_config/proto.rs | 285 +++++++++++++++ datafusion/datasource/src/source.rs | 36 ++ datafusion/proto/src/physical_plan/from_proto.rs | 112 +----- datafusion/proto/src/physical_plan/mod.rs | 382 +++++++++++++++++++++ datafusion/proto/src/physical_plan/to_proto.rs | 89 +---- 8 files changed, 879 insertions(+), 183 deletions(-) diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index b78ac616de..f09447b694 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,8 +34,10 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] -# Enables protobuf conversions for datasource types and serialization hooks. -# Off by default so consumers that never serialize plans pay nothing. +# Enables protobuf conversions for datasource types, source serialization hooks, +# and the shared `FileScanConfig` <-> proto conversion. Off by default so +# consumers that never serialize plans pay nothing. Mirrors the `proto` feature +# on `datafusion-physical-plan`. proto = [ "dep:datafusion-proto-models", "datafusion-physical-plan/proto", diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 07460b2369..691bb314b7 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -351,6 +351,31 @@ pub trait FileSource: Any + Send + Sync { fn schema_adapter_factory(&self) -> Option<Arc<dyn SchemaAdapterFactory>> { None } + + /// Serialize this file source into a full [`PhysicalPlanNode`] (a + /// `DataSourceExec` wrapping the `FileScanConfig`), if it knows how. + /// + /// `base` is the shared [`FileScanConfig`] this source is wrapped in; the + /// format-agnostic parts (file groups, schema, statistics, ordering, + /// projection, …) are encoded via + /// [`FileScanConfig::try_to_proto`](crate::file_scan_config::FileScanConfig::try_to_proto), + /// and the concrete source appends its format-specific fields (e.g. CSV + /// delimiter/quote) around it. + /// + /// * `Ok(None)` (the default) — this source has no proto hook yet; the + /// caller falls back to the central downcast chain in `datafusion-proto`. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`FileScanConfig`]: crate::file_scan_config::FileScanConfig + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _base: &FileScanConfig, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + Ok(None) + } } impl dyn FileSource { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index d1dd3c11fc..766da2f8f7 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -20,6 +20,13 @@ pub(crate) mod sort_pushdown; +/// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature. +/// Attaches inherent `try_to_proto` / `try_from_proto` / +/// `parse_table_schema_from_proto` helpers to [`FileScanConfig`] used by every +/// file source's `try_to_proto` hook. +#[cfg(feature = "proto")] +mod proto; + use crate::file_groups::FileGroup; use crate::{ PartitionedFile, display::FileGroupsDisplay, file::FileSource, @@ -1176,6 +1183,18 @@ impl DataSource for FileScanConfig { Some(Arc::new(SharedWorkSource::from_config(self)) as Arc<dyn Any + Send + Sync>) } + + /// Serialize this file scan by delegating to the concrete + /// [`FileSource`]'s + /// [`try_to_proto`](crate::file::FileSource::try_to_proto) hook, passing + /// `self` as the shared spine it needs to emit the base config. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + self.file_source().try_to_proto(self, ctx) + } } impl FileScanConfig { @@ -1569,12 +1588,18 @@ mod tests { use datafusion_common::{Result, assert_batches_eq, internal_err}; use datafusion_execution::TaskContext; use datafusion_expr::SortExpr; + #[cfg(feature = "proto")] + use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::create_physical_sort_expr; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::execution_plan::collect; + #[cfg(feature = "proto")] + use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx}; + #[cfg(feature = "proto")] + use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use futures::FutureExt as _; use futures::StreamExt as _; use futures::stream; @@ -1633,6 +1658,108 @@ mod tests { } } + #[cfg(feature = "proto")] + #[derive(Clone)] + struct ProtoHookSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + } + + #[cfg(feature = "proto")] + impl ProtoHookSource { + fn new(table_schema: TableSchema) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + } + } + } + + #[cfg(feature = "proto")] + impl FileSource for ProtoHookSource { + fn create_file_opener( + &self, + _object_store: Arc<dyn ObjectStore>, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result<Arc<dyn crate::file_stream::FileOpener>> { + internal_err!("not needed for proto delegation test") + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> { + Arc::new(self.clone()) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "proto-hook-test" + } + + fn try_to_proto( + &self, + _base: &FileScanConfig, + _ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<PhysicalPlanNode>> { + Ok(Some(PhysicalPlanNode::default())) + } + } + + #[cfg(feature = "proto")] + struct UnusedPlanEncoder; + + #[cfg(feature = "proto")] + impl ExecutionPlanEncode for UnusedPlanEncoder { + fn encode_plan( + &self, + _plan: &Arc<dyn ExecutionPlan>, + ) -> Result<PhysicalPlanNode> { + internal_err!("not needed for proto delegation test") + } + + fn encode_expr(&self, _expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udf(&self, _udf: &ScalarUDF) -> Result<Option<Vec<u8>>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result<Option<Vec<u8>>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udwf(&self, _udwf: &WindowUDF) -> Result<Option<Vec<u8>>> { + internal_err!("not needed for proto delegation test") + } + } + + #[cfg(feature = "proto")] + #[test] + fn data_source_exec_delegates_proto_to_file_source() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema))); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .build(); + let exec = DataSourceExec::from_data_source(config); + let encoder = UnusedPlanEncoder; + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert_eq!(exec.try_to_proto(&ctx)?, Some(PhysicalPlanNode::default())); + Ok(()) + } + #[test] fn physical_plan_config_no_projection_tab_cols_as_field() { let file_schema = aggr_test_schema(); diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs new file mode 100644 index 0000000000..d7135173c8 --- /dev/null +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -0,0 +1,285 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared serialization of the format-agnostic [`FileScanConfig`] spine. +//! +//! This is the relocated body of `datafusion-proto`'s +//! `serialize_file_scan_config` / `parse_protobuf_file_scan_config`, ported to +//! ride the +//! [`ExecutionPlanEncodeCtx`](datafusion_physical_plan::proto::ExecutionPlanEncodeCtx) / +//! [`ExecutionPlanDecodeCtx`](datafusion_physical_plan::proto::ExecutionPlanDecodeCtx) +//! instead of the raw `PhysicalExtensionCodec` + +//! `PhysicalProtoConverterExtension`. Every +//! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its +//! `*ScanExecNode` around [`FileScanConfig::try_to_proto`] and decodes with +//! [`FileScanConfig::try_from_proto`], keeping a single copy of the shared +//! wire logic. The wire format is byte-for-byte identical to the old central +//! serializer. +//! +//! Child physical expressions (sort orderings, hash/range partitioning, and +//! projection expressions) are (de)serialized through `ctx.encode_expr` / +//! `ctx.decode_expr`; `Schema`, `Statistics`, `Constraints`, and `ScalarValue` +//! go through `datafusion-proto-common`. Nothing here needs the raw codec. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; +use datafusion_physical_expr::{LexOrdering, Partitioning}; +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; +use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; +use datafusion_proto_models::protobuf; + +use crate::file::FileSource; +use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use crate::table_schema::TableSchema; + +impl FileScanConfig { + /// Serialize the shared, format-agnostic part of a file scan into a + /// [`protobuf::FileScanExecConf`]. + /// + /// Each concrete [`FileSource::try_to_proto`] + /// wraps the returned value in its own `*ScanExecNode`. Byte-compatible with + /// the former `serialize_file_scan_config` in `datafusion-proto`. + pub fn try_to_proto( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result<protobuf::FileScanExecConf> { + let file_groups = self + .file_groups + .iter() + .map(TryInto::try_into) + .collect::<Result<Vec<_>>>()?; + + let mut output_ordering = vec![]; + for order in &self.output_ordering { + let nodes = sort_exprs_try_to_proto(order.iter(), &ctx.expr_ctx())?; + output_ordering.push(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes: nodes, + }); + } + + let output_partitioning = self + .output_partitioning + .as_ref() + .map(|partitioning| partitioning.try_to_proto(&ctx.expr_ctx())) + .transpose()?; + + // Fields must be added to the schema so that they can persist in the + // protobuf, and then removed from the schema in `try_from_proto`. + let mut fields = self + .file_schema() + .fields() + .iter() + .cloned() + .collect::<Vec<_>>(); + fields.extend(self.table_partition_cols().iter().cloned()); + let schema = + Schema::new(fields).with_metadata(self.file_schema().metadata.clone()); + + let projection_exprs = self + .file_source() + .projection() + .as_ref() + .map(|projection_exprs| { + Ok::<_, DataFusionError>(protobuf::ProjectionExprs { + projections: projection_exprs + .iter() + .map(|expr| { + Ok(protobuf::ProjectionExpr { + alias: expr.alias.to_string(), + expr: Some(ctx.encode_expr(&expr.expr)?), + }) + }) + .collect::<Result<Vec<_>>>()?, + }) + }) + .transpose()?; + + Ok(protobuf::FileScanExecConf { + file_groups, + statistics: Some((&self.statistics()).into()), + limit: self.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), + projection: vec![], + schema: Some((&schema).try_into()?), + table_partition_cols: self + .table_partition_cols() + .iter() + .map(|x| x.name().clone()) + .collect::<Vec<_>>(), + object_store_url: self.object_store_url.to_string(), + output_ordering, + constraints: Some(self.constraints.clone().into()), + batch_size: self.batch_size.map(|s| s as u64), + projection_exprs, + output_partitioning, + }) + } + + /// Reconstruct a [`FileScanConfig`] from a [`protobuf::FileScanExecConf`] + /// and a `file_source` the caller has already rebuilt (typically from the + /// table schema via [`FileScanConfig::parse_table_schema_from_proto`]). + /// + /// Byte-compatible with the former `parse_protobuf_file_scan_config`. + pub fn try_from_proto( + conf: &protobuf::FileScanExecConf, + ctx: &ExecutionPlanDecodeCtx<'_>, + file_source: Arc<dyn FileSource>, + ) -> Result<FileScanConfig> { + let schema = parse_file_scan_schema(conf)?; + + let constraints = conf + .constraints + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'constraints'" + ) + })? + .try_into()?; + let statistics = conf + .statistics + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'statistics'" + ) + })? + .try_into()?; + + let file_groups = conf + .file_groups + .iter() + .map(TryInto::try_into) + .collect::<Result<Vec<_>>>()?; + + let object_store_url = match conf.object_store_url.is_empty() { + false => ObjectStoreUrl::parse(&conf.object_store_url)?, + true => ObjectStoreUrl::local_filesystem(), + }; + + let mut output_ordering = vec![]; + for node_collection in &conf.output_ordering { + let sort_exprs = sort_exprs_try_from_proto( + &node_collection.physical_sort_expr_nodes, + &ctx.expr_ctx(&schema), + )?; + output_ordering.extend(LexOrdering::new(sort_exprs)); + } + + let output_partitioning = conf + .output_partitioning + .as_ref() + .map(|partitioning| { + Partitioning::try_from_proto(partitioning, &ctx.expr_ctx(&schema)) + }) + .transpose()? + .flatten(); + + // Parse projection expressions if present and apply to the file source. + let file_source = if let Some(proto_projection_exprs) = &conf.projection_exprs { + let projection_exprs: Vec<ProjectionExpr> = proto_projection_exprs + .projections + .iter() + .map(|proto_expr| { + let expr = ctx.decode_expr( + proto_expr.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("ProjectionExpr missing expr field") + })?, + &schema, + )?; + Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) + }) + .collect::<Result<Vec<_>>>()?; + + let projection_exprs = ProjectionExprs::new(projection_exprs); + + file_source + .try_pushdown_projection(&projection_exprs)? + .unwrap_or(file_source) + } else { + file_source + }; + + let config_builder = FileScanConfigBuilder::new(object_store_url, file_source) + .with_file_groups(file_groups) + .with_constraints(constraints) + .with_statistics(statistics) + .with_limit(conf.limit.as_ref().map(|sl| sl.limit as usize)) + .with_output_ordering(output_ordering) + .with_output_partitioning(output_partitioning) + .with_batch_size(conf.batch_size.map(|s| s as usize)); + Ok(config_builder.build()) + } + + /// Parse a [`TableSchema`] (file schema + partition columns) from a + /// [`protobuf::FileScanExecConf`]. File sources use this to rebuild their + /// concrete source before calling [`FileScanConfig::try_from_proto`]. + /// + /// Byte-compatible with the former `parse_table_schema_from_proto`. + pub fn parse_table_schema_from_proto( + conf: &protobuf::FileScanExecConf, + ) -> Result<TableSchema> { + let schema = parse_file_scan_schema(conf)?; + + // Reacquire the partition column types from the schema before removing + // them below. + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|col| Ok(Arc::new(schema.field_with_name(col)?.clone()))) + .collect::<Result<Vec<_>>>()?; + + // Remove partition columns from the schema after recreating + // table_partition_cols because the partition columns are not in the + // file. They are present to allow the partition column types to be + // reconstructed after serde. + let file_schema = Arc::new( + Schema::new( + schema + .fields() + .iter() + .filter(|field| !table_partition_cols.contains(field)) + .cloned() + .collect::<Vec<_>>(), + ) + .with_metadata(schema.metadata.clone()), + ); + + Ok(TableSchema::builder(file_schema) + .with_table_partition_cols(table_partition_cols) + .build()) + } +} + +/// Parse the full (file + partition columns) schema off the base conf. +fn parse_file_scan_schema(conf: &protobuf::FileScanExecConf) -> Result<Arc<Schema>> { + let schema: Schema = conf + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'schema'" + ) + })? + .try_into()?; + Ok(Arc::new(schema)) +} diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index c280470bb0..1fd5f865c4 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -264,6 +264,30 @@ pub trait DataSource: Any + Send + Sync + Debug { fn open_with_args(&self, args: OpenArgs) -> Result<SendableRecordBatchStream> { self.open(args.partition, args.context) } + + /// Serialize this data source to a full [`PhysicalPlanNode`] (a + /// `DataSourceExec` wrapping this source), if it knows how. + /// + /// This is the `DataSource` analog of + /// [`ExecutionPlan::try_to_proto`]. + /// [`DataSourceExec::try_to_proto`](crate::source::DataSourceExec) delegates + /// to this hook, which for file scans forwards to + /// [`FileSource::try_to_proto`] + /// through the shared [`FileScanConfig`] + /// spine. + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller falls + /// back to the central downcast chain in `datafusion-proto`. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + Ok(None) + } } /// Arguments for [`DataSource::open_with_args`] @@ -553,6 +577,18 @@ impl ExecutionPlan for DataSourceExec { new_exec.execution_state = Arc::new(OnceLock::new()); Ok(Arc::new(new_exec)) } + + /// Delegates serialization to the wrapped [`DataSource`]. For file scans the + /// concrete [`FileSource`] emits the node via its + /// own `try_to_proto` hook, keeping the format-specific wire logic in the + /// format crate. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + self.data_source().try_to_proto(ctx) + } } impl DataSourceExec { diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index ade6ea183b..645854295b 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -26,34 +26,33 @@ use arrow::ipc::reader::StreamReader; use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err}; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::{FileRange, PartitionedFile, TableSchema}; use datafusion_datasource_csv::file_format::CsvSink; use datafusion_datasource_json::file_format::JsonSink; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::ParquetSink; -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; -use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::{ - HigherOrderFunctionExpr, LexOrdering, PhysicalSortExpr, ScalarFunctionExpr, + HigherOrderFunctionExpr, PhysicalSortExpr, ScalarFunctionExpr, }; use datafusion_physical_plan::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, }; use datafusion_physical_plan::joins::HashExpr; +use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use datafusion_proto_common::common::proto_error; use super::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - PhysicalProtoConverterExtension, + ConverterPlanDecoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; use crate::convert::TryFromProto; use crate::protobuf::physical_expr_node::ExprType; @@ -436,33 +435,7 @@ pub fn parse_protobuf_file_scan_schema( pub fn parse_table_schema_from_proto( proto: &protobuf::FileScanExecConf, ) -> Result<TableSchema> { - let schema: Arc<Schema> = parse_protobuf_file_scan_schema(proto)?; - - // Reacquire the partition column types from the schema before removing them below. - let table_partition_cols = proto - .table_partition_cols - .iter() - .map(|col| Ok(Arc::new(schema.field_with_name(col)?.clone()))) - .collect::<Result<Vec<_>>>()?; - - // Remove partition columns from the schema after recreating table_partition_cols - // because the partition columns are not in the file. They are present to allow - // the partition column types to be reconstructed after serde. - let file_schema = Arc::new( - Schema::new( - schema - .fields() - .iter() - .filter(|field| !table_partition_cols.contains(field)) - .cloned() - .collect::<Vec<_>>(), - ) - .with_metadata(schema.metadata.clone()), - ); - - Ok(TableSchema::builder(file_schema) - .with_table_partition_cols(table_partition_cols) - .build()) + FileScanConfig::parse_table_schema_from_proto(proto) } pub fn parse_protobuf_file_scan_config( @@ -471,76 +444,15 @@ pub fn parse_protobuf_file_scan_config( proto_converter: &dyn PhysicalProtoConverterExtension, file_source: Arc<dyn FileSource>, ) -> Result<FileScanConfig> { - let schema: Arc<Schema> = parse_protobuf_file_scan_schema(proto)?; - - let constraints = convert_required!(proto.constraints)?; - let statistics = convert_required!(proto.statistics)?; - - let file_groups = proto - .file_groups - .iter() - .map(FileGroup::try_from_proto) - .collect::<Result<Vec<_>, _>>()?; - - let object_store_url = match proto.object_store_url.is_empty() { - false => ObjectStoreUrl::parse(&proto.object_store_url)?, - true => ObjectStoreUrl::local_filesystem(), - }; - - let mut output_ordering = vec![]; - for node_collection in &proto.output_ordering { - let sort_exprs = parse_physical_sort_exprs( - &node_collection.physical_sort_expr_nodes, - ctx, - &schema, - proto_converter, - )?; - output_ordering.extend(LexOrdering::new(sort_exprs)); - } - let output_partitioning = parse_protobuf_partitioning( - proto.output_partitioning.as_ref(), + let decoder = ConverterPlanDecoder { ctx, - &schema, proto_converter, - )?; - - // Parse projection expressions if present and apply to file source - let file_source = if let Some(proto_projection_exprs) = &proto.projection_exprs { - let projection_exprs: Vec<ProjectionExpr> = proto_projection_exprs - .projections - .iter() - .map(|proto_expr| { - let expr = proto_converter.proto_to_physical_expr( - proto_expr.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!("ProjectionExpr missing expr field") - })?, - &schema, - ctx, - )?; - Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) - }) - .collect::<Result<Vec<_>>>()?; - - let projection_exprs = ProjectionExprs::new(projection_exprs); - - // Apply projection to file source - file_source - .try_pushdown_projection(&projection_exprs)? - .unwrap_or(file_source) - } else { - file_source }; - - let config = FileScanConfigBuilder::new(object_store_url, file_source) - .with_file_groups(file_groups) - .with_constraints(constraints) - .with_statistics(statistics) - .with_limit(proto.limit.as_ref().map(|sl| sl.limit as usize)) - .with_output_ordering(output_ordering) - .with_output_partitioning(output_partitioning) - .with_batch_size(proto.batch_size.map(|s| s as usize)) - .build(); - Ok(config) + FileScanConfig::try_from_proto( + proto, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) } pub fn parse_record_batches(buf: &[u8]) -> Result<Vec<RecordBatch>> { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index bb17f9dbca..7e162bf954 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -119,6 +119,388 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { ) } +#[cfg(test)] +mod file_scan_config_serde { + use super::*; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; + use datafusion_datasource::file::FileSource; + use datafusion_datasource::file_groups::FileGroup; + use datafusion_datasource::file_stream::FileOpener; + use datafusion_datasource::{PartitionedFile, TableSchema}; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::projection::{ + ProjectionExpr as FileProjectionExpr, ProjectionExprs as FileProjectionExprs, + }; + use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use object_store::ObjectStore; + + #[derive(Clone)] + struct SerdeTestSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + projection: Option<FileProjectionExprs>, + } + + impl SerdeTestSource { + fn new( + table_schema: TableSchema, + projection: Option<FileProjectionExprs>, + ) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + projection, + } + } + } + + impl FileSource for SerdeTestSource { + fn create_file_opener( + &self, + _object_store: Arc<dyn ObjectStore>, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result<Arc<dyn FileOpener>> { + internal_err!("not needed for FileScanConfig serde tests") + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc<dyn FileSource> { + Arc::new(self.clone()) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "serde-test" + } + + fn try_pushdown_projection( + &self, + projection: &FileProjectionExprs, + ) -> Result<Option<Arc<dyn FileSource>>> { + Ok(Some(Arc::new(Self { + projection: Some(projection.clone()), + ..self.clone() + }))) + } + + fn projection(&self) -> Option<&FileProjectionExprs> { + self.projection.as_ref() + } + } + + fn populated_projection() -> FileProjectionExprs { + FileProjectionExprs::new(vec![FileProjectionExpr::new( + Arc::new(Column::new("value", 0)), + "projected_value", + )]) + } + + fn test_config(output_partitioning: Option<Partitioning>) -> FileScanConfig { + test_config_with_projection(output_partitioning, Some(populated_projection())) + } + + fn test_config_with_projection( + output_partitioning: Option<Partitioning>, + projection: Option<FileProjectionExprs>, + ) -> FileScanConfig { + let file_schema = Arc::new( + Schema::new(vec![ + Field::new("value", DataType::Int32, false), + Field::new("label", DataType::Utf8, true), + ]) + .with_metadata(HashMap::from([( + "serde_test_key".to_string(), + "serde_test_value".to_string(), + )])), + ); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Utf8, + false, + ))]) + .build(); + let table_statistics = Statistics::new_unknown(table_schema.table_schema()); + let source = Arc::new(SerdeTestSource::new(table_schema, projection)); + let first_file = PartitionedFile::new("data/part=a/file.arrow", 1024) + .with_partition_values(vec![ScalarValue::Utf8(Some("a".to_string()))]) + .with_range(10, 900) + .with_arrow_schema(Arc::clone(&file_schema)) + .with_statistics(Arc::new(table_statistics.clone())); + let second_file = PartitionedFile::new("data/part=b/file.arrow", 2048) + .with_partition_values(vec![ScalarValue::Utf8(Some("b".to_string()))]); + let third_file = PartitionedFile::new("data/part=c/file.arrow", 4096) + .with_partition_values(vec![ScalarValue::Utf8(Some("c".to_string()))]) + .with_arrow_schema(Arc::clone(&file_schema)); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("value", 0), + ))]) + .expect("single expression ordering"); + + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file_groups(vec![ + FileGroup::new(vec![first_file, second_file]), + FileGroup::new(vec![third_file]), + ]) + .with_constraints(Constraints::new_unverified(vec![Constraint::PrimaryKey( + vec![0], + )])) + .with_statistics(table_statistics) + .with_limit(Some(17)) + .with_batch_size(Some(256)) + .with_output_ordering(vec![ordering]) + .with_output_partitioning(output_partitioning) + .build() + } + + fn hash_partitioning() -> Partitioning { + Partitioning::Hash(vec![Arc::new(Column::new("value", 0))], 3) + } + + fn range_partitioning() -> Partitioning { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("value", 0), + ))]) + .expect("single expression ordering"); + Partitioning::Range(RangePartitioning::new( + ordering, + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )) + } + + fn decode_source(conf: &protobuf::FileScanExecConf) -> Result<Arc<dyn FileSource>> { + Ok(Arc::new(SerdeTestSource::new( + FileScanConfig::parse_table_schema_from_proto(conf)?, + None, + ))) + } + + struct FileScanSerdeHarness { + codec: DefaultPhysicalExtensionCodec, + converter: DefaultPhysicalProtoConverter, + task_ctx: TaskContext, + } + + impl FileScanSerdeHarness { + fn new() -> Self { + Self { + codec: DefaultPhysicalExtensionCodec {}, + converter: DefaultPhysicalProtoConverter {}, + task_ctx: TaskContext::default(), + } + } + + fn encode(&self, config: &FileScanConfig) -> Result<protobuf::FileScanExecConf> { + let encoder = ConverterPlanEncoder { + codec: &self.codec, + proto_converter: &self.converter, + }; + config.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) + } + + fn decode(&self, conf: &protobuf::FileScanExecConf) -> Result<FileScanConfig> { + self.decode_with_source(conf, decode_source(conf)?) + } + + fn decode_with_source( + &self, + conf: &protobuf::FileScanExecConf, + file_source: Arc<dyn FileSource>, + ) -> Result<FileScanConfig> { + let physical_decode_ctx = + PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); + let decoder = ConverterPlanDecoder { + ctx: &physical_decode_ctx, + proto_converter: &self.converter, + }; + FileScanConfig::try_from_proto( + conf, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) + } + } + + #[test] + fn new_file_scan_config_serde_roundtrips_all_partitioning_variants() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + for config in [ + test_config(None), + test_config(Some(Partitioning::RoundRobinBatch(2))), + test_config(Some(hash_partitioning())), + test_config(Some(range_partitioning())), + test_config(Some(Partitioning::UnknownPartitioning(4))), + ] { + let encoded = serde.encode(&config)?; + let reencoded = serde.encode(&serde.decode(&encoded)?)?; + assert_eq!(reencoded.output_partitioning, encoded.output_partitioning); + } + + Ok(()) + } + + #[test] + fn new_file_scan_config_serde_preserves_complete_fixture() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let config = test_config(None); + let decoded = serde.decode(&serde.encode(&config)?)?; + + assert_eq!(decoded.constraints, config.constraints); + assert_eq!( + decoded.file_schema().metadata, + config.file_schema().metadata + ); + assert_eq!(decoded.file_groups.len(), 2); + assert_eq!(decoded.file_groups[0].len(), 2); + assert_eq!(decoded.file_groups[1].len(), 1); + assert!(decoded.file_groups[0].files()[0].arrow_schema.is_some()); + assert!(decoded.file_groups[0].files()[1].arrow_schema.is_none()); + + Ok(()) + } + + #[test] + fn new_file_scan_config_serde_preserves_projection_presence() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + let absent = serde.encode(&test_config_with_projection(None, None))?; + assert!(absent.projection_exprs.is_none()); + assert!(serde.decode(&absent)?.file_source().projection().is_none()); + + let empty = serde.encode(&test_config_with_projection( + None, + Some(FileProjectionExprs::new(vec![])), + ))?; + assert!( + empty + .projection_exprs + .as_ref() + .is_some_and(|projection| projection.projections.is_empty()) + ); + assert!( + serde + .decode(&empty)? + .file_source() + .projection() + .is_some_and(|projection| projection.as_ref().is_empty()) + ); + + Ok(()) + } + + #[test] + fn new_file_scan_config_decode_rejects_malformed_required_fields() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let valid = serde.encode(&test_config(None))?; + let file_source = decode_source(&valid)?; + + for (field, malformed) in [ + ( + "schema", + protobuf::FileScanExecConf { + schema: None, + ..valid.clone() + }, + ), + ( + "constraints", + protobuf::FileScanExecConf { + constraints: None, + ..valid.clone() + }, + ), + ( + "statistics", + protobuf::FileScanExecConf { + statistics: None, + ..valid.clone() + }, + ), + ] { + let err = serde + .decode_with_source(&malformed, Arc::clone(&file_source)) + .expect_err("missing required field must fail"); + assert!(err.to_string().contains(field), "unexpected error: {err}"); + } + + let mut missing_projection_expr = valid.clone(); + missing_projection_expr + .projection_exprs + .as_mut() + .expect("test config has projection expressions") + .projections[0] + .expr = None; + let err = serde + .decode_with_source(&missing_projection_expr, file_source) + .expect_err("missing projection expression must fail"); + assert!( + err.to_string() + .contains("ProjectionExpr missing expr field"), + "unexpected error: {err}" + ); + + Ok(()) + } + + #[test] + fn new_file_scan_config_decode_rejects_invalid_range_ordering() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let mut proto = serde.encode(&test_config(Some(range_partitioning())))?; + + let mut duplicate_ordering = proto.clone(); + let range = match duplicate_ordering + .output_partitioning + .as_mut() + .and_then(|p| p.partition_method.as_mut()) + { + Some(protobuf::partitioning::PartitionMethod::Range(range)) => range, + other => panic!("expected range partitioning, got {other:?}"), + }; + range.sort_expr.push(range.sort_expr[0].clone()); + + let err = serde + .decode(&duplicate_ordering) + .expect_err("duplicate range ordering must fail"); + assert!( + err.to_string().contains("duplicate expressions"), + "unexpected error: {err}" + ); + + let range = match proto + .output_partitioning + .as_mut() + .and_then(|p| p.partition_method.as_mut()) + { + Some(protobuf::partitioning::PartitionMethod::Range(range)) => range, + other => panic!("expected range partitioning, got {other:?}"), + }; + range.sort_expr.clear(); + + let err = serde + .decode(&proto) + .expect_err("empty range ordering must fail"); + assert!( + err.to_string().contains("requires non-empty ordering"), + "unexpected error: {err}" + ); + + Ok(()) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 2d4aa72ff0..c8a7ea383a 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -34,18 +34,18 @@ use datafusion_expr::WindowFrame; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; use datafusion_physical_expr::{HigherOrderFunctionExpr, ScalarFunctionExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use datafusion_physical_plan::proto::ExecutionPlanEncodeCtx; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use super::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + ConverterPlanEncoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalProtoConverterExtension, encode_human_display_alias, }; use crate::convert::TryFromProto; use crate::protobuf::{ - self, PhysicalSortExprNode, PhysicalSortExprNodeCollection, - physical_aggregate_expr_node, physical_window_expr_node, + self, PhysicalSortExprNode, physical_aggregate_expr_node, physical_window_expr_node, }; #[expect(clippy::needless_pass_by_value)] @@ -407,84 +407,11 @@ pub fn serialize_file_scan_config( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<protobuf::FileScanExecConf> { - let file_groups = conf - .file_groups - .iter() - .map(TryInto::try_into) - .collect::<Result<Vec<_>, _>>()?; - - let mut output_orderings = vec![]; - for order in &conf.output_ordering { - let ordering = - serialize_physical_sort_exprs(order.to_vec(), codec, proto_converter)?; - output_orderings.push(ordering) - } - let output_partitioning = conf - .output_partitioning - .as_ref() - .map(|partitioning| serialize_partitioning(partitioning, codec, proto_converter)) - .transpose()?; - - // Fields must be added to the schema so that they can persist in the protobuf, - // and then they are to be removed from the schema in `parse_protobuf_file_scan_config` - let mut fields = conf - .file_schema() - .fields() - .iter() - .cloned() - .collect::<Vec<_>>(); - fields.extend(conf.table_partition_cols().iter().cloned()); - - let schema = Arc::new( - Schema::new(fields.clone()).with_metadata(conf.file_schema().metadata.clone()), - ); - - let projection_exprs = conf - .file_source - .projection() - .as_ref() - .map(|projection_exprs| { - let projections = projection_exprs.iter().cloned().collect::<Vec<_>>(); - Ok::<_, DataFusionError>(protobuf::ProjectionExprs { - projections: projections - .into_iter() - .map(|expr| { - Ok(protobuf::ProjectionExpr { - alias: expr.alias.to_string(), - expr: Some( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - ), - }) - }) - .collect::<Result<Vec<_>>>()?, - }) - }) - .transpose()?; - - Ok(protobuf::FileScanExecConf { - file_groups, - statistics: Some((&conf.statistics()).into()), - limit: conf.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), - projection: vec![], - schema: Some(schema.as_ref().try_into()?), - table_partition_cols: conf - .table_partition_cols() - .iter() - .map(|x| x.name().clone()) - .collect::<Vec<_>>(), - object_store_url: conf.object_store_url.to_string(), - output_ordering: output_orderings - .into_iter() - .map(|e| PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: e, - }) - .collect::<Vec<_>>(), - constraints: Some(conf.constraints.clone().into()), - batch_size: conf.batch_size.map(|s| s as u64), - projection_exprs, - output_partitioning, - }) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + conf.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) } pub fn serialize_maybe_filter( --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
