Add task_count to context

Author: gabotechs

src/execution_plans/partition_isolator.rs

@@ -1,4 +1,3 @@
-use crate::StageExec;
 use crate::distributed_physical_optimizer_rule::limit_tasks_err;
 use crate::execution_plans::DistributedTaskContext;
 use datafusion::common::{exec_err, plan_err};
@@ -191,12 +190,14 @@ impl ExecutionPlan for PartitionIsolatorExec {
         };
 
         let task_context = DistributedTaskContext::from_ctx(&context);
-        let stage = StageExec::from_ctx(&context)?;
 
         let input_partitions = self_ready.input.output_partitioning().partition_count();
 
-        let partition_group =
-            Self::partition_group(input_partitions, task_context.task_index, stage.tasks.len());
+        let partition_group = Self::partition_group(
+            input_partitions,
+            task_context.task_index,
+            task_context.task_count,
+        );
 
         // if our partition group is [7,8,9] and we are asked for parittion 1,
         // then look up that index in our group and execute that partition, in this

src/execution_plans/stage.rs

@@ -149,13 +149,10 @@ pub struct ExecutionTask {
 #[derive(Debug, Clone, Default)]
 pub struct DistributedTaskContext {
     pub task_index: usize,
+    pub task_count: usize,
 }
 
 impl DistributedTaskContext {
-    pub fn new(task_index: usize) -> Self {
-        Self { task_index }
-    }
-
     pub fn from_ctx(ctx: &Arc<TaskContext>) -> Arc<Self> {
         ctx.session_config()
             .get_extension::<Self>()
@@ -335,7 +332,10 @@ impl ExecutionPlan for StageExec {
             .session_config()
             .clone()
             .with_extension(assigned_stage.clone())
-            .with_extension(Arc::new(DistributedTaskContext { task_index: 0 }));
+            .with_extension(Arc::new(DistributedTaskContext {
+                task_index: 0,
+                task_count: 1,
+            }));
 
         let new_ctx =
             SessionContext::new_with_config_rt(config, context.runtime_env().clone()).task_ctx();

src/flight_service/do_get.rs

@@ -113,9 +113,10 @@ impl ArrowFlightEndpoint {
         let cfg = session_state.config_mut();
         cfg.set_extension(Arc::clone(&stage));
         cfg.set_extension(Arc::new(ContextGrpcMetadata(metadata.into_headers())));
-        cfg.set_extension(Arc::new(DistributedTaskContext::new(
-            doget.target_task_index as usize,
-        )));
+        cfg.set_extension(Arc::new(DistributedTaskContext {
+            task_index: doget.target_task_index as usize,
+            task_count: stage.tasks.len(),
+        }));
 
         let partition_count = stage.plan.properties().partitioning.partition_count();
         let target_partition = doget.target_partition as usize;