Repository: apache/datafusion
Pull request #16081: Set `TrackConsumersPool` as default in datafusion-cli
Author: login: ding-young
Labels: name: documentation
description: Improvements or additions to documentation
Created: 2025-05-18T03:21:55Z | Merged: 2025-05-22T06:25:36Z

## Which issue does this PR close?
Part of #16065 
<!--
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.
-->


## Rationale for this change
Currently, datafusion-cli does not provide an option to use `TrackConsumersPool` as memory pool. It would be useful to set `TrackConsumersPool` as default so that users can view top memory consumers when running out of memory. 
 
<!--
 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?
This pr adds a new option `top-memory-consumers` in datafusion-cli. Thereby, user can set `TrackConsumersPool { inner: <mem-pool-type>, top: <top-memory-consumers>}` as memory pool. If `top-memory-consumers` is set to 0, datafusion-cli will use `mem-pool-type` not wrapped with `TrackConsumersPool`. 

<!--
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, via locally. 
<!--
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.
-->
I updated the docs. (`usage.md`) 

Linked issue #123: Implement physical plan for EXISTS subquery
*Note*: migrated from original JIRA: https://issues.apache.org/jira/browse/ARROW-10819

The TPC-H queries include use of the EXISTS which is used to test for the existence of any record in a subquery. For example:

```sql
and *exists* (
    select
        *
    from
        lineitem
    where
        l_orderkey = o_orderkey
        and l_commitdate < l_receiptdate
)
```
Linked issue comments: user: login: alamb
created_at: 2021-04-26T13:23:43Z
body: Comment from Andy Grove(andygrove) @ 2020-12-31T19:42:45.132+0000:
<pre>The example given here is a correlated subquery that can be translated into a join.

Here is a random stackoverflow discussion on this for reference (I have not reviewed it)

https://stackoverflow.com/questions/1772609/procedurally-transform-subquery-into-join</pre>
author_association: CONTRIBUTOR
user: login: alamb
created_at: 2022-10-21T19:21:30Z
body: In case anyone is curious -- we support correlated versions of these queries (via a join) but if there is no correlation (not super useful) we do not

```sql
❯ create table foo as select * from (values (1), (2), (NULL)) as sql
;
0 rows in set. Query took 0.022 seconds.
3 rows in set. Query took 0.007 seconds.
❯ create table bar as select * from (values (1), (2), (NULL)) as sql;
0 rows in set. Query took 0.000 seconds.
❯ select * from foo where exists (select column1 from bar);
NotImplemented("Physical plan does not support logical expression EXISTS (<subquery>)")
❯ select * from foo where exists (select column1 from bar where foo.column1 = bar.column1);
+---------+
| column1 |
+---------+
| 2       |
| 1       |
+---------+
```
author_association: CONTRIBUTOR
user: login: logan-keede
created_at: 2025-02-09T19:53:56Z
body: 
> In case anyone is curious -- we support correlated versions of these queries (via a join) but if there is no correlation (not super useful) we do not
> 
> ❯ create table foo as select * from (values (1), (2), (NULL)) as sql
> ;
> 0 rows in set. Query took 0.022 seconds.
> 3 rows in set. Query took 0.007 seconds.
> ❯ create table bar as select * from (values (1), (2), (NULL)) as sql;
> 0 rows in set. Query took 0.000 seconds.
> ❯ select * from foo where exists (select column1 from bar);
> NotImplemented("Physical plan does not support logical expression EXISTS (<subquery>)")
> ❯ select * from foo where exists (select column1 from bar where foo.column1 = bar.column1);
> +---------+
> | column1 |
> +---------+
> | 2       |
> | 1       |
> +---------+

