From 82dbeda52fc935ee4bcb5602545cf56e83443e06 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 09:22:26 -0500 Subject: [PATCH 1/4] refactor(proto): move PartitionedFile / FileGroup serde into datafusion-datasource The protobuf conversions for the file-scan leaf types live in `datafusion-proto` today, as `TryFromProto` impls, because that is historically the only crate that could name both sides. That forces any other crate needing them (e.g. a `FileSource` serializing its own scan config, per #23494) to reimplement the same wire logic, which can then drift from the central serializer. Put the single copy next to the types that own it, in a new `datafusion_datasource::proto` module behind a new `proto` feature: - `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 existing callers are unaffected and the two sides cannot disagree. The wire format is unchanged. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 1 + datafusion/datasource/Cargo.toml | 5 + datafusion/datasource/src/mod.rs | 4 + datafusion/datasource/src/proto.rs | 218 ++++++++++++++++++ datafusion/proto/Cargo.toml | 2 +- .../proto/src/physical_plan/from_proto.rs | 52 +---- .../proto/src/physical_plan/to_proto.rs | 40 +--- 7 files changed, 245 insertions(+), 77 deletions(-) create mode 100644 datafusion/datasource/src/proto.rs diff --git a/Cargo.lock b/Cargo.lock index 451ff70b1cd48..9cd5d12d5bca4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1944,6 +1944,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..459ca436f365d 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,6 +34,10 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] +# Enables the protobuf conversions for the file-scan leaf types owned by this +# crate (`FileRange`, `PartitionedFile`, `FileGroup`). Off by default so +# consumers that never serialize plans pay nothing. +proto = ["dep:datafusion-proto-models"] [dependencies] arrow = { workspace = true } @@ -56,6 +60,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/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/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/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 7aa6376313c96..66a3929bf048c 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::{ DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, @@ -632,64 +629,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() { - // The wire format carries statistics for the full table schema (file + partition - // columns), so assign directly — `with_statistics` would append the partition - // column stats a second time. - pf.statistics = Some(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) } } @@ -812,6 +775,9 @@ impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprD mod tests { use super::*; use arrow::datatypes::{DataType, Field, Schema}; + use chrono::{TimeZone, Utc}; + use object_store::ObjectMeta; + use object_store::path::Path; #[test] fn partitioned_file_path_roundtrip_percent_encoded() { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index e13923dbb9519..15f3716ca61a1 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 9b2e67a8b8aef7e0be92bd1a48f79f41ef758282 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:32:48 -0500 Subject: [PATCH 2/4] refactor(proto): express the leaf conversions as TryFrom MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `TryFromProto` exists because `datafusion-proto` owns neither side of the conversions it hosts: with both the DataFusion type and the prost message foreign to it, `impl TryFrom for X` is rejected by the orphan rule, so a local trait was needed to say the same thing. Moving a conversion into the crate that owns the DataFusion type removes that constraint, and `&T` is `#[fundamental]`, so both directions are expressible with the standard trait: impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile // ok impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile // ok Use it for the three types this PR moves, instead of inherent `try_to_proto` / `try_from_proto` methods. The hook naming stays for conversions that need an encode/decode context (plans, expressions, scan configs) — the standard trait cannot carry that second argument, and the distinction now tells a reader whether a conversion recurses. The slice form is the one exception: `&[PartitionedFile]` is not a type this crate owns, so `protobuf::FileGroup`'s slice conversion stays a `TryFromProto` shim. Callers inside DataFusion go through `FileGroup`, whose impl lives next to the type. `datafusion-proto`'s `TryFromProto` impls remain as delegating shims, so downstream callers are unaffected. --- datafusion/datasource/src/proto.rs | 87 +++++++++++-------- .../proto/src/physical_plan/from_proto.rs | 14 +-- .../proto/src/physical_plan/to_proto.rs | 19 ++-- 3 files changed, 70 insertions(+), 50 deletions(-) diff --git a/datafusion/datasource/src/proto.rs b/datafusion/datasource/src/proto.rs index 2796aa224c5de..a79d5d4648ae8 100644 --- a/datafusion/datasource/src/proto.rs +++ b/datafusion/datasource/src/proto.rs @@ -24,7 +24,12 @@ //! 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`. +//! field is plain data or goes through `datafusion-proto-common`. That is why +//! they are plain [`TryFrom`] impls rather than the `try_to_proto(ctx)` / +//! `try_from_proto(node, ctx)` hooks used for plans, expressions and scan +//! configs: the standard trait can express a conversion that takes nothing but +//! the value, and the orphan rule allows it here because one side of each +//! conversion is a type this crate owns. use std::sync::Arc; @@ -37,17 +42,21 @@ 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 { +impl TryFrom<&FileRange> for protobuf::FileRange { + type Error = DataFusionError; + + fn try_from(range: &FileRange) -> Result { Ok(protobuf::FileRange { - start: self.start, - end: self.end, + start: range.start, + end: range.end, }) } +} - /// Reconstruct a [`FileRange`] from its protobuf representation. - pub fn try_from_proto(range: &protobuf::FileRange) -> Result { +impl TryFrom<&protobuf::FileRange> for FileRange { + type Error = DataFusionError; + + fn try_from(range: &protobuf::FileRange) -> Result { Ok(FileRange { start: range.start, end: range.end, @@ -55,40 +64,40 @@ impl FileRange { } } -impl PartitionedFile { - /// Serialize this file into its protobuf representation. - pub fn try_to_proto(&self) -> Result { - let last_modified = self.object_meta.last_modified; +impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile { + type Error = DataFusionError; + + fn try_from(file: &PartitionedFile) -> Result { + let last_modified = file.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: file .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, + path: file.object_meta.location.as_ref().to_owned(), + size: file.object_meta.size, last_modified_ns, - partition_values: self + partition_values: file .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()), + range: file.range.as_ref().map(TryInto::try_into).transpose()?, + statistics: file.statistics.as_ref().map(|s| s.as_ref().into()), }) } +} + +impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { + type Error = DataFusionError; - /// Reconstruct a [`PartitionedFile`] from its protobuf representation. - pub fn try_from_proto(file: &protobuf::PartitionedFile) -> Result { + fn try_from(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}") @@ -110,7 +119,7 @@ impl PartitionedFile { )); } if let Some(range) = file.range.as_ref() { - let range = FileRange::try_from_proto(range)?; + let range = FileRange::try_from(range)?; pf = pf.with_range(range.start, range.end); } if let Some(proto_stats) = file.statistics.as_ref() { @@ -120,25 +129,29 @@ impl PartitionedFile { } } -impl FileGroup { - /// Serialize this group into its protobuf representation. - pub fn try_to_proto(&self) -> Result { +impl TryFrom<&FileGroup> for protobuf::FileGroup { + type Error = DataFusionError; + + fn try_from(group: &FileGroup) -> Result { Ok(protobuf::FileGroup { - files: self + files: group .files() .iter() - .map(PartitionedFile::try_to_proto) + .map(TryInto::try_into) .collect::>>()?, }) } +} + +impl TryFrom<&protobuf::FileGroup> for FileGroup { + type Error = DataFusionError; - /// Reconstruct a [`FileGroup`] from its protobuf representation. - pub fn try_from_proto(group: &protobuf::FileGroup) -> Result { + fn try_from(group: &protobuf::FileGroup) -> Result { Ok(FileGroup::new( group .files .iter() - .map(PartitionedFile::try_from_proto) + .map(TryInto::try_into) .collect::>>()?, )) } @@ -166,7 +179,8 @@ mod tests { .with_arrow_schema(Arc::clone(&schema)) .with_statistics(Arc::new(Statistics::new_unknown(&schema))); - let decoded = PartitionedFile::try_from_proto(&pf.try_to_proto()?)?; + let encoded = protobuf::PartitionedFile::try_from(&pf)?; + let decoded = PartitionedFile::try_from(&encoded)?; assert_eq!(decoded.object_meta.location, pf.object_meta.location); assert_eq!(decoded.object_meta.size, pf.object_meta.size); @@ -192,7 +206,7 @@ mod tests { ..Default::default() }; - let err = PartitionedFile::try_from_proto(&proto).unwrap_err(); + let err = PartitionedFile::try_from(&proto).unwrap_err(); assert!( err.to_string().contains("Invalid object_store path"), "unexpected error: {err}" @@ -206,7 +220,8 @@ mod tests { PartitionedFile::new("b.parquet", 2), ]); - let decoded = FileGroup::try_from_proto(&group.try_to_proto()?)?; + let encoded = protobuf::FileGroup::try_from(&group)?; + let decoded = FileGroup::try_from(&encoded)?; assert_eq!(decoded.len(), 2); assert_eq!( diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 66a3929bf048c..34ad8c7a62fc7 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -629,30 +629,30 @@ pub fn parse_record_batches(buf: &[u8]) -> Result> { Ok(batches) } -/// Thin shim over [`PartitionedFile::try_from_proto`], which owns the wire logic. +/// Thin shim over `TryFrom<&protobuf::PartitionedFile>`, which owns the wire logic. impl TryFromProto<&protobuf::PartitionedFile> for PartitionedFile { type Error = DataFusionError; fn try_from_proto(val: &protobuf::PartitionedFile) -> Result { - PartitionedFile::try_from_proto(val) + PartitionedFile::try_from(val) } } -/// Thin shim over [`FileRange::try_from_proto`], which owns the wire logic. +/// Thin shim over `TryFrom<&protobuf::FileRange>`, which owns the wire logic. impl TryFromProto<&protobuf::FileRange> for FileRange { type Error = DataFusionError; fn try_from_proto(value: &protobuf::FileRange) -> Result { - FileRange::try_from_proto(value) + FileRange::try_from(value) } } -/// Thin shim over [`FileGroup::try_from_proto`], which owns the wire logic. +/// Thin shim over `TryFrom<&protobuf::FileGroup>`, which owns the wire logic. impl TryFromProto<&protobuf::FileGroup> for FileGroup { type Error = DataFusionError; fn try_from_proto(val: &protobuf::FileGroup) -> Result { - FileGroup::try_from_proto(val) + FileGroup::try_from(val) } } @@ -697,7 +697,7 @@ impl TryFromProto<&protobuf::FileSinkConfig> for FileSinkConfig { let file_group = FileGroup::new( conf.file_groups .iter() - .map(PartitionedFile::try_from_proto) + .map(TryInto::try_into) .collect::>>()?, ); let table_paths = conf diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 15f3716ca61a1..5189972f0e200 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -424,25 +424,30 @@ fn serialize_range_split_point( }) } -/// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. +/// Thin shim over `TryFrom<&PartitionedFile>`, which owns the wire logic. impl TryFromProto<&PartitionedFile> for protobuf::PartitionedFile { type Error = DataFusionError; fn try_from_proto(pf: &PartitionedFile) -> Result { - pf.try_to_proto() + pf.try_into() } } -/// Thin shim over [`FileRange::try_to_proto`], which owns the wire logic. +/// Thin shim over `TryFrom<&FileRange>`, which owns the wire logic. impl TryFromProto<&FileRange> for protobuf::FileRange { type Error = DataFusionError; fn try_from_proto(value: &FileRange) -> Result { - value.try_to_proto() + value.try_into() } } -/// Thin shim over [`PartitionedFile::try_to_proto`], which owns the wire logic. +/// Thin shim over `TryFrom<&PartitionedFile>`, which owns the wire logic. +/// +/// The slice form cannot be a `TryFrom` impl: the orphan rule only accepts a +/// type this crate owns, and `&[PartitionedFile]` is not one (`&FileGroup` is, +/// hence the impl next to the type). Callers inside DataFusion go through +/// `FileGroup`; this stays for downstream users of the published signature. impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { type Error = DataFusionError; @@ -450,7 +455,7 @@ impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { Ok(protobuf::FileGroup { files: gr .iter() - .map(PartitionedFile::try_to_proto) + .map(TryInto::try_into) .collect::>>()?, }) } @@ -464,7 +469,7 @@ pub fn serialize_file_scan_config( let file_groups = conf .file_groups .iter() - .map(|p| protobuf::FileGroup::try_from_proto(p.files())) + .map(TryInto::try_into) .collect::, _>>()?; let mut output_orderings = vec![]; From 8213bf901917739c6a7d1e0727dbe9ec824357ec Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:03:46 -0500 Subject: [PATCH 3/4] refactor(proto): keep the datasource proto module private The module holds nothing but `TryFrom` impls, and trait impls are registered globally by coherence, so they reach every downstream crate whether or not the module that writes them is reachable by path. `pub` therefore adds no capability: it only publishes an item-less module page and a path we would then owe compatibility to. If a later conversion in this family needs a named helper that `TryFrom` cannot express (an encode/decode context, per #23494), exporting the module again is additive. Co-Authored-By: Claude Opus 5 (1M context) --- datafusion/datasource/src/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index e718064c819da..e415b3e48a02a 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -44,7 +44,7 @@ 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; +mod proto; pub mod schema_adapter; pub mod sink; pub mod source; From 53cb606de5a01b065d655796296ecaa9298b9e48 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Thu, 30 Jul 2026 13:01:11 -0500 Subject: [PATCH 4/4] refactor(proto): carry #23999 into the moved PartitionedFile decode #23999 landed in `datafusion-proto` while this branch was open, fixing the duplicate-partition-statistics decode that #23998 described. The moved copy in `datafusion-datasource` predates it, so replaying the move on top of current main would have silently reverted the fix. Apply it to the copy that now owns the wire logic, so the delegating shim behaves exactly as `main` does today. Upstream's regression test, `partitioned_file_statistics_roundtrip_with_partition_values`, exercises the shim and therefore pins the moved impl. The round-trip test here asserted only that statistics survived decode, because on the old base they did not survive it faithfully. That caveat is gone: it now asserts the decoded statistics equal the originals, and that they span the full table schema. Co-Authored-By: Claude Opus 5 (1M context) --- datafusion/datasource/src/proto.rs | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/datafusion/datasource/src/proto.rs b/datafusion/datasource/src/proto.rs index a79d5d4648ae8..cf48a461655c7 100644 --- a/datafusion/datasource/src/proto.rs +++ b/datafusion/datasource/src/proto.rs @@ -123,7 +123,10 @@ impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { 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()?)); + // The wire format carries statistics for the full table schema (file + partition + // columns), so assign directly — `with_statistics` would append the partition + // column stats a second time. + pf.statistics = Some(Arc::new(proto_stats.try_into()?)); } Ok(pf) } @@ -191,11 +194,13 @@ mod tests { 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()); + // Statistics span the full table schema (file columns followed by one + // entry per partition column), and survive the round trip intact. + assert_eq!( + pf.statistics.as_ref().unwrap().column_statistics.len(), + schema.fields().len() + pf.partition_values.len() + ); + assert_eq!(decoded.statistics, pf.statistics); Ok(()) }