From 68144f7ba7fd75f8deab25342aa34174979c7e9c Mon Sep 17 00:00:00 2001 From: Joy Bestourous Date: Fri, 21 Aug 2026 17:24:01 -0400 Subject: [PATCH] serde for metadata_size_hint --- datafusion/common/src/scalar/mod.rs | 2 +- datafusion/proto/proto/datafusion.proto | 1 + datafusion/proto/src/generated/pbjson.rs | 22 +++++++++ datafusion/proto/src/generated/prost.rs | 2 + .../proto/src/physical_plan/from_proto.rs | 2 +- .../proto/src/physical_plan/to_proto.rs | 1 + .../tests/cases/roundtrip_physical_plan.rs | 48 +++++++++++++++++++ 7 files changed, 76 insertions(+), 2 deletions(-) diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index 9059ae07e648f..85635b3383b7f 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -37,7 +37,7 @@ use crate::cast::{ as_decimal128_array, as_decimal256_array, as_dictionary_array, as_fixed_size_binary_array, as_fixed_size_list_array, }; -use crate::error::{DataFusionError, Result, _exec_err, _internal_err, _not_impl_err}; +use crate::error::{_exec_err, _internal_err, _not_impl_err, DataFusionError, Result}; use crate::hash_utils::create_hashes; use crate::utils::SingleRowListArrayBuilder; use arrow::array::{ diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto/proto/datafusion.proto index 1cdfe6d216e32..018c0e481dcb3 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto/proto/datafusion.proto @@ -1254,6 +1254,7 @@ message PartitionedFile { repeated datafusion_common.ScalarValue partition_values = 4; FileRange range = 5; datafusion_common.Statistics statistics = 6; + optional uint64 metadata_size_hint = 7; } message FileRange { diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto/src/generated/pbjson.rs index 6e09e9a797ea0..d70db307992db 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto/src/generated/pbjson.rs @@ -13064,6 +13064,9 @@ impl serde::Serialize for PartitionedFile { if self.statistics.is_some() { len += 1; } + if self.metadata_size_hint.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PartitionedFile", len)?; if !self.path.is_empty() { struct_ser.serialize_field("path", &self.path)?; @@ -13087,6 +13090,11 @@ impl serde::Serialize for PartitionedFile { if let Some(v) = self.statistics.as_ref() { struct_ser.serialize_field("statistics", v)?; } + if let Some(v) = self.metadata_size_hint.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("metadataSizeHint", ToString::to_string(&v).as_str())?; + } struct_ser.end() } } @@ -13105,6 +13113,8 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { "partitionValues", "range", "statistics", + "metadata_size_hint", + "metadataSizeHint", ]; #[allow(clippy::enum_variant_names)] @@ -13115,6 +13125,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { PartitionValues, Range, Statistics, + MetadataSizeHint, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -13142,6 +13153,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { "partitionValues" | "partition_values" => Ok(GeneratedField::PartitionValues), "range" => Ok(GeneratedField::Range), "statistics" => Ok(GeneratedField::Statistics), + "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -13167,6 +13179,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { let mut partition_values__ = None; let mut range__ = None; let mut statistics__ = None; + let mut metadata_size_hint__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Path => { @@ -13209,6 +13222,14 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { } statistics__ = map_.next_value()?; } + GeneratedField::MetadataSizeHint => { + if metadata_size_hint__.is_some() { + return Err(serde::de::Error::duplicate_field("metadataSizeHint")); + } + metadata_size_hint__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } } } Ok(PartitionedFile { @@ -13218,6 +13239,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { partition_values: partition_values__.unwrap_or_default(), range: range__, statistics: statistics__, + metadata_size_hint: metadata_size_hint__, }) } } diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto/src/generated/prost.rs index f5ec45da48f2a..81c5f738e6080 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto/src/generated/prost.rs @@ -1879,6 +1879,8 @@ pub struct PartitionedFile { pub range: ::core::option::Option, #[prost(message, optional, tag = "6")] pub statistics: ::core::option::Option, + #[prost(uint64, optional, tag = "7")] + pub metadata_size_hint: ::core::option::Option, } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct FileRange { diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 34fb5bb6ddc19..fa04e122c4812 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -568,7 +568,7 @@ impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { range: val.range.as_ref().map(|v| v.try_into()).transpose()?, statistics: val.statistics.as_ref().map(|v| v.try_into()).transpose()?, extensions: None, - metadata_size_hint: None, + metadata_size_hint: val.metadata_size_hint.map(|v| v as usize), }) } } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index 3f67842fe625c..fb6d3d4434fc9 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -450,6 +450,7 @@ impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile { .collect::, _>>()?, range: pf.range.as_ref().map(|r| r.try_into()).transpose()?, statistics: pf.statistics.as_ref().map(|s| s.into()), + metadata_size_hint: pf.metadata_size_hint.map(|v| v as u64), }) } } diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index a8ecb2d0749e6..124c43dcf4393 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -50,6 +50,7 @@ use datafusion::datasource::physical_plan::{ wrap_partition_type_in_dict, wrap_partition_value_in_dict, FileScanConfig, FileSinkConfig, ParquetSource, }; +use datafusion::datasource::source::DataSourceExec; use datafusion::execution::FunctionRegistry; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::functions_window::nth_value::nth_value_udwf; @@ -682,6 +683,53 @@ fn roundtrip_sort_preserve_partitioning() -> Result<()> { )) } +#[test] +fn roundtrip_partitioned_file_metadata_size_hint() -> Result<()> { + let mut file_group = + PartitionedFile::new("/path/to/part=0/file.parquet".to_string(), 1024) + .with_metadata_size_hint(123); + + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let source = Arc::new(ParquetSource::default()); + let scan_config = FileScanConfig { + object_store_url: ObjectStoreUrl::local_filesystem(), + file_groups: vec![vec![file_group]], + constraints: Constraints::empty(), + statistics: Statistics::new_unknown(&schema), + file_schema: schema, + projection: Some(vec![0, 1]), + limit: None, + table_partition_cols: vec![Field::new( + "part".to_string(), + wrap_partition_type_in_dict(DataType::Int16), + false, + )], + output_ordering: vec![], + file_compression_type: FileCompressionType::UNCOMPRESSED, + new_lines_in_values: false, + source, + }; + + let ctx = SessionContext::new(); + let reconstructed_plan = roundtrip_test_and_return( + scan_config.build(), + &ctx, + &DefaultPhysicalExtensionCodec {}, + )?; + let exec = reconstructed_plan + .as_any() + .downcast_ref::() + .expect("DataSourceExec not found"); + let scan_config = exec + .source() + .as_any() + .downcast_ref::() + .expect("FileScanConfig not found"); + assert_eq!(scan_config.file_groups[0][0].metadata_size_hint, Some(123)); + Ok(()) +} + #[test] fn roundtrip_coalesce_with_fetch() -> Result<()> { let field_a = Field::new("a", DataType::Boolean, false);