```sql

> explain select * from foo where exists (select column1 from bar);
+---------------+-----------------------------------------------------+
| plan_type     | plan                                                |
+---------------+-----------------------------------------------------+
| logical_plan  | LeftSemi Join:                                      |
|               |   TableScan: foo projection=[column1]               |
|               |   SubqueryAlias: __correlated_sq_1                  |
|               |     TableScan: bar projection=[]                    |
| physical_plan | NestedLoopJoinExec: join_type=RightSemi             |
|               |   DataSourceExec: partitions=1, partition_sizes=[1] |
|               |   DataSourceExec: partitions=1, partition_sizes=[1] |
|               |                                                     |
+---------------+-----------------------------------------------------+
2 row(s) fetched. 
Elapsed 0.007 seconds.

> select * from foo where exists (select column1 from bar);
+---------+
| column1 |
+---------+
| 1       |
| 2       |
| NULL    |
+---------+
3 row(s) fetched. 
Elapsed 0.006 seconds.
```


I get this on `datafusion 45.0.0`, perhaps this issue can be closed or at least updated if there is a case where we still don't support EXISTS in physical plan.

author_association: CONTRIBUTOR
user: login: alamb
created_at: 2025-02-10T15:04:31Z
body: Looks good -- thanks for checking @logan-keede  -- we can open a new issue if we find another hole
author_association: CONTRIBUTOR

Diff:
diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs
index 16ddf2a6a5f84..091b177d0819a 100644
--- a/datafusion-cli/src/main.rs
+++ b/datafusion-cli/src/main.rs
@@ -17,13 +17,16 @@
 
 use std::collections::HashMap;
 use std::env;
+use std::num::NonZeroUsize;
 use std::path::Path;
 use std::process::ExitCode;
 use std::sync::{Arc, LazyLock};
 
 use datafusion::error::{DataFusionError, Result};
 use datafusion::execution::context::SessionConfig;
-use datafusion::execution::memory_pool::{FairSpillPool, GreedyMemoryPool, MemoryPool};
+use datafusion::execution::memory_pool::{
+    FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
+};
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::execution::DiskManager;
 use datafusion::prelude::SessionContext;
@@ -118,6 +121,13 @@ struct Args {
     )]
     mem_pool_type: PoolType,
 
+    #[clap(
+        long,
+        help = "The number of top memory consumers to display when query fails due to memory exhaustion. To disable memory consumer tracking, set this value to 0",
+        default_value = "3"
+    )]
+    top_memory_consumers: usize,
+
     #[clap(
         long,
         help = "The max number of rows to display for 'Table' format\n[possible values: numbers(0/10/...), inf(no limit)]",
