Skip to content

Commit cbebc6f

Browse files
Fix missing field partitioned_by_file_group in serialization (apache#22365)
I'm not super versed in the serialization machinery involved here, please review carefully. ## Which issue does this PR close? - Closes apache#22363. ## Rationale for this change The partitioned_by_file_group field was introduced in apache#21351 and apache#21342 but not added to the protobuf schema, breaking `datafusion-distributed`. ## What changes are included in this PR? - Add optional `bool partitioned_by_file_group = 14` to `FileScanExecConf` in `datafusion.proto` - Serialize the field in `to_proto.rs` - Deserialize the field in `from_proto.rs` - Regenerate prost/pbjson code ## Are these changes tested? Yes, added roundtrip_parquet_exec_partitioned_by_file_group test. ## Are there any user-facing changes? No
1 parent 971cf99 commit cbebc6f

6 files changed

Lines changed: 72 additions & 0 deletions

File tree

datafusion/proto-models/proto/datafusion.proto

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1120,6 +1120,7 @@ message FileScanExecConf {
11201120
optional uint64 batch_size = 12;
11211121

11221122
optional ProjectionExprs projection_exprs = 13;
1123+
optional bool partitioned_by_file_group = 14;
11231124
}
11241125

11251126
message ParquetScanExecNode {

datafusion/proto-models/src/generated/pbjson.rs

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6848,6 +6848,9 @@ impl serde::Serialize for FileScanExecConf {
68486848
if self.projection_exprs.is_some() {
68496849
len += 1;
68506850
}
6851+
if self.partitioned_by_file_group.is_some() {
6852+
len += 1;
6853+
}
68516854
let mut struct_ser = serializer.serialize_struct("datafusion.FileScanExecConf", len)?;
68526855
if !self.file_groups.is_empty() {
68536856
struct_ser.serialize_field("fileGroups", &self.file_groups)?;
@@ -6884,6 +6887,9 @@ impl serde::Serialize for FileScanExecConf {
68846887
if let Some(v) = self.projection_exprs.as_ref() {
68856888
struct_ser.serialize_field("projectionExprs", v)?;
68866889
}
6890+
if let Some(v) = self.partitioned_by_file_group.as_ref() {
6891+
struct_ser.serialize_field("partitionedByFileGroup", v)?;
6892+
}
68876893
struct_ser.end()
68886894
}
68896895
}
@@ -6911,6 +6917,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf {
69116917
"batchSize",
69126918
"projection_exprs",
69136919
"projectionExprs",
6920+
"partitioned_by_file_group",
6921+
"partitionedByFileGroup",
69146922
];
69156923

69166924
#[allow(clippy::enum_variant_names)]
@@ -6926,6 +6934,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf {
69266934
Constraints,
69276935
BatchSize,
69286936
ProjectionExprs,
6937+
PartitionedByFileGroup,
69296938
}
69306939
impl<'de> serde::Deserialize<'de> for GeneratedField {
69316940
fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -6958,6 +6967,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf {
69586967
"constraints" => Ok(GeneratedField::Constraints),
69596968
"batchSize" | "batch_size" => Ok(GeneratedField::BatchSize),
69606969
"projectionExprs" | "projection_exprs" => Ok(GeneratedField::ProjectionExprs),
6970+
"partitionedByFileGroup" | "partitioned_by_file_group" => Ok(GeneratedField::PartitionedByFileGroup),
69616971
_ => Err(serde::de::Error::unknown_field(value, FIELDS)),
69626972
}
69636973
}
@@ -6988,6 +6998,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf {
69886998
let mut constraints__ = None;
69896999
let mut batch_size__ = None;
69907000
let mut projection_exprs__ = None;
7001+
let mut partitioned_by_file_group__ = None;
69917002
while let Some(k) = map_.next_key()? {
69927003
match k {
69937004
GeneratedField::FileGroups => {
@@ -7061,6 +7072,12 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf {
70617072
}
70627073
projection_exprs__ = map_.next_value()?;
70637074
}
7075+
GeneratedField::PartitionedByFileGroup => {
7076+
if partitioned_by_file_group__.is_some() {
7077+
return Err(serde::de::Error::duplicate_field("partitionedByFileGroup"));
7078+
}
7079+
partitioned_by_file_group__ = map_.next_value()?;
7080+
}
70647081
}
70657082
}
70667083
Ok(FileScanExecConf {
@@ -7075,6 +7092,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf {
70757092
constraints: constraints__,
70767093
batch_size: batch_size__,
70777094
projection_exprs: projection_exprs__,
7095+
partitioned_by_file_group: partitioned_by_file_group__,
70787096
})
70797097
}
70807098
}

datafusion/proto-models/src/generated/prost.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1677,6 +1677,8 @@ pub struct FileScanExecConf {
16771677
pub batch_size: ::core::option::Option<u64>,
16781678
#[prost(message, optional, tag = "13")]
16791679
pub projection_exprs: ::core::option::Option<ProjectionExprs>,
1680+
#[prost(bool, optional, tag = "14")]
1681+
pub partitioned_by_file_group: ::core::option::Option<bool>,
16801682
}
16811683
#[derive(Clone, PartialEq, ::prost::Message)]
16821684
pub struct ParquetScanExecNode {

datafusion/proto/src/physical_plan/from_proto.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -711,6 +711,7 @@ pub fn parse_protobuf_file_scan_config(
711711
.with_limit(proto.limit.as_ref().map(|sl| sl.limit as usize))
712712
.with_output_ordering(output_ordering)
713713
.with_batch_size(proto.batch_size.map(|s| s as usize))
714+
.with_partitioned_by_file_group(proto.partitioned_by_file_group.unwrap_or(false))
714715
.build();
715716
Ok(config)
716717
}

datafusion/proto/src/physical_plan/to_proto.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,7 @@ pub fn serialize_file_scan_config(
757757
constraints: Some(conf.constraints.clone().into()),
758758
batch_size: conf.batch_size.map(|s| s as u64),
759759
projection_exprs,
760+
partitioned_by_file_group: Some(conf.partitioned_by_file_group),
760761
})
761762
}
762763

datafusion/proto/tests/cases/roundtrip_physical_plan.rs

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4067,5 +4067,54 @@ fn test_custom_node_with_dynamic_filter_dedup_roundtrip() -> Result<()> {
40674067
// rewrite can reconstruct the remapped form on the other side.
40684068
assert_dynamic_filters_equal(deser_custom_df, deser_filter_df);
40694069
assert_dynamic_filter_update_is_visible(deser_custom_df, deser_filter_df)?;
4070+
4071+
Ok(())
4072+
}
4073+
4074+
#[test]
4075+
fn roundtrip_parquet_exec_partitioned_by_file_group() -> Result<()> {
4076+
use datafusion::datasource::physical_plan::FileScanConfig;
4077+
4078+
let file_schema =
4079+
Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
4080+
let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema)));
4081+
let scan_config =
4082+
FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source)
4083+
.with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new(
4084+
"/path/to/file.parquet".to_string(),
4085+
1024,
4086+
)])])
4087+
.with_partitioned_by_file_group(true)
4088+
.build();
4089+
4090+
assert!(scan_config.partitioned_by_file_group);
4091+
4092+
let exec_plan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(scan_config);
4093+
4094+
let ctx = SessionContext::new();
4095+
let codec = DefaultPhysicalExtensionCodec {};
4096+
let proto_converter = DefaultPhysicalProtoConverter {};
4097+
let bytes = physical_plan_to_bytes_with_proto_converter(
4098+
Arc::clone(&exec_plan),
4099+
&codec,
4100+
&proto_converter,
4101+
)?;
4102+
let result_plan = physical_plan_from_bytes_with_proto_converter(
4103+
bytes.as_ref(),
4104+
ctx.task_ctx().as_ref(),
4105+
&codec,
4106+
&proto_converter,
4107+
)?;
4108+
4109+
let data_source_exec = result_plan
4110+
.downcast_ref::<DataSourceExec>()
4111+
.expect("Expected DataSourceExec");
4112+
let file_scan_config = data_source_exec
4113+
.data_source()
4114+
.downcast_ref::<FileScanConfig>()
4115+
.expect("Expected FileScanConfig");
4116+
4117+
assert!(file_scan_config.partitioned_by_file_group);
4118+
40704119
Ok(())
40714120
}

0 commit comments

Comments
 (0)