Add memory profiling support to DataFusion CLI and memory pool metrics #17021
Conversation
github-actions
bot
added
documentation
…y profiling and tracker fields in SessionStateBuilder
…blank lines - Fixed the spelling of "Apche" to "Apache" in the license comment. - Removed unnecessary blank lines in `no_grouping.rs`, `symmetric_hash_join.rs`, and `sort.rs`.
…porting" This reverts commit 1258405.
…hensive reporting" This reverts commit 8dbdfa0.
…aFusion" This reverts commit 7db0933.
kosiew
changed the title
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thanks @kosiew -- I am sorry for the slow reviews, but I am really having trouble reviewing this PR -- I don't know why there are more unrelated changes that showed up here the changes to the CLI testing
Are you by any chance using a AI coding tool? If so, perhaps you can spend some time reviewing the PRs a bit more yourself and spit them into smaller PRs for review?
Some of the other changes in this PR are nice (like the change to tracked pool, potentially) but they seem unrelated to printing the contents of the memory
I am sorry for another round of feedback (this is like the 3rd time and there are already 280 commits 🤯 )
In general, I think we would have an easier time reviewing PRs (and get them merged faster) if you could keep them smaller and more focused
| let lower_stdout = stdout.to_lowercase(); | ||
| let lower_stderr = stderr.to_lowercase(); | ||
|
|
||
| let has_planning_error = lower_stdout.contains("failed to coerce arguments") |
There was a problem hiding this comment.
I don't understand any of these changes -- why would we accept arrow cast errors?
There was a problem hiding this comment.
🤦 This seemed to be from tests related to another branch related to cast_column for struct support.
I am very sorry that I did not notice the unrelated change creeping into this PR.
| /// Returns a string with peak usage, cumulative allocations, and totals per | ||
| /// operator category. The caller is responsible for printing the returned | ||
| /// string if desired. | ||
| pub fn format_metrics(metrics: &[ConsumerMemoryMetrics]) -> String { |
There was a problem hiding this comment.
I think this formatting stuff should be put in the CLI code for now (not the main repo)
| ) -> Result<DataFrame, DataFusionError>; | ||
|
|
||
| /// Return true if memory profiling is enabled. | ||
| fn memory_profiling(&self) -> bool { |
There was a problem hiding this comment.
I think the memory profiling flag would best be stored on PrintOptions, similarly to the quiet mode flag (that suppresses execution time printing). Then you would not need to introduce so much new code and a new trait
There was a problem hiding this comment.
That would require print_options to be mut because we can
toggle memory_profiling
It was moved out of print_options after a comment that print_options should not be mut for memory profiling.
| @@ -1,9 +1,9 @@ | |||
| --- | |||
| source: tests/cli_integration.rs | |||
| source: datafusion-cli/tests/cli_integration.rs | |||
There was a problem hiding this comment.
Why did you change these files?
There was a problem hiding this comment.
These unrelated changes have been cleaned up.
kosiew
changed the title
… documentation Updated the default value for the `--top-memory-consumers` option from 3 to 5 in the CLI usage documentation to reflect the correct configuration.
…acking Added detailed metrics collection for memory consumption in the memory profiling example. This includes peak memory and cumulative allocation bytes for each operator or execution stage, along with a summarized output by operator category.
kosiew
changed the title
Resolved conflicts in memory pool implementation: - Kept cumulative memory tracking feature in TrackedConsumer - Updated test snapshots to match current error message format - All conflicts resolved in favor of HEAD branch (memory-16904)
|
Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days. |
4 participants
Which issue does this PR close?
Rationale for this change
Currently, the DataFusion CLI has limited visibility into query memory usage. Users cannot easily track peak memory, cumulative allocations, or operator-level memory breakdowns. This PR introduces memory profiling features, providing actionable insights for debugging, tuning, and capacity planning. It also strengthens error reporting with clearer context about memory consumers.
What changes are included in this PR?
CLI Enhancements:
\memory_profilingto toggle profiling during CLI sessions.Core Enhancements:
TrackedPooltrait for memory pools with consumer tracking.ConsumerMemoryMetricsandformat_metricsfor reporting.MemoryReservationto track peak usage and provide human-readable summaries.Examples and Documentation:
memory_profiling.rsexample demonstrating programmatic usage.Tests and Snapshots:
Are these changes tested?
Yes:
TrackConsumersPooland memory metrics.cli_memory_auto_report,cli_memory_disable_stops_report).Are there any user-facing changes?
Yes:
\memory_profiling(toggle profiling on/off).--top-memory-consumers N(default now 5, previously 3).memory_profiling.rs.No breaking API changes; all new features are additive.