chore: Make query stage / shuffle code easier to understand (#54)

Author: andygrove

datafusion_ray/context.py

@@ -50,7 +50,7 @@ def execute_query_stage(
 
     # if the query stage has a single output partition then we need to execute for the output
     # partition, otherwise we need to execute in parallel for each input partition
-    concurrency = stage.get_input_partition_count()
+    concurrency = stage.get_execution_partition_count()
     output_partitions_count = stage.get_output_partition_count()
     if output_partitions_count == 1:
         # reduce stage
@@ -159,5 +159,6 @@ def plan(self, execution_plan: Any) -> List[pa.RecordBatch]:
         )
         _, partitions = ray.get(future)
         # assert len(partitions) == 1, len(partitions)
-        result_set = ray.get(partitions[0])
-        return result_set
+        record_batches = ray.get(partitions[0])
+        # filter out empty batches
+        return [batch for batch in record_batches if batch.num_rows > 0]

src/planner.rs

@@ -399,7 +399,7 @@ mod test {
             let query_stage = graph.query_stages.get(&id).unwrap();
             output.push_str(&format!(
                 "Query Stage #{id} ({} -> {}):\n{}\n",
-                query_stage.get_input_partition_count(),
+                query_stage.get_execution_partition_count(),
                 query_stage.get_output_partition_count(),
                 displayable(query_stage.plan.as_ref()).indent(false)
             ));

src/query_stage.rs

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::context::serialize_execution_plan;
-use crate::shuffle::{ShuffleCodec, ShuffleReaderExec};
+use crate::shuffle::{ShuffleCodec, ShuffleReaderExec, ShuffleWriterExec};
 use datafusion::error::Result;
 use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, Partitioning};
 use datafusion::prelude::SessionContext;
@@ -60,8 +60,8 @@ impl PyQueryStage {
         self.stage.get_child_stage_ids()
     }
 
-    pub fn get_input_partition_count(&self) -> usize {
-        self.stage.get_input_partition_count()
+    pub fn get_execution_partition_count(&self) -> usize {
+        self.stage.get_execution_partition_count()
     }
 
     pub fn get_output_partition_count(&self) -> usize {
@@ -75,16 +75,6 @@ pub struct QueryStage {
     pub plan: Arc<dyn ExecutionPlan>,
 }
 
-fn _get_output_partition_count(plan: &dyn ExecutionPlan) -> usize {
-    // UnknownPartitioning and HashPartitioning with empty expressions will
-    // both return 1 partition.
-    match plan.properties().output_partitioning() {
-        Partitioning::UnknownPartitioning(_) => 1,
-        Partitioning::Hash(expr, _) if expr.is_empty() => 1,
-        p => p.partition_count(),
-    }
-}
-
 impl QueryStage {
     pub fn new(id: usize, plan: Arc<dyn ExecutionPlan>) -> Self {
         Self { id, plan }
@@ -96,21 +86,27 @@ impl QueryStage {
         ids
     }
 
-    /// Get the input partition count. This is the same as the number of concurrent tasks
-    /// when we schedule this query stage for execution
-    pub fn get_input_partition_count(&self) -> usize {
-        if self.plan.children().is_empty() {
-            // leaf node (file scan)
-            self.plan.output_partitioning().partition_count()
+    /// Get the number of partitions that can be executed in parallel
+    pub fn get_execution_partition_count(&self) -> usize {
+        if let Some(shuffle) = self.plan.as_any().downcast_ref::<ShuffleWriterExec>() {
+            // use the partitioning of the input to the shuffle write because we are
+            // really executing that and then using the shuffle writer to repartition
+            // the output
+            shuffle.input_plan.output_partitioning().partition_count()
         } else {
-            self.plan.children()[0]
-                .output_partitioning()
-                .partition_count()
+            // for any other plan, use its output partitioning
+            self.plan.output_partitioning().partition_count()
         }
     }
 
     pub fn get_output_partition_count(&self) -> usize {
-        _get_output_partition_count(self.plan.as_ref())
+        // UnknownPartitioning and HashPartitioning with empty expressions will
+        // both return 1 partition.
+        match self.plan.properties().output_partitioning() {
+            Partitioning::UnknownPartitioning(_) => 1,
+            Partitioning::Hash(expr, _) if expr.is_empty() => 1,
+            p => p.partition_count(),
+        }
     }
 }
 

src/shuffle/codec.rs

@@ -102,7 +102,7 @@ impl PhysicalExtensionCodec for ShuffleCodec {
             };
             PlanType::ShuffleReader(reader)
         } else if let Some(writer) = node.as_any().downcast_ref::<ShuffleWriterExec>() {
-            let plan = PhysicalPlanNode::try_from_physical_plan(writer.plan.clone(), self)?;
+            let plan = PhysicalPlanNode::try_from_physical_plan(writer.input_plan.clone(), self)?;
             let partitioning =
                 encode_partitioning_scheme(writer.properties().output_partitioning())?;
             let writer = ShuffleWriterExecNode {

src/shuffle/writer.rs

@@ -47,7 +47,7 @@ use std::sync::Arc;
 #[derive(Debug)]
 pub struct ShuffleWriterExec {
     pub stage_id: usize,
-    pub(crate) plan: Arc<dyn ExecutionPlan>,
+    pub(crate) input_plan: Arc<dyn ExecutionPlan>,
     /// Output partitioning
     properties: PlanProperties,
     /// Directory to write shuffle files from
@@ -84,7 +84,7 @@ impl ShuffleWriterExec {
 
         Self {
             stage_id,
-            plan,
+            input_plan: plan,
             properties,
             shuffle_dir: shuffle_dir.to_string(),
             metrics: ExecutionPlanMetricsSet::new(),
@@ -98,11 +98,11 @@ impl ExecutionPlan for ShuffleWriterExec {
     }
 
     fn schema(&self) -> SchemaRef {
-        self.plan.schema()
+        self.input_plan.schema()
     }
 
     fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        vec![&self.plan]
+        vec![&self.input_plan]
     }
 
     fn with_new_children(
@@ -122,7 +122,7 @@ impl ExecutionPlan for ShuffleWriterExec {
             self.stage_id
         );
 
-        let mut stream = self.plan.execute(input_partition, context)?;
+        let mut stream = self.input_plan.execute(input_partition, context)?;
         let write_time =
             MetricBuilder::new(&self.metrics).subset_time("write_time", input_partition);
         let repart_time =

testdata/expected-plans/q1.txt

@@ -42,7 +42,7 @@ ShuffleWriterExec(stage_id=1, output_partitioning=Hash([Column { name: "l_return
         CoalesceBatchesExec: target_batch_size=8192
           ShuffleReaderExec(stage_id=0, input_partitioning=Hash([Column { name: "l_returnflag", index: 0 }, Column { name: "l_linestatus", index: 1 }], 2))
 
-Query Stage #2 (2 -> 1):
+Query Stage #2 (1 -> 1):
 SortPreservingMergeExec: [l_returnflag@0 ASC NULLS LAST, l_linestatus@1 ASC NULLS LAST]
   ShuffleReaderExec(stage_id=1, input_partitioning=Hash([Column { name: "l_returnflag", index: 0 }, Column { name: "l_linestatus", index: 1 }], 2))
 

testdata/expected-plans/q10.txt

@@ -117,7 +117,7 @@ ShuffleWriterExec(stage_id=7, output_partitioning=Hash([Column { name: "c_custke
         CoalesceBatchesExec: target_batch_size=8192
           ShuffleReaderExec(stage_id=6, input_partitioning=Hash([Column { name: "c_custkey", index: 0 }, Column { name: "c_name", index: 1 }, Column { name: "c_acctbal", index: 2 }, Column { name: "c_phone", index: 3 }, Column { name: "n_name", index: 4 }, Column { name: "c_address", index: 5 }, Column { name: "c_comment", index: 6 }], 2))
 
-Query Stage #8 (2 -> 1):
+Query Stage #8 (1 -> 1):
 SortPreservingMergeExec: [revenue@2 DESC], fetch=20
   ShuffleReaderExec(stage_id=7, input_partitioning=Hash([Column { name: "c_custkey", index: 0 }, Column { name: "c_name", index: 1 }, Column { name: "c_acctbal", index: 3 }, Column { name: "c_phone", index: 6 }, Column { name: "n_name", index: 4 }, Column { name: "c_address", index: 5 }, Column { name: "c_comment", index: 7 }], 2))
 

testdata/expected-plans/q11.txt

@@ -167,7 +167,7 @@ ShuffleWriterExec(stage_id=10, output_partitioning=Hash([Column { name: "ps_part
           CoalesceBatchesExec: target_batch_size=8192
             ShuffleReaderExec(stage_id=9, input_partitioning=Hash([Column { name: "ps_partkey", index: 0 }], 2))
 
-Query Stage #11 (2 -> 1):
+Query Stage #11 (1 -> 1):
 SortPreservingMergeExec: [value@1 DESC]
   ShuffleReaderExec(stage_id=10, input_partitioning=Hash([Column { name: "ps_partkey", index: 0 }], 2))
 

testdata/expected-plans/q12.txt

@@ -65,7 +65,7 @@ ShuffleWriterExec(stage_id=3, output_partitioning=Hash([Column { name: "l_shipmo
         CoalesceBatchesExec: target_batch_size=8192
           ShuffleReaderExec(stage_id=2, input_partitioning=Hash([Column { name: "l_shipmode", index: 0 }], 2))
 
-Query Stage #4 (2 -> 1):
+Query Stage #4 (1 -> 1):
 SortPreservingMergeExec: [l_shipmode@0 ASC NULLS LAST]
   ShuffleReaderExec(stage_id=3, input_partitioning=Hash([Column { name: "l_shipmode", index: 0 }], 2))
 

testdata/expected-plans/q13.txt

@@ -70,7 +70,7 @@ ShuffleWriterExec(stage_id=3, output_partitioning=Hash([Column { name: "c_count"
         CoalesceBatchesExec: target_batch_size=8192
           ShuffleReaderExec(stage_id=2, input_partitioning=Hash([Column { name: "c_count", index: 0 }], 2))
 
-Query Stage #4 (2 -> 1):
+Query Stage #4 (1 -> 1):
 SortPreservingMergeExec: [custdist@1 DESC, c_count@0 DESC]
   ShuffleReaderExec(stage_id=3, input_partitioning=Hash([Column { name: "c_count", index: 0 }], 2))