From 7e415c6f128ff79354f2f52955804a5dff3b4ef3 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 18 Jul 2026 12:23:28 +0530 Subject: [PATCH 1/2] 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 | 534 ++++++++++++++++++ datafusion/datasource/src/source.rs | 36 ++ datafusion/proto/Cargo.toml | 2 +- datafusion/proto/src/physical_plan/mod.rs | 447 +++++++++++++++ 8 files changed, 1179 insertions(+), 1 deletion(-) create mode 100644 datafusion/datasource/src/file_scan_config/proto.rs diff --git a/Cargo.lock b/Cargo.lock index f00c931f15032..2e6a622941699 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1936,6 +1936,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 91caabeee6a41..4223a60e645da 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, @@ -1167,6 +1173,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 { @@ -1558,12 +1576,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; @@ -1622,6 +1646,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..82264fe29bc07 --- /dev/null +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -0,0 +1,534 @@ +// 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, output_partitioning_from_partition_fields, +}; +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, + // Partition grouping is now encoded in `output_partitioning`; this + // legacy wire field is left unset (readers rely on + // `output_partitioning`). + partitioned_by_file_group: None, + 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)?; + let output_partitioning = match output_partitioning { + Some(output_partitioning) => Some(output_partitioning), + None if conf.partitioned_by_file_group.unwrap_or(false) => { + // Backward compatibility: older serialized plans used only + // `partitioned_by_file_group` to declare scan output partitioning. + let table_schema = Self::parse_table_schema_from_proto(conf)?; + output_partitioning_from_partition_fields( + &schema, + table_schema.table_partition_cols(), + file_groups.len(), + ) + } + None => None, + }; + + // 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 e5f8aa072db71..0275d6a78b42c 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -178,6 +178,453 @@ mod tests { ); } + 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_preserves_legacy_partitioning() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let mut proto = serde.encode(&test_config(None))?; + proto.partitioned_by_file_group = Some(true); + proto.output_partitioning = None; + + let decoded = serde.decode(&proto)?; + + match decoded.output_partitioning { + Some(Partitioning::Hash(exprs, partition_count)) => { + assert_eq!(partition_count, 2); + let column = exprs[0] + .downcast_ref::() + .expect("partition expression should be a column"); + assert_eq!(column.name(), "part"); + assert_eq!(column.index(), 2); + } + other => panic!("expected legacy hash partitioning, got {other:?}"), + } + + 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 bc9a8e5c0e09595bbc52ea77a0a614904a765616 Mon Sep 17 00:00:00 2001 From: Kumar Ujjawal Date: Sat, 18 Jul 2026 12:52:52 +0530 Subject: [PATCH 2/2] 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 82264fe29bc07..8ec6162239dda 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