add partitioning scheme for unresolved shuffle and shuffle reader exec

ballista/core/proto/ballista.proto

@@ -50,14 +50,15 @@ message ShuffleWriterExecNode {
 message UnresolvedShuffleExecNode {
   uint32 stage_id = 1;
   datafusion_common.Schema schema = 2;
-  uint32 output_partition_count = 4;
+  datafusion.Partitioning partitioning = 5;
 }
 
 message ShuffleReaderExecNode {
   repeated ShuffleReaderPartition partition = 1;
   datafusion_common.Schema schema = 2;
   // The stage to read from
   uint32 stage_id = 3;
+  datafusion.Partitioning partitioning = 4;
 }
 
 message ShuffleReaderPartition {

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -74,12 +74,11 @@ impl ShuffleReaderExec {
         stage_id: usize,
         partition: Vec<Vec<PartitionLocation>>,
         schema: SchemaRef,
+        partitioning: Partitioning,
     ) -> Result<Self> {
         let properties = PlanProperties::new(
             datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
-            // TODO partitioning may be known and could be populated here
-            // see https://github.com/apache/arrow-datafusion/issues/758
-            Partitioning::UnknownPartitioning(partition.len()),
+            partitioning,
             datafusion::physical_plan::ExecutionMode::Bounded,
         );
         Ok(Self {
@@ -134,6 +133,7 @@ impl ExecutionPlan for ShuffleReaderExec {
             self.stage_id,
             self.partition.clone(),
             self.schema.clone(),
+            self.properties().output_partitioning().clone(),
         )?))
     }
 
@@ -553,6 +553,7 @@ mod tests {
             input_stage_id,
             vec![partitions],
             Arc::new(schema),
+            Partitioning::UnknownPartitioning(4),
         )?;
         let mut stream = shuffle_reader_exec.execute(0, task_ctx)?;
         let batches = utils::collect_stream(&mut stream).await;

ballista/core/src/execution_plans/unresolved_shuffle.rs

@@ -46,22 +46,16 @@ pub struct UnresolvedShuffleExec {
 
 impl UnresolvedShuffleExec {
     /// Create a new UnresolvedShuffleExec
-    pub fn new(
-        stage_id: usize,
-        schema: SchemaRef,
-        output_partition_count: usize,
-    ) -> Self {
+    pub fn new(stage_id: usize, schema: SchemaRef, partitioning: Partitioning) -> Self {
         let properties = PlanProperties::new(
             datafusion::physical_expr::EquivalenceProperties::new(schema.clone()),
-            // TODO the output partition is known and should be populated here!
-            // see https://github.com/apache/arrow-datafusion/issues/758
-            Partitioning::UnknownPartitioning(output_partition_count),
+            partitioning,
             datafusion::physical_plan::ExecutionMode::Bounded,
         );
         Self {
             stage_id,
             schema,
-            output_partition_count,
+            output_partition_count: properties.partitioning.partition_count(),
             properties,
         }
     }
@@ -75,7 +69,11 @@ impl DisplayAs for UnresolvedShuffleExec {
     ) -> std::fmt::Result {
         match t {
             DisplayFormatType::Default | DisplayFormatType::Verbose => {
-                write!(f, "UnresolvedShuffleExec")
+                write!(
+                    f,
+                    "UnresolvedShuffleExec: {:?}",
+                    self.properties().output_partitioning()
+                )
             }
         }
     }

ballista/core/src/serde/generated/ballista.rs

@@ -42,8 +42,8 @@ pub struct UnresolvedShuffleExecNode {
     pub stage_id: u32,
     #[prost(message, optional, tag = "2")]
     pub schema: ::core::option::Option<::datafusion_proto_common::Schema>,
-    #[prost(uint32, tag = "4")]
-    pub output_partition_count: u32,
+    #[prost(message, optional, tag = "5")]
+    pub partitioning: ::core::option::Option<::datafusion_proto::protobuf::Partitioning>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ShuffleReaderExecNode {
@@ -54,6 +54,8 @@ pub struct ShuffleReaderExecNode {
     /// The stage to read from
     #[prost(uint32, tag = "3")]
     pub stage_id: u32,
+    #[prost(message, optional, tag = "4")]
+    pub partitioning: ::core::option::Option<::datafusion_proto::protobuf::Partitioning>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ShuffleReaderPartition {

ballista/core/src/serde/mod.rs

@@ -21,6 +21,7 @@
 use crate::{error::BallistaError, serde::scheduler::Action as BallistaAction};
 
 use arrow_flight::sql::ProstMessageExt;
+use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::{DataFusionError, Result};
 use datafusion::execution::FunctionRegistry;
 use datafusion::physical_plan::{ExecutionPlan, Partitioning};
@@ -29,6 +30,8 @@ use datafusion_proto::logical_plan::file_formats::{
     JsonLogicalExtensionCodec, ParquetLogicalExtensionCodec,
 };
 use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning;
+use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning;
+use datafusion_proto::physical_plan::to_proto::serialize_partitioning;
 use datafusion_proto::protobuf::proto_error;
 use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
 use datafusion_proto::{
@@ -291,8 +294,11 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                 )?))
             }
             PhysicalPlanType::ShuffleReader(shuffle_reader) => {
+                let default_codec =
+                    datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec {};
                 let stage_id = shuffle_reader.stage_id as usize;
-                let schema = Arc::new(convert_required!(shuffle_reader.schema)?);
+                let schema: SchemaRef =
+                    Arc::new(convert_required!(shuffle_reader.schema)?);
                 let partition_location: Vec<Vec<PartitionLocation>> = shuffle_reader
                     .partition
                     .iter()
@@ -309,16 +315,39 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                             .collect::<Result<Vec<_>, _>>()
                     })
                     .collect::<Result<Vec<_>, DataFusionError>>()?;
-                let shuffle_reader =
-                    ShuffleReaderExec::try_new(stage_id, partition_location, schema)?;
+                let partitioning = parse_protobuf_partitioning(
+                    shuffle_reader.partitioning.as_ref(),
+                    registry,
+                    schema.as_ref(),
+                    &default_codec,
+                )?;
+                let partitioning = partitioning
+                    .ok_or_else(|| proto_error("missing required partitioning field"))?;
+                let shuffle_reader = ShuffleReaderExec::try_new(
+                    stage_id,
+                    partition_location,
+                    schema,
+                    partitioning,
+                )?;
                 Ok(Arc::new(shuffle_reader))
             }
             PhysicalPlanType::UnresolvedShuffle(unresolved_shuffle) => {
-                let schema = Arc::new(convert_required!(unresolved_shuffle.schema)?);
+                let default_codec =
+                    datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec {};
+                let schema: SchemaRef =
+                    Arc::new(convert_required!(unresolved_shuffle.schema)?);
+                let partitioning = parse_protobuf_partitioning(
+                    unresolved_shuffle.partitioning.as_ref(),
+                    registry,
+                    schema.as_ref(),
+                    &default_codec,
+                )?;
+                let partitioning = partitioning
+                    .ok_or_else(|| proto_error("missing required partitioning field"))?;
                 Ok(Arc::new(UnresolvedShuffleExec::new(
                     unresolved_shuffle.stage_id as usize,
                     schema,
-                    unresolved_shuffle.output_partition_count as usize,
+                    partitioning,
                 )))
             }
         }
@@ -387,12 +416,17 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
                         .collect::<Result<Vec<_>, _>>()?,
                 });
             }
