Implement tree explain for PartialSortExec (#15066)

* Implement tree explain for PartialSortExec

* simplify

* fix

* remove square brackets

Author: irenjj

datafusion/physical-plan/src/sorts/partial_sort.rs

@@ -226,10 +226,15 @@ impl DisplayAs for PartialSortExec {
                     None => write!(f, "PartialSortExec: expr=[{}], common_prefix_length=[{common_prefix_length}]", self.expr),
                 }
             }
-            DisplayFormatType::TreeRender => {
-                // TODO: collect info
-                write!(f, "")
-            }
+            DisplayFormatType::TreeRender => match self.fetch {
+                Some(fetch) => {
+                    writeln!(f, "limit={fetch}")?;
+                    writeln!(f, "sort keys={}", self.expr)
+                }
+                None => {
+                    writeln!(f, "sort keys={}", self.expr)
+                }
+            },
         }
     }
 }

datafusion/sqllogictest/test_files/explain_tree.slt

@@ -82,7 +82,18 @@ CREATE EXTERNAL TABLE table5
 STORED AS ARROW
 LOCATION 'test_files/scratch/explain_tree/table5.arrow';
 
-
+statement ok
+CREATE UNBOUNDED EXTERNAL TABLE annotated_data_infinite2 (
+  a0 INTEGER,
+  a INTEGER,
+  b INTEGER,
+  c INTEGER,
+  d INTEGER
+)
+STORED AS CSV
+WITH ORDER (a ASC, b ASC, c ASC)
+LOCATION '../core/tests/data/window_2.csv'
+OPTIONS ('format.has_header' 'true');
 
 ######## Begin Queries ########
 
@@ -528,6 +539,52 @@ physical_plan
 17)│       format: arrow       │
 18)└───────────────────────────┘
 
+# Query with PartialSortExec.
+query TT
+EXPLAIN SELECT *
+FROM annotated_data_infinite2
+ORDER BY a, b, d;
+----
+logical_plan
+01)Sort: annotated_data_infinite2.a ASC NULLS LAST, annotated_data_infinite2.b ASC NULLS LAST, annotated_data_infinite2.d ASC NULLS LAST
+02)--TableScan: annotated_data_infinite2 projection=[a0, a, b, c, d]
+physical_plan
+01)┌───────────────────────────┐
+02)│      PartialSortExec      │
+03)│    --------------------   │
+04)│         sort keys:        │
+05)│  a@1 ASC NULLS LAST, b@2  │
+06)│     ASC NULLS LAST, d@4   │
+07)│       ASC NULLS LAST      │
+08)└─────────────┬─────────────┘
+09)┌─────────────┴─────────────┐
+10)│     StreamingTableExec    │
+11)└───────────────────────────┘
+
+query TT
+EXPLAIN SELECT *
+FROM annotated_data_infinite2
+ORDER BY a, b, d
+LIMIT 50;
+----
+logical_plan
+01)Sort: annotated_data_infinite2.a ASC NULLS LAST, annotated_data_infinite2.b ASC NULLS LAST, annotated_data_infinite2.d ASC NULLS LAST, fetch=50
+02)--TableScan: annotated_data_infinite2 projection=[a0, a, b, c, d]
+physical_plan
+01)┌───────────────────────────┐
+02)│      PartialSortExec      │
+03)│    --------------------   │
+04)│         limit: 50         │
+05)│                           │
+06)│         sort keys:        │
+07)│  a@1 ASC NULLS LAST, b@2  │
+08)│     ASC NULLS LAST, d@4   │
+09)│       ASC NULLS LAST      │
+10)└─────────────┬─────────────┘
+11)┌─────────────┴─────────────┐
+12)│     StreamingTableExec    │
+13)└───────────────────────────┘
+
 # Query with hash join.
 query TT
 explain select * from table1 inner join table2 on table1.int_col = table2.int_col and table1.string_col = table2.string_col;
@@ -616,7 +673,6 @@ physical_plan
 34)│        format: csv        │
 35)└───────────────────────────┘
 
-
 # cleanup
 statement ok
 drop table table1;