You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "kosiew (via GitHub)" <gi...@apache.org> on 2025/08/03 10:54:24 UTC
[PR] Add On-Demand Memory Profiling Support to DataFusion CLI [datafusion]
kosiew opened a new pull request, #17021:
URL: https://github.com/apache/datafusion/pull/17021
## Which issue does this PR close?
* Closes #9246
## Rationale for this change
Understanding and optimizing memory usage is critical for query performance and stability, especially for large or complex queries. This change introduces an experimental but powerful feature to enable fine-grained, on-demand memory profiling in the DataFusion CLI, helping users better analyze and debug memory-intensive workloads.
## What changes are included in this PR?
* Added `\memory_profiling` command in the CLI with `enable`, `disable`, and `show` subcommands.
* Integrated `MemoryProfilingHandle` and enhanced memory tracking into `SessionContext`.
* Implemented `EnhancedMemoryReport` for detailed categorized memory usage analysis.
* CLI and SQL command execution pipelines now toggle memory tracking based on user input.
* Added test coverage, including snapshot tests and integration-level memory profiling tests.
* Introduced a memory profiling example in `datafusion-examples` demonstrating real-world usage.
* Updated documentation and user guide to reflect the new command and usage.
* Added new `MemoryProfilingMode` config option in `SessionConfig`.
## Are these changes tested?
Yes, the following test coverage is provided:
* Unit tests for `EnhancedMemoryReport` and `MemoryProfilingHandle`
* Integration tests validating correct toggling and memory capture
* CLI snapshot tests verifying `\memory_profiling enable/show` output
* Example usage in `examples/memory_profiling.rs`
## Are there any user-facing changes?
Yes:
* New CLI command: `\memory_profiling [enable|disable|show]`
* Memory usage report with categorized operators and summary printed to stdout
* Documentation and user guide updated
These changes are additive, opt-in, and backward compatible. The memory profiling remains disabled by default and must be explicitly enabled per query.
---
*NOTE:* The memory profiling feature is currently experimental and should be used for analysis and debugging purposes. It may not yet capture all memory allocations.
--
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] Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2294374390
##########
datafusion-cli/tests/cli_integration.rs:
##########
@@ -352,10 +401,34 @@ fn test_backtrace_output(#[case] query: &str) {
let stderr = String::from_utf8_lossy(&output.stderr);
let combined_output = format!("{}{}", stdout, stderr);
- // Assert that the output includes literal 'backtrace'
+ // Accept either a printed backtrace or a readable error message.
+ // Some builds may not include backtrace support in the binary; in that
+ // case the CLI prints a clear planning error message instead. Verify one
+ // of these is present to avoid a fragile test.
+ let has_backtrace = combined_output.to_lowercase().contains("backtrace");
+ let lower_stdout = stdout.to_lowercase();
+ let lower_stderr = stderr.to_lowercase();
+
+ let has_planning_error = lower_stdout.contains("failed to coerce arguments")
Review Comment:
I don't understand any of these changes -- why would we accept arrow cast errors?
##########
datafusion/execution/src/memory_pool/metrics.rs:
##########
@@ -0,0 +1,90 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Memory usage metrics for query execution.
+
+use super::{human_readable_size, ConsumerMemoryMetrics};
+use std::{collections::BTreeMap, fmt::Write};
+/// Format summary of memory usage metrics.
+///
+/// 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 {
Review Comment:
I think this formatting stuff should be put in the CLI code for now (not the main repo)
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -52,6 +52,19 @@ pub trait CliSessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError>;
+
+ /// Return true if memory profiling is enabled.
+ fn memory_profiling(&self) -> bool {
Review 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
https://github.com/apache/datafusion/blob/df45d186d34f2ac131d64e4a068d9f39b35e99c7/datafusion-cli/src/print_options.rs#L73-L72
##########
datafusion-cli/tests/snapshots/aws_options.snap:
##########
@@ -1,9 +1,9 @@
---
-source: tests/cli_integration.rs
+source: datafusion-cli/tests/cli_integration.rs
Review Comment:
Why did you change these files?
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "rluvaton (via GitHub)" <gi...@apache.org>.
rluvaton commented on PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#issuecomment-3157497138
Can you please attach how this will look like for the user?
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#issuecomment-3151073315
FYI @ding-young @2010YOUY01 and @rluvaton who I think have been working on memory usage recently
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2290689291
##########
datafusion-cli/tests/snapshots/aws_options.snap:
##########
@@ -12,6 +12,11 @@ exit_code: 0
0 row(s) fetched.
[ELAPSED]
+Peak memory usage: 0.0 B
Review Comment:
Amended to
memory_profiling off by default
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284378398
##########
datafusion-cli/src/command.rs:
##########
@@ -206,11 +243,31 @@ impl FromStr for Command {
Self::OutputFormat(Some(subcommand.to_string()))
}
("pset", None) => Self::OutputFormat(None),
+ ("memory_profiling", sub) => {
Review Comment:
Good idea!
Implemented
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2290666767
##########
datafusion-cli/src/command.rs:
##########
@@ -206,11 +240,30 @@ impl FromStr for Command {
Self::OutputFormat(Some(subcommand.to_string()))
}
("pset", None) => Self::OutputFormat(None),
+ ("memory_profiling", sub) => {
+ let sub = match sub {
+ Some(s) => Some(s.parse::<MemoryProfilingCommand>().map_err(|_| ())?),
+ None => None,
+ };
+ Self::MemoryProfiling(sub)
+ }
_ => return Err(()),
})
}
}
+impl FromStr for MemoryProfilingCommand {
Review Comment:
Removed
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284281614
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +156,128 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ tracked_memory_pool: RwLock<Option<Arc<dyn TrackedPool>>>,
+ top_memory_consumers: usize,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ Self {
+ ctx,
+ memory_profiling: AtomicBool::new(false),
+ base_memory_pool,
+ tracked_memory_pool: RwLock::new(None),
+ top_memory_consumers,
+ }
+ }
+}
+
+#[async_trait::async_trait]
+impl CliSessionContext for ReplSessionContext {
Review Comment:
still required for casting SessionContext to dyn CliSessionContext.
--
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] Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#issuecomment-3494534719
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.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2290688291
##########
datafusion/core/src/execution/session_state.rs:
##########
@@ -16,69 +16,76 @@
// under the License.
//! [`SessionState`]: information required to run queries in a session
-
-use std::any::Any;
-use std::collections::hash_map::Entry;
-use std::collections::{HashMap, HashSet};
-use std::fmt::Debug;
-use std::sync::Arc;
-
-use crate::catalog::{CatalogProviderList, SchemaProvider, TableProviderFactory};
-use crate::datasource::cte_worktable::CteWorkTable;
-use crate::datasource::file_format::{format_as_file_type, FileFormatFactory};
-use crate::datasource::provider_as_source;
-use crate::execution::context::{EmptySerializerRegistry, FunctionFactory, QueryPlanner};
-use crate::execution::SessionStateDefaults;
-use crate::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
+use crate::{
Review Comment:
Reverted
##########
datafusion-cli/README.md:
##########
@@ -30,3 +30,31 @@ DataFusion CLI (`datafusion-cli`) is a small command line utility that runs SQL
## Where can I find more information?
See the [`datafusion-cli` documentation](https://datafusion.apache.org/user-guide/cli/index.html) for further information.
+
+## Memory Profiling
Review Comment:
Moved
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2290686071
##########
datafusion-cli/src/main.rs:
##########
@@ -174,27 +176,13 @@ async fn main_inner() -> Result<()> {
let session_config = get_session_config(&args)?;
let mut rt_builder = RuntimeEnvBuilder::new();
- // set memory pool size
if let Some(memory_limit) = args.memory_limit {
// set memory pool type
let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
- 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(),
- )),
+ PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
Review Comment:
No.
✅ Previous Behavior (Before Change)
If --top-memory-consumers > 0, a TrackConsumersPool wraps the actual memory pool.
Otherwise, just the base pool (FairSpillPool or GreedyMemoryPool) is used.
No tracking if --top-memory-consumers == 0.
✅ New Behavior (After Change)
The base_memory_pool is always constructed first, regardless of tracking.
If --top-memory-consumers > 0, a TrackConsumersPool is wrapped around base_memory_pool and passed to the runtime.
It is also assigned to a tracked_pool: Option<Arc<dyn TrackedPool>> and passed to ReplSessionContext.
Additionally, disable_tracking() is explicitly called to ensure it's initially not profiling until enabled.
If --top-memory-consumers == 0, tracking is skipped, and the base memory pool is used as-is.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#issuecomment-3160334167
hi @rluvaton ,
I included
- datafusion-cli/README.md - has a brief example of how to run this from datafusion-cli
- datafusion-examples/examples/memory_profiling.rs - which shows how to run it in a .rs file.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2279058579
##########
datafusion-cli/README.md:
##########
@@ -30,3 +30,33 @@ DataFusion CLI (`datafusion-cli`) is a small command line utility that runs SQL
## Where can I find more information?
See the [`datafusion-cli` documentation](https://datafusion.apache.org/user-guide/cli/index.html) for further information.
+
+## Memory Profiling
+
+> **Tip:** Memory profiling requires the tracked pool. Start the CLI with `--top-memory-consumers N` (N≥1), or profiling will report no metrics. By default, CLI starts with --top-memory-consumers 5.
+
+Enable memory tracking for the next query and display the report afterwards:
+
+```text
+> \memory_profiling enable
+Memory profiling enabled
+> SELECT v % 100 AS group_key, COUNT(*) AS cnt, SUM(v) AS sum_v FROM generate_series(1,100000) AS t(v) GROUP BY group_key ORDER BY group_key;
+
++-----------+------+----------+
+| group_key | cnt | sum_v |
++-----------+------+----------+
+| 0 | 1000 | 50050000 |
+| 1 | 1000 | 49951000 |
+| 2 | 1000 | 49952000 |
+...
+
+\memory_profiling show
Review Comment:
Thanks @kosiew
In terms of user experience, I feel like `\memory_profiling show` is redundant
Specifically, I think if I as a user enabled memory profiling I would expect that all queries I ran would produce a memory report until it was disabled.
So that would mean the report would always be shown
##########
datafusion-cli/src/main.rs:
##########
@@ -174,27 +176,15 @@ async fn main_inner() -> Result<()> {
let session_config = get_session_config(&args)?;
let mut rt_builder = RuntimeEnvBuilder::new();
- // set memory pool size
+ let mut base_pool: Option<Arc<dyn MemoryPool>> = None;
Review Comment:
Does this disable memory tracking by default?
I was worried that that will be a regression when a query OOMs it won't show the top consumers anymore
However, I tried it locally and it seems to still work:
```shell
andrewlamb@Andrews-MacBook-Pro-3:~/Software/datafusion$ cargo run -p datafusion-cli -- -m 10m
warning: unused variable: `options`
--> datafusion/datasource-parquet/src/reader.rs:247:9
|
247 | options: Option<&'a ArrowReaderOptions>,
| ^^^^^^^ help: if this is intentional, prefix it with an underscore: `_options`
|
= note: `#[warn(unused_variables)]` on by default
warning: `datafusion-datasource-parquet` (lib) generated 1 warning
Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.14s
Running `target/debug/datafusion-cli -m 10m`
DataFusion CLI v49.0.0
> select distinct "ADDRESS1", "ADDRESS2" from '/Users/andrewlamb/Downloads/f94d0c87-8798-4bf6-9c98-8d89971e2539.parquet';
Resources exhausted: Additional allocation failed with top memory consumers (across reservations) as:
GroupedHashAggregateStream[11] ()#27(can spill: true) consumed 4.7 MB,
GroupedHashAggregateStream[0] ()#15(can spill: true) consumed 4.5 MB,
GroupedHashAggregateStream[2] ()#18(can spill: true) consumed 0.0 B,
GroupedHashAggregateStream[10] ()#26(can spill: true) consumed 0.0 B,
RepartitionExec[5]#38(can spill: false) consumed 0.0 B.
Error: Failed to allocate additional 1065.0 KB for RepartitionExec[0] with 0.0 B already allocated for this reservation - 844.8 KB remain available for the total pool
```
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -15,19 +15,69 @@
// specific language governing permissions and limitations
// under the License.
-use std::sync::Arc;
+use std::num::NonZeroUsize;
+use std::sync::{
+ atomic::{AtomicBool, Ordering},
+ Arc, RwLock,
+};
use datafusion::{
dataframe::DataFrame,
error::DataFusionError,
- execution::{context::SessionState, TaskContext},
+ execution::{
+ context::SessionState,
+ memory_pool::{
+ MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
+ TrackConsumersPool, TrackedPool,
+ },
+ runtime_env::RuntimeEnvBuilder,
+ session_state::SessionStateBuilder,
+ TaskContext,
+ },
logical_expr::LogicalPlan,
prelude::SessionContext,
};
use object_store::ObjectStore;
use crate::object_storage::{AwsOptions, GcpOptions};
+#[derive(Debug)]
Review Comment:
As I understand this code, I think it makes an `Arc<dyn MemoryPool>` implement the MemoryPool API directly
However since we can already copy an `Arc` I don't undertand why this code is necessary so I played around with it.
Here is an alternate suggestion that I think is more general -- basically automatically derive MemoryPool for any `Arc<dyn MemoryPool>`:
- https://github.com/kosiew/datafusion/pull/30
##########
datafusion-cli/src/command.rs:
##########
@@ -206,11 +243,31 @@ impl FromStr for Command {
Self::OutputFormat(Some(subcommand.to_string()))
}
("pset", None) => Self::OutputFormat(None),
+ ("memory_profiling", sub) => {
Review Comment:
What do you think about supporting `\memory_profling` without an argument that works like a toggle to enable/disable printing?
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +156,128 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ tracked_memory_pool: RwLock<Option<Arc<dyn TrackedPool>>>,
+ top_memory_consumers: usize,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ Self {
+ ctx,
+ memory_profiling: AtomicBool::new(false),
+ base_memory_pool,
+ tracked_memory_pool: RwLock::new(None),
+ top_memory_consumers,
+ }
+ }
+}
+
+#[async_trait::async_trait]
+impl CliSessionContext for ReplSessionContext {
+ fn task_ctx(&self) -> Arc<TaskContext> {
+ self.ctx.task_ctx()
+ }
+
+ fn session_state(&self) -> SessionState {
+ self.ctx.state()
+ }
+
+ fn register_object_store(
+ &self,
+ url: &url::Url,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Option<Arc<dyn ObjectStore + 'static>> {
+ self.ctx.register_object_store(url, object_store)
+ }
+
+ fn register_table_options_extension_from_scheme(&self, scheme: &str) {
+ match scheme {
+ // For Amazon S3 or Alibaba Cloud OSS
+ "s3" | "oss" | "cos" => self
+ .ctx
+ .register_table_options_extension(AwsOptions::default()),
+ // For Google Cloud Storage
+ "gs" | "gcs" => self
+ .ctx
+ .register_table_options_extension(GcpOptions::default()),
+ // For unsupported schemes, do nothing:
+ _ => {}
+ }
+ }
+
+ async fn execute_logical_plan(
+ &self,
+ plan: LogicalPlan,
+ ) -> Result<DataFrame, DataFusionError> {
+ self.ctx.execute_logical_plan(plan).await
+ }
+
+ fn memory_profiling(&self) -> bool {
+ self.memory_profiling.load(Ordering::Relaxed)
+ }
+
+ fn set_memory_profiling(&self, enable: bool) {
+ if enable {
+ if self.top_memory_consumers == 0 {
+ return;
+ }
+ if self.memory_profiling.swap(true, Ordering::Relaxed) {
+ return;
+ }
+ let tracked = Arc::new(TrackConsumersPool::new(
Review Comment:
One thing I don't understand about this design is that the current datafusion-cli always uses a TrackConsumersPool, but this PR seems to install / uninstall a tracked pool bsed on the `memory_consumers` setting
Would it be possible to *always* use the same TrackConsumersPool that is used today, and have `\memory_profiing` control when that information was displayed?
I think the PR would get a lot simpler then -- it woudl be a matter of adding logic to display the TrackConsumersPool that was already present
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +156,128 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ tracked_memory_pool: RwLock<Option<Arc<dyn TrackedPool>>>,
+ top_memory_consumers: usize,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ Self {
+ ctx,
+ memory_profiling: AtomicBool::new(false),
+ base_memory_pool,
+ tracked_memory_pool: RwLock::new(None),
+ top_memory_consumers,
+ }
+ }
+}
+
+#[async_trait::async_trait]
+impl CliSessionContext for ReplSessionContext {
Review Comment:
does this mean we can remove the impl for normal `SessionContext`?
https://github.com/apache/datafusion/blob/e4e1cee82a734502ef79186e1bb16dddc237443a/datafusion-cli/src/cli_context.rs#L121-L120
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284281614
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +156,128 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ tracked_memory_pool: RwLock<Option<Arc<dyn TrackedPool>>>,
+ top_memory_consumers: usize,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ Self {
+ ctx,
+ memory_profiling: AtomicBool::new(false),
+ base_memory_pool,
+ tracked_memory_pool: RwLock::new(None),
+ top_memory_consumers,
+ }
+ }
+}
+
+#[async_trait::async_trait]
+impl CliSessionContext for ReplSessionContext {
Review Comment:
It is still required for casting SessionContext to dyn CliSessionContext.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284282422
##########
datafusion-cli/src/main.rs:
##########
@@ -174,27 +176,15 @@ async fn main_inner() -> Result<()> {
let session_config = get_session_config(&args)?;
let mut rt_builder = RuntimeEnvBuilder::new();
- // set memory pool size
+ let mut base_pool: Option<Arc<dyn MemoryPool>> = None;
Review Comment:
Amended to not require
`mut base_pool`
--
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] DRAFT Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2298185451
##########
datafusion/execution/src/memory_pool/metrics.rs:
##########
@@ -0,0 +1,90 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Memory usage metrics for query execution.
+
+use super::{human_readable_size, ConsumerMemoryMetrics};
+use std::{collections::BTreeMap, fmt::Write};
+/// Format summary of memory usage metrics.
+///
+/// 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 {
Review Comment:
Moved
--
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] Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2297004893
##########
datafusion-cli/tests/cli_integration.rs:
##########
@@ -352,10 +401,34 @@ fn test_backtrace_output(#[case] query: &str) {
let stderr = String::from_utf8_lossy(&output.stderr);
let combined_output = format!("{}{}", stdout, stderr);
- // Assert that the output includes literal 'backtrace'
+ // Accept either a printed backtrace or a readable error message.
+ // Some builds may not include backtrace support in the binary; in that
+ // case the CLI prints a clear planning error message instead. Verify one
+ // of these is present to avoid a fragile test.
+ let has_backtrace = combined_output.to_lowercase().contains("backtrace");
+ let lower_stdout = stdout.to_lowercase();
+ let lower_stderr = stderr.to_lowercase();
+
+ let has_planning_error = lower_stdout.contains("failed to coerce arguments")
Review Comment:
<img width="1512" height="400" alt="image" src="https://github.com/user-attachments/assets/71776aa2-97dc-474b-97f3-7fde8b56fcba" />
🤦 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.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284285161
##########
datafusion-cli/README.md:
##########
@@ -30,3 +30,33 @@ DataFusion CLI (`datafusion-cli`) is a small command line utility that runs SQL
## Where can I find more information?
See the [`datafusion-cli` documentation](https://datafusion.apache.org/user-guide/cli/index.html) for further information.
+
+## Memory Profiling
+
+> **Tip:** Memory profiling requires the tracked pool. Start the CLI with `--top-memory-consumers N` (N≥1), or profiling will report no metrics. By default, CLI starts with --top-memory-consumers 5.
+
+Enable memory tracking for the next query and display the report afterwards:
+
+```text
+> \memory_profiling enable
+Memory profiling enabled
+> SELECT v % 100 AS group_key, COUNT(*) AS cnt, SUM(v) AS sum_v FROM generate_series(1,100000) AS t(v) GROUP BY group_key ORDER BY group_key;
+
++-----------+------+----------+
+| group_key | cnt | sum_v |
++-----------+------+----------+
+| 0 | 1000 | 50050000 |
+| 1 | 1000 | 49951000 |
+| 2 | 1000 | 49952000 |
+...
+
+\memory_profiling show
Review Comment:
Added automatic printing of memory report
and removed `\memory_profiling show `
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284282422
##########
datafusion-cli/src/main.rs:
##########
@@ -174,27 +176,15 @@ async fn main_inner() -> Result<()> {
let session_config = get_session_config(&args)?;
let mut rt_builder = RuntimeEnvBuilder::new();
- // set memory pool size
+ let mut base_pool: Option<Arc<dyn MemoryPool>> = None;
Review Comment:
Simplified and removed
`mut base_pool`
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2279144398
##########
datafusion-cli/README.md:
##########
@@ -30,3 +30,33 @@ DataFusion CLI (`datafusion-cli`) is a small command line utility that runs SQL
## Where can I find more information?
See the [`datafusion-cli` documentation](https://datafusion.apache.org/user-guide/cli/index.html) for further information.
+
+## Memory Profiling
+
+> **Tip:** Memory profiling requires the tracked pool. Start the CLI with `--top-memory-consumers N` (N≥1), or profiling will report no metrics. By default, CLI starts with --top-memory-consumers 5.
+
+Enable memory tracking for the next query and display the report afterwards:
+
+```text
+> \memory_profiling enable
+Memory profiling enabled
+> SELECT v % 100 AS group_key, COUNT(*) AS cnt, SUM(v) AS sum_v FROM generate_series(1,100000) AS t(v) GROUP BY group_key ORDER BY group_key;
+
++-----------+------+----------+
+| group_key | cnt | sum_v |
++-----------+------+----------+
+| 0 | 1000 | 50050000 |
+| 1 | 1000 | 49951000 |
+| 2 | 1000 | 49952000 |
+...
+
+\memory_profiling show
Review Comment:
I think probably we can include this memory report inside `EXPLAIN ANALYZE`, and remove this `\memory_profiling show` command
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2289199544
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -1373,7 +1373,8 @@ impl DataFrame {
pub async fn collect(self) -> Result<Vec<RecordBatch>> {
let task_ctx = Arc::new(self.task_ctx());
let plan = self.create_physical_plan().await?;
- collect(plan, task_ctx).await
+ let batches = collect(plan, task_ctx).await?;
Review Comment:
what is the purpose of these changes?
##########
datafusion-cli/tests/snapshots/aws_options.snap:
##########
@@ -12,6 +12,11 @@ exit_code: 0
0 row(s) fetched.
[ELAPSED]
+Peak memory usage: 0.0 B
Review Comment:
this change shows the memory profiling is enabled by default. I think memory profiling would be more useful if we had it off by default and then enabled when \profiling was on
##########
datafusion/core/src/execution/session_state.rs:
##########
@@ -16,69 +16,76 @@
// under the License.
//! [`SessionState`]: information required to run queries in a session
-
-use std::any::Any;
-use std::collections::hash_map::Entry;
-use std::collections::{HashMap, HashSet};
-use std::fmt::Debug;
-use std::sync::Arc;
-
-use crate::catalog::{CatalogProviderList, SchemaProvider, TableProviderFactory};
-use crate::datasource::cte_worktable::CteWorkTable;
-use crate::datasource::file_format::{format_as_file_type, FileFormatFactory};
-use crate::datasource::provider_as_source;
-use crate::execution::context::{EmptySerializerRegistry, FunctionFactory, QueryPlanner};
-use crate::execution::SessionStateDefaults;
-use crate::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
+use crate::{
Review Comment:
I think these changes lead to uncessary churn and make the code base more inconsistent -- can we please pull them out into their own PR if you want to refactor the `use` statements?
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +116,120 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ tracked_memory_pool: Option<Arc<dyn TrackedPool>>,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ let tracked_memory_pool = if top_memory_consumers > 0 {
Review Comment:
this seems like it is redundant with the top level pool that is installed with the CLI (which also may be a tracked pool)
##########
datafusion-cli/README.md:
##########
@@ -30,3 +30,31 @@ DataFusion CLI (`datafusion-cli`) is a small command line utility that runs SQL
## Where can I find more information?
See the [`datafusion-cli` documentation](https://datafusion.apache.org/user-guide/cli/index.html) for further information.
+
+## Memory Profiling
Review Comment:
This seems like it should be in the CLI documentation rather than this README
##########
datafusion-cli/src/command.rs:
##########
@@ -206,11 +240,30 @@ impl FromStr for Command {
Self::OutputFormat(Some(subcommand.to_string()))
}
("pset", None) => Self::OutputFormat(None),
+ ("memory_profiling", sub) => {
+ let sub = match sub {
+ Some(s) => Some(s.parse::<MemoryProfilingCommand>().map_err(|_| ())?),
+ None => None,
+ };
+ Self::MemoryProfiling(sub)
+ }
_ => return Err(()),
})
}
}
+impl FromStr for MemoryProfilingCommand {
Review Comment:
this eems a bit left over from when there was a command that took a value
##########
datafusion-cli/src/main.rs:
##########
@@ -174,27 +176,13 @@ async fn main_inner() -> Result<()> {
let session_config = get_session_config(&args)?;
let mut rt_builder = RuntimeEnvBuilder::new();
- // set memory pool size
if let Some(memory_limit) = args.memory_limit {
// set memory pool type
let pool: Arc<dyn MemoryPool> = match args.mem_pool_type {
- 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(),
- )),
+ PoolType::Fair => Arc::new(FairSpillPool::new(memory_limit)),
Review Comment:
if we change this doesn't it mean that we will break the existing behavior of showing hte top consumers when memory profiling is not enabled?
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2284283586
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +156,128 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ tracked_memory_pool: RwLock<Option<Arc<dyn TrackedPool>>>,
+ top_memory_consumers: usize,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ Self {
+ ctx,
+ memory_profiling: AtomicBool::new(false),
+ base_memory_pool,
+ tracked_memory_pool: RwLock::new(None),
+ top_memory_consumers,
+ }
+ }
+}
+
+#[async_trait::async_trait]
+impl CliSessionContext for ReplSessionContext {
+ fn task_ctx(&self) -> Arc<TaskContext> {
+ self.ctx.task_ctx()
+ }
+
+ fn session_state(&self) -> SessionState {
+ self.ctx.state()
+ }
+
+ fn register_object_store(
+ &self,
+ url: &url::Url,
+ object_store: Arc<dyn ObjectStore>,
+ ) -> Option<Arc<dyn ObjectStore + 'static>> {
+ self.ctx.register_object_store(url, object_store)
+ }
+
+ fn register_table_options_extension_from_scheme(&self, scheme: &str) {
+ match scheme {
+ // For Amazon S3 or Alibaba Cloud OSS
+ "s3" | "oss" | "cos" => self
+ .ctx
+ .register_table_options_extension(AwsOptions::default()),
+ // For Google Cloud Storage
+ "gs" | "gcs" => self
+ .ctx
+ .register_table_options_extension(GcpOptions::default()),
+ // For unsupported schemes, do nothing:
+ _ => {}
+ }
+ }
+
+ async fn execute_logical_plan(
+ &self,
+ plan: LogicalPlan,
+ ) -> Result<DataFrame, DataFusionError> {
+ self.ctx.execute_logical_plan(plan).await
+ }
+
+ fn memory_profiling(&self) -> bool {
+ self.memory_profiling.load(Ordering::Relaxed)
+ }
+
+ fn set_memory_profiling(&self, enable: bool) {
+ if enable {
+ if self.top_memory_consumers == 0 {
+ return;
+ }
+ if self.memory_profiling.swap(true, Ordering::Relaxed) {
+ return;
+ }
+ let tracked = Arc::new(TrackConsumersPool::new(
Review Comment:
Simplified as suggested.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2289924264
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -1373,7 +1373,8 @@ impl DataFrame {
pub async fn collect(self) -> Result<Vec<RecordBatch>> {
let task_ctx = Arc::new(self.task_ctx());
let plan = self.create_physical_plan().await?;
- collect(plan, task_ctx).await
+ let batches = collect(plan, task_ctx).await?;
Review Comment:
Oops, it was a change made during testing. I'll revert it.
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2290687827
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -93,6 +116,120 @@ impl CliSessionContext for SessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError> {
- self.execute_logical_plan(plan).await
+ SessionContext::execute_logical_plan(self, plan).await
+ }
+}
+
+/// Session context used by the CLI with memory profiling support.
+pub struct ReplSessionContext {
+ ctx: SessionContext,
+ memory_profiling: AtomicBool,
+ tracked_memory_pool: Option<Arc<dyn TrackedPool>>,
+}
+
+impl ReplSessionContext {
+ pub fn new(
+ ctx: SessionContext,
+ base_memory_pool: Arc<dyn MemoryPool>,
+ top_memory_consumers: usize,
+ ) -> Self {
+ let tracked_memory_pool = if top_memory_consumers > 0 {
Review Comment:
Changed implementation
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2263831579
##########
datafusion-cli/src/exec.rs:
##########
@@ -227,9 +227,19 @@ pub(super) async fn exec_and_print(
let statements = DFParser::parse_sql_with_dialect(&sql, dialect.as_ref())?;
for statement in statements {
+ let _mem_handle = if print_options.memory_profiling {
+ // RAII guard: dropping the handle disables profiling after execution
+ Some(ctx.enable_memory_profiling())
+ } else {
+ None
+ };
StatementExecutor::new(statement)
.execute(ctx, print_options)
.await?;
+ // disable after each statement
Review Comment:
If you are going to use a RAAI style thing I would expect the `Drop` impl to disable profiling 🤔
##########
datafusion-cli/src/print_options.rs:
##########
@@ -73,6 +73,7 @@ pub struct PrintOptions {
pub quiet: bool,
pub maxrows: MaxRows,
pub color: bool,
+ pub memory_profiling: bool,
Review Comment:
memory profiling seems like something that should be attached to the session context, as it is not related to formatting output
##########
datafusion-cli/src/exec.rs:
##########
@@ -54,7 +54,7 @@ use tokio::signal;
pub async fn exec_from_commands(
ctx: &dyn CliSessionContext,
commands: Vec<String>,
- print_options: &PrintOptions,
+ print_options: &mut PrintOptions,
Review Comment:
It seems strange that now PrintOptions are mutable - why would executing a command change the options?
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -90,6 +90,244 @@ use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use parking_lot::RwLock;
use url::Url;
+/// Memory profiling report for a query.
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct MemoryReport {
+ metrics: std::collections::HashMap<String, usize>,
+}
+
+impl MemoryReport {
+ /// Create a new [`MemoryReport`] from the provided metrics.
+ pub fn new(metrics: std::collections::HashMap<String, usize>) -> Self {
+ Self { metrics }
+ }
+
+ /// Returns `true` if the report contains no metrics.
+ pub fn is_empty(&self) -> bool {
+ self.metrics.is_empty()
+ }
+
+ /// Number of tracked operators in the report.
+ pub fn len(&self) -> usize {
+ self.metrics.len()
+ }
+
+ /// Consume the report and return the underlying metrics.
+ pub fn into_inner(self) -> std::collections::HashMap<String, usize> {
+ self.metrics
+ }
+}
+
+impl std::ops::Deref for MemoryReport {
+ type Target = std::collections::HashMap<String, usize>;
+
+ fn deref(&self) -> &Self::Target {
+ &self.metrics
+ }
+}
+
+impl From<std::collections::HashMap<String, usize>> for MemoryReport {
+ fn from(metrics: std::collections::HashMap<String, usize>) -> Self {
+ Self::new(metrics)
+ }
+}
+// Implement IntoIterator for &MemoryReport to allow iterating over &report
+impl<'a> IntoIterator for &'a MemoryReport {
+ type Item = (&'a String, &'a usize);
+ type IntoIter = std::collections::hash_map::Iter<'a, String, usize>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ self.metrics.iter()
+ }
+}
+/// Enhanced memory profiling report with categorization and analysis
+#[derive(Debug)]
+pub struct EnhancedMemoryReport {
Review Comment:
I wonder why you chose to put this in the session/execution context and not in the existing memory pool https://docs.rs/datafusion/latest/datafusion/execution/memory_pool/index.html
In particular, I think this code has a lot of similarity to https://docs.rs/datafusion/latest/datafusion/execution/memory_pool/struct.TrackConsumersPool.html
Perhaps we can modify that pool or add another wrapper pool for DataFusion CLI that saves the maximum consumers, etc
##########
datafusion/core/src/execution/context/mod.rs:
##########
@@ -90,6 +90,244 @@ use chrono::{DateTime, Utc};
use object_store::ObjectStore;
use parking_lot::RwLock;
use url::Url;
+/// Memory profiling report for a query.
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct MemoryReport {
+ metrics: std::collections::HashMap<String, usize>,
Review Comment:
can we please document what this hash map contains (is it memory consumer names?)
##########
datafusion-cli/src/exec.rs:
##########
@@ -227,9 +227,19 @@ pub(super) async fn exec_and_print(
let statements = DFParser::parse_sql_with_dialect(&sql, dialect.as_ref())?;
for statement in statements {
+ let _mem_handle = if print_options.memory_profiling {
+ // RAII guard: dropping the handle disables profiling after execution
Review Comment:
I don't understand why this is named with a leading underscore given it is referenced below
--
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] Add Memory Profiling Support to DataFusion CLI [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2279414621
##########
datafusion-cli/README.md:
##########
@@ -30,3 +30,33 @@ DataFusion CLI (`datafusion-cli`) is a small command line utility that runs SQL
## Where can I find more information?
See the [`datafusion-cli` documentation](https://datafusion.apache.org/user-guide/cli/index.html) for further information.
+
+## Memory Profiling
+
+> **Tip:** Memory profiling requires the tracked pool. Start the CLI with `--top-memory-consumers N` (N≥1), or profiling will report no metrics. By default, CLI starts with --top-memory-consumers 5.
+
+Enable memory tracking for the next query and display the report afterwards:
+
+```text
+> \memory_profiling enable
+Memory profiling enabled
+> SELECT v % 100 AS group_key, COUNT(*) AS cnt, SUM(v) AS sum_v FROM generate_series(1,100000) AS t(v) GROUP BY group_key ORDER BY group_key;
+
++-----------+------+----------+
+| group_key | cnt | sum_v |
++-----------+------+----------+
+| 0 | 1000 | 50050000 |
+| 1 | 1000 | 49951000 |
+| 2 | 1000 | 49952000 |
+...
+
+\memory_profiling show
Review Comment:
In order to do that we might have to change the core of DataFusion (specifically build the reporting into EXPLAIN ANALYZE)
I think it is worth keeping things modular and reporting outside datafusion's core at first and keep it in datafusion-cli, and if it proves to be a popular feature we can then consider moving it into the core
--
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] Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2297716757
##########
datafusion-cli/tests/snapshots/aws_options.snap:
##########
@@ -1,9 +1,9 @@
---
-source: tests/cli_integration.rs
+source: datafusion-cli/tests/cli_integration.rs
Review Comment:
These unrelated changes have been cleaned up.
--
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] Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "kosiew (via GitHub)" <gi...@apache.org>.
kosiew commented on code in PR #17021:
URL: https://github.com/apache/datafusion/pull/17021#discussion_r2297533222
##########
datafusion-cli/src/cli_context.rs:
##########
@@ -52,6 +52,19 @@ pub trait CliSessionContext {
&self,
plan: LogicalPlan,
) -> Result<DataFrame, DataFusionError>;
+
+ /// Return true if memory profiling is enabled.
+ fn memory_profiling(&self) -> bool {
Review 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.
--
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] Add memory profiling support to DataFusion CLI and memory pool metrics [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] closed pull request #17021: Add memory profiling support to DataFusion CLI and memory pool metrics
URL: https://github.com/apache/datafusion/pull/17021
--
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