+            let default_codec =
+                datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec {};
+            let partitioning =
+                serialize_partitioning(&exec.properties().partitioning, &default_codec)?;
             let proto = protobuf::BallistaPhysicalPlanNode {
                 physical_plan_type: Some(PhysicalPlanType::ShuffleReader(
                     protobuf::ShuffleReaderExecNode {
                         stage_id,
                         partition,
                         schema: Some(exec.schema().as_ref().try_into()?),
+                        partitioning: Some(partitioning),
                     },
                 )),
             };
@@ -404,12 +438,16 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec {
 
             Ok(())
         } else if let Some(exec) = node.as_any().downcast_ref::<UnresolvedShuffleExec>() {
+            let default_codec =
+                datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec {};
+            let partitioning =
+                serialize_partitioning(&exec.properties().partitioning, &default_codec)?;
             let proto = protobuf::BallistaPhysicalPlanNode {
                 physical_plan_type: Some(PhysicalPlanType::UnresolvedShuffle(
                     protobuf::UnresolvedShuffleExecNode {
                         stage_id: exec.stage_id as u32,
                         schema: Some(exec.schema().as_ref().try_into()?),
-                        output_partition_count: exec.output_partition_count as u32,
+                        partitioning: Some(partitioning),
                     },
                 )),
             };

ballista/scheduler/src/planner.rs

@@ -168,10 +168,7 @@ fn create_unresolved_shuffle(
     Arc::new(UnresolvedShuffleExec::new(
         shuffle_writer.stage_id(),
         shuffle_writer.schema(),
-        shuffle_writer
-            .properties()
-            .output_partitioning()
-            .partition_count(),
+        shuffle_writer.properties().output_partitioning().clone(),
     ))
 }
 
@@ -239,6 +236,10 @@ pub fn remove_unresolved_shuffles(
                 unresolved_shuffle.stage_id,
                 relevant_locations,
                 unresolved_shuffle.schema().clone(),
+                unresolved_shuffle
+                    .properties()
+                    .output_partitioning()
+                    .clone(),
             )?))
         } else {
             new_children.push(remove_unresolved_shuffles(
@@ -259,16 +260,12 @@ pub fn rollback_resolved_shuffles(
     let mut new_children: Vec<Arc<dyn ExecutionPlan>> = vec![];
     for child in stage.children() {
         if let Some(shuffle_reader) = child.as_any().downcast_ref::<ShuffleReaderExec>() {
-            let output_partition_count = shuffle_reader
-                .properties()
-                .output_partitioning()
-                .partition_count();
             let stage_id = shuffle_reader.stage_id;
 
             let unresolved_shuffle = Arc::new(UnresolvedShuffleExec::new(
                 stage_id,
                 shuffle_reader.schema(),
-                output_partition_count,
+                shuffle_reader.properties().partitioning.clone(),
             ));
             new_children.push(unresolved_shuffle);
         } else {