From 14963726fb8988c487c93f2e05f9a9237163b9c1 Mon Sep 17 00:00:00 2001
From: zebsme <zeb.yeung@gmail.com>
Date: Tue, 18 Mar 2025 00:00:05 +0800
Subject: [PATCH] Implement tree explain for UnionExec

---
 datafusion/physical-plan/src/union.rs         |  5 +---
 .../sqllogictest/test_files/explain_tree.slt  | 30 +++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs
index 9315e9005b74..2b666093f29e 100644
--- a/datafusion/physical-plan/src/union.rs
+++ b/datafusion/physical-plan/src/union.rs
@@ -157,10 +157,7 @@ impl DisplayAs for UnionExec {
             DisplayFormatType::Default | DisplayFormatType::Verbose => {
                 write!(f, "UnionExec")
             }
-            DisplayFormatType::TreeRender => {
-                // TODO: collect info
-                write!(f, "")
-            }
+            DisplayFormatType::TreeRender => Ok(()),
         }
     }
 }
diff --git a/datafusion/sqllogictest/test_files/explain_tree.slt b/datafusion/sqllogictest/test_files/explain_tree.slt
index b9f4074a450d..4774b8ddace6 100644
--- a/datafusion/sqllogictest/test_files/explain_tree.slt
+++ b/datafusion/sqllogictest/test_files/explain_tree.slt
@@ -1541,6 +1541,32 @@ physical_plan
 71)│          rows: 1          ││          rows: 1          │
 72)└───────────────────────────┘└───────────────────────────┘
 
+# Test explain tree for UnionExec
+query TT
+EXPLAIN
+SELECT id, name FROM t1
+UNION ALL
+SELECT id, name FROM t2;
+----
+physical_plan
+01)┌───────────────────────────┐
+02)│         UnionExec         ├──────────────┐
+03)└─────────────┬─────────────┘              │
+04)┌─────────────┴─────────────┐┌─────────────┴─────────────┐
+05)│       DataSourceExec      ││       ProjectionExec      │
+06)│    --------------------   ││    --------------------   │
+07)│        bytes: 1320        ││   id: CAST(id AS Int32)   │
+08)│       format: memory      ││         name: name        │
+09)│          rows: 1          ││                           │
+10)└───────────────────────────┘└─────────────┬─────────────┘
+11)-----------------------------┌─────────────┴─────────────┐
+12)-----------------------------│       DataSourceExec      │
+13)-----------------------------│    --------------------   │
+14)-----------------------------│        bytes: 1312        │
+15)-----------------------------│       format: memory      │
+16)-----------------------------│          rows: 1          │
+17)-----------------------------└───────────────────────────┘
+
 # cleanup
 statement ok
 drop table t1;
@@ -2000,6 +2026,10 @@ physical_plan
 58)│          rows: 1          │
 59)└───────────────────────────┘
 
+# clean up
+statement ok
+drop table t1;
+
 # Test explain tree for LazyMemoryExec
 query TT
 EXPLAIN SELECT * FROM generate_series(1, 100)