From 412cb4084e52335c8d6eac2d32b82a6b4f25b148 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 18 Jul 2026 12:23:28 +0530 Subject: [PATCH 1/8] Add DataSource/FileSource proto hooks and FileScanConfig serde --- Cargo.lock | 1 + datafusion/datasource/Cargo.toml | 9 + datafusion/datasource/src/file.rs | 25 + .../datasource/src/file_scan_config/mod.rs | 126 +++++ .../datasource/src/file_scan_config/proto.rs | 514 ++++++++++++++++++ datafusion/datasource/src/source.rs | 36 ++ datafusion/proto/Cargo.toml | 2 +- datafusion/proto/src/physical_plan/mod.rs | 424 +++++++++++++++ 8 files changed, 1136 insertions(+), 1 deletion(-) create mode 100644 datafusion/datasource/src/file_scan_config/proto.rs diff --git a/Cargo.lock b/Cargo.lock index a41734e064d4c..fa1b8adcbd9f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1942,6 +1942,7 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 2ac42ed900095..77b10a171f3fc 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,6 +34,14 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] +# Enables `DataSource::try_to_proto` / `FileSource::try_to_proto` 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", +] [dependencies] arrow = { workspace = true } @@ -56,6 +64,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } flate2 = { workspace = true, optional = true } futures = { workspace = true } diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 07460b23694b7..44fdf8dc16d74 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> { 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::to_proto_conf`](crate::file_scan_config::FileScanConfig::to_proto_conf), + /// 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> { + 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 962df06302386..36f2273576659 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -20,6 +20,12 @@ pub(crate) mod sort_pushdown; +/// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature. +/// Attaches inherent `to_proto_conf` / `from_proto_conf` / `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, @@ -1175,6 +1181,18 @@ impl DataSource for FileScanConfig { Some(Arc::new(SharedWorkSource::from_config(self)) as Arc) } + + /// 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> { + self.file_source().try_to_proto(self, ctx) + } } impl FileScanConfig { @@ -1566,12 +1584,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; @@ -1630,6 +1654,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, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + 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 { + 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> { + Ok(Some(PhysicalPlanNode::default())) + } + } + + #[cfg(feature = "proto")] + struct UnusedPlanEncoder; + + #[cfg(feature = "proto")] + impl ExecutionPlanEncode for UnusedPlanEncoder { + fn encode_plan( + &self, + _plan: &Arc, + ) -> Result { + internal_err!("not needed for proto delegation test") + } + + fn encode_expr(&self, _expr: &Arc) -> Result { + internal_err!("not needed for proto delegation test") + } + + fn encode_udf(&self, _udf: &ScalarUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udwf(&self, _udwf: &WindowUDF) -> Result>> { + 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 0000000000000..5d7ce3efbfc8b --- /dev/null +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -0,0 +1,514 @@ +// 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`] / [`ExecutionPlanDecodeCtx`] instead of +//! the raw `PhysicalExtensionCodec` + `PhysicalProtoConverterExtension`. Every +//! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its +//! `*ScanExecNode` around [`FileScanConfig::to_proto_conf`] and decodes with +//! [`FileScanConfig::from_proto_conf`], 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::compute::SortOptions; +use arrow::datatypes::Schema; +use chrono::{TimeZone, Utc}; +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, PhysicalSortExpr, RangePartitioning, SplitPoint, +}; +use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; +use datafusion_proto_models::protobuf; +use object_store::ObjectMeta; +use object_store::path::Path; + +use crate::PartitionedFile; +use crate::file::FileSource; +use crate::file_groups::FileGroup; +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 to_proto_conf( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result { + let file_groups = self + .file_groups + .iter() + .map(file_group_to_proto) + .collect::>>()?; + + // Sort orderings: only the child expressions need the ctx; the + // asc/nulls_first wrapping is plain data inlined into a + // `PhysicalSortExprNode` (same shape as `sorts/sort.rs`). + let mut output_ordering = vec![]; + for order in &self.output_ordering { + let nodes = order + .iter() + .map(|sort_expr| { + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), + asc: !sort_expr.options.descending, + nulls_first: sort_expr.options.nulls_first, + }) + }) + .collect::>>()?; + output_ordering.push(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes: nodes, + }); + } + + let output_partitioning = self + .output_partitioning + .as_ref() + .map(|p| partitioning_to_proto(p, ctx)) + .transpose()?; + + // Fields must be added to the schema so that they can persist in the + // protobuf, and then removed from the schema in `from_proto_conf`. + let mut fields = self + .file_schema() + .fields() + .iter() + .cloned() + .collect::>(); + 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::>>()?, + }) + }) + .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::>(), + 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 from_proto_conf( + conf: &protobuf::FileScanExecConf, + ctx: &ExecutionPlanDecodeCtx<'_>, + file_source: Arc, + ) -> Result { + 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(file_group_from_proto) + .collect::>>()?; + + 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 = parse_sort_exprs( + &node_collection.physical_sort_expr_nodes, + ctx, + &schema, + )?; + output_ordering.extend(LexOrdering::new(sort_exprs)); + } + + let output_partitioning = + partitioning_from_proto(conf.output_partitioning.as_ref(), ctx, &schema)?; + + // 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 = 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::>>()?; + + 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::from_proto_conf`]. + /// + /// Byte-compatible with the former `parse_table_schema_from_proto`. + pub fn parse_table_schema_from_proto( + conf: &protobuf::FileScanExecConf, + ) -> Result { + 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::>>()?; + + // 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::>(), + ) + .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> { + 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)) +} + +fn parse_sort_exprs( + nodes: &[protobuf::PhysicalSortExprNode], + ctx: &ExecutionPlanDecodeCtx<'_>, + schema: &Schema, +) -> Result> { + nodes + .iter() + .map(|sort_expr| { + let expr = sort_expr.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("Unexpected empty physical expression") + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr, schema)?, + options: SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect() +} + +/// Inlined equivalent of `datafusion-proto`'s `serialize_partitioning`. Only +/// child physical expressions and `ScalarValue`s need the ctx; the +/// `protobuf::Partitioning` wrapping is built directly here. +fn partitioning_to_proto( + partitioning: &Partitioning, + ctx: &ExecutionPlanEncodeCtx<'_>, +) -> Result { + let partition_method = match partitioning { + Partitioning::RoundRobinBatch(n) => { + protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) + } + Partitioning::Hash(exprs, n) => { + let hash_expr = ctx.encode_expressions(exprs)?; + protobuf::partitioning::PartitionMethod::Hash( + protobuf::PhysicalHashRepartition { + hash_expr, + partition_count: *n as u64, + }, + ) + } + Partitioning::Range(range) => { + let sort_expr = range + .ordering() + .iter() + .map(|sort_expr| { + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), + asc: !sort_expr.options.descending, + nulls_first: sort_expr.options.nulls_first, + }) + }) + .collect::>>()?; + let split_point = range + .split_points() + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>()?; + protobuf::partitioning::PartitionMethod::Range( + protobuf::PhysicalRangePartitioning { + sort_expr, + split_point, + }, + ) + } + Partitioning::UnknownPartitioning(n) => { + protobuf::partitioning::PartitionMethod::Unknown(*n as u64) + } + }; + Ok(protobuf::Partitioning { + partition_method: Some(partition_method), + }) +} + +/// Inlined equivalent of `datafusion-proto`'s `parse_protobuf_partitioning`. +fn partitioning_from_proto( + partitioning: Option<&protobuf::Partitioning>, + ctx: &ExecutionPlanDecodeCtx<'_>, + schema: &Schema, +) -> Result> { + let Some(partitioning) = partitioning else { + return Ok(None); + }; + let Some(partition_method) = partitioning.partition_method.as_ref() else { + return Ok(None); + }; + let partitioning = match partition_method { + protobuf::partitioning::PartitionMethod::RoundRobin(n) => { + Partitioning::RoundRobinBatch(*n as usize) + } + protobuf::partitioning::PartitionMethod::Hash(hash) => { + let exprs = hash + .hash_expr + .iter() + .map(|expr| ctx.decode_expr(expr, schema)) + .collect::>>()?; + Partitioning::Hash(exprs, hash.partition_count as usize) + } + protobuf::partitioning::PartitionMethod::Unknown(n) => { + Partitioning::UnknownPartitioning(*n as usize) + } + protobuf::partitioning::PartitionMethod::Range(range) => { + let sort_exprs = parse_sort_exprs(&range.sort_expr, ctx, schema)?; + let sort_expr_count = sort_exprs.len(); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!("Range partitioning requires non-empty ordering") + })?; + if ordering.len() != sort_expr_count { + return Err(internal_datafusion_err!( + "Range partitioning ordering must not contain duplicate expressions" + )); + } + let split_points = range + .split_point + .iter() + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| { + datafusion_common::ScalarValue::try_from(value) + .map_err(Into::into) + }) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>()?; + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) + } + }; + Ok(Some(partitioning)) +} + +fn file_group_to_proto(group: &FileGroup) -> Result { + Ok(protobuf::FileGroup { + files: group + .files() + .iter() + .map(partitioned_file_to_proto) + .collect::>>()?, + }) +} + +fn file_group_from_proto(group: &protobuf::FileGroup) -> Result { + let files = group + .files + .iter() + .map(partitioned_file_from_proto) + .collect::>>()?; + Ok(FileGroup::new(files)) +} + +fn partitioned_file_to_proto(pf: &PartitionedFile) -> Result { + let last_modified = pf.object_meta.last_modified; + let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { + DataFusionError::Plan(format!( + "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" + )) + })? as u64; + Ok(protobuf::PartitionedFile { + arrow_schema: pf + .arrow_schema + .as_ref() + .map(|s| s.as_ref().try_into()) + .transpose()?, + path: pf.object_meta.location.as_ref().to_owned(), + size: pf.object_meta.size, + last_modified_ns, + partition_values: pf + .partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + range: pf.range.as_ref().map(|range| protobuf::FileRange { + start: range.start, + end: range.end, + }), + statistics: pf.statistics.as_ref().map(|s| s.as_ref().into()), + }) +} + +fn partitioned_file_from_proto( + val: &protobuf::PartitionedFile, +) -> Result { + let mut pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse(val.path.as_str()) + .map_err(|e| internal_datafusion_err!("Invalid object_store path: {e}"))?, + last_modified: Utc.timestamp_nanos(val.last_modified_ns as i64), + size: val.size, + e_tag: None, + version: None, + }) + .with_partition_values( + val.partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + ); + if let Some(proto_schema) = val.arrow_schema.as_ref() { + pf = pf.with_arrow_schema(Arc::new( + proto_schema.try_into().map_err(DataFusionError::from)?, + )); + } + if let Some(range) = val.range.as_ref() { + pf = pf.with_range(range.start, range.end); + } + if let Some(proto_stats) = val.statistics.as_ref() { + pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); + } + Ok(pf) +} diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index c280470bb0d0b..1fd5f865c45ab 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 { 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> { + 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> { + self.data_source().try_to_proto(ctx) + } } impl DataSourceExec { diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index cfff8a949418a..037be27769f4d 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -54,7 +54,7 @@ chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true } -datafusion-datasource = { workspace = true } +datafusion-datasource = { workspace = true, features = ["proto"] } datafusion-datasource-arrow = { workspace = true } datafusion-datasource-avro = { workspace = true, optional = true } datafusion-datasource-csv = { workspace = true } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 7ee173cb36868..94c18527bd1a1 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -125,6 +125,430 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { mod tests { use super::*; + 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, RangePartitioning, SplitPoint, + }; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use object_store::ObjectStore; + + #[derive(Clone)] + struct SerdeTestSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + projection: Option, + } + + impl SerdeTestSource { + fn new( + table_schema: TableSchema, + projection: Option, + ) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + projection, + } + } + } + + impl FileSource for SerdeTestSource { + fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + 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 { + 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>> { + 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) -> FileScanConfig { + test_config_with_projection(output_partitioning, Some(populated_projection())) + } + + fn test_config_with_projection( + output_partitioning: Option, + projection: Option, + ) -> 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> { + 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 { + let encoder = ConverterPlanEncoder { + codec: &self.codec, + proto_converter: &self.converter, + }; + config.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) + } + + fn legacy_encode( + &self, + config: &FileScanConfig, + ) -> Result { + serialize_file_scan_config(config, &self.codec, &self.converter) + } + + fn decode( + &self, + conf: &protobuf::FileScanExecConf, + ) -> Result { + self.decode_with_source(conf, decode_source(conf)?) + } + + fn decode_with_source( + &self, + conf: &protobuf::FileScanExecConf, + file_source: Arc, + ) -> Result { + let physical_decode_ctx = + PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); + let decoder = ConverterPlanDecoder { + ctx: &physical_decode_ctx, + proto_converter: &self.converter, + }; + FileScanConfig::from_proto_conf( + conf, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) + } + + fn legacy_decode( + &self, + conf: &protobuf::FileScanExecConf, + ) -> Result { + let physical_decode_ctx = + PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); + parse_protobuf_file_scan_config( + conf, + &physical_decode_ctx, + &self.converter, + decode_source(conf)?, + ) + } + } + + #[test] + fn new_file_scan_config_serde_matches_legacy_wire_format() -> 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))), + test_config_with_projection(None, None), + test_config_with_projection(None, Some(FileProjectionExprs::new(vec![]))), + ] { + let legacy = serde.legacy_encode(&config)?; + let migrated = serde.encode(&config)?; + + assert_eq!(migrated, legacy); + assert_eq!(migrated.encode_to_vec(), legacy.encode_to_vec()); + + let migrated_reencoded = serde.encode(&serde.decode(&migrated)?)?; + let legacy_reencoded = + serde.legacy_encode(&serde.legacy_decode(&legacy)?)?; + assert_eq!(migrated_reencoded, legacy_reencoded); + assert_eq!( + migrated_reencoded.encode_to_vec(), + legacy_reencoded.encode_to_vec() + ); + } + + 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(()) + } + } + /// Unit tests for the bytes-only function serde exposed on /// [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] and backed by /// [`ConverterPlanEncoder`] / [`ConverterPlanDecoder`]. Function-carrying From cc98f1b8b04505531a6f66770327572d2bfbcb14 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 18 Jul 2026 12:52:52 +0530 Subject: [PATCH 2/8] Fix FileScanConfig proto rustdoc links --- datafusion/datasource/src/file_scan_config/proto.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 5d7ce3efbfc8b..399cfe81396f0 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -19,8 +19,11 @@ //! //! This is the relocated body of `datafusion-proto`'s //! `serialize_file_scan_config` / `parse_protobuf_file_scan_config`, ported to -//! ride the [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] instead of -//! the raw `PhysicalExtensionCodec` + `PhysicalProtoConverterExtension`. Every +//! 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::to_proto_conf`] and decodes with //! [`FileScanConfig::from_proto_conf`], keeping a single copy of the shared From c083115772ae22c5edc8bc0a4c4b8fc831a2f125 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Wed, 22 Jul 2026 11:53:38 +0530 Subject: [PATCH 3/8] Delegate FileScanConfig proto wrappers --- .../proto/src/physical_plan/from_proto.rs | 109 ++---------------- datafusion/proto/src/physical_plan/mod.rs | 43 +------ .../proto/src/physical_plan/to_proto.rs | 89 ++------------ 3 files changed, 24 insertions(+), 217 deletions(-) diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index d5a1e0efac6b6..dd88646054d95 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -29,7 +29,7 @@ use datafusion_common::{ }; 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, ListingTableUrl, PartitionedFile, TableSchema}; use datafusion_datasource_csv::file_format::CsvSink; @@ -41,7 +41,6 @@ use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; use datafusion_expr::dml::InsertOp; 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, @@ -51,6 +50,7 @@ use datafusion_physical_plan::expressions::{ 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, RangePartitioning, SplitPoint, WindowExpr, @@ -60,8 +60,8 @@ use object_store::ObjectMeta; use object_store::path::Path; use super::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - PhysicalProtoConverterExtension, + ConverterPlanDecoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; use crate::convert::TryFromProto; use crate::protobuf::physical_expr_node::ExprType; @@ -513,33 +513,7 @@ pub fn parse_protobuf_file_scan_schema( pub fn parse_table_schema_from_proto( proto: &protobuf::FileScanExecConf, ) -> Result { - let schema: Arc = 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::>>()?; - - // 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::>(), - ) - .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( @@ -548,76 +522,15 @@ pub fn parse_protobuf_file_scan_config( proto_converter: &dyn PhysicalProtoConverterExtension, file_source: Arc, ) -> Result { - let schema: Arc = 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::, _>>()?; - - 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 = 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::>>()?; - - 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::from_proto_conf( + proto, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) } pub fn parse_record_batches(buf: &[u8]) -> Result> { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 94c18527bd1a1..9c05a12b9e83b 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -320,13 +320,6 @@ mod tests { config.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) } - fn legacy_encode( - &self, - config: &FileScanConfig, - ) -> Result { - serialize_file_scan_config(config, &self.codec, &self.converter) - } - fn decode( &self, conf: &protobuf::FileScanExecConf, @@ -351,24 +344,11 @@ mod tests { file_source, ) } - - fn legacy_decode( - &self, - conf: &protobuf::FileScanExecConf, - ) -> Result { - let physical_decode_ctx = - PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); - parse_protobuf_file_scan_config( - conf, - &physical_decode_ctx, - &self.converter, - decode_source(conf)?, - ) - } } #[test] - fn new_file_scan_config_serde_matches_legacy_wire_format() -> Result<()> { + fn new_file_scan_config_serde_roundtrips_all_partitioning_variants() -> Result<()> + { let serde = FileScanSerdeHarness::new(); for config in [ @@ -377,23 +357,10 @@ mod tests { test_config(Some(hash_partitioning())), test_config(Some(range_partitioning())), test_config(Some(Partitioning::UnknownPartitioning(4))), - test_config_with_projection(None, None), - test_config_with_projection(None, Some(FileProjectionExprs::new(vec![]))), ] { - let legacy = serde.legacy_encode(&config)?; - let migrated = serde.encode(&config)?; - - assert_eq!(migrated, legacy); - assert_eq!(migrated.encode_to_vec(), legacy.encode_to_vec()); - - let migrated_reencoded = serde.encode(&serde.decode(&migrated)?)?; - let legacy_reencoded = - serde.legacy_encode(&serde.legacy_decode(&legacy)?)?; - assert_eq!(migrated_reencoded, legacy_reencoded); - assert_eq!( - migrated_reencoded.encode_to_vec(), - legacy_reencoded.encode_to_vec() - ); + let encoded = serde.encode(&config)?; + let reencoded = serde.encode(&serde.decode(&encoded)?)?; + assert_eq!(reencoded.output_partitioning, encoded.output_partitioning); } Ok(()) diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index e13923dbb9519..56cacabccc859 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -34,6 +34,7 @@ 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::{ @@ -41,13 +42,12 @@ use datafusion_physical_plan::{ }; 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)] @@ -487,84 +487,11 @@ pub fn serialize_file_scan_config( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let file_groups = conf - .file_groups - .iter() - .map(|p| protobuf::FileGroup::try_from_proto(p.files())) - .collect::, _>>()?; - - 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::>(); - 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::>(); - 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::>>()?, - }) - }) - .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::>(), - object_store_url: conf.object_store_url.to_string(), - output_ordering: output_orderings - .into_iter() - .map(|e| PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: e, - }) - .collect::>(), - 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.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) } pub fn serialize_maybe_filter( From a6d029166a96e6920cfc4b84d95efaa3fbd0133f Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 25 Jul 2026 10:16:36 +0530 Subject: [PATCH 4/8] Address FileScanConfig proto review feedback --- .../datasource/src/file_scan_config/mod.rs | 2 +- .../datasource/src/file_scan_config/proto.rs | 6 +- datafusion/proto/src/physical_plan/mod.rs | 683 +++++++++--------- 3 files changed, 342 insertions(+), 349 deletions(-) diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 36f2273576659..f9535d81417dd 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -24,7 +24,7 @@ pub(crate) mod sort_pushdown; /// Attaches inherent `to_proto_conf` / `from_proto_conf` / `parse_table_schema_from_proto` /// helpers to [`FileScanConfig`] used by every file source's `try_to_proto` hook. #[cfg(feature = "proto")] -mod proto; +pub(crate) mod proto; use crate::file_groups::FileGroup; use crate::{ diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 399cfe81396f0..89a764ad143aa 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -456,7 +456,9 @@ fn file_group_from_proto(group: &protobuf::FileGroup) -> Result { Ok(FileGroup::new(files)) } -fn partitioned_file_to_proto(pf: &PartitionedFile) -> Result { +pub(crate) fn partitioned_file_to_proto( + pf: &PartitionedFile, +) -> Result { let last_modified = pf.object_meta.last_modified; let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { DataFusionError::Plan(format!( @@ -485,7 +487,7 @@ fn partitioned_file_to_proto(pf: &PartitionedFile) -> Result Result { let mut pf = PartitionedFile::new_from_meta(ObjectMeta { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 9c05a12b9e83b..b70275e84fc2d 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -122,399 +122,390 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { } #[cfg(test)] -mod tests { +mod file_scan_config_serde { use super::*; - - 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, RangePartitioning, SplitPoint, - }; - use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; - use object_store::ObjectStore; - - #[derive(Clone)] - struct SerdeTestSource { - metrics: ExecutionPlanMetricsSet, + 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, RangePartitioning, SplitPoint, + }; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use object_store::ObjectStore; + + #[derive(Clone)] + struct SerdeTestSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + projection: Option, + } + + impl SerdeTestSource { + fn new( table_schema: TableSchema, projection: Option, - } - - impl SerdeTestSource { - fn new( - table_schema: TableSchema, - projection: Option, - ) -> Self { - Self { - metrics: ExecutionPlanMetricsSet::new(), - table_schema, - projection, - } + ) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + projection, } } + } - impl FileSource for SerdeTestSource { - fn create_file_opener( - &self, - _object_store: Arc, - _base_config: &FileScanConfig, - _partition: usize, - ) -> Result> { - 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 { - 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>> { - Ok(Some(Arc::new(Self { - projection: Some(projection.clone()), - ..self.clone() - }))) - } + impl FileSource for SerdeTestSource { + fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + internal_err!("not needed for FileScanConfig serde tests") + } - fn projection(&self) -> Option<&FileProjectionExprs> { - self.projection.as_ref() - } + fn table_schema(&self) -> &TableSchema { + &self.table_schema } - fn populated_projection() -> FileProjectionExprs { - FileProjectionExprs::new(vec![FileProjectionExpr::new( - Arc::new(Column::new("value", 0)), - "projected_value", - )]) + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) } - fn test_config(output_partitioning: Option) -> FileScanConfig { - test_config_with_projection(output_partitioning, Some(populated_projection())) + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics } - fn test_config_with_projection( - output_partitioning: Option, - projection: Option, - ) -> 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 file_type(&self) -> &str { + "serde-test" } - fn decode_source( - conf: &protobuf::FileScanExecConf, - ) -> Result> { - Ok(Arc::new(SerdeTestSource::new( - FileScanConfig::parse_table_schema_from_proto(conf)?, - None, - ))) + fn try_pushdown_projection( + &self, + projection: &FileProjectionExprs, + ) -> Result>> { + Ok(Some(Arc::new(Self { + projection: Some(projection.clone()), + ..self.clone() + }))) } - struct FileScanSerdeHarness { - codec: DefaultPhysicalExtensionCodec, - converter: DefaultPhysicalProtoConverter, - task_ctx: TaskContext, + fn projection(&self) -> Option<&FileProjectionExprs> { + self.projection.as_ref() } + } - impl FileScanSerdeHarness { - fn new() -> Self { - Self { - codec: DefaultPhysicalExtensionCodec {}, - converter: DefaultPhysicalProtoConverter {}, - task_ctx: TaskContext::default(), - } - } + fn populated_projection() -> FileProjectionExprs { + FileProjectionExprs::new(vec![FileProjectionExpr::new( + Arc::new(Column::new("value", 0)), + "projected_value", + )]) + } - fn encode( - &self, - config: &FileScanConfig, - ) -> Result { - let encoder = ConverterPlanEncoder { - codec: &self.codec, - proto_converter: &self.converter, - }; - config.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) - } + fn test_config(output_partitioning: Option) -> FileScanConfig { + test_config_with_projection(output_partitioning, Some(populated_projection())) + } - fn decode( - &self, - conf: &protobuf::FileScanExecConf, - ) -> Result { - self.decode_with_source(conf, decode_source(conf)?) - } + fn test_config_with_projection( + output_partitioning: Option, + projection: Option, + ) -> 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> { + Ok(Arc::new(SerdeTestSource::new( + FileScanConfig::parse_table_schema_from_proto(conf)?, + None, + ))) + } - fn decode_with_source( - &self, - conf: &protobuf::FileScanExecConf, - file_source: Arc, - ) -> Result { - let physical_decode_ctx = - PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); - let decoder = ConverterPlanDecoder { - ctx: &physical_decode_ctx, - proto_converter: &self.converter, - }; - FileScanConfig::from_proto_conf( - conf, - &ExecutionPlanDecodeCtx::new(&decoder), - file_source, - ) + struct FileScanSerdeHarness { + codec: DefaultPhysicalExtensionCodec, + converter: DefaultPhysicalProtoConverter, + task_ctx: TaskContext, + } + + impl FileScanSerdeHarness { + fn new() -> Self { + Self { + codec: DefaultPhysicalExtensionCodec {}, + converter: DefaultPhysicalProtoConverter {}, + task_ctx: TaskContext::default(), } } - #[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); - } + fn encode(&self, config: &FileScanConfig) -> Result { + let encoder = ConverterPlanEncoder { + codec: &self.codec, + proto_converter: &self.converter, + }; + config.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) + } - Ok(()) + fn decode(&self, conf: &protobuf::FileScanExecConf) -> Result { + self.decode_with_source(conf, decode_source(conf)?) } - #[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)?)?; + fn decode_with_source( + &self, + conf: &protobuf::FileScanExecConf, + file_source: Arc, + ) -> Result { + let physical_decode_ctx = + PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); + let decoder = ConverterPlanDecoder { + ctx: &physical_decode_ctx, + proto_converter: &self.converter, + }; + FileScanConfig::from_proto_conf( + conf, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) + } + } - 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()); + #[test] + fn new_file_scan_config_serde_roundtrips_all_partitioning_variants() -> Result<()> { + let serde = FileScanSerdeHarness::new(); - Ok(()) + 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); } - #[test] - fn new_file_scan_config_serde_preserves_projection_presence() -> Result<()> { - let serde = FileScanSerdeHarness::new(); + Ok(()) + } - 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()); + #[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)?)?; - 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()) - ); + 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(()) - } + Ok(()) + } + #[test] + fn new_file_scan_config_serde_preserves_projection_presence() -> Result<()> { + let serde = FileScanSerdeHarness::new(); - #[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 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 mut missing_projection_expr = valid.clone(); - missing_projection_expr + let empty = serde.encode(&test_config_with_projection( + None, + Some(FileProjectionExprs::new(vec![])), + ))?; + assert!( + empty .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}" - ); + .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(()) - } + 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()); + #[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(&duplicate_ordering) - .expect_err("duplicate range ordering must fail"); - assert!( - err.to_string().contains("duplicate expressions"), - "unexpected error: {err}" - ); + .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}" + ); - 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(); + Ok(()) + } - let err = serde - .decode(&proto) - .expect_err("empty range ordering must fail"); - assert!( - err.to_string().contains("requires non-empty ordering"), - "unexpected error: {err}" - ); + #[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())))?; - Ok(()) - } + 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::*; /// Unit tests for the bytes-only function serde exposed on /// [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] and backed by From 1213e29bdfa5a4fc11f7fe9f4036a7ac38322408 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:54:34 -0500 Subject: [PATCH 5/8] Move PartitionedFile/FileGroup proto conversions into datafusion-datasource `FileScanConfig`'s new proto module carried private copies of the `PartitionedFile` / `FileGroup` wire logic that already existed as `TryFromProto` impls in `datafusion-proto`, so the two could drift. Put the single copy next to the types that own it, in a new `datafusion_datasource::proto` module (feature `proto`): - `FileRange::try_to_proto` / `try_from_proto` - `PartitionedFile::try_to_proto` / `try_from_proto` - `FileGroup::try_to_proto` / `try_from_proto` `datafusion-proto`'s `TryFromProto` impls for these types become one-line shims delegating to the above, so the central serializer and the per-source `try_to_proto` hooks cannot disagree, and the private `partitioned_file_*` / `file_group_*` helpers in `file_scan_config/proto.rs` are gone. Making them inherent methods (rather than moving `TryFromProto` into `datafusion-proto-models`) keeps `TryFromProto` local to `datafusion-proto`, which its other 45 impls rely on for the orphan rule. Co-Authored-By: Claude Opus 5 --- .../datasource/src/file_scan_config/proto.rs | 89 +------ datafusion/datasource/src/mod.rs | 4 + datafusion/datasource/src/proto.rs | 218 ++++++++++++++++++ .../proto/src/physical_plan/from_proto.rs | 50 +--- .../proto/src/physical_plan/to_proto.rs | 40 +--- 5 files changed, 241 insertions(+), 160 deletions(-) create mode 100644 datafusion/datasource/src/proto.rs diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 89a764ad143aa..88908f8a9fef9 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -39,7 +39,6 @@ use std::sync::Arc; use arrow::compute::SortOptions; use arrow::datatypes::Schema; -use chrono::{TimeZone, Utc}; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; @@ -48,10 +47,7 @@ use datafusion_physical_expr::{ }; use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; use datafusion_proto_models::protobuf; -use object_store::ObjectMeta; -use object_store::path::Path; -use crate::PartitionedFile; use crate::file::FileSource; use crate::file_groups::FileGroup; use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -71,7 +67,7 @@ impl FileScanConfig { let file_groups = self .file_groups .iter() - .map(file_group_to_proto) + .map(FileGroup::try_to_proto) .collect::>>()?; // Sort orderings: only the child expressions need the ctx; the @@ -185,7 +181,7 @@ impl FileScanConfig { let file_groups = conf .file_groups .iter() - .map(file_group_from_proto) + .map(FileGroup::try_from_proto) .collect::>>()?; let object_store_url = match conf.object_store_url.is_empty() { @@ -436,84 +432,3 @@ fn partitioning_from_proto( }; Ok(Some(partitioning)) } - -fn file_group_to_proto(group: &FileGroup) -> Result { - Ok(protobuf::FileGroup { - files: group - .files() - .iter() - .map(partitioned_file_to_proto) - .collect::>>()?, - }) -} - -fn file_group_from_proto(group: &protobuf::FileGroup) -> Result { - let files = group - .files - .iter() - .map(partitioned_file_from_proto) - .collect::>>()?; - Ok(FileGroup::new(files)) -} - -pub(crate) fn partitioned_file_to_proto( - pf: &PartitionedFile, -) -> Result { - let last_modified = pf.object_meta.last_modified; - let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" - )) - })? as u64; - Ok(protobuf::PartitionedFile { - arrow_schema: pf - .arrow_schema - .as_ref() - .map(|s| s.as_ref().try_into()) - .transpose()?, - path: pf.object_meta.location.as_ref().to_owned(), - size: pf.object_meta.size, - last_modified_ns, - partition_values: pf - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: pf.range.as_ref().map(|range| protobuf::FileRange { - start: range.start, - end: range.end, - }), - statistics: pf.statistics.as_ref().map(|s| s.as_ref().into()), - }) -} - -pub(crate) fn partitioned_file_from_proto( - val: &protobuf::PartitionedFile, -) -> Result { - let mut pf = PartitionedFile::new_from_meta(ObjectMeta { - location: Path::parse(val.path.as_str()) - .map_err(|e| internal_datafusion_err!("Invalid object_store path: {e}"))?, - last_modified: Utc.timestamp_nanos(val.last_modified_ns as i64), - size: val.size, - e_tag: None, - version: None, - }) - .with_partition_values( - val.partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - ); - if let Some(proto_schema) = val.arrow_schema.as_ref() { - pf = pf.with_arrow_schema(Arc::new( - proto_schema.try_into().map_err(DataFusionError::from)?, - )); - } - if let Some(range) = val.range.as_ref() { - pf = pf.with_range(range.start, range.end); - } - if let Some(proto_stats) = val.statistics.as_ref() { - pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); - } - Ok(pf) -} diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 7c8cae337f1eb..e718064c819da 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -41,6 +41,10 @@ pub mod file_stream; pub mod memory; pub mod morsel; pub mod projection; +/// Protobuf conversions for [`FileRange`], [`PartitionedFile`] and +/// [`FileGroup`](crate::file_groups::FileGroup), gated on the `proto` feature. +#[cfg(feature = "proto")] +pub mod proto; pub mod schema_adapter; pub mod sink; pub mod source; diff --git a/datafusion/datasource/src/proto.rs b/datafusion/datasource/src/proto.rs new file mode 100644 index 0000000000000..2796aa224c5de --- /dev/null +++ b/datafusion/datasource/src/proto.rs @@ -0,0 +1,218 @@ +// 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. + +//! Protobuf conversions for the file-scan leaf types owned by this crate: +//! [`FileRange`], [`PartitionedFile`] and [`FileGroup`]. +//! +//! These are the single copy of that wire logic. `datafusion-proto`'s +//! `TryFromProto` implementations for the same types are thin shims that +//! delegate here, so the format cannot drift between the central serializer and +//! the per-source `try_to_proto` hooks. +//! +//! None of these conversions need a codec or an encode/decode context: every +//! field is plain data or goes through `datafusion-proto-common`. + +use std::sync::Arc; + +use chrono::{TimeZone, Utc}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_proto_models::protobuf; +use object_store::ObjectMeta; +use object_store::path::Path; + +use crate::file_groups::FileGroup; +use crate::{FileRange, PartitionedFile}; + +impl FileRange { + /// Serialize this range into its protobuf representation. + pub fn try_to_proto(&self) -> Result { + Ok(protobuf::FileRange { + start: self.start, + end: self.end, + }) + } + + /// Reconstruct a [`FileRange`] from its protobuf representation. + pub fn try_from_proto(range: &protobuf::FileRange) -> Result { + Ok(FileRange { + start: range.start, + end: range.end, + }) + } +} + +impl PartitionedFile { + /// Serialize this file into its protobuf representation. + pub fn try_to_proto(&self) -> Result { + let last_modified = self.object_meta.last_modified; + let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { + DataFusionError::Plan(format!( + "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" + )) + })? as u64; + Ok(protobuf::PartitionedFile { + arrow_schema: self + .arrow_schema + .as_ref() + .map(|s| s.as_ref().try_into()) + .transpose()?, + path: self.object_meta.location.as_ref().to_owned(), + size: self.object_meta.size, + last_modified_ns, + partition_values: self + .partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + range: self + .range + .as_ref() + .map(FileRange::try_to_proto) + .transpose()?, + statistics: self.statistics.as_ref().map(|s| s.as_ref().into()), + }) + } + + /// Reconstruct a [`PartitionedFile`] from its protobuf representation. + pub fn try_from_proto(file: &protobuf::PartitionedFile) -> Result { + let mut pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse(file.path.as_str()).map_err(|e| { + internal_datafusion_err!("Invalid object_store path: {e}") + })?, + last_modified: Utc.timestamp_nanos(file.last_modified_ns as i64), + size: file.size, + e_tag: None, + version: None, + }) + .with_partition_values( + file.partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + ); + if let Some(proto_schema) = file.arrow_schema.as_ref() { + pf = pf.with_arrow_schema(Arc::new( + proto_schema.try_into().map_err(DataFusionError::from)?, + )); + } + if let Some(range) = file.range.as_ref() { + let range = FileRange::try_from_proto(range)?; + pf = pf.with_range(range.start, range.end); + } + if let Some(proto_stats) = file.statistics.as_ref() { + pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); + } + Ok(pf) + } +} + +impl FileGroup { + /// Serialize this group into its protobuf representation. + pub fn try_to_proto(&self) -> Result { + Ok(protobuf::FileGroup { + files: self + .files() + .iter() + .map(PartitionedFile::try_to_proto) + .collect::>>()?, + }) + } + + /// Reconstruct a [`FileGroup`] from its protobuf representation. + pub fn try_from_proto(group: &protobuf::FileGroup) -> Result { + Ok(FileGroup::new( + group + .files + .iter() + .map(PartitionedFile::try_from_proto) + .collect::>>()?, + )) + } +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ScalarValue, Statistics}; + + use super::*; + + #[test] + fn partitioned_file_roundtrip_preserves_all_fields() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse("foo/bar.parquet")?, + last_modified: Utc.timestamp_nanos(1_000_000_000), + size: 1234, + e_tag: None, + version: None, + }) + .with_partition_values(vec![ScalarValue::from("2024-01-01")]) + .with_range(10, 20) + .with_arrow_schema(Arc::clone(&schema)) + .with_statistics(Arc::new(Statistics::new_unknown(&schema))); + + let decoded = PartitionedFile::try_from_proto(&pf.try_to_proto()?)?; + + assert_eq!(decoded.object_meta.location, pf.object_meta.location); + assert_eq!(decoded.object_meta.size, pf.object_meta.size); + assert_eq!( + decoded.object_meta.last_modified, + pf.object_meta.last_modified + ); + assert_eq!(decoded.partition_values, pf.partition_values); + assert_eq!(decoded.range, pf.range); + assert_eq!(decoded.arrow_schema.as_deref(), Some(schema.as_ref())); + // Statistics on `PartitionedFile` span the full table schema, and + // `PartitionedFile::with_statistics` re-derives the partition column + // entries, so the decoded statistics are not compared field by field + // here; this is pre-existing behavior of the wire format. + assert!(decoded.statistics.is_some()); + Ok(()) + } + + #[test] + fn partitioned_file_from_proto_rejects_invalid_path() { + let proto = protobuf::PartitionedFile { + path: "foo//bar.parquet".to_string(), + ..Default::default() + }; + + let err = PartitionedFile::try_from_proto(&proto).unwrap_err(); + assert!( + err.to_string().contains("Invalid object_store path"), + "unexpected error: {err}" + ); + } + + #[test] + fn file_group_roundtrip() -> Result<()> { + let group = FileGroup::new(vec![ + PartitionedFile::new("a.parquet", 1), + PartitionedFile::new("b.parquet", 2), + ]); + + let decoded = FileGroup::try_from_proto(&group.try_to_proto()?)?; + + assert_eq!(decoded.len(), 2); + assert_eq!( + decoded.files()[1].object_meta.location, + group.files()[1].object_meta.location + ); + Ok(()) + } +} diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index dd88646054d95..f87fb1f3e1580 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -23,7 +23,6 @@ use arrow::array::RecordBatch; use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; -use chrono::{TimeZone, Utc}; use datafusion_common::{ DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, }; @@ -56,8 +55,6 @@ use datafusion_physical_plan::{ Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, WindowExpr, }; use datafusion_proto_common::common::proto_error; -use object_store::ObjectMeta; -use object_store::path::Path; use super::{ ConverterPlanDecoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -545,61 +542,30 @@ pub fn parse_record_batches(buf: &[u8]) -> Result> { Ok(batches) } +/// Thin shim over [`PartitionedFile::try_from_proto`], which owns the wire logic. impl TryFromProto<&protobuf::PartitionedFile> for PartitionedFile { type Error = DataFusionError; fn try_from_proto(val: &protobuf::PartitionedFile) -> Result { - let mut pf = PartitionedFile::new_from_meta(ObjectMeta { - location: Path::parse(val.path.as_str()) - .map_err(|e| proto_error(format!("Invalid object_store path: {e}")))?, - last_modified: Utc.timestamp_nanos(val.last_modified_ns as i64), - size: val.size, - e_tag: None, - version: None, - }) - .with_partition_values( - val.partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - ); - if let Some(proto_schema) = val.arrow_schema.as_ref() { - pf = pf.with_arrow_schema(Arc::new( - proto_schema.try_into().map_err(DataFusionError::from)?, - )); - } - if let Some(range) = val.range.as_ref() { - let file_range = FileRange::try_from_proto(range)?; - pf = pf.with_range(file_range.start, file_range.end); - } - if let Some(proto_stats) = val.statistics.as_ref() { - pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); - } - Ok(pf) + PartitionedFile::try_from_proto(val) } } +/// Thin shim over [`FileRange::try_from_proto`], which owns the wire logic. impl TryFromProto<&protobuf::FileRange> for FileRange { type Error = DataFusionError; fn try_from_proto(value: &protobuf::FileRange) -> Result { - Ok(FileRange { - start: value.start, - end: value.end, - }) + FileRange::try_from_proto(value) } } +/// Thin shim over [`FileGroup::try_from_proto`], which owns the wire logic. impl TryFromProto<&protobuf::FileGroup> for FileGroup { type Error = DataFusionError; fn try_from_proto(val: &protobuf::FileGroup) -> Result { - let files = val - .files - .iter() - .map(PartitionedFile::try_from_proto) - .collect::, _>>()?; - Ok(FileGroup::new(files)) + FileGroup::try_from_proto(val) } } @@ -721,6 +687,10 @@ impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprD #[cfg(test)] mod tests { + use chrono::{TimeZone, Utc}; + use object_store::ObjectMeta; + use object_store::path::Path; + use super::*; #[test] diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 56cacabccc859..e6628d511e4c5 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -424,51 +424,25 @@ fn serialize_range_split_point( }) } +/// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. impl TryFromProto<&PartitionedFile> for protobuf::PartitionedFile { type Error = DataFusionError; fn try_from_proto(pf: &PartitionedFile) -> Result { - let last_modified = pf.object_meta.last_modified; - let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" - )) - })? as u64; - Ok(protobuf::PartitionedFile { - arrow_schema: pf - .arrow_schema - .as_ref() - .map(|s| s.as_ref().try_into()) - .transpose()?, - path: pf.object_meta.location.as_ref().to_owned(), - size: pf.object_meta.size, - last_modified_ns, - partition_values: pf - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: pf - .range - .as_ref() - .map(protobuf::FileRange::try_from_proto) - .transpose()?, - statistics: pf.statistics.as_ref().map(|s| s.as_ref().into()), - }) + pf.try_to_proto() } } +/// Thin shim over [`FileRange::try_to_proto`], which owns the wire logic. impl TryFromProto<&FileRange> for protobuf::FileRange { type Error = DataFusionError; fn try_from_proto(value: &FileRange) -> Result { - Ok(protobuf::FileRange { - start: value.start, - end: value.end, - }) + value.try_to_proto() } } +/// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { type Error = DataFusionError; @@ -476,8 +450,8 @@ impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { Ok(protobuf::FileGroup { files: gr .iter() - .map(protobuf::PartitionedFile::try_from_proto) - .collect::, _>>()?, + .map(PartitionedFile::try_to_proto) + .collect::>>()?, }) } } From d38ebef1fad99cd22daae1b2036f7b21744951cc Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:57:25 -0500 Subject: [PATCH 6/8] Rename FileScanConfig proto helpers to try_to_proto / try_from_proto Matches the naming used by the `ExecutionPlan` / `DataSource` / `FileSource` hooks, as raised in review. `FileScanConfig` also implements `DataSource::try_to_proto`, so the inherent encode method shadows the trait method for callers holding a concrete `FileScanConfig` (they differ in return type, so a mistaken call is a compile error, not silent misbehavior, and `dyn DataSource` callers are unaffected). Documented on the method; the decode side has no such overlap. Co-Authored-By: Claude Opus 5 --- datafusion/datasource/src/file.rs | 2 +- .../datasource/src/file_scan_config/mod.rs | 2 +- .../datasource/src/file_scan_config/proto.rs | 19 +++++++++++++------ .../proto/src/physical_plan/from_proto.rs | 2 +- datafusion/proto/src/physical_plan/mod.rs | 4 ++-- .../proto/src/physical_plan/to_proto.rs | 2 +- 6 files changed, 19 insertions(+), 12 deletions(-) diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 44fdf8dc16d74..691bb314b7c03 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -358,7 +358,7 @@ pub trait FileSource: Any + Send + Sync { /// `base` is the shared [`FileScanConfig`] this source is wrapped in; the /// format-agnostic parts (file groups, schema, statistics, ordering, /// projection, …) are encoded via - /// [`FileScanConfig::to_proto_conf`](crate::file_scan_config::FileScanConfig::to_proto_conf), + /// [`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. /// diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index f9535d81417dd..d02ecff914396 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -21,7 +21,7 @@ pub(crate) mod sort_pushdown; /// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature. -/// Attaches inherent `to_proto_conf` / `from_proto_conf` / `parse_table_schema_from_proto` +/// 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")] pub(crate) mod proto; diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 88908f8a9fef9..9a119790bba48 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -25,8 +25,8 @@ //! instead of the raw `PhysicalExtensionCodec` + //! `PhysicalProtoConverterExtension`. Every //! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its -//! `*ScanExecNode` around [`FileScanConfig::to_proto_conf`] and decodes with -//! [`FileScanConfig::from_proto_conf`], keeping a single copy of the shared +//! `*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. //! @@ -60,7 +60,14 @@ impl FileScanConfig { /// 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 to_proto_conf( + /// + /// Note this inherent method shadows + /// [`DataSource::try_to_proto`](crate::source::DataSource::try_to_proto) for + /// callers holding a concrete `FileScanConfig`: that one serializes the whole + /// `DataSourceExec` node (by delegating to the file source), this one only the + /// shared `FileScanExecConf` spine. Spell the trait method + /// `DataSource::try_to_proto(&config, ctx)` when you want the plan node. + pub fn try_to_proto( &self, ctx: &ExecutionPlanEncodeCtx<'_>, ) -> Result { @@ -97,7 +104,7 @@ impl FileScanConfig { .transpose()?; // Fields must be added to the schema so that they can persist in the - // protobuf, and then removed from the schema in `from_proto_conf`. + // protobuf, and then removed from the schema in `try_from_proto`. let mut fields = self .file_schema() .fields() @@ -152,7 +159,7 @@ impl FileScanConfig { /// table schema via [`FileScanConfig::parse_table_schema_from_proto`]). /// /// Byte-compatible with the former `parse_protobuf_file_scan_config`. - pub fn from_proto_conf( + pub fn try_from_proto( conf: &protobuf::FileScanExecConf, ctx: &ExecutionPlanDecodeCtx<'_>, file_source: Arc, @@ -240,7 +247,7 @@ impl FileScanConfig { /// Parse a [`TableSchema`] (file schema + partition columns) from a /// [`protobuf::FileScanExecConf`]. File sources use this to rebuild their - /// concrete source before calling [`FileScanConfig::from_proto_conf`]. + /// 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( diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index f87fb1f3e1580..b8db8c7f398e4 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -523,7 +523,7 @@ pub fn parse_protobuf_file_scan_config( ctx, proto_converter, }; - FileScanConfig::from_proto_conf( + FileScanConfig::try_from_proto( proto, &ExecutionPlanDecodeCtx::new(&decoder), file_source, diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index b70275e84fc2d..7d724c4c23937 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -309,7 +309,7 @@ mod file_scan_config_serde { codec: &self.codec, proto_converter: &self.converter, }; - config.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) + config.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) } fn decode(&self, conf: &protobuf::FileScanExecConf) -> Result { @@ -327,7 +327,7 @@ mod file_scan_config_serde { ctx: &physical_decode_ctx, proto_converter: &self.converter, }; - FileScanConfig::from_proto_conf( + FileScanConfig::try_from_proto( conf, &ExecutionPlanDecodeCtx::new(&decoder), file_source, diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index e6628d511e4c5..612661dba8c50 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -465,7 +465,7 @@ pub fn serialize_file_scan_config( codec, proto_converter, }; - conf.to_proto_conf(&ExecutionPlanEncodeCtx::new(&encoder)) + conf.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) } pub fn serialize_maybe_filter( From bf930284b51f2a2696b869942e2d2355e5d1e4a2 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:31:52 -0500 Subject: [PATCH 7/8] Deprecate datafusion-proto's parse_table_schema_from_proto wrapper `datafusion_proto::physical_plan::from_proto::parse_table_schema_from_proto` is now a pure delegation to `FileScanConfig::parse_table_schema_from_proto`, so mark it deprecated (as flagged in review) and point the in-crate callers at the `FileScanConfig` method directly. Co-Authored-By: Claude Opus 5 --- .../proto/src/physical_plan/from_proto.rs | 4 ++++ datafusion/proto/src/physical_plan/mod.rs | 18 +++++++++------- .../library-user-guide/upgrading/55.0.0.md | 21 +++++++++++++++++++ 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index b8db8c7f398e4..2ec77c00596c7 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -507,6 +507,10 @@ pub fn parse_protobuf_file_scan_schema( } /// Parses a TableSchema from protobuf, extracting the file schema and partition columns +#[deprecated( + since = "55.0.0", + note = "use `FileScanConfig::parse_table_schema_from_proto` instead" +)] pub fn parse_table_schema_from_proto( proto: &protobuf::FileScanExecConf, ) -> Result { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 7d724c4c23937..77e05cdc40109 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -100,7 +100,7 @@ use crate::convert::TryFromProto; use crate::convert_required; use crate::physical_plan::from_proto::{ parse_physical_expr_with_converter, parse_physical_sort_exprs, - parse_protobuf_file_scan_config, parse_record_batches, parse_table_schema_from_proto, + parse_protobuf_file_scan_config, parse_record_batches, }; use crate::physical_plan::to_proto::{ serialize_file_scan_config, serialize_physical_expr_with_converter, @@ -1381,8 +1381,9 @@ pub trait PhysicalPlanNodeExt: Sized { }; // Parse table schema with partition columns - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; + let table_schema = FileScanConfig::parse_table_schema_from_proto( + scan.base_conf.as_ref().unwrap(), + )?; let csv_options = CsvOptions { has_header: Some(scan.has_header), @@ -1416,7 +1417,7 @@ pub trait PhysicalPlanNodeExt: Sized { proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { let base_conf = scan.base_conf.as_ref().unwrap(); - let table_schema = parse_table_schema_from_proto(base_conf)?; + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; let scan_conf = parse_protobuf_file_scan_config( base_conf, ctx, @@ -1435,7 +1436,7 @@ pub trait PhysicalPlanNodeExt: Sized { 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 table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; let scan_conf = parse_protobuf_file_scan_config( base_conf, ctx, @@ -1490,7 +1491,7 @@ pub trait PhysicalPlanNodeExt: Sized { } // Parse table schema with partition columns - let table_schema = parse_table_schema_from_proto(base_conf)?; + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; let object_store_url = match base_conf.object_store_url.is_empty() { false => ObjectStoreUrl::parse(&base_conf.object_store_url)?, true => ObjectStoreUrl::local_filesystem(), @@ -1537,8 +1538,9 @@ pub trait PhysicalPlanNodeExt: Sized { ) -> Result> { #[cfg(feature = "avro")] { - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; + let table_schema = FileScanConfig::parse_table_schema_from_proto( + scan.base_conf.as_ref().unwrap(), + )?; let conf = parse_protobuf_file_scan_config( scan.base_conf.as_ref().unwrap(), ctx, diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index 6097c8dc717df..a4218b592e4c2 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -894,6 +894,27 @@ let plan = deserialize_bytes(&proto_bytes)?; See [PR #23827](https://github.com/apache/datafusion/pull/23827) for details. +### `parse_table_schema_from_proto` is deprecated + +`datafusion_proto::physical_plan::from_proto::parse_table_schema_from_proto` is +deprecated in favor of the equivalent method on `FileScanConfig`, which lives in +`datafusion-datasource` alongside the rest of the shared file-scan protobuf +conversions: + +```rust,ignore +// Before +let table_schema = parse_table_schema_from_proto(base_conf)?; + +// After +let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; +``` + +`PartitionedFile`, `FileGroup` and `FileRange` also gained inherent +`try_to_proto` / `try_from_proto` methods (in `datafusion_datasource::proto`, +behind the `proto` feature) which now own that wire logic; +`datafusion-proto`'s `TryFromProto` implementations for those types delegate to +them and keep working unchanged. + ### `MSRV` updated to 1.94.0 The Minimum Supported Rust Version (MSRV) has been updated to [`1.94.0`]. From 262b402eedd6cc6558a71f893917fc2bb50d43f8 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 08:35:00 -0500 Subject: [PATCH 8/8] Put Partitioning / PhysicalSortExpr proto conversion on the types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The new `FileScanConfig` proto module carried its own copy of the `Partitioning` and `PhysicalSortExprNode` wire logic, which already existed inline in `RepartitionExec`'s hook and in `datafusion-proto`'s `serialize_partitioning` / `parse_protobuf_partitioning` — three copies. Apply the same principle as the `PartitionedFile` commit and put the single copy next to the types that own it, taking the expression-level context (`datafusion-physical-expr{,-common}` already have the `proto` feature): - `PhysicalSortExpr::try_to_proto` / `try_from_proto` (physical-expr-common) - `Partitioning::try_to_proto` / `try_from_proto` (physical-expr) To let plan hooks reach them, `ExecutionPlanEncodeCtx` / `ExecutionPlanDecodeCtx` now back the expression-level contexts and hand one out via `expr_ctx()`. `FileScanConfig`, `RepartitionExec` and `datafusion-proto`'s central serializer all route through the type methods, which also retires `serialize_range_partitioning`, `serialize_range_split_point`, `parse_protobuf_range_partitioning` and `parse_protobuf_range_split_point`. Wire format is unchanged. Behavior differences: out-of-range partition counts now error instead of wrapping or panicking on `unwrap`, and a missing sort-expression child reports which field is missing rather than "Unexpected empty physical expression". Co-Authored-By: Claude Opus 5 --- .../datasource/src/file_scan_config/proto.rs | 197 +++--------------- .../physical-expr-common/src/sort_expr.rs | 43 ++++ datafusion/physical-expr/src/partitioning.rs | 140 +++++++++++++ datafusion/physical-plan/src/proto.rs | 47 +++++ .../physical-plan/src/repartition/mod.rs | 133 +----------- .../proto/src/physical_plan/from_proto.rs | 95 ++------- .../proto/src/physical_plan/to_proto.rs | 74 +------ 7 files changed, 288 insertions(+), 441 deletions(-) diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs index 9a119790bba48..aac2364b5f030 100644 --- a/datafusion/datasource/src/file_scan_config/proto.rs +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -30,21 +30,21 @@ //! 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. +//! Sort orderings and output partitioning ride the shared +//! `ctx.encode_sort_exprs` / `ctx.encode_partitioning` helpers (and their decode +//! siblings); projection expressions go through `ctx.encode_expr` / +//! `ctx.decode_expr`; file groups through +//! [`FileGroup::try_to_proto`](crate::file_groups::FileGroup::try_to_proto); +//! and `Schema`, `Statistics`, `Constraints` and `ScalarValue` through +//! `datafusion-proto-common`. Nothing here needs the raw codec. use std::sync::Arc; -use arrow::compute::SortOptions; 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, PhysicalSortExpr, RangePartitioning, SplitPoint, -}; +use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalSortExpr}; use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; use datafusion_proto_models::protobuf; @@ -77,30 +77,21 @@ impl FileScanConfig { .map(FileGroup::try_to_proto) .collect::>>()?; - // Sort orderings: only the child expressions need the ctx; the - // asc/nulls_first wrapping is plain data inlined into a - // `PhysicalSortExprNode` (same shape as `sorts/sort.rs`). + let expr_ctx = ctx.expr_ctx(); let mut output_ordering = vec![]; for order in &self.output_ordering { - let nodes = order - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; output_ordering.push(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: nodes, + physical_sort_expr_nodes: order + .iter() + .map(|sort_expr| sort_expr.try_to_proto(&expr_ctx)) + .collect::>>()?, }); } let output_partitioning = self .output_partitioning .as_ref() - .map(|p| partitioning_to_proto(p, ctx)) + .map(|partitioning| partitioning.try_to_proto(&expr_ctx)) .transpose()?; // Fields must be added to the schema so that they can persist in the @@ -196,18 +187,23 @@ impl FileScanConfig { true => ObjectStoreUrl::local_filesystem(), }; + let expr_ctx = ctx.expr_ctx(&schema); let mut output_ordering = vec![]; for node_collection in &conf.output_ordering { - let sort_exprs = parse_sort_exprs( - &node_collection.physical_sort_expr_nodes, - ctx, - &schema, - )?; + let sort_exprs = node_collection + .physical_sort_expr_nodes + .iter() + .map(|sort_expr| PhysicalSortExpr::try_from_proto(sort_expr, &expr_ctx)) + .collect::>>()?; output_ordering.extend(LexOrdering::new(sort_exprs)); } - let output_partitioning = - partitioning_from_proto(conf.output_partitioning.as_ref(), ctx, &schema)?; + let output_partitioning = conf + .output_partitioning + .as_ref() + .map(|partitioning| Partitioning::try_from_proto(partitioning, &expr_ctx)) + .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 { @@ -298,144 +294,3 @@ fn parse_file_scan_schema(conf: &protobuf::FileScanExecConf) -> Result, - schema: &Schema, -) -> Result> { - nodes - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!("Unexpected empty physical expression") - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, schema)?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect() -} - -/// Inlined equivalent of `datafusion-proto`'s `serialize_partitioning`. Only -/// child physical expressions and `ScalarValue`s need the ctx; the -/// `protobuf::Partitioning` wrapping is built directly here. -fn partitioning_to_proto( - partitioning: &Partitioning, - ctx: &ExecutionPlanEncodeCtx<'_>, -) -> Result { - let partition_method = match partitioning { - Partitioning::RoundRobinBatch(n) => { - protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) - } - Partitioning::Hash(exprs, n) => { - let hash_expr = ctx.encode_expressions(exprs)?; - protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr, - partition_count: *n as u64, - }, - ) - } - Partitioning::Range(range) => { - let sort_expr = range - .ordering() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; - let split_point = range - .split_points() - .iter() - .map(|split_point| { - let value = split_point - .values() - .iter() - .map(|value| value.try_into().map_err(Into::into)) - .collect::>>()?; - Ok(protobuf::PhysicalRangeSplitPoint { value }) - }) - .collect::>>()?; - protobuf::partitioning::PartitionMethod::Range( - protobuf::PhysicalRangePartitioning { - sort_expr, - split_point, - }, - ) - } - Partitioning::UnknownPartitioning(n) => { - protobuf::partitioning::PartitionMethod::Unknown(*n as u64) - } - }; - Ok(protobuf::Partitioning { - partition_method: Some(partition_method), - }) -} - -/// Inlined equivalent of `datafusion-proto`'s `parse_protobuf_partitioning`. -fn partitioning_from_proto( - partitioning: Option<&protobuf::Partitioning>, - ctx: &ExecutionPlanDecodeCtx<'_>, - schema: &Schema, -) -> Result> { - let Some(partitioning) = partitioning else { - return Ok(None); - }; - let Some(partition_method) = partitioning.partition_method.as_ref() else { - return Ok(None); - }; - let partitioning = match partition_method { - protobuf::partitioning::PartitionMethod::RoundRobin(n) => { - Partitioning::RoundRobinBatch(*n as usize) - } - protobuf::partitioning::PartitionMethod::Hash(hash) => { - let exprs = hash - .hash_expr - .iter() - .map(|expr| ctx.decode_expr(expr, schema)) - .collect::>>()?; - Partitioning::Hash(exprs, hash.partition_count as usize) - } - protobuf::partitioning::PartitionMethod::Unknown(n) => { - Partitioning::UnknownPartitioning(*n as usize) - } - protobuf::partitioning::PartitionMethod::Range(range) => { - let sort_exprs = parse_sort_exprs(&range.sort_expr, ctx, schema)?; - let sort_expr_count = sort_exprs.len(); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - internal_datafusion_err!("Range partitioning requires non-empty ordering") - })?; - if ordering.len() != sort_expr_count { - return Err(internal_datafusion_err!( - "Range partitioning ordering must not contain duplicate expressions" - )); - } - let split_points = range - .split_point - .iter() - .map(|split_point| { - let values = split_point - .value - .iter() - .map(|value| { - datafusion_common::ScalarValue::try_from(value) - .map_err(Into::into) - }) - .collect::>>()?; - Ok(SplitPoint::new(values)) - }) - .collect::>>()?; - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) - } - }; - Ok(Some(partitioning)) -} diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 84ffb92eaa600..82f7a6e5fc33f 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -183,6 +183,49 @@ impl PhysicalSortExpr { } } +/// Protobuf conversions for [`PhysicalSortExpr`]. +/// +/// This is the flat [`PhysicalSortExprNode`] representation used wherever the +/// wire format stores an ordering (scan output orderings, range partitioning, +/// window frames, …). It is *not* the `PhysicalExprNode::Sort` wrapping that +/// `SortExec` uses for its own `expr` field. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +impl PhysicalSortExpr { + /// Serialize this sort expression, encoding its child expression through + /// `ctx`. + pub fn try_to_proto( + &self, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result { + Ok(datafusion_proto_models::protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_child(&self.expr)?)), + asc: !self.options.descending, + nulls_first: self.options.nulls_first, + }) + } + + /// Reconstruct a [`PhysicalSortExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalSortExprNode, + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result { + let expr = ctx.decode_required_expression( + node.expr.as_deref(), + "PhysicalSortExpr", + "expr", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + } +} + impl PartialEq for PhysicalSortExpr { fn eq(&self, other: &Self) -> bool { self.options == other.options && self.expr.eq(&other.expr) diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index 59d36c4efc1bb..9a24658411300 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -515,6 +515,146 @@ impl Partitioning { } } +/// Protobuf conversions for [`Partitioning`]. +/// +/// Child expressions (hash keys, range orderings) and `ScalarValue` split +/// points are (de)serialized through the expression-level context, so this is +/// the single copy of the partitioning wire format: `RepartitionExec`, +/// `FileScanConfig` and `datafusion-proto`'s central serializer all route +/// through it. +/// +/// [`protobuf::Partitioning`]: datafusion_proto_models::protobuf::Partitioning +#[cfg(feature = "proto")] +impl Partitioning { + /// Serialize this partitioning into its protobuf representation. + pub fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result { + use datafusion_proto_models::protobuf; + + let partition_method = match self { + Partitioning::RoundRobinBatch(n) => { + protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) + } + Partitioning::Hash(exprs, n) => { + protobuf::partitioning::PartitionMethod::Hash( + protobuf::PhysicalHashRepartition { + hash_expr: ctx.encode_children_expressions(exprs)?, + partition_count: *n as u64, + }, + ) + } + Partitioning::Range(range) => { + let sort_expr = range + .ordering() + .iter() + .map(|sort_expr| sort_expr.try_to_proto(ctx)) + .collect::>>()?; + let split_point = range + .split_points() + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>()?; + protobuf::partitioning::PartitionMethod::Range( + protobuf::PhysicalRangePartitioning { + sort_expr, + split_point, + }, + ) + } + Partitioning::UnknownPartitioning(n) => { + protobuf::partitioning::PartitionMethod::Unknown(*n as u64) + } + }; + Ok(protobuf::Partitioning { + partition_method: Some(partition_method), + }) + } + + /// Reconstruct a [`Partitioning`] from its protobuf representation. + /// + /// Returns `Ok(None)` when the message carries no `partition_method`, which + /// the wire format uses to mean "no output partitioning declared"; callers + /// for which it is required should turn that into their own error. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::Partitioning, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::{ScalarValue, internal_datafusion_err, internal_err}; + use datafusion_proto_models::protobuf; + + let Some(partition_method) = node.partition_method.as_ref() else { + return Ok(None); + }; + let partitioning = match partition_method { + protobuf::partitioning::PartitionMethod::RoundRobin(n) => { + Partitioning::RoundRobinBatch(partition_count(*n)?) + } + protobuf::partitioning::PartitionMethod::Hash(hash) => { + let exprs = hash + .hash_expr + .iter() + .map(|expr| ctx.decode(expr)) + .collect::>>()?; + Partitioning::Hash(exprs, partition_count(hash.partition_count)?) + } + protobuf::partitioning::PartitionMethod::Unknown(n) => { + Partitioning::UnknownPartitioning(partition_count(*n)?) + } + protobuf::partitioning::PartitionMethod::Range(range) => { + let sort_exprs = range + .sort_expr + .iter() + .map(|sort_expr| PhysicalSortExpr::try_from_proto(sort_expr, ctx)) + .collect::>>()?; + let sort_expr_count = sort_exprs.len(); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!( + "Range partitioning requires non-empty ordering" + ) + })?; + if ordering.len() != sort_expr_count { + return internal_err!( + "Range partitioning ordering must not contain duplicate expressions" + ); + } + let split_points = range + .split_point + .iter() + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>()?; + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) + } + }; + Ok(Some(partitioning)) + } +} + +/// Narrow a wire partition count to `usize`. +#[cfg(feature = "proto")] +fn partition_count(count: u64) -> Result { + usize::try_from(count).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Partition count {count} exceeds usize::MAX" + ) + }) +} + impl PartialEq for Partitioning { fn eq(&self, other: &Partitioning) -> bool { match (self, other) { diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs index 1731203f6c767..5883f1d1fec81 100644 --- a/datafusion/physical-plan/src/proto.rs +++ b/datafusion/physical-plan/src/proto.rs @@ -64,6 +64,12 @@ use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::TaskContext; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, +}; +use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, +}; use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use crate::ExecutionPlan; @@ -190,6 +196,25 @@ impl<'a> ExecutionPlanEncodeCtx<'a> { pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { self.encoder.encode_udwf(udwf) } + + /// An expression-level encode context backed by this plan context. + /// + /// Lets a plan hand `ctx` to expression-level conversions that own their own + /// wire logic — e.g. + /// [`Partitioning::try_to_proto`](datafusion_physical_expr::Partitioning::try_to_proto) + /// and + /// [`PhysicalSortExpr::try_to_proto`](datafusion_physical_expr::PhysicalSortExpr::try_to_proto). + pub fn expr_ctx(&self) -> PhysicalExprEncodeCtx<'_> { + PhysicalExprEncodeCtx::new(self) + } +} + +/// Lets [`ExecutionPlanEncodeCtx`] back an [`PhysicalExprEncodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprEncode for ExecutionPlanEncodeCtx<'_> { + fn encode(&self, expr: &Arc) -> Result { + self.encode_expr(expr) + } } /// Context handed to a plan's `try_from_proto` associated function. @@ -286,6 +311,28 @@ impl<'a> ExecutionPlanDecodeCtx<'a> { ) -> Result> { self.decoder.decode_udwf(name, payload) } + + /// An expression-level decode context backed by this plan context, bound to + /// `input_schema`. + /// + /// The decode counterpart of + /// [`ExecutionPlanEncodeCtx::expr_ctx`], for calling conversions such as + /// [`Partitioning::try_from_proto`](datafusion_physical_expr::Partitioning::try_from_proto). + pub fn expr_ctx<'s>(&'s self, input_schema: &'s Schema) -> PhysicalExprDecodeCtx<'s> { + PhysicalExprDecodeCtx::new(input_schema, self) + } +} + +/// Lets [`ExecutionPlanDecodeCtx`] back a [`PhysicalExprDecodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> { + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.decode_expr(node, schema) + } } /// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 3473aad9b3fc0..873f35fd6aed9 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1713,64 +1713,14 @@ impl ExecutionPlan for RepartitionExec { let input = ctx.encode_child(self.input())?; - // Keep the existing protobuf wire representation unchanged. - let partition_method = match self.partitioning() { - Partitioning::RoundRobinBatch(n) => { - protobuf::partitioning::PartitionMethod::RoundRobin(*n as u64) - } - Partitioning::Hash(exprs, n) => { - let hash_expr = ctx.encode_expressions(exprs)?; - protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr, - partition_count: *n as u64, - }, - ) - } - Partitioning::Range(range) => { - let sort_expr = range - .ordering() - .iter() - .map(|sort_expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), - asc: !sort_expr.options.descending, - nulls_first: sort_expr.options.nulls_first, - }) - }) - .collect::>>()?; - let split_point = range - .split_points() - .iter() - .map(|split_point| { - let value = split_point - .values() - .iter() - .map(|value| value.try_into().map_err(Into::into)) - .collect::>>()?; - Ok(protobuf::PhysicalRangeSplitPoint { value }) - }) - .collect::>>()?; - protobuf::partitioning::PartitionMethod::Range( - protobuf::PhysicalRangePartitioning { - sort_expr, - split_point, - }, - ) - } - Partitioning::UnknownPartitioning(n) => { - protobuf::partitioning::PartitionMethod::Unknown(*n as u64) - } - }; + let partitioning = self.partitioning().try_to_proto(&ctx.expr_ctx())?; Ok(Some(protobuf::PhysicalPlanNode { physical_plan_type: Some( protobuf::physical_plan_node::PhysicalPlanType::Repartition(Box::new( protobuf::RepartitionExecNode { input: Some(Box::new(input)), - partitioning: Some(protobuf::Partitioning { - partition_method: Some(partition_method), - }), + partitioning: Some(partitioning), preserve_order: self.preserve_order(), }, )), @@ -1800,84 +1750,23 @@ impl RepartitionExec { )?; let input_schema = input.schema(); - let partition_method = repart + let partitioning = repart .partitioning .as_ref() - .and_then(|p| p.partition_method.as_ref()) + .map(|partitioning| { + Partitioning::try_from_proto( + partitioning, + &ctx.expr_ctx(input_schema.as_ref()), + ) + }) + .transpose()? + .flatten() .ok_or_else(|| { datafusion_common::internal_datafusion_err!( "RepartitionExec is missing required field 'partitioning'" ) })?; - let partitioning = match partition_method { - protobuf::partitioning::PartitionMethod::RoundRobin(n) => { - Partitioning::RoundRobinBatch(*n as usize) - } - protobuf::partitioning::PartitionMethod::Hash(hash) => { - let exprs = hash - .hash_expr - .iter() - .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) - .collect::>>()?; - let partition_count = - usize::try_from(hash.partition_count).map_err(|_| { - datafusion_common::internal_datafusion_err!( - "Hash partition count {} exceeds usize::MAX", - hash.partition_count - ) - })?; - Partitioning::Hash(exprs, partition_count) - } - protobuf::partitioning::PartitionMethod::Unknown(n) => { - Partitioning::UnknownPartitioning(*n as usize) - } - protobuf::partitioning::PartitionMethod::Range(range) => { - let sort_exprs = range - .sort_expr - .iter() - .map(|sort_expr| { - let expr = sort_expr.expr.as_ref().ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "Unexpected empty physical expression" - ) - })?; - Ok(PhysicalSortExpr { - expr: ctx.decode_expr(expr, input_schema.as_ref())?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - }) - .collect::>>()?; - let sort_expr_count = sort_exprs.len(); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - datafusion_common::internal_datafusion_err!( - "Range partitioning requires non-empty ordering" - ) - })?; - if ordering.len() != sort_expr_count { - return datafusion_common::internal_err!( - "Range partitioning ordering must not contain duplicate expressions" - ); - } - let split_points = range - .split_point - .iter() - .map(|split_point| { - let values = split_point - .value - .iter() - .map(|value| ScalarValue::try_from(value).map_err(Into::into)) - .collect::>>()?; - Ok(SplitPoint::new(values)) - }) - .collect::>>()?; - Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) - } - }; - let mut repart_exec = RepartitionExec::try_new(input, partitioning)?; if repart.preserve_order { repart_exec = repart_exec.with_preserve_order(); diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 2ec77c00596c7..fcefe402b175a 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -23,9 +23,7 @@ use arrow::array::RecordBatch; use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; -use datafusion_common::{ - DataFusionError, Result, ScalarValue, internal_datafusion_err, not_impl_err, -}; +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; @@ -42,7 +40,7 @@ use datafusion_expr::dml::InsertOp; use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; 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, @@ -51,9 +49,7 @@ use datafusion_physical_plan::expressions::{ 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, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use datafusion_proto_common::common::proto_error; use super::{ @@ -423,83 +419,20 @@ pub fn parse_protobuf_partitioning( input_schema: &Schema, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - match partitioning { - Some(protobuf::Partitioning { partition_method }) => match partition_method { - Some(protobuf::partitioning::PartitionMethod::RoundRobin( - partition_count, - )) => Ok(Some(Partitioning::RoundRobinBatch( - *partition_count as usize, - ))), - Some(protobuf::partitioning::PartitionMethod::Hash(hash_repartition)) => { - parse_protobuf_hash_partitioning( - Some(hash_repartition), - ctx, - input_schema, - proto_converter, - ) - } - Some(protobuf::partitioning::PartitionMethod::Range(range_partitioning)) => { - Ok(Some(parse_protobuf_range_partitioning( - range_partitioning, - ctx, - input_schema, - proto_converter, - )?)) - } - Some(protobuf::partitioning::PartitionMethod::Unknown(partition_count)) => { - Ok(Some(Partitioning::UnknownPartitioning( - *partition_count as usize, - ))) - } - None => Ok(None), - }, - None => Ok(None), - } -} - -fn parse_protobuf_range_partitioning( - range_partitioning: &protobuf::PhysicalRangePartitioning, - ctx: &PhysicalPlanDecodeContext<'_>, - input_schema: &Schema, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - let sort_exprs = parse_physical_sort_exprs( - &range_partitioning.sort_expr, + let decoder = ConverterDecoder { ctx, - input_schema, proto_converter, - )?; - let sort_expr_count = sort_exprs.len(); - let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { - internal_datafusion_err!("Range partitioning requires non-empty ordering") - })?; - if ordering.len() != sort_expr_count { - return Err(internal_datafusion_err!( - "Range partitioning ordering must not contain duplicate expressions" - )); - } - let split_points = range_partitioning - .split_point - .iter() - .map(parse_protobuf_range_split_point) - .collect::>()?; - Ok(Partitioning::Range(RangePartitioning::try_new( - ordering, - split_points, - )?)) -} - -fn parse_protobuf_range_split_point( - split_point: &protobuf::PhysicalRangeSplitPoint, -) -> Result { - let values = split_point - .value - .iter() - .map(|value| ScalarValue::try_from(value).map_err(Into::into)) - .collect::>()?; - Ok(SplitPoint::new(values)) + }; + let decode_ctx = + datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::new( + input_schema, + &decoder, + ); + partitioning + .map(|partitioning| Partitioning::try_from_proto(partitioning, &decode_ctx)) + .transpose() + .map(Option::flatten) } - pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result> { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 612661dba8c50..04252176d2845 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -37,9 +37,7 @@ 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, RangePartitioning, SplitPoint, WindowExpr, -}; +use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use super::{ ConverterPlanEncoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, @@ -358,72 +356,14 @@ pub fn serialize_partitioning( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let serialized_partitioning = match partitioning { - Partitioning::RoundRobinBatch(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::RoundRobin( - *partition_count as u64, - )), - }, - Partitioning::Hash(exprs, partition_count) => { - let serialized_exprs = - serialize_physical_exprs(exprs, codec, proto_converter)?; - protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr: serialized_exprs, - partition_count: *partition_count as u64, - }, - )), - } - } - Partitioning::Range(range) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Range( - serialize_range_partitioning(range, codec, proto_converter)?, - )), - }, - Partitioning::UnknownPartitioning(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Unknown( - *partition_count as u64, - )), - }, + let encoder = ConverterEncoder { + codec, + proto_converter, }; - Ok(serialized_partitioning) -} - -fn serialize_range_partitioning( - range: &RangePartitioning, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - Ok(protobuf::PhysicalRangePartitioning { - sort_expr: serialize_physical_sort_exprs( - range.ordering().iter().cloned(), - codec, - proto_converter, - )?, - split_point: range - .split_points() - .iter() - .map(serialize_range_split_point) - .collect::>()?, - }) -} - -fn serialize_range_split_point( - split_point: &SplitPoint, -) -> Result { - Ok(protobuf::PhysicalRangeSplitPoint { - value: split_point - .values() - .iter() - .map(|value| { - TryInto::::try_into(value) - .map_err(Into::into) - }) - .collect::>()?, - }) + partitioning.try_to_proto( + &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx::new(&encoder), + ) } - /// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. impl TryFromProto<&PartitionedFile> for protobuf::PartitionedFile { type Error = DataFusionError;