add error check

apache · Sep 18, 2019 · 56edbc5 · 56edbc5
1 parent 07c31d8
commit 56edbc5
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/rust/datafusion/src/execution/context.rs b/rust/datafusion/src/execution/context.rs
@@ -287,10 +287,17 @@ impl ExecutionContext {
                 common::collect(it)
             }
             _ => {
+                // merge into a single partition
                 let plan = MergeExec::new(plan.schema().clone(), partitions);
                 let partitions = plan.partitions()?;
-                let it = partitions[0].execute()?;
-                common::collect(it)
+                if partitions.len() == 1 {
+                    common::collect(partitions[0].execute()?)
+                } else {
+                    Err(ExecutionError::InternalError(format!(
+                        "MergeExec returned {} partitions",
+                        partitions.len()
+                    )))
+                }
             }
         }
     }