refactor(proto): move PartitionedFile / FileGroup serde into datafusion-datasource

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 }

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")]
+mod proto;
 pub mod schema_adapter;
 pub mod sink;
 pub mod source;

datafusion/datasource/src/proto.rs

@@ -0,0 +1,238 @@
+// 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`. 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;
+
+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 TryFrom<&FileRange> for protobuf::FileRange {
+    type Error = DataFusionError;
+
+    fn try_from(range: &FileRange) -> Result<Self> {
+        Ok(protobuf::FileRange {
+            start: range.start,
+            end: range.end,
+        })
+    }
+}
+
+impl TryFrom<&protobuf::FileRange> for FileRange {
+    type Error = DataFusionError;
+
+    fn try_from(range: &protobuf::FileRange) -> Result<Self> {
+        Ok(FileRange {
+            start: range.start,
+            end: range.end,
+        })
+    }
+}
+
+impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile {
+    type Error = DataFusionError;
+
+    fn try_from(file: &PartitionedFile) -> Result<Self> {
+        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: file
+                .arrow_schema
+                .as_ref()
+                .map(|s| s.as_ref().try_into())
+                .transpose()?,
+            path: file.object_meta.location.as_ref().to_owned(),
+            size: file.object_meta.size,
+            last_modified_ns,
+            partition_values: file
+                .partition_values
+                .iter()
+                .map(|v| v.try_into())
+                .collect::<Result<Vec<_>, _>>()?,
+            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;
+
+    fn try_from(file: &protobuf::PartitionedFile) -> Result<Self> {
+        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::<Result<Vec<_>, _>>()?,
+        );
+        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(range)?;
+            pf = pf.with_range(range.start, range.end);
+        }
+        if let Some(proto_stats) = file.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)
+    }
+}
+
+impl TryFrom<&FileGroup> for protobuf::FileGroup {
+    type Error = DataFusionError;
+
+    fn try_from(group: &FileGroup) -> Result<Self> {
+        Ok(protobuf::FileGroup {
+            files: group
+                .files()
+                .iter()
+                .map(TryInto::try_into)
+                .collect::<Result<Vec<_>>>()?,
+        })
+    }
+}
+
+impl TryFrom<&protobuf::FileGroup> for FileGroup {
+    type Error = DataFusionError;
+
+    fn try_from(group: &protobuf::FileGroup) -> Result<Self> {
+        Ok(FileGroup::new(
+            group
+                .files
+                .iter()
+                .map(TryInto::try_into)
+                .collect::<Result<Vec<_>>>()?,
+        ))
+    }
+}
+
+#[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 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);
+        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 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(())
+    }
+
+    #[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).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 encoded = protobuf::FileGroup::try_from(&group)?;
+        let decoded = FileGroup::try_from(&encoded)?;
+
+        assert_eq!(decoded.len(), 2);
+        assert_eq!(
+            decoded.files()[1].object_meta.location,
+            group.files()[1].object_meta.location
+        );
+        Ok(())
+    }
+}

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 }

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<Vec<RecordBatch>> {
     Ok(batches)
 }
 
+/// 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<Self, Self::Error> {
-        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::<Result<Vec<_>, _>>()?,
-        );
-        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(val)
     }
 }
 
+/// 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<Self, Self::Error> {
-        Ok(FileRange {
-            start: value.start,
-            end: value.end,
-        })
+        FileRange::try_from(value)
     }
 }
 
+/// 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<Self, Self::Error> {
-        let files = val
-            .files
-            .iter()
-            .map(PartitionedFile::try_from_proto)
-            .collect::<Result<Vec<_>, _>>()?;
-        Ok(FileGroup::new(files))
+        FileGroup::try_from(val)
     }
 }
 
@@ -734,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::<Result<Vec<_>>>()?,
         );
         let table_paths = conf
@@ -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() {