Skip to content

Commit 833e501

Browse files
buraksennadriangb
andauthored
fix(proto): prevent duplicate partition statistics on roundtrip (#23999)
## Which issue does this PR close? - Closes #23998 ## Rationale for this change Decoding a `PartitionedFile` re-appended partition-column statistics already present in protobuf. ## What changes are included in this PR? Assign decoded statistics directly. ## Are these changes tested? Yes new roundtrip test for this ## Are there any user-facing changes? no --------- Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>
1 parent 541caab commit 833e501

1 file changed

Lines changed: 27 additions & 3 deletions

File tree

datafusion/proto/src/physical_plan/from_proto.rs

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -660,7 +660,10 @@ impl TryFromProto<&protobuf::PartitionedFile> for PartitionedFile {
660660
pf = pf.with_range(file_range.start, file_range.end);
661661
}
662662
if let Some(proto_stats) = val.statistics.as_ref() {
663-
pf = pf.with_statistics(Arc::new(proto_stats.try_into()?));
663+
// The wire format carries statistics for the full table schema (file + partition
664+
// columns), so assign directly — `with_statistics` would append the partition
665+
// column stats a second time.
666+
pf.statistics = Some(Arc::new(proto_stats.try_into()?));
664667
}
665668
Ok(pf)
666669
}
@@ -807,8 +810,8 @@ impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprD
807810

808811
#[cfg(test)]
809812
mod tests {
810-
811813
use super::*;
814+
use arrow::datatypes::{DataType, Field, Schema};
812815

813816
#[test]
814817
fn partitioned_file_path_roundtrip_percent_encoded() {
@@ -833,7 +836,6 @@ mod tests {
833836

834837
#[test]
835838
fn partitioned_file_arrow_schema_roundtrip() {
836-
use arrow::datatypes::{DataType, Field, Schema};
837839
use std::collections::HashMap;
838840

839841
let arrow_schema = Arc::new(Schema::new_with_metadata(
@@ -858,6 +860,28 @@ mod tests {
858860
);
859861
}
860862

863+
#[test]
864+
fn partitioned_file_statistics_roundtrip_with_partition_values() {
865+
use datafusion_common::Statistics;
866+
let file_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
867+
let pf = PartitionedFile::new("foo/bar.parquet", 1234)
868+
.with_partition_values(vec![ScalarValue::from("2024-01-01")])
869+
.with_statistics(Arc::new(Statistics::new_unknown(&file_schema)));
870+
871+
// `statistics` covers the full table schema: file columns followed by one
872+
// entry per partition column.
873+
let expected_len = file_schema.fields().len() + pf.partition_values.len();
874+
assert_eq!(
875+
pf.statistics.as_ref().unwrap().column_statistics.len(),
876+
expected_len
877+
);
878+
879+
let proto = protobuf::PartitionedFile::try_from_proto(&pf).unwrap();
880+
let decoded = PartitionedFile::try_from_proto(&proto).unwrap();
881+
882+
assert_eq!(decoded.statistics, pf.statistics);
883+
}
884+
861885
#[test]
862886
fn partitioned_file_from_proto_invalid_path() {
863887
let proto = protobuf::PartitionedFile {

0 commit comments

Comments
 (0)