@@ -169,9 +179,22 @@ async fn main_inner() -> Result<()> {
     if let Some(memory_limit) = args.memory_limit {
         // set memory pool type
         let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
-            PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
-            PoolType::Greedy => Arc::new(GreedyMemoryPool::new(memory_limit)),
+            PoolType::Fair if args.top_memory_consumers == 0 => {
+                Arc::new(FairSpillPool::new(memory_limit))
+            }
+            PoolType::Fair => Arc::new(TrackConsumersPool::new(
+                FairSpillPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
+            PoolType::Greedy if args.top_memory_consumers == 0 => {
+                Arc::new(GreedyMemoryPool::new(memory_limit))
+            }
+            PoolType::Greedy => Arc::new(TrackConsumersPool::new(
+                GreedyMemoryPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
         };
+
         rt_builder = rt_builder.with_memory_pool(pool)
     }
 
diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs
index bdab38026fcf3..fb2f08157f674 100644
--- a/datafusion-cli/tests/cli_integration.rs
+++ b/datafusion-cli/tests/cli_integration.rs
@@ -122,6 +122,42 @@ fn test_cli_format<'a>(#[case] format: &'a str) {
     assert_cmd_snapshot!(cmd);
 }
 
+#[rstest]
+#[case("no_track", ["--top-memory-consumers", "0"])]
+#[case("top2", ["--top-memory-consumers", "2"])]
+#[case("top3_default", [])]
+#[test]
+fn test_cli_top_memory_consumers<'a>(
+    #[case] snapshot_name: &str,
+    #[case] top_memory_consumers: impl IntoIterator<Item = &'a str>,
+) {
+    let mut settings = make_settings();
+
+    settings.set_snapshot_suffix(snapshot_name);
+
+    settings.add_filter(
+        r"[^\s]+\#\d+\(can spill: (true|false)\) consumed .*?B",
+        "Consumer(can spill: bool) consumed XB",
+    );
+    settings.add_filter(
+        r"Error: Failed to allocate additional .*? for .*? with .*? already allocated for this reservation - .*? remain available for the total pool",
+        "Error: Failed to allocate ",
+    );
+    settings.add_filter(
+        r"Resources exhausted: Failed to allocate additional .*? for .*? with .*? already allocated for this reservation - .*? remain available for the total pool",
+        "Resources exhausted: Failed to allocate",
+    );
+
+    let _bound = settings.bind_to_scope();
+
+    let mut cmd = cli();
+    let sql = "select * from generate_series(1,500000) as t1(v1) order by v1;";
+    cmd.args(["--memory-limit", "10M", "--command", sql]);
+    cmd.args(top_memory_consumers);
+
+    assert_cmd_snapshot!(cmd);
+}
+
 #[tokio::test]
 async fn test_cli() {
     if env::var("TEST_STORAGE_INTEGRATION").is_err() {
diff --git a/datafusion-cli/tests/snapshots/cli_top_memory_consumers@no_track.snap b/datafusion-cli/tests/snapshots/cli_top_memory_consumers@no_track.snap
new file mode 100644
index 0000000000000..89b646a531f8b
--- /dev/null
+++ b/datafusion-cli/tests/snapshots/cli_top_memory_consumers@no_track.snap
@@ -0,0 +1,21 @@
+---
+source: datafusion-cli/tests/cli_integration.rs
+info:
+  program: datafusion-cli
+  args:
+    - "--memory-limit"
+    - 10M
+    - "--command"
+    - "select * from generate_series(1,500000) as t1(v1) order by v1;"
+    - "--top-memory-consumers"
+    - "0"
+---
+success: false
+exit_code: 1
+----- stdout -----
+[CLI_VERSION]
+Error: Not enough memory to continue external sort. Consider increasing the memory limit, or decreasing sort_spill_reservation_bytes
+caused by
+Resources exhausted: Failed to allocate
+
+----- stderr -----
diff --git a/datafusion-cli/tests/snapshots/cli_top_memory_consumers@top2.snap b/datafusion-cli/tests/snapshots/cli_top_memory_consumers@top2.snap
new file mode 100644
index 0000000000000..ed925a6f64613
--- /dev/null
+++ b/datafusion-cli/tests/snapshots/cli_top_memory_consumers@top2.snap
@@ -0,0 +1,24 @@
+---
+source: datafusion-cli/tests/cli_integration.rs
+info:
+  program: datafusion-cli
+  args:
+    - "--memory-limit"
+    - 10M
+    - "--command"
+    - "select * from generate_series(1,500000) as t1(v1) order by v1;"
+    - "--top-memory-consumers"
+    - "2"
+---
+success: false
+exit_code: 1
+----- stdout -----
+[CLI_VERSION]
+Error: Not enough memory to continue external sort. Consider increasing the memory limit, or decreasing sort_spill_reservation_bytes
+caused by
+Resources exhausted: Additional allocation failed with top memory consumers (across reservations) as:
+  Consumer(can spill: bool) consumed XB,
+  Consumer(can spill: bool) consumed XB.
+Error: Failed to allocate 
+
+----- stderr -----
diff --git a/datafusion-cli/tests/snapshots/cli_top_memory_consumers@top3_default.snap b/datafusion-cli/tests/snapshots/cli_top_memory_consumers@top3_default.snap
new file mode 100644
index 0000000000000..f35e3b117178f
--- /dev/null
+++ b/datafusion-cli/tests/snapshots/cli_top_memory_consumers@top3_default.snap
@@ -0,0 +1,23 @@
+---
+source: datafusion-cli/tests/cli_integration.rs
+info:
+  program: datafusion-cli
+  args:
+    - "--memory-limit"
+    - 10M
+    - "--command"
+    - "select * from generate_series(1,500000) as t1(v1) order by v1;"
+---
+success: false
+exit_code: 1
+----- stdout -----
+[CLI_VERSION]
+Error: Not enough memory to continue external sort. Consider increasing the memory limit, or decreasing sort_spill_reservation_bytes
+caused by
+Resources exhausted: Additional allocation failed with top memory consumers (across reservations) as:
+  Consumer(can spill: bool) consumed XB,
+  Consumer(can spill: bool) consumed XB,
+  Consumer(can spill: bool) consumed XB.
+Error: Failed to allocate 
+
+----- stderr -----
diff --git a/docs/source/user-guide/cli/usage.md b/docs/source/user-guide/cli/usage.md
index 68b09d3199840..13f0e7cff175d 100644
--- a/docs/source/user-guide/cli/usage.md
+++ b/docs/source/user-guide/cli/usage.md
@@ -57,6 +57,9 @@ OPTIONS:
         --mem-pool-type <MEM_POOL_TYPE>
             Specify the memory pool type 'greedy' or 'fair', default to 'greedy'
 
+        --top-memory-consumers <TOP_MEMORY_CONSUMERS>
+            The number of top memory consumers to display when query fails due to memory exhaustion. To disable memory consumer tracking, set this value to 0 [default: 3]
+
     -d, --disk-limit <DISK_LIMIT>
             Available disk space for spilling queries (e.g. '10g'), default to None (uses DataFusion's default value of '100g')
 


Reviews:
user: login: jfahne
body: LGTM. Thanks for adding the test.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T03:04:09Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR sets TrackConsumersPool as the default memory pool in datafusion-cli by introducing a new CLI option --top-memory-consumers.  
- Adds the --top-memory-consumers option to control tracking of top memory consumers.  
- Extends integration tests and updates snapshots to validate the new memory pool behavior.  
- Updates CLI documentation to include the new option.

### Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

| File                                                            | Description                                                  |
| ---------------------------------------------------------------- | ------------------------------------------------------------ |
| docs/source/user-guide/cli/usage.md                              | Added documentation for the new --top-memory-consumers option. |
| datafusion-cli/tests/snapshots/cli_top_memory_consumers@top3_default.snap, cli_top_memory_consumers@top2.snap, cli_top_memory_consumers@no_track.snap | Updated snapshots to reflect error messages with and without memory consumer tracking. |
| datafusion-cli/tests/cli_integration.rs                          | Added integration tests to exercise the new --top-memory-consumers behavior. |
| datafusion-cli/src/main.rs                                       | Updated CLI argument parsing and memory pool initialization to support TrackConsumersPool. |




state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T04:54:52Z
user: login: 2010YOUY01
body: Thank you @ding-young and @jfahne 

I think the reviews from the bot are also worth addressing before merging.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T05:01:36Z
user: login: ding-young
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T16:40:47Z
user: login: ding-young
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T16:41:38Z

Review comments:
diff_hunk: @@ -169,9 +179,22 @@ async fn main_inner() -> Result<()> {
     if let Some(memory_limit) = args.memory_limit {
         // set memory pool type
         let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
-            PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
-            PoolType::Greedy => Arc::new(GreedyMemoryPool::new(memory_limit)),
+            PoolType::Fair if args.top_memory_consumers == 0 => {
+                Arc::new(FairSpillPool::new(memory_limit))
+            }
+            PoolType::Fair => Arc::new(TrackConsumersPool::new(
+                FairSpillPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
+            PoolType::Greedy if args.top_memory_consumers == 0 => {
+                Arc::new(GreedyMemoryPool::new(memory_limit))
+            }
+            PoolType::Greedy => Arc::new(TrackConsumersPool::new(
+                GreedyMemoryPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
         };
path: datafusion-cli/src/main.rs
user: login: Copilot
body: [nitpick] The logic for wrapping Fair and Greedy pools with TrackConsumersPool is duplicated. Consider refactoring this common functionality into a helper function to reduce code duplication.
created_at: 2025-05-21T04:54:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -57,6 +57,9 @@ OPTIONS:
         --mem-pool-type <MEM_POOL_TYPE>
             Specify the memory pool type 'greedy' or 'fair', default to 'greedy'
 
+        --top-memory-consumers <TOP_MEMORY_CONSUMERS>
path: docs/source/user-guide/cli/usage.md
user: login: Copilot
body: [nitpick] It would be helpful to document that the default value is 3, ensuring users know that setting it to 0 disables memory consumer tracking.
created_at: 2025-05-21T04:54:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -122,6 +122,42 @@ fn test_cli_format<'a>(#[case] format: &'a str) {
     assert_cmd_snapshot!(cmd);
 }
 
+#[rstest]
+#[case("no_track", ["--top-memory-consumers", "0"])]
+#[case("top2", ["--top-memory-consumers", "2"])]
+#[case("top3_default", [])]
+#[test]
+fn test_cli_top_memory_consumers<'a>(
+    #[case] snapshot_name: &str,
+    #[case] top_memory_consumers: impl IntoIterator<Item = &'a str>,
+) {
+    let mut settings = make_settings();
+
+    settings.set_snapshot_suffix(snapshot_name);
+
+    settings.add_filter(
path: datafusion-cli/tests/cli_integration.rs
user: login: 2010YOUY01
body: TIL👍🏼
created_at: 2025-05-21T04:58:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -169,9 +179,22 @@ async fn main_inner() -> Result<()> {
     if let Some(memory_limit) = args.memory_limit {
         // set memory pool type
         let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
-            PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
-            PoolType::Greedy => Arc::new(GreedyMemoryPool::new(memory_limit)),
+            PoolType::Fair if args.top_memory_consumers == 0 => {
+                Arc::new(FairSpillPool::new(memory_limit))
+            }
+            PoolType::Fair => Arc::new(TrackConsumersPool::new(
+                FairSpillPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
+            PoolType::Greedy if args.top_memory_consumers == 0 => {
+                Arc::new(GreedyMemoryPool::new(memory_limit))
+            }
+            PoolType::Greedy => Arc::new(TrackConsumersPool::new(
+                GreedyMemoryPool::new(memory_limit),
+                NonZeroUsize::new(args.top_memory_consumers).unwrap(),
+            )),
         };
path: datafusion-cli/src/main.rs
user: login: ding-young
body: @2010YOUY01 I ended up keeping the match statement as is since `TrackConsumersPool::new()` requires a concrete type implementing `MemoryPool`, and we can't wrap a `dyn MemoryPool` directly. :cry: 
created_at: 2025-05-21T16:40:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -57,6 +57,9 @@ OPTIONS:
         --mem-pool-type <MEM_POOL_TYPE>
             Specify the memory pool type 'greedy' or 'fair', default to 'greedy'
 
+        --top-memory-consumers <TOP_MEMORY_CONSUMERS>
+            The number of top memory consumers to display when query fails due to memory exhaustion. To disable memory consumer tracking, set this value to 0 [default: 3]
path: docs/source/user-guide/cli/usage.md
user: login: ding-young
body: Description updated
created_at: 2025-05-21T16:41:38Z
author_association: CONTRIBUTOR

PR comments:
user: login: jfahne
created_at: 2025-05-19T03:09:43Z
body: Ideally there would be a basic unit test to document/verify the functionality of the configuration option. I believe it would be added here ([link to cli_integration.rs](https://github.com/apache/datafusion/blob/main/datafusion-cli/tests/cli_integration.rs)).
author_association: CONTRIBUTOR
user: login: 2010YOUY01
created_at: 2025-05-19T09:22:27Z
body: I think it's good to go after adding the test.
author_association: CONTRIBUTOR
user: login: ding-young
created_at: 2025-05-19T10:02:34Z
body: @jfahne @2010YOUY01 Thank you for review! I'll add the test
author_association: CONTRIBUTOR
user: login: ding-young
created_at: 2025-05-19T13:52:23Z
body: I think the order of top memory consumers can even change :sweat: I'll update the snapshot tests.  
author_association: CONTRIBUTOR
user: login: ding-young
created_at: 2025-05-21T05:09:13Z
body: Thank you @2010YOUY01  @jfahne . I'll address them by today. 
author_association: CONTRIBUTOR