Repository: apache/datafusion (rust)

Issue:
Implement `tree` explain for `LocalLimitExec`
### Is your feature request related to a problem or challenge?

- Part of https://github.com/apache/datafusion/issues/14914

@irenjj  added a new `tree` explain mode in https://github.com/apache/datafusion/pull/14677. Now we need to add support for different types of operators. 

Here is an example of how to see the new explain plans:
```sql
set datafusion.explain.format = 'tree';
create table foo(x int, y int) as values (1,2), (3,4);
explain select * from foo where x = 4;
+---------------+------------------------------------+
| plan_type     | plan                               |
+---------------+------------------------------------+
| logical_plan  | Filter: foo.x = Int32(4)           |
|               |   TableScan: foo projection=[x, y] |
| physical_plan | ┌───────────────────────────┐      |
|               | │    CoalesceBatchesExec    │      |
|               | └─────────────┬─────────────┘      |
|               | ┌─────────────┴─────────────┐      |
|               | │         FilterExec        │      |
|               | └─────────────┬─────────────┘      |
|               | ┌─────────────┴─────────────┐      |
|               | │       DataSourceExec      │      |
|               | │    --------------------   │      |
|               | │    partition_sizes: [1]   │      |
|               | │       partitions: 1       │      |
|               | └───────────────────────────┘      |
|               |                                    |
+---------------+------------------------------------+
```


### Describe the solution you'd like

Add `tree` format to the ExecutionPlan specified in the subject of this ticket

Note that the tree mode should have only the most relevant information for users to understand what plan is being run. Detailed iformation

The process goes like:
1. Add the relevant code to the operator  
2. Add / update explain_tree.slt test to show the new code in action

The relevant code looks like this
```rust
            DisplayFormatType::TreeRender => {
                // TODO: collect info
                write!(f, "")?;
            }
```


You can run the tests like
```shell
cargo test --test sqllogictests -- explain_tree
```

You can update the test like this:
```shell
cargo test --test sqllogictests -- explain_tree --complete
```



### Describe alternatives you've considered

Here is an example PR that shows what is needed
- https://github.com/apache/datafusion/pull/15001

### Additional context

Since this is well specified and has several examples and will give exposure to DataFusion code and testing, I have marked this as a good first issue

Pull request:
Implement tree explain for `LocalLimitExec`
## Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->

- Closes #15025 .

## Rationale for this change

<!--
 Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
 Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
-->

## What changes are included in this PR?

<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->

## Are these changes tested?
Yes

<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code

If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->

## Are there any user-facing changes?

<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->

<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->


Patch:
diff --git a/datafusion/core/src/datasource/file_format/arrow.rs b/datafusion/core/src/datasource/file_format/arrow.rs
index 6835d9a6d..71d8d0026 100644
--- a/datafusion/core/src/datasource/file_format/arrow.rs
+++ b/datafusion/core/src/datasource/file_format/arrow.rs
@@ -302,8 +302,8 @@ impl DisplayAs for ArrowFileSink {
                 write!(f, ")")
             }
             DisplayFormatType::TreeRender => {
-                // TODO: collect info
-                write!(f, "")
+                writeln!(f, "format: arrow")?;
+                write!(f, "file={}", &self.config.original_url)
             }
         }
     }
diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs
index e2c7d1eca..3b71593b3 100644
--- a/datafusion/core/src/datasource/file_format/parquet.rs
+++ b/datafusion/core/src/datasource/file_format/parquet.rs
@@ -1373,6 +1373,7 @@ mod tests {
         let object_store_url = ObjectStoreUrl::local_filesystem();
 
         let file_sink_config = FileSinkConfig {
+            original_url: String::default(),
             object_store_url: object_store_url.clone(),
             file_groups: vec![PartitionedFile::new("/tmp".to_string(), 1)],
             table_paths: vec![ListingTableUrl::parse(table_path)?],
@@ -1458,6 +1459,7 @@ mod tests {
 
         // set file config to include partitioning on field_a
         let file_sink_config = FileSinkConfig {
+            original_url: String::default(),
             object_store_url: object_store_url.clone(),
             file_groups: vec![PartitionedFile::new("/tmp".to_string(), 1)],
             table_paths: vec![ListingTableUrl::parse("file:///")?],
@@ -1541,6 +1543,7 @@ mod tests {
             let object_store_url = ObjectStoreUrl::local_filesystem();
 
             let file_sink_config = FileSinkConfig {
+                original_url: String::default(),
                 object_store_url: object_store_url.clone(),
                 file_groups: vec![PartitionedFile::new("/tmp".to_string(), 1)],
                 table_paths: vec![ListingTableUrl::parse("file:///")?],
diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs
index 4d7762784..21b35bac2 100644
--- a/datafusion/core/src/datasource/listing/table.rs
+++ b/datafusion/core/src/datasource/listing/table.rs
@@ -1037,6 +1037,7 @@ impl TableProvider for ListingTable {
 
         // Sink related option, apart from format
         let config = FileSinkConfig {
+            original_url: String::default(),
             object_store_url: self.table_paths()[0].object_store(),
             table_paths: self.table_paths().clone(),
             file_groups,
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
index 6aff9280f..170f85af7 100644
--- a/datafusion/core/src/physical_planner.rs
+++ b/datafusion/core/src/physical_planner.rs
@@ -502,6 +502,7 @@ impl DefaultPhysicalPlanner {
                 partition_by,
                 options: source_option_tuples,
             }) => {
+                let original_url = output_url.clone();
                 let input_exec = children.one()?;
                 let parsed_url = ListingTableUrl::parse(output_url)?;
                 let object_store_url = parsed_url.object_store();
@@ -531,6 +532,7 @@ impl DefaultPhysicalPlanner {
 
                 // Set file sink related options
                 let config = FileSinkConfig {
+                    original_url,
                     object_store_url,
                     table_paths: vec![parsed_url],
                     file_groups: vec![],
diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs
index cc1b63cfc..522cb12db 100644
--- a/datafusion/datasource-csv/src/file_format.rs
+++ b/datafusion/datasource-csv/src/file_format.rs
@@ -666,10 +666,8 @@ impl DisplayAs for CsvSink {
                 write!(f, ")")
             }
             DisplayFormatType::TreeRender => {
-                if !self.config.file_groups.is_empty() {
-                    FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?;
-                }
-                Ok(())
+                writeln!(f, "format: csv")?;
+                write!(f, "file={}", &self.config.original_url)
             }
         }
     }
diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs
index bec3a524f..9b6d5925f 100644
--- a/datafusion/datasource-json/src/file_format.rs
+++ b/datafusion/datasource-json/src/file_format.rs
@@ -325,10 +325,8 @@ impl DisplayAs for JsonSink {
                 write!(f, ")")
             }
             DisplayFormatType::TreeRender => {
-                if !self.config.file_groups.is_empty() {
-                    FileGroupDisplay(&self.config.file_groups).fmt_as(t, f)?;
-                }
-                Ok(())
+                writeln!(f, "format: json")?;
+                write!(f, "file={}", &self.config.original_url)
             }
         }
     }
diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs
index 6087f930d..279c9d210 100644
--- a/datafusion/datasource/src/file_sink_config.rs
+++ b/datafusion/datasource/src/file_sink_config.rs
@@ -86,6 +86,8 @@ pub trait FileSink: DataSink {
 /// The base configurations to provide when creating a physical plan for
 /// writing to any given file format.
 pub struct FileSinkConfig {
+    /// The unresolved URL specified by the user
+    pub original_url: String,
     /// Object store URL, used to get an ObjectStore instance
     pub object_store_url: ObjectStoreUrl,
     /// A vector of [`PartitionedFile`] structs, each representing a file partition
diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs
index b9464e3a8..89cf47a6d 100644
--- a/datafusion/physical-plan/src/limit.rs
+++ b/datafusion/physical-plan/src/limit.rs
@@ -268,8 +268,7 @@ impl DisplayAs for LocalLimitExec {
                 write!(f, "LocalLimitExec: fetch={}", self.fetch)
             }
             DisplayFormatType::TreeRender => {
-                // TODO: collect info
-                write!(f, "")
+                write!(f, "limit={}", self.fetch)
             }
         }
     }
diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs
index 6331b7fb3..0bf9fdb63 100644
--- a/datafusion/proto/src/physical_plan/from_proto.rs
+++ b/datafusion/proto/src/physical_plan/from_proto.rs
@@ -657,6 +657,7 @@ impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig {
             protobuf::InsertOp::Replace => InsertOp::Replace,
         };
         Ok(Self {
+            original_url: String::default(),
             object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?,
             file_groups,
             table_paths,


Test patch:
diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
index b5bfef99a..21622390c 100644
--- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
+++ b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs
@@ -1281,6 +1281,7 @@ fn roundtrip_json_sink() -> Result<()> {
     let input = Arc::new(PlaceholderRowExec::new(schema.clone()));
 
     let file_sink_config = FileSinkConfig {
+        original_url: String::default(),
         object_store_url: ObjectStoreUrl::local_filesystem(),
         file_groups: vec![PartitionedFile::new("/tmp".to_string(), 1)],
         table_paths: vec![ListingTableUrl::parse("file:///")?],
@@ -1317,6 +1318,7 @@ fn roundtrip_csv_sink() -> Result<()> {
     let input = Arc::new(PlaceholderRowExec::new(schema.clone()));
 
     let file_sink_config = FileSinkConfig {
+        original_url: String::default(),
         object_store_url: ObjectStoreUrl::local_filesystem(),
         file_groups: vec![PartitionedFile::new("/tmp".to_string(), 1)],
         table_paths: vec![ListingTableUrl::parse("file:///")?],
@@ -1372,6 +1374,7 @@ fn roundtrip_parquet_sink() -> Result<()> {
     let input = Arc::new(PlaceholderRowExec::new(schema.clone()));
 
     let file_sink_config = FileSinkConfig {
+        original_url: String::default(),
         object_store_url: ObjectStoreUrl::local_filesystem(),
         file_groups: vec![PartitionedFile::new("/tmp".to_string(), 1)],
         table_paths: vec![ListingTableUrl::parse("file:///")?],
diff --git a/datafusion/sqllogictest/test_files/explain_tree.slt b/datafusion/sqllogictest/test_files/explain_tree.slt
index 99190122f..83d2edee7 100644
--- a/datafusion/sqllogictest/test_files/explain_tree.slt
+++ b/datafusion/sqllogictest/test_files/explain_tree.slt
@@ -133,6 +133,11 @@ AS SELECT
 FROM
   hashjoin_datatype_table_t2_source
 
+statement ok
+CREATE TABLE limit_table AS
+SELECT * FROM table1
+UNION ALL SELECT * FROM table1
+
 ######## Begin Queries ########
 
 # Filter
@@ -246,6 +251,21 @@ physical_plan
 11)│        format: csv        │
 12)└───────────────────────────┘
 
+query TT
+explain SELECT * FROM limit_table LIMIT 10;
+----
+physical_plan
+01)┌───────────────────────────┐
+02)│   CoalescePartitionsExec  │
+03)└─────────────┬─────────────┘
+04)┌─────────────┴─────────────┐
+05)│       DataSourceExec      │
+06)│    --------------------   │
+07)│        bytes: 3120        │
+08)│       format: memory      │
+09)│          rows: 2          │
+10)└───────────────────────────┘
+
 # 2 Joins
 query TT
 explain SELECT table1.string_col, table2.date_col FROM table1 JOIN table2 ON table1.int_col = table2.int_col;
@@ -1827,14 +1847,21 @@ TO 'test_files/scratch/explain_tree/1.json';
 physical_plan
 01)┌───────────────────────────┐
 02)│        DataSinkExec       │
