From 84b8f584ca3c96c64d379592126a8af3b2d4ad45 Mon Sep 17 00:00:00 2001
From: Rob Tandy <rob.tandy@datadoghq.com>
Date: Wed, 17 Sep 2025 14:02:57 -0400
Subject: [PATCH] fix bug in graphviz for determining output partitions

---
 src/execution_plans/stage.rs | 21 +--------------------
 1 file changed, 1 insertion(+), 20 deletions(-)

diff --git a/src/execution_plans/stage.rs b/src/execution_plans/stage.rs
index 18ce8f7..d99f3a9 100644
--- a/src/execution_plans/stage.rs
+++ b/src/execution_plans/stage.rs
@@ -570,26 +570,7 @@ pub fn display_plan(
         if let Some(parent) = maybe_parent {
             let output_partitions = plan.output_partitioning().partition_count();
 
-            let partitions = if !distributed {
-                #[deny(clippy::if_same_then_else)]
-                output_partitions
-            } else if plan
-                .as_any()
-                .downcast_ref::<ArrowFlightReadExec>()
-                .is_some()
-                || plan
-                    .as_any()
-                    .downcast_ref::<PartitionIsolatorExec>()
-                    .is_some()
-            {
-                output_partitions
-            } else if let Some(child) = plan.children().first() {
-                child.output_partitioning().partition_count()
-            } else {
-                output_partitions
-            };
-
-            for i in 0..partitions {
+            for i in 0..output_partitions {
                 let mut style = "";
                 if plan
                     .as_any()