fix

Author: andygrove

src/query_stage.rs

@@ -99,14 +99,13 @@ impl QueryStage {
     /// 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.as_any().is::<ShuffleWriterExec>() {
-            // most query stages represent a shuffle write
+        if self.plan.children().is_empty() {
+            // leaf node (file scan)
+            self.plan.output_partitioning().partition_count()
+        } else {
             self.plan.children()[0]
                 .output_partitioning()
                 .partition_count()
-        } else {
-            // probably the final query stage
-            self.plan.output_partitioning().partition_count()
         }
     }