andygrove commented on code in PR #56:
URL: https://github.com/apache/datafusion-ray/pull/56#discussion_r1885769257
##########
src/query_stage.rs:
##########
@@ -99,14 +100,23 @@ 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.children().is_empty() {
- // leaf node (file scan)
- self.plan.output_partitioning().partition_count()
- } else {
- self.plan.children()[0]
- .output_partitioning()
- .partition_count()
+ let mut output_partition_counts = HashSet::new();
+
+ for child in self.plan.children() {
+
output_partition_counts.insert(child.output_partitioning().partition_count());
+ if output_partition_counts.len() > 1 {
+ panic!(
+ "Children plan of {:#?} have a distinct output
partitioning partition count",
+ self.plan
+ );
+ }
}
Review Comment:
We can do this check without creating additional data structures:
```rust
let partition_count =
self.plan.children()[0].output_partitioning().partition_count();
for i in 1..self.plan.children().len() {
if
self.plan.children()[i].output_partitioning().partition_count() !=
partition_count {
return Err(...)
}
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
