diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index b22fad9c0ebe4..fdb33c4fd2607 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1369,10 +1369,16 @@ message JoinOn { message EmptyExecNode { datafusion_common.Schema schema = 1; + // Number of output partitions. Absent (0) means a single partition, so that + // plans encoded before this field existed decode to the previous default. + uint32 partitions = 2; } message PlaceholderRowExecNode { datafusion_common.Schema schema = 1; + // Number of output partitions. Absent (0) means a single partition, so that + // plans encoded before this field existed decode to the previous default. + uint32 partitions = 2; } message ProjectionExecNode { diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index c334eac2f53e9..050d780d0860c 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -5857,10 +5857,16 @@ impl serde::Serialize for EmptyExecNode { if self.schema.is_some() { len += 1; } + if self.partitions != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.EmptyExecNode", len)?; if let Some(v) = self.schema.as_ref() { struct_ser.serialize_field("schema", v)?; } + if self.partitions != 0 { + struct_ser.serialize_field("partitions", &self.partitions)?; + } struct_ser.end() } } @@ -5872,11 +5878,13 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { { const FIELDS: &[&str] = &[ "schema", + "partitions", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Schema, + Partitions, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -5899,6 +5907,7 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { { match value { "schema" => Ok(GeneratedField::Schema), + "partitions" => Ok(GeneratedField::Partitions), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -5919,6 +5928,7 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { V: serde::de::MapAccess<'de>, { let mut schema__ = None; + let mut partitions__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Schema => { @@ -5927,10 +5937,19 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { } schema__ = map_.next_value()?; } + GeneratedField::Partitions => { + if partitions__.is_some() { + return Err(serde::de::Error::duplicate_field("partitions")); + } + partitions__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } } } Ok(EmptyExecNode { schema: schema__, + partitions: partitions__.unwrap_or_default(), }) } } @@ -22125,10 +22144,16 @@ impl serde::Serialize for PlaceholderRowExecNode { if self.schema.is_some() { len += 1; } + if self.partitions != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PlaceholderRowExecNode", len)?; if let Some(v) = self.schema.as_ref() { struct_ser.serialize_field("schema", v)?; } + if self.partitions != 0 { + struct_ser.serialize_field("partitions", &self.partitions)?; + } struct_ser.end() } } @@ -22140,11 +22165,13 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { { const FIELDS: &[&str] = &[ "schema", + "partitions", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Schema, + Partitions, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -22167,6 +22194,7 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { { match value { "schema" => Ok(GeneratedField::Schema), + "partitions" => Ok(GeneratedField::Partitions), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -22187,6 +22215,7 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { V: serde::de::MapAccess<'de>, { let mut schema__ = None; + let mut partitions__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Schema => { @@ -22195,10 +22224,19 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { } schema__ = map_.next_value()?; } + GeneratedField::Partitions => { + if partitions__.is_some() { + return Err(serde::de::Error::duplicate_field("partitions")); + } + partitions__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } } } Ok(PlaceholderRowExecNode { schema: schema__, + partitions: partitions__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index db51edfd5d9c2..7fb3f1575240f 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2068,11 +2068,19 @@ pub struct JoinOn { pub struct EmptyExecNode { #[prost(message, optional, tag = "1")] pub schema: ::core::option::Option, + /// Number of output partitions. Absent (0) means a single partition, so that + /// plans encoded before this field existed decode to the previous default. + #[prost(uint32, tag = "2")] + pub partitions: u32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PlaceholderRowExecNode { #[prost(message, optional, tag = "1")] pub schema: ::core::option::Option, + /// Number of output partitions. Absent (0) means a single partition, so that + /// plans encoded before this field existed decode to the previous default. + #[prost(uint32, tag = "2")] + pub partitions: u32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ProjectionExecNode { diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 0744a94dcebd1..f3d14ec53f394 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1710,7 +1710,10 @@ pub trait PhysicalPlanNodeExt: Sized { _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { let schema = Arc::new(convert_required!(empty.schema)?); - Ok(Arc::new(EmptyExec::new(schema))) + // A zero (absent) partition count comes from a plan encoded before the + // field existed, which always meant a single partition. + let partitions = empty.partitions.max(1) as usize; + Ok(Arc::new(EmptyExec::new(schema).with_partitions(partitions))) } fn try_into_placeholder_row_physical_plan( @@ -1719,7 +1722,12 @@ pub trait PhysicalPlanNodeExt: Sized { _ctx: &PhysicalPlanDecodeContext<'_>, ) -> Result> { let schema = Arc::new(convert_required!(placeholder.schema)?); - Ok(Arc::new(PlaceholderRowExec::new(schema))) + // A zero (absent) partition count comes from a plan encoded before the + // field existed, which always meant a single partition. + let partitions = placeholder.partitions.max(1) as usize; + Ok(Arc::new( + PlaceholderRowExec::new(schema).with_partitions(partitions), + )) } fn try_into_sort_physical_plan( @@ -3015,6 +3023,8 @@ pub trait PhysicalPlanNodeExt: Sized { Ok(protobuf::PhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::Empty(protobuf::EmptyExecNode { schema: Some(schema), + partitions: empty.properties().output_partitioning().partition_count() + as u32, })), }) } @@ -3028,6 +3038,8 @@ pub trait PhysicalPlanNodeExt: Sized { physical_plan_type: Some(PhysicalPlanType::PlaceholderRow( protobuf::PlaceholderRowExecNode { schema: Some(schema), + partitions: empty.properties().output_partitioning().partition_count() + as u32, }, )), }) diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs index 6ede6fc0e9ae3..1dc45803028eb 100644 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs @@ -88,9 +88,9 @@ use datafusion::physical_plan::windows::{ create_udwf_window_expr, }; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PhysicalExpr, PlanProperties, RangePartitioning, SendableRecordBatchStream, - SplitPoint, Statistics, displayable, + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, + Partitioning, PhysicalExpr, PlanProperties, RangePartitioning, + SendableRecordBatchStream, SplitPoint, Statistics, displayable, }; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion::scalar::ScalarValue; @@ -235,6 +235,58 @@ fn roundtrip_empty() -> Result<()> { roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty())))) } +#[test] +fn roundtrip_empty_with_partitions() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let plan = Arc::new(EmptyExec::new(Arc::new(Schema::empty())).with_partitions(4)); + let plan = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + assert_eq!(plan.output_partitioning().partition_count(), 4); + Ok(()) +} + +#[test] +fn roundtrip_placeholder_row_with_partitions() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let plan = + Arc::new(PlaceholderRowExec::new(Arc::new(Schema::empty())).with_partitions(4)); + let plan = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + assert_eq!(plan.output_partitioning().partition_count(), 4); + Ok(()) +} + +/// Plans encoded before `partitions` was added carry no value for it, which +/// decodes as zero and must be treated as the previous default of one. +#[test] +fn decode_empty_and_placeholder_row_without_partitions() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let schema: protobuf::Schema = (&Schema::empty()).try_into()?; + + for physical_plan_type in [ + protobuf::physical_plan_node::PhysicalPlanType::Empty(protobuf::EmptyExecNode { + schema: Some(schema.clone()), + partitions: 0, + }), + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( + protobuf::PlaceholderRowExecNode { + schema: Some(schema.clone()), + partitions: 0, + }, + ), + ] { + let node = protobuf::PhysicalPlanNode { + physical_plan_type: Some(physical_plan_type), + }; + let plan = node.try_into_physical_plan(ctx.task_ctx().as_ref(), &codec)?; + assert_eq!(plan.output_partitioning().partition_count(), 1); + } + Ok(()) +} + #[derive(Debug)] struct DowncastDelegatingExec { inner: Arc, diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md index f7d2745549c99..d92748082af34 100644 --- a/docs/source/library-user-guide/upgrading/55.0.0.md +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -591,3 +591,36 @@ metadata and schema fingerprint match. - Pass the current schema fingerprint to `CachedFileMetadata::is_valid_for`. See [PR #23201](https://github.com/apache/datafusion/pull/23201) for details. + +### `EmptyExecNode` and `PlaceholderRowExecNode` gained a `partitions` field + +The generated protobuf structs `EmptyExecNode` and `PlaceholderRowExecNode` +encoded only a schema, so the partition count set by `EmptyExec::with_partitions` +was silently dropped when a physical plan was serialized and deserialized: a plan +that reported `n` partitions before encoding reported `1` after. Both messages now +carry a `partitions` field that round-trips the count. + +**Who is affected:** + +- Users constructing `EmptyExecNode` or `PlaceholderRowExecNode` with an + exhaustive struct literal. + +**Migration guide:** + +Set the new field, or fill it from `Default`: + +```rust,ignore +// Before +EmptyExecNode { schema: Some(schema) } + +// After +EmptyExecNode { schema: Some(schema), partitions: 4 } +// or +EmptyExecNode { schema: Some(schema), ..Default::default() } +``` + +The wire format stays compatible in both directions. Plans encoded before this +field existed decode as a single partition, the previous default, and plans +encoded after it add a field that older readers ignore. + +See [PR #23643](https://github.com/apache/datafusion/pull/23643) for details.