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..e415b3e48a02a 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")] +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..cf48a461655c7 --- /dev/null +++ b/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 { + 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 { + Ok(FileRange { + start: range.start, + end: range.end, + }) + } +} + +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: 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::, _>>()?, + 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 { + 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(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 { + Ok(protobuf::FileGroup { + files: group + .files() + .iter() + .map(TryInto::try_into) + .collect::>>()?, + }) + } +} + +impl TryFrom<&protobuf::FileGroup> for FileGroup { + type Error = DataFusionError; + + fn try_from(group: &protobuf::FileGroup) -> Result { + Ok(FileGroup::new( + group + .files + .iter() + .map(TryInto::try_into) + .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 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(()) + } +} 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..34ad8c7a62fc7 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 `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 { - 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(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 { - 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 { - let files = val - .files - .iter() - .map(PartitionedFile::try_from_proto) - .collect::, _>>()?; - 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::>>()?, ); 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() { diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index e13923dbb9519..5189972f0e200 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -424,51 +424,30 @@ fn serialize_range_split_point( }) } +/// 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 { - 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_into() } } +/// 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 { - Ok(protobuf::FileRange { - start: value.start, - end: value.end, - }) + value.try_into() } } +/// 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; @@ -476,8 +455,8 @@ impl TryFromProto<&[PartitionedFile]> for protobuf::FileGroup { Ok(protobuf::FileGroup { files: gr .iter() - .map(protobuf::PartitionedFile::try_from_proto) - .collect::, _>>()?, + .map(TryInto::try_into) + .collect::>>()?, }) } } @@ -490,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![];