-03)└─────────────┬─────────────┘
-04)┌─────────────┴─────────────┐
-05)│       DataSourceExec      │
-06)│    --------------------   │
-07)│        bytes: 2672        │
-08)│       format: memory      │
-09)│          rows: 1          │
-10)└───────────────────────────┘
+03)│    --------------------   │
+04)│           file:           │
+05)│     test_files/scratch    │
+06)│      /explain_tree/1      │
+07)│           .json           │
+08)│                           │
+09)│        format: json       │
+10)└─────────────┬─────────────┘
+11)┌─────────────┴─────────────┐
+12)│       DataSourceExec      │
+13)│    --------------------   │
+14)│        bytes: 2672        │
+15)│       format: memory      │
+16)│          rows: 1          │
+17)└───────────────────────────┘
 
 query TT
 explain COPY (VALUES (1, 'foo', 1, '2023-01-01'), (2, 'bar', 2, '2023-01-02'), (3, 'baz', 3, '2023-01-03'))
@@ -1843,14 +1870,45 @@ TO 'test_files/scratch/explain_tree/2.csv';
 physical_plan
 01)┌───────────────────────────┐
 02)│        DataSinkExec       │
-03)└─────────────┬─────────────┘
-04)┌─────────────┴─────────────┐
-05)│       DataSourceExec      │
-06)│    --------------------   │
-07)│        bytes: 2672        │
-08)│       format: memory      │
-09)│          rows: 1          │
-10)└───────────────────────────┘
+03)│    --------------------   │
+04)│           file:           │
+05)│     test_files/scratch    │
+06)│      /explain_tree/2      │
+07)│            .csv           │
+08)│                           │
+09)│        format: csv        │
+10)└─────────────┬─────────────┘
+11)┌─────────────┴─────────────┐
+12)│       DataSourceExec      │
+13)│    --------------------   │
+14)│        bytes: 2672        │
+15)│       format: memory      │
+16)│          rows: 1          │
+17)└───────────────────────────┘
+
+query TT
+explain COPY (VALUES (1, 'foo', 1, '2023-01-01'), (2, 'bar', 2, '2023-01-02'), (3, 'baz', 3, '2023-01-03'))
+TO 'test_files/scratch/explain_tree/3.arrow';
+----
+physical_plan
+01)┌───────────────────────────┐
+02)│        DataSinkExec       │
+03)│    --------------------   │
+04)│           file:           │
+05)│     test_files/scratch    │
+06)│      /explain_tree/3      │
+07)│           .arrow          │
+08)│                           │
+09)│       format: arrow       │
+10)└─────────────┬─────────────┘
+11)┌─────────────┴─────────────┐
+12)│       DataSourceExec      │
+13)│    --------------------   │
+14)│        bytes: 2672        │
+15)│       format: memory      │
+16)│          rows: 1          │
+17)└───────────────────────────┘
+
 
 # Test explain tree rendering for CoalesceBatchesExec with limit
 statement ok