You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "ding-young (via GitHub)" <gi...@apache.org> on 2025/05/18 03:21:58 UTC

[PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

ding-young opened a new pull request, #16081:
URL: https://github.com/apache/datafusion/pull/16081

   ## 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`) 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#issuecomment-2890284245

   I think it's good to go after adding the test.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#issuecomment-2896605817

   Thank you @2010YOUY01  @jfahne . I'll address them by today. 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "jfahne (via GitHub)" <gi...@apache.org>.
jfahne commented on PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#issuecomment-2889486250

   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)).


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on code in PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#discussion_r2100742234


##########
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]

Review Comment:
   Description updated



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#discussion_r2099329043


##########
datafusion-cli/src/main.rs:
##########
@@ -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(),
+            )),
         };

Review Comment:
   [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.



##########
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>

Review Comment:
   [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.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on code in PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#discussion_r2100740833


##########
datafusion-cli/src/main.rs:
##########
@@ -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(),
+            )),
         };

Review Comment:
   @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: 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#issuecomment-2890419199

   @jfahne @2010YOUY01 Thank you for review! I'll add the test


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 merged PR #16081:
URL: https://github.com/apache/datafusion/pull/16081


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "ding-young (via GitHub)" <gi...@apache.org>.
ding-young commented on PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#issuecomment-2891122963

   I think the order of top memory consumers can even change :sweat: I'll modify the snapshot tests soon.  


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org


Re: [PR] Set `TrackConsumersPool` as default in datafusion-cli [datafusion]

Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #16081:
URL: https://github.com/apache/datafusion/pull/16081#discussion_r2099336841


##########
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(

Review Comment:
   TIL👍🏼



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org