Improve stage encoding size by removing unrelated partition location

[milenkovicm]

Is your feature request related to a problem or challenge? Please describe what you are trying to do.

ShuffleReaderExec holds information about each partition (PartitionLocation) even this information is not required for a specific task (partition). This will make plan very big when partition number is big, making serialistaion hitting grpc message size (#1853 (comment))

Currently, once serialised stage plan is cached (and reused for each task), this functionallity could be disabled using disable-stage-plan-cache scheduler feature, but even disabled it will still encode information about all partitions

Describe the solution you'd like

Implement stage encoding logic which would take into account actual partition and remove unrelated data

Describe alternatives you've considered

Additional context

This work is orthogonal to #1853

Scheduler feature disable-stage-plan-cache might disable plan caching but current implementation will not take into account actual partition information,

I had a quick experiment with

diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml
index f882d26f..3a60d897 100644
--- a/ballista/scheduler/Cargo.toml
+++ b/ballista/scheduler/Cargo.toml
@@ -34,7 +34,7 @@ required-features = ["build-binary"]
 
 [features]
 build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary"]
-default = ["build-binary", "rest-api"]
+default = ["build-binary", "rest-api", "disable-stage-plan-cache"]
 # job info can cache stage plans, in some cases where
 # task plans can be re-computed, cache behavior may need to be disabled.
 disable-stage-plan-cache = []
diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs
index f04c10e6..255aa27e 100644
--- a/ballista/scheduler/src/state/task_manager.rs
+++ b/ballista/scheduler/src/state/task_manager.rs
@@ -167,6 +167,7 @@ impl JobInfoCache {
     fn encode_stage_plan<U: AsExecutionPlan>(
         &mut self,
         stage_id: usize,
+        partition_id: usize,
         plan: &Arc<dyn ExecutionPlan>,
         codec: &dyn PhysicalExtensionCodec,
     ) -> Result<Vec<u8>> {
@@ -186,13 +187,57 @@ impl JobInfoCache {
     fn encode_stage_plan<U: AsExecutionPlan>(
         &mut self,
         _stage_id: usize,
+        partition_id: usize,
         plan: &Arc<dyn ExecutionPlan>,
         codec: &dyn PhysicalExtensionCodec,
     ) -> Result<Vec<u8>> {
+        use ballista_core::execution_plans::ShuffleReaderExec;
+        use datafusion::common::tree_node::{Transformed, TreeNode};
+
+        let transformed = plan.clone().transform_up(|node| {
+            if let Some(reader) = node.as_any().downcast_ref::<ShuffleReaderExec>() {
+                if reader.broadcast {
+                    return Ok(Transformed::no(node));
+                }
+                let new_partition: Vec<Vec<_>> = reader
+                    .partition
+                    .iter()
+                    .enumerate()
+                    .map(|(i, locs)| {
+                        if i == partition_id || locs.len() <= 1 {
+                            locs.clone()
+                        } else {
+                            vec![]
+                        }
+                    })
+                    .collect();
+                let new_reader = if let Some(coalesce) = reader.coalesce.clone() {
+                    ShuffleReaderExec::try_new_coalesced(
+                        reader.stage_id,
+                        new_partition,
+                        coalesce,
+                        reader.schema(),
+                        reader.properties().output_partitioning().clone(),
+                    )?
+                } else {
+                    ShuffleReaderExec::try_new(
+                        reader.stage_id,
+                        new_partition,
+                        reader.schema(),
+                        reader.properties().output_partitioning().clone(),
+                    )?
+                };
+                Ok(Transformed::yes(
+                    Arc::new(new_reader) as Arc<dyn ExecutionPlan>
+                ))
+            } else {
+                Ok(Transformed::no(node))
+            }
+        })?;
+
         let mut plan_buf: Vec<u8> = vec![];
-        let plan_proto = U::try_from_physical_plan(plan.clone(), codec)?;
+        let plan_proto = U::try_from_physical_plan(transformed.data, codec)?;
         plan_proto.try_encode(&mut plan_buf)?;
-
         Ok(plan_buf)
     }
 }
@@ -656,6 +701,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
         if let Some(mut job_info) = self.active_job_cache.get_mut(&job_id) {
             let plan = job_info.encode_stage_plan::<PhysicalPlanNode>(
                 stage_id,
+                task.partition.partition_id,
                 &task.plan,
                 self.codec.physical_extension_codec(),
             )?;
@@ -733,6 +779,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
             if let Some(mut job_info) = self.active_job_cache.get_mut(&job_id) {
                 let plan = job_info.encode_stage_plan::<PhysicalPlanNode>(
                     stage_id,
+                    task.partition.partition_id,
                     &task.plan,
                     self.codec.physical_extension_codec(),
                 )?;

and there is no encoding exception as seen in #1853 (comment) but for some reason it returns 0 result, which i believe its due to last stage for which, my quick change, does not work

[milenkovicm]

cc @coderfender interested ?

[coderfender]

Sure !

[milenkovicm]

the code is a quick hack to illustrate the point, it might not be correct

[milenkovicm]

Cc @andygrove

This issue might be relevant to TPCH SF1K