diff --git a/Cargo.lock b/Cargo.lock
index 47aa178a57f..9b895626f1b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -14412,6 +14412,7 @@ dependencies = [
"common-time",
"common-wal",
"datafusion-common",
+ "datafusion-execution",
"datafusion-expr",
"datafusion-physical-plan",
"datatypes",
diff --git a/Cargo.toml b/Cargo.toml
index e2f91a54336..3cb4e77fb77 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -133,6 +133,7 @@ dashmap = "6.1"
datafusion = "=55.1.0"
datafusion-common = "=55.1.0"
datafusion-datasource = "=55.1.0"
+datafusion-execution = "=55.1.0"
datafusion-expr = "=55.1.0"
datafusion-expr-common = "=55.1.0"
datafusion-functions = "=55.1.0"
@@ -368,6 +369,7 @@ substrait = { git = "https://github.com/GreptimeTeam/substrait-rs.git", rev = "9
datafusion = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
datafusion-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
datafusion-datasource = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
+datafusion-execution = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
datafusion-expr = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
datafusion-expr-common = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
datafusion-functions = { git = "https://github.com/GreptimeTeam/datafusion.git", rev = "e3221b8f74b15c7192621576e5d1e354a6cf097a" }
diff --git a/config/config.md b/config/config.md
index 82e98fe96ec..8663ff1cc28 100644
--- a/config/config.md
+++ b/config/config.md
@@ -142,7 +142,7 @@
| `procedure.max_running_procedures` | Integer | `128` | Max running procedures.
The maximum number of procedures that can be running at the same time.
If the number of running procedures exceeds this limit, the procedure will be rejected. |
| `query` | -- | -- | The query engine options. |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. |
-| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. |
+| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join,
and internal series merges during table scans).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
Projected Parquet row groups and streamed scan batches use scan_memory_limit separately. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). |
| `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. |
| `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". |
@@ -210,10 +210,11 @@
| `region_engine.mito.enable_refill_cache_on_read` | Bool | `true` | Enable refilling cache on read operations (default: true).
When disabled, cache refilling on read won't happen. |
| `region_engine.mito.manifest_cache_size` | String | `256MB` | Capacity for manifest cache (default: 256MB). |
| `region_engine.mito.sst_write_buffer_size` | String | `8MB` | Buffer size for SST writing. |
-| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently. |
+| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently when the memory limit is disabled
or a scan includes extension sources not covered by row-group accounting. |
| `region_engine.mito.allow_stale_entries` | Bool | `false` | Whether to allow stale WAL entries read during replay. |
-| `region_engine.mito.scan_memory_limit` | String | `unlimited` | Memory limit for table scans across all queries.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit. |
-| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Controls what happens when a scan cannot get memory immediately.
"fail" (default) fails fast and is the recommended option for most users.
"wait" / "wait()" waits for memory to become available. This is mainly
for advanced tuning in bursty workloads where temporary contention is common and
higher latency is acceptable.
"wait" means "wait(10s)", not unlimited waiting. |
+| `region_engine.mito.scan_memory_limit` | String | `25%` | Memory limit for table scans across all queries and regions.
Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit and restores file-count checks. |
+| `region_engine.mito.scan_memory_mode` | String | `basic` | Memory accounting mode: "basic" (default) counts projected Parquet columns'
uncompressed row-group sizes while decoding, including files from older versions.
"extended" also counts cumulative streamed batches until stream drop.
Limited reservations round up to KB; unlimited accounting is byte-exact.
Internal series merges use query.memory_pool_size in both modes, sharing capacity
with other query operators on the executing node. Neither budget measures total process memory. |
+| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Scan memory exhaustion always fails immediately.
Legacy "wait" / "wait()" settings are deprecated. They are accepted
for upgrade compatibility, but startup logs a warning and treats them as "fail".
Change existing wait settings to "fail"; scans no longer wait for quota to free up.
This policy applies only to scan reservations. Series merges use the query pool
and fail immediately when that pool is exhausted. |
| `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). |
| `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.
Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. |
| `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. |
@@ -387,7 +388,7 @@
| `query` | -- | -- | The query engine options. |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. |
| `query.allow_query_fallback` | Bool | `false` | Whether to allow query fallback when push down optimize fails.
Default to false, meaning when push down optimize failed, return error msg |
-| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans (only applies to datanodes). |
+| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
This pool applies to operators on this frontend. Datanodes use their own query
pools for local operators and internal series merges during table scans. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). |
| `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. |
| `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". |
@@ -618,7 +619,7 @@
| `wal.accept_missing_history` | Bool | `false` | Whether a region opens although the WAL objects under the prefix do not hold its history.
An existing region fails to open when the prefix held no object when the WAL opened, or when the flushed entry id of the region names an object above every object of the prefix that holds entries. Both happen after `storage_provider`, or the bucket or root of the provider it names, changed without the WAL objects being copied, or after the objects were deleted.
To move the WAL, stop the instance, copy every object under the prefix to the new location and then change the configuration; the regions open with all their writes.
When the old WAL is lost, set this to `true` for one start: the regions open without their writes that were not flushed, which are lost, a warning names each region, and `greptime_logstore_object_store_wal_missing_history_accepted_total` counts them. Turn it off again once the instance has written to the WAL after that start; until then a later start refuses a region whose flushed entry id names an object above every object of the prefix that holds entries.
**It's only used when the provider is `experimental_object_store`**. |
| `query` | -- | -- | The query engine options. |
| `query.parallelism` | Integer | `0` | Parallelism of the query engine.
Default to 0, which means the number of CPU cores. |
-| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by table scans. |
+| `query.memory_pool_size` | String | 0 | Memory pool size for query execution operators (aggregation, sorting, join,
and internal series merges during table scans).
Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
When this limit is reached, queries will fail with ResourceExhausted error.
Projected Parquet row groups and streamed scan batches use scan_memory_limit separately. |
| `query.experimental_memory_pool_policy` | String | `greedy` | Experimental memory pool allocation policy:
- "greedy" (default): first-come-first-served allocation; preserves current behavior.
- "fair": divides available memory among spillable operators and may spill earlier.
Only effective when `memory_pool_size` is bounded (>0). |
| `query.experimental_spill_mode` | String | `default` | Spill mode:
- "default": preserve DataFusion built-in OS temp directory (default).
- "custom": explicitly configure spill path, quota, and compression.
- "disabled": explicitly disable disk spilling.
Set this to "custom" before using the path/quota/compression keys below. |
| `query.experimental_spill_path` | String | Unset | Spill directory path. Ignored unless mode is "custom". |
@@ -687,10 +688,11 @@
| `region_engine.mito.enable_refill_cache_on_read` | Bool | `true` | Enable refilling cache on read operations (default: true).
When disabled, cache refilling on read won't happen. |
| `region_engine.mito.manifest_cache_size` | String | `256MB` | Capacity for manifest cache (default: 256MB). |
| `region_engine.mito.sst_write_buffer_size` | String | `8MB` | Buffer size for SST writing. |
-| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently. |
+| `region_engine.mito.max_concurrent_scan_files` | Integer | `384` | Maximum number of SST files to scan concurrently when the memory limit is disabled
or a scan includes extension sources not covered by row-group accounting. |
| `region_engine.mito.allow_stale_entries` | Bool | `false` | Whether to allow stale WAL entries read during replay. |
-| `region_engine.mito.scan_memory_limit` | String | `unlimited` | Memory limit for table scans across all queries.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit. |
-| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Controls what happens when a scan cannot get memory immediately.
"fail" (default) fails fast and is the recommended option for most users.
"wait" / "wait()" waits for memory to become available. This is mainly
for advanced tuning in bursty workloads where temporary contention is common and
higher latency is acceptable.
"wait" means "wait(10s)", not unlimited waiting. |
+| `region_engine.mito.scan_memory_limit` | String | `25%` | Memory limit for table scans across all queries and regions.
Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available.
Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 or "unlimited" disables the limit and restores file-count checks. |
+| `region_engine.mito.scan_memory_mode` | String | `basic` | Memory accounting mode: "basic" (default) counts projected Parquet columns'
uncompressed row-group sizes while decoding, including files from older versions.
"extended" also counts cumulative streamed batches until stream drop.
Limited reservations round up to KB; unlimited accounting is byte-exact.
Internal series merges use query.memory_pool_size in both modes, sharing capacity
with other query operators on the executing node. Neither budget measures total process memory. |
+| `region_engine.mito.scan_memory_on_exhausted` | String | `fail` | Scan memory exhaustion always fails immediately.
Legacy "wait" / "wait()" settings are deprecated. They are accepted
for upgrade compatibility, but startup logs a warning and treats them as "fail".
Change existing wait settings to "fail"; scans no longer wait for quota to free up.
This policy applies only to scan reservations. Series merges use the query pool
and fail immediately when that pool is exhausted. |
| `region_engine.mito.min_compaction_interval` | String | `0m` | Minimum time interval between two compactions.
To align with the old behavior, the default value is 0 (no restrictions). |
| `region_engine.mito.schedule_compaction_after_edit` | Bool | `true` | Whether to allow to schedule a compaction after a successful region edit.
Setting this to "true" is a necessary but not sufficient condition for scheduling compaction after a region edit.
Other constraints, such as "min_compaction_interval", may still prevent compaction from being scheduled.
Setting this to "false", however, guarantees that compaction will not be scheduled after a region edit. |
| `region_engine.mito.default_flat_format` | Bool | `true` | Whether to enable flat format as the default SST format. |
diff --git a/config/datanode.example.toml b/config/datanode.example.toml
index d5f8e1e9be0..4b8eb41242e 100644
--- a/config/datanode.example.toml
+++ b/config/datanode.example.toml
@@ -337,11 +337,12 @@ overwrite_entry_start_id = false
## Default to 0, which means the number of CPU cores.
parallelism = 0
-## Memory pool size for query execution operators (aggregation, sorting, join).
+## Memory pool size for query execution operators (aggregation, sorting, join,
+## and internal series merges during table scans).
## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
## When this limit is reached, queries will fail with ResourceExhausted error.
-## NOTE: This does NOT limit memory used by table scans.
+## Projected Parquet row groups and streamed scan batches use scan_memory_limit separately.
## @toml2docs:none-default="0"
memory_pool_size = "50%"
@@ -665,22 +666,33 @@ manifest_cache_size = "256MB"
## Buffer size for SST writing.
sst_write_buffer_size = "8MB"
-## Maximum number of SST files to scan concurrently.
+## Maximum number of SST files to scan concurrently when the memory limit is disabled
+## or a scan includes extension sources not covered by row-group accounting.
max_concurrent_scan_files = 384
## Whether to allow stale WAL entries read during replay.
allow_stale_entries = false
-## Memory limit for table scans across all queries.
+## Memory limit for table scans across all queries and regions.
+## Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available.
## Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
-## Setting it to 0 or "unlimited" disables the limit.
-scan_memory_limit = "unlimited"
-## Controls what happens when a scan cannot get memory immediately.
-## "fail" (default) fails fast and is the recommended option for most users.
-## "wait" / "wait()" waits for memory to become available. This is mainly
-## for advanced tuning in bursty workloads where temporary contention is common and
-## higher latency is acceptable.
-## "wait" means "wait(10s)", not unlimited waiting.
+## Setting it to 0 or "unlimited" disables the limit and restores file-count checks.
+scan_memory_limit = "25%"
+
+## Memory accounting mode: "basic" (default) counts projected Parquet columns'
+## uncompressed row-group sizes while decoding, including files from older versions.
+## "extended" also counts cumulative streamed batches until stream drop.
+## Limited reservations round up to KB; unlimited accounting is byte-exact.
+## Internal series merges use query.memory_pool_size in both modes, sharing capacity
+## with other query operators on the executing node. Neither budget measures total process memory.
+scan_memory_mode = "basic"
+
+## Scan memory exhaustion always fails immediately.
+## Legacy "wait" / "wait()" settings are deprecated. They are accepted
+## for upgrade compatibility, but startup logs a warning and treats them as "fail".
+## Change existing wait settings to "fail"; scans no longer wait for quota to free up.
+## This policy applies only to scan reservations. Series merges use the query pool
+## and fail immediately when that pool is exhausted.
scan_memory_on_exhausted = "fail"
## Minimum time interval between two compactions.
diff --git a/config/frontend.example.toml b/config/frontend.example.toml
index 24aef8bd796..97a939c9829 100644
--- a/config/frontend.example.toml
+++ b/config/frontend.example.toml
@@ -370,7 +370,8 @@ allow_query_fallback = false
## Supports absolute size (e.g., "4GB", "8GB") or percentage of system memory (e.g., "30%").
## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
## When this limit is reached, queries will fail with ResourceExhausted error.
-## NOTE: This does NOT limit memory used by table scans (only applies to datanodes).
+## This pool applies to operators on this frontend. Datanodes use their own query
+## pools for local operators and internal series merges during table scans.
## @toml2docs:none-default="0"
memory_pool_size = "50%"
diff --git a/config/standalone.example.toml b/config/standalone.example.toml
index 6fba1d91398..e16210d3644 100644
--- a/config/standalone.example.toml
+++ b/config/standalone.example.toml
@@ -548,11 +548,12 @@ max_running_procedures = 128
## Default to 0, which means the number of CPU cores.
parallelism = 0
-## Memory pool size for query execution operators (aggregation, sorting, join).
+## Memory pool size for query execution operators (aggregation, sorting, join,
+## and internal series merges during table scans).
## Supports absolute size (e.g., "2GB", "4GB") or percentage of system memory (e.g., "20%").
## Setting it to 0 disables the limit. Defaults to 0 (unlimited); this example sets it to `50%`.
## When this limit is reached, queries will fail with ResourceExhausted error.
-## NOTE: This does NOT limit memory used by table scans.
+## Projected Parquet row groups and streamed scan batches use scan_memory_limit separately.
## @toml2docs:none-default="0"
memory_pool_size = "50%"
@@ -872,22 +873,33 @@ manifest_cache_size = "256MB"
## Buffer size for SST writing.
sst_write_buffer_size = "8MB"
-## Maximum number of SST files to scan concurrently.
+## Maximum number of SST files to scan concurrently when the memory limit is disabled
+## or a scan includes extension sources not covered by row-group accounting.
max_concurrent_scan_files = 384
## Whether to allow stale WAL entries read during replay.
allow_stale_entries = false
-## Memory limit for table scans across all queries.
+## Memory limit for table scans across all queries and regions.
+## Defaults to 25% of cgroup memory, or host memory when no cgroup limit is available.
## Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "20%").
-## Setting it to 0 or "unlimited" disables the limit.
-scan_memory_limit = "unlimited"
-## Controls what happens when a scan cannot get memory immediately.
-## "fail" (default) fails fast and is the recommended option for most users.
-## "wait" / "wait()" waits for memory to become available. This is mainly
-## for advanced tuning in bursty workloads where temporary contention is common and
-## higher latency is acceptable.
-## "wait" means "wait(10s)", not unlimited waiting.
+## Setting it to 0 or "unlimited" disables the limit and restores file-count checks.
+scan_memory_limit = "25%"
+
+## Memory accounting mode: "basic" (default) counts projected Parquet columns'
+## uncompressed row-group sizes while decoding, including files from older versions.
+## "extended" also counts cumulative streamed batches until stream drop.
+## Limited reservations round up to KB; unlimited accounting is byte-exact.
+## Internal series merges use query.memory_pool_size in both modes, sharing capacity
+## with other query operators on the executing node. Neither budget measures total process memory.
+scan_memory_mode = "basic"
+
+## Scan memory exhaustion always fails immediately.
+## Legacy "wait" / "wait()" settings are deprecated. They are accepted
+## for upgrade compatibility, but startup logs a warning and treats them as "fail".
+## Change existing wait settings to "fail"; scans no longer wait for quota to free up.
+## This policy applies only to scan reservations. Series merges use the query pool
+## and fail immediately when that pool is exhausted.
scan_memory_on_exhausted = "fail"
## Minimum time interval between two compactions.
diff --git a/src/cmd/src/datanode/scanbench.rs b/src/cmd/src/datanode/scanbench.rs
index 9fbde16fb7b..667a5f54be1 100644
--- a/src/cmd/src/datanode/scanbench.rs
+++ b/src/cmd/src/datanode/scanbench.rs
@@ -1017,6 +1017,7 @@ impl ScanbenchCommand {
let num_partitions = scanner.properties().partitions.len();
let ctx = QueryScanContext {
explain_verbose: collect_scanner_explain,
+ ..Default::default()
};
let metrics_set = ExecutionPlanMetricsSet::new();
diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs
index 00cdafd9b8a..2781183c0c8 100644
--- a/src/cmd/tests/load_config_test.rs
+++ b/src/cmd/tests/load_config_test.rs
@@ -201,7 +201,8 @@ fn test_load_datanode_example_config() {
auto_flush_interval: Duration::from_secs(10 * 60),
default_region_write_buffer_size: ReadableSize::mb(0),
write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)),
- scan_memory_limit: MemoryLimit::Unlimited,
+ scan_memory_limit: MemoryLimit::Percentage(25),
+ scan_memory_mode: mito2::config::ScanMemoryMode::Basic,
..Default::default()
}),
RegionEngineConfig::File(FileEngineConfig {}),
@@ -440,7 +441,8 @@ fn test_load_standalone_example_config() {
auto_flush_interval: Duration::from_secs(10 * 60),
default_region_write_buffer_size: ReadableSize::mb(0),
write_cache_ttl: Some(Duration::from_secs(60 * 60 * 8)),
- scan_memory_limit: MemoryLimit::Unlimited,
+ scan_memory_limit: MemoryLimit::Percentage(25),
+ scan_memory_mode: mito2::config::ScanMemoryMode::Basic,
..Default::default()
}),
RegionEngineConfig::File(FileEngineConfig {}),
diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs
index 0ccb8a4c28e..a3455b366dd 100644
--- a/src/common/recordbatch/src/lib.rs
+++ b/src/common/recordbatch/src/lib.rs
@@ -17,21 +17,18 @@ pub mod cursor;
pub mod error;
pub mod ext;
pub mod filter;
+mod memory;
pub mod recordbatch;
pub mod util;
use std::fmt::{self, Write};
-use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use adapter::RecordBatchMetrics;
use arc_swap::ArcSwapOption;
use common_base::readable_size::ReadableSize;
-use common_error::ext::BoxedError;
-use common_memory_manager::{
- MemoryGuard, MemoryManager, MemoryMetrics, OnExhaustedPolicy, PermitGranularity,
-};
+use common_memory_manager::{MemoryManager, MemoryMetrics, PermitGranularity};
use common_telemetry::tracing::Span;
pub use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
use datatypes::arrow::array::{Array, ArrayRef, AsArray, StringBuilder};
@@ -53,9 +50,10 @@ use error::Result;
use futures::task::{Context, Poll};
use futures::{Stream, TryStreamExt};
pub use recordbatch::RecordBatch;
-use snafu::{IntoError, ResultExt, ensure};
+use snafu::{ResultExt, ensure};
use crate::error::{ArrowComputeSnafu, NewDfRecordBatchSnafu};
+pub use crate::memory::QueryMemoryReservation;
pub trait RecordBatchStream: Stream- > {
fn name(&self) -> &str {
@@ -595,14 +593,14 @@ impl> + Unpin> Stream for RecordBatchStream
}
}
-/// Memory tracker for RecordBatch streams. Clone to share the same limit across queries.
-///
-/// Each stream acquires quota independently from this tracker.
+/// Memory tracker for RecordBatch streams and reader reservations. Clone to share
+/// the same limit across queries. Limited accounting rounds up to KB; unlimited
+/// accounting is byte-exact. Exhaustion fails immediately. DataFusion operators use
+/// the separate query runtime pool.
#[derive(Clone)]
pub struct QueryMemoryTracker {
manager: MemoryManager,
metrics: CallbackMemoryMetrics,
- on_exhausted_policy: OnExhaustedPolicy,
}
impl fmt::Debug for QueryMemoryTracker {
@@ -610,7 +608,6 @@ impl fmt::Debug for QueryMemoryTracker {
f.debug_struct("QueryMemoryTracker")
.field("current", &self.current())
.field("limit", &self.limit())
- .field("on_exhausted_policy", &self.on_exhausted_policy)
.field("on_update", &self.metrics.has_on_update())
.field("on_exhausted", &self.metrics.has_on_exhausted())
.field("on_rejected", &self.metrics.has_on_rejected())
@@ -620,13 +617,9 @@ impl fmt::Debug for QueryMemoryTracker {
impl QueryMemoryTracker {
/// Create a builder for a query memory tracker.
- pub fn builder(
- limit: usize,
- on_exhausted_policy: OnExhaustedPolicy,
- ) -> QueryMemoryTrackerBuilder {
+ pub fn builder(limit: usize) -> QueryMemoryTrackerBuilder {
QueryMemoryTrackerBuilder {
limit,
- on_exhausted_policy,
on_update: None,
on_exhausted: None,
on_reject: None,
@@ -636,7 +629,10 @@ impl QueryMemoryTracker {
fn new_stream_tracker(&self) -> StreamMemoryTracker {
StreamMemoryTracker {
tracker: self.clone(),
- guard: self.manager.try_acquire(0).unwrap(),
+ guard: QueryMemoryReservation {
+ // A zero-byte acquisition always succeeds; this manager is never closed.
+ inner: self.manager.try_acquire(0).unwrap(),
+ },
tracked_bytes: 0,
}
}
@@ -645,10 +641,22 @@ impl QueryMemoryTracker {
self.manager.used_bytes() as usize
}
- fn limit(&self) -> usize {
+ /// Returns the memory limit in bytes, or zero when unlimited.
+ pub fn limit(&self) -> usize {
self.manager.limit_bytes() as usize
}
+ /// Reserves bytes until the returned guard is dropped, failing immediately if
+ /// the shared budget cannot accommodate the reservation.
+ pub fn reserve(&self, bytes: usize) -> Result {
+ let mut tracker = self.new_stream_tracker();
+ if let Err(error) = tracker.try_track(bytes) {
+ tracker.inc_rejected();
+ return Err(error);
+ }
+ Ok(tracker.guard)
+ }
+
fn reject_error(
&self,
current: usize,
@@ -660,7 +668,7 @@ impl QueryMemoryTracker {
"{} requested, {} used globally ({}%), {} used by this stream, hard limit: {}",
ReadableSize(additional as u64),
ReadableSize(current as u64),
- (current * 100).checked_div(limit).unwrap_or(0),
+ current.saturating_mul(100).checked_div(limit).unwrap_or(0),
ReadableSize(stream_tracked as u64),
ReadableSize(limit as u64)
);
@@ -675,7 +683,6 @@ impl QueryMemoryTracker {
/// Builder for constructing a [`QueryMemoryTracker`] with optional callbacks.
pub struct QueryMemoryTrackerBuilder {
limit: usize,
- on_exhausted_policy: OnExhaustedPolicy,
on_update: Option,
on_exhausted: Option,
on_reject: Option,
@@ -687,7 +694,7 @@ impl QueryMemoryTrackerBuilder {
///
/// # Note
/// The callback is called after both successful `track()` and stream drop.
- /// Usage is exact in unlimited mode and 1KB-aligned in limited mode.
+ /// Usage is byte-exact in unlimited mode; otherwise it is 1KB-aligned.
pub fn on_update(mut self, on_update: F) -> Self
where
F: Fn(usize) + Send + Sync + 'static,
@@ -700,7 +707,6 @@ impl QueryMemoryTrackerBuilder {
///
/// # Note
/// This is called when the non-blocking allocation fast path fails.
- /// Requests using `OnExhaustedPolicy::Wait` may still succeed after waiting.
/// It is never called when `limit == 0` (unlimited mode).
pub fn on_exhausted(mut self, on_exhausted: F) -> Self
where
@@ -728,22 +734,16 @@ impl QueryMemoryTrackerBuilder {
metrics.clone(),
);
- QueryMemoryTracker {
- manager,
- metrics,
- on_exhausted_policy: self.on_exhausted_policy,
- }
+ QueryMemoryTracker { manager, metrics }
}
}
struct StreamMemoryTracker {
tracker: QueryMemoryTracker,
- guard: MemoryGuard,
+ guard: QueryMemoryReservation,
tracked_bytes: usize,
}
-type MemoryAcquireResult = std::result::Result<(), common_memory_manager::Error>;
-
impl StreamMemoryTracker {
fn inc_rejected(&self) {
self.tracker.inc_rejected();
@@ -758,51 +758,13 @@ impl StreamMemoryTracker {
}
}
- async fn track_with_policy(mut self, additional: usize) -> (Self, MemoryAcquireResult) {
- let result = self
- .guard
- .acquire_additional_with_policy(additional as u64, self.tracker.on_exhausted_policy)
- .await;
- if result.is_ok() {
- self.tracked_bytes = self.tracked_bytes.saturating_add(additional);
- }
- (self, result)
- }
-
fn reject_error(&self, additional: usize) -> error::Error {
let current = self.tracker.current();
self.tracker
.reject_error(current, additional, self.tracked_bytes)
}
-
- fn wait_error(&self, additional: usize, source: common_memory_manager::Error) -> error::Error {
- match source {
- common_memory_manager::Error::MemoryLimitExceeded { .. } => {
- self.reject_error(additional)
- }
- common_memory_manager::Error::MemoryAcquireTimeout { waited, .. } => {
- let current = self.tracker.current();
- let limit = self.tracker.limit();
- let msg = format!(
- "timed out waiting {:?} for {}, {} used globally ({}%), {} used by this stream, hard limit: {}",
- waited,
- ReadableSize(additional as u64),
- ReadableSize(current as u64),
- (current * 100).checked_div(limit).unwrap_or(0),
- ReadableSize(self.tracked_bytes as u64),
- ReadableSize(limit as u64)
- );
- error::ExceedMemoryLimitSnafu { msg }.build()
- }
- error => error::ExternalSnafu.into_error(BoxedError::new(error)),
- }
- }
}
-type PendingTrackFuture = Pin<
- Box + Send>,
->;
-
#[derive(Clone)]
struct CallbackMemoryMetrics {
inner: Arc,
@@ -871,99 +833,24 @@ impl MemoryMetrics for CallbackMemoryMetrics {
/// A wrapper stream that tracks memory usage of RecordBatches.
pub struct MemoryTrackedStream {
inner: SendableRecordBatchStream,
- tracker: Option,
- // Waiting stores a batch that has already been pulled from the inner stream but has not yet
- // acquired additional quota. This keeps `poll_next()` non-blocking and allows bounded waits,
- // at the cost of temporarily holding one untracked batch per blocked stream in memory.
- waiting: Option,
+ tracker: StreamMemoryTracker,
}
impl MemoryTrackedStream {
pub fn new(inner: SendableRecordBatchStream, tracker: QueryMemoryTracker) -> Self {
Self {
inner,
- tracker: Some(tracker.new_stream_tracker()),
- waiting: None,
+ tracker: tracker.new_stream_tracker(),
}
}
- fn ready_tracker_mut(&mut self) -> &mut StreamMemoryTracker {
- debug_assert!(
- self.waiting.is_none(),
- "a ready tracker must not coexist with a waiting future"
- );
- self.tracker.as_mut().unwrap()
- }
-
- fn enter_waiting(&mut self, batch: RecordBatch, additional: usize) {
- debug_assert!(
- self.waiting.is_none(),
- "enter_waiting should only be called from the ready state"
- );
- debug_assert!(
- self.tracker.is_some(),
- "enter_waiting requires a tracker in the ready state"
- );
- let tracker = self.tracker.take().unwrap();
- self.waiting = Some(Self::start_waiting(tracker, batch, additional));
- }
-
- fn start_waiting(
- tracker: StreamMemoryTracker,
- batch: RecordBatch,
- additional: usize,
- ) -> PendingTrackFuture {
- Box::pin(async move {
- let (tracker, result) = tracker.track_with_policy(additional).await;
- (tracker, batch, additional, result)
- })
- }
-
- fn poll_waiting(&mut self, cx: &mut Context<'_>) -> Poll>> {
- let future = self.waiting.as_mut().unwrap();
- match future.as_mut().poll(cx) {
- Poll::Ready((tracker, batch, additional, result)) => {
- let output = match result {
- Ok(()) => Ok(batch),
- Err(error) => {
- tracker.inc_rejected();
- Err(tracker.wait_error(additional, error))
- }
- };
- self.waiting = None;
- self.tracker = Some(tracker);
- Poll::Ready(Some(output))
- }
- Poll::Pending => Poll::Pending,
- }
- }
-
- fn poll_batch(
- &mut self,
- batch: RecordBatch,
- cx: &mut Context<'_>,
- ) -> Poll>> {
+ fn track_batch(&mut self, batch: RecordBatch) -> Result {
let additional = batch.logical_slice_memory_size();
- let tracker = self.ready_tracker_mut();
-
- if let Err(error) = tracker.try_track(additional) {
- match tracker.tracker.on_exhausted_policy {
- OnExhaustedPolicy::Fail => {
- tracker.inc_rejected();
- return Poll::Ready(Some(Err(error)));
- }
- // `Wait` is a deliberate tradeoff: the batch has already been materialized, so we
- // keep it in memory while waiting for quota instead of failing immediately. Under
- // contention, real memory usage can therefore exceed `scan_memory_limit` by up to
- // one buffered batch per blocked stream.
- OnExhaustedPolicy::Wait { .. } => {
- self.enter_waiting(batch, additional);
- return self.poll_waiting(cx);
- }
- }
+ if let Err(error) = self.tracker.try_track(additional) {
+ self.tracker.inc_rejected();
+ return Err(error);
}
-
- Poll::Ready(Some(Ok(batch)))
+ Ok(batch)
}
}
@@ -971,12 +858,8 @@ impl Stream for MemoryTrackedStream {
type Item = Result;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> {
- if self.waiting.is_some() {
- return self.poll_waiting(cx);
- }
-
match Pin::new(&mut self.inner).poll_next(cx) {
- Poll::Ready(Some(Ok(batch))) => self.poll_batch(batch, cx),
+ Poll::Ready(Some(Ok(batch))) => Poll::Ready(Some(self.track_batch(batch))),
Poll::Ready(Some(Err(error))) => Poll::Ready(Some(Err(error))),
Poll::Ready(None) => Poll::Ready(None),
Poll::Pending => Poll::Pending,
@@ -1006,9 +889,8 @@ impl RecordBatchStream for MemoryTrackedStream {
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
- use std::time::Duration;
- use common_memory_manager::{OnExhaustedPolicy, PermitGranularity};
+ use common_memory_manager::PermitGranularity;
use datatypes::arrow::array::{
DictionaryArray, Int32Array, ListArray, StringArray, UInt32Array,
};
@@ -1020,7 +902,6 @@ mod tests {
use datatypes::schema::{ColumnSchema, Schema};
use datatypes::vectors::{BooleanVector, Int32Vector, StringVector};
use futures::StreamExt;
- use tokio::time::{sleep, timeout};
use super::*;
@@ -1041,6 +922,51 @@ mod tests {
as usize
}
+ #[test]
+ fn test_reader_reservation_rejects_and_releases_capacity() {
+ let tracker = QueryMemoryTracker::builder(2048).build();
+ let first = tracker.reserve(1025).unwrap();
+ assert_eq!(2048, tracker.current());
+ assert!(tracker.reserve(1).is_err());
+ assert_eq!(2048, tracker.current());
+ drop(first);
+ assert_eq!(0, tracker.current());
+
+ assert!(tracker.reserve(2049).is_err());
+ assert_eq!(0, tracker.current());
+ let guard = tracker.reserve(2048).unwrap();
+ assert_eq!(2048, tracker.current());
+ drop(guard);
+ assert_eq!(0, tracker.current());
+ }
+
+ #[tokio::test]
+ async fn test_reader_and_batch_reservations_share_scan_budget() {
+ let batch = large_string_batch(512);
+ let batch_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size());
+ let tracker = QueryMemoryTracker::builder(batch_bytes + 1024).build();
+ let reader = tracker.reserve(1).unwrap();
+ let mut stream = MemoryTrackedStream::new(
+ RecordBatches::try_new(batch.schema.clone(), vec![batch.clone(), batch])
+ .unwrap()
+ .as_stream(),
+ tracker.clone(),
+ );
+ stream.next().await.unwrap().unwrap();
+ assert_eq!(tracker.limit(), tracker.current());
+ assert!(tracker.reserve(1).is_err());
+ assert!(stream.next().await.unwrap().is_err());
+ drop(reader);
+ assert_eq!(batch_bytes, tracker.current());
+ drop(stream);
+ assert_eq!(0, tracker.current());
+ let unlimited = QueryMemoryTracker::builder(0).build();
+ let guard = unlimited.reserve(17).unwrap();
+ assert_eq!(17, unlimited.current());
+ drop(guard);
+ assert_eq!(0, unlimited.current());
+ }
+
#[tokio::test]
async fn test_memory_tracked_stream_charges_logical_slice_size() {
let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
@@ -1057,7 +983,7 @@ mod tests {
.unwrap();
let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size());
assert!(expected_bytes < aligned_tracked_bytes(batch.buffer_memory_size()));
- let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail).build();
+ let tracker = QueryMemoryTracker::builder(MB).build();
let mut stream = MemoryTrackedStream::new(
RecordBatches::try_new(batch.schema.clone(), vec![batch])
.unwrap()
@@ -1282,8 +1208,7 @@ mod tests {
#[test]
fn test_query_memory_tracker_basic() {
- let tracker =
- Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build());
+ let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
let mut stream1 = tracker.new_stream_tracker();
assert!(stream1.try_track(5 * MB).is_ok());
@@ -1300,8 +1225,7 @@ mod tests {
#[test]
fn test_query_memory_tracker_shared_global_limit() {
- let tracker =
- Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build());
+ let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
let mut stream1 = tracker.new_stream_tracker();
let mut stream2 = tracker.new_stream_tracker();
@@ -1325,8 +1249,7 @@ mod tests {
#[test]
fn test_query_memory_tracker_hard_limit() {
- let tracker =
- Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build());
+ let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
let mut stream = tracker.new_stream_tracker();
assert!(stream.try_track(9 * MB).is_ok());
@@ -1347,7 +1270,7 @@ mod tests {
#[test]
fn test_query_memory_tracker_unlimited() {
- let tracker = Arc::new(QueryMemoryTracker::builder(0, OnExhaustedPolicy::Fail).build());
+ let tracker = Arc::new(QueryMemoryTracker::builder(0).build());
let mut stream = tracker.new_stream_tracker();
assert!(stream.try_track(10 * MB).is_ok());
@@ -1358,8 +1281,7 @@ mod tests {
#[test]
fn test_query_memory_tracker_rounds_to_kilobytes() {
- let tracker =
- Arc::new(QueryMemoryTracker::builder(10 * MB, OnExhaustedPolicy::Fail).build());
+ let tracker = Arc::new(QueryMemoryTracker::builder(10 * MB).build());
let mut stream = tracker.new_stream_tracker();
assert!(stream.try_track(1_537).is_ok());
@@ -1370,110 +1292,12 @@ mod tests {
}
#[tokio::test]
- async fn test_memory_tracked_stream_waits_for_capacity() {
+ async fn test_memory_tracked_stream_rejects_immediately() {
let exhausted = Arc::new(AtomicUsize::new(0));
let rejected = Arc::new(AtomicUsize::new(0));
let exhausted_counter = exhausted.clone();
let rejected_counter = rejected.clone();
- let tracker = QueryMemoryTracker::builder(
- MB,
- OnExhaustedPolicy::Wait {
- timeout: Duration::from_millis(200),
- },
- )
- .on_exhausted(move || {
- exhausted_counter.fetch_add(1, Ordering::Relaxed);
- })
- .on_reject(move || {
- rejected_counter.fetch_add(1, Ordering::Relaxed);
- })
- .build();
- let batch = large_string_batch(700 * 1024);
- let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size());
-
- let mut stream1 = MemoryTrackedStream::new(
- RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()])
- .unwrap()
- .as_stream(),
- tracker.clone(),
- );
- let first = stream1.next().await.unwrap().unwrap();
- assert_eq!(first.num_rows(), 1);
- assert_eq!(tracker.current(), expected_bytes);
-
- let stream2 = MemoryTrackedStream::new(
- RecordBatches::try_new(batch.schema.clone(), vec![batch])
- .unwrap()
- .as_stream(),
- tracker.clone(),
- );
- let waiter = tokio::spawn(async move {
- let mut stream2 = stream2;
- stream2.next().await.unwrap()
- });
-
- sleep(Duration::from_millis(50)).await;
- assert!(!waiter.is_finished());
-
- drop(stream1);
- let second = waiter.await.unwrap().unwrap();
- assert_eq!(second.num_rows(), 1);
- assert_eq!(exhausted.load(Ordering::Relaxed), 1);
- assert_eq!(rejected.load(Ordering::Relaxed), 0);
- }
-
- #[tokio::test]
- async fn test_memory_tracked_stream_wait_times_out() {
- let exhausted = Arc::new(AtomicUsize::new(0));
- let rejected = Arc::new(AtomicUsize::new(0));
- let exhausted_counter = exhausted.clone();
- let rejected_counter = rejected.clone();
- let tracker = QueryMemoryTracker::builder(
- MB,
- OnExhaustedPolicy::Wait {
- timeout: Duration::from_millis(50),
- },
- )
- .on_exhausted(move || {
- exhausted_counter.fetch_add(1, Ordering::Relaxed);
- })
- .on_reject(move || {
- rejected_counter.fetch_add(1, Ordering::Relaxed);
- })
- .build();
- let batch = large_string_batch(700 * 1024);
-
- let mut stream1 = MemoryTrackedStream::new(
- RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()])
- .unwrap()
- .as_stream(),
- tracker.clone(),
- );
- let first = stream1.next().await.unwrap().unwrap();
- assert_eq!(first.num_rows(), 1);
-
- let mut stream2 = MemoryTrackedStream::new(
- RecordBatches::try_new(batch.schema.clone(), vec![batch])
- .unwrap()
- .as_stream(),
- tracker,
- );
- let result = timeout(Duration::from_secs(1), stream2.next())
- .await
- .unwrap();
- let error = result.unwrap().unwrap_err();
- assert!(error.to_string().contains("timed out waiting"));
- assert_eq!(exhausted.load(Ordering::Relaxed), 1);
- assert_eq!(rejected.load(Ordering::Relaxed), 1);
- }
-
- #[tokio::test]
- async fn test_memory_tracked_stream_fail_policy_rejects_immediately() {
- let exhausted = Arc::new(AtomicUsize::new(0));
- let rejected = Arc::new(AtomicUsize::new(0));
- let exhausted_counter = exhausted.clone();
- let rejected_counter = rejected.clone();
- let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail)
+ let tracker = QueryMemoryTracker::builder(MB)
.on_exhausted(move || {
exhausted_counter.fetch_add(1, Ordering::Relaxed);
})
@@ -1498,8 +1322,10 @@ mod tests {
.as_stream(),
tracker,
);
- let result = stream2.next().await.unwrap();
- assert!(result.is_err());
+ assert!(matches!(
+ futures::poll!(stream2.next()),
+ Poll::Ready(Some(Err(_)))
+ ));
assert_eq!(exhausted.load(Ordering::Relaxed), 1);
assert_eq!(rejected.load(Ordering::Relaxed), 1);
}
diff --git a/src/common/recordbatch/src/memory.rs b/src/common/recordbatch/src/memory.rs
new file mode 100644
index 00000000000..402599f6de4
--- /dev/null
+++ b/src/common/recordbatch/src/memory.rs
@@ -0,0 +1,31 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed 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.
+
+//! Scan memory reservations held for reader lifetimes.
+
+use common_memory_manager::MemoryGuard;
+
+use crate::CallbackMemoryMetrics;
+
+/// Reserved scan memory, released automatically on drop.
+#[derive(Debug)]
+pub struct QueryMemoryReservation {
+ pub(crate) inner: MemoryGuard,
+}
+
+impl QueryMemoryReservation {
+ pub(crate) fn try_acquire_additional(&mut self, bytes: u64) -> bool {
+ self.inner.try_acquire_additional(bytes)
+ }
+}
diff --git a/src/mito2/AGENTS.md b/src/mito2/AGENTS.md
index f70aa622a27..08600cd00e4 100644
--- a/src/mito2/AGENTS.md
+++ b/src/mito2/AGENTS.md
@@ -51,6 +51,10 @@ immutable `Version` (`region/version.rs`) → scans memtables and Parquet SSTs
(`sst/parquet.rs`) → merges (`read/`) and dedups by sequence → projected,
filtered `RecordBatch` stream.
+Internal series merges use the query runtime pool forwarded through
+`QueryScanContext`; reader and output-batch accounting uses the separate scan
+memory tracker. Direct scanner callers default to an unbounded operator pool.
+
## Public surface
- Entry: `MitoEngine` in `src/mito2/src/engine.rs`, built via `MitoEngineBuilder`.
diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs
index c0fa74281ca..98b4ba4be42 100644
--- a/src/mito2/src/config.rs
+++ b/src/mito2/src/config.rs
@@ -51,11 +51,22 @@ const PAGE_CACHE_SIZE_FACTOR: u64 = 8;
/// Use `1/INDEX_CREATE_MEM_THRESHOLD_FACTOR` of OS memory size as mem threshold for creating index
const INDEX_CREATE_MEM_THRESHOLD_FACTOR: u64 = 16;
+/// Memory accounted against the shared scan limit.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum ScanMemoryMode {
+ /// Count projected Parquet columns' uncompressed row-group sizes.
+ #[default]
+ Basic,
+ /// Also count cumulative streamed batches.
+ Extended,
+}
+
/// Fetch option timeout
pub(crate) const FETCH_OPTION_TIMEOUT: Duration = Duration::from_secs(3);
/// Configuration for [MitoEngine](crate::engine::MitoEngine).
-/// Before using the config, make sure to call `MitoConfig::validate()` to check if the config is valid.
+/// Before using the config, call [`MitoConfig::sanitize()`] to validate and normalize it.
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
#[serde(default)]
pub struct MitoConfig {
@@ -170,18 +181,22 @@ pub struct MitoConfig {
/// Buffer size for SST writing.
// TODO(#9240): Support per-object-store write buffer sizes for backend-specific tuning.
pub sst_write_buffer_size: ReadableSize,
- /// Maximum number of SST files to scan concurrently (default 384).
+ /// Maximum number of SST files to scan concurrently when memory accounting
+ /// does not cover the scan or the limit is disabled (default 384).
pub max_concurrent_scan_files: usize,
/// Whether to allow stale entries read during replay.
pub allow_stale_entries: bool,
- /// Memory limit for table scans across all queries.
+ /// Memory limit for table scans across all queries (default: 25% of system memory).
/// Setting it to 0 or "unlimited" disables the limit.
/// Supports absolute size (e.g., "2GB") or percentage of system memory (e.g., "50%").
pub scan_memory_limit: MemoryLimit,
- /// Behavior when scan memory tracking cannot acquire memory from the budget.
- /// `wait` means `wait(10s)`, not unlimited waiting.
- /// Defaults to [`OnExhaustedPolicy::Fail`], which intentionally differs from
- /// [`OnExhaustedPolicy::default()`].
+ /// Scan memory accounting: basic counts projected Parquet row groups;
+ /// extended also counts cumulative streamed batches. Limited reservations round up to KB.
+ /// Internal series merges use the query runtime pool in both modes.
+ pub scan_memory_mode: ScanMemoryMode,
+ /// Scan memory exhaustion always fails immediately (default: `fail`).
+ /// Legacy `wait` and `wait()` values are accepted but deprecated:
+ /// sanitization warns and replaces them with [`OnExhaustedPolicy::Fail`].
pub scan_memory_on_exhausted: OnExhaustedPolicy,
/// Index configs.
@@ -253,7 +268,8 @@ impl Default for MitoConfig {
sst_write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE,
max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES,
allow_stale_entries: false,
- scan_memory_limit: MemoryLimit::default(),
+ scan_memory_limit: MemoryLimit::Percentage(25),
+ scan_memory_mode: ScanMemoryMode::Basic,
scan_memory_on_exhausted: OnExhaustedPolicy::Fail,
index: IndexConfig::default(),
inverted_index: InvertedIndexConfig::default(),
@@ -281,6 +297,18 @@ impl MitoConfig {
///
/// Returns an error if there is a configuration that unable to sanitize.
pub fn sanitize(&mut self, data_home: &str) -> Result<()> {
+ if matches!(
+ self.scan_memory_on_exhausted,
+ OnExhaustedPolicy::Wait { .. }
+ ) {
+ warn!(
+ "scan_memory_on_exhausted = {:?} is deprecated and is treated as fail; \
+ scan memory exhaustion now fails immediately. Set scan_memory_on_exhausted = \"fail\"",
+ self.scan_memory_on_exhausted
+ );
+ self.scan_memory_on_exhausted = OnExhaustedPolicy::Fail;
+ }
+
if self.experimental_enable_series_index {
snafu::ensure!(
self.experimental_series_index_max_size.as_bytes() >= 1024,
@@ -419,6 +447,39 @@ impl MitoConfig {
mod tests {
use super::*;
+ #[test]
+ fn test_scan_memory_policy_upgrade() {
+ for policy in [
+ None,
+ Some("fail"),
+ Some("skip"),
+ Some("wait"),
+ Some("wait(30s)"),
+ ] {
+ let input = policy
+ .map(|policy| format!("scan_memory_on_exhausted = {policy:?}"))
+ .unwrap_or_default();
+ let mut config: MitoConfig = toml::from_str(&input).unwrap();
+ let compaction_policy = config.experimental_compaction_on_exhausted;
+ config.sanitize("/data").unwrap();
+ assert_eq!(OnExhaustedPolicy::Fail, config.scan_memory_on_exhausted);
+ assert_eq!(
+ compaction_policy,
+ config.experimental_compaction_on_exhausted
+ );
+ let serialized = toml::to_string(&config).unwrap();
+ assert!(serialized.contains("scan_memory_on_exhausted = \"fail\""));
+ config.sanitize("/data").unwrap();
+ assert_eq!(serialized, toml::to_string(&config).unwrap());
+ }
+ for policy in ["unknown", "wait(invalid)", "wait()"] {
+ assert!(
+ toml::from_str::(&format!("scan_memory_on_exhausted = {policy:?}"))
+ .is_err()
+ );
+ }
+ }
+
#[test]
fn test_default_auto_flush_interval() {
assert_eq!(
diff --git a/src/mito2/src/engine.rs b/src/mito2/src/engine.rs
index 3036c61e7ff..cdfc78ff0c4 100644
--- a/src/mito2/src/engine.rs
+++ b/src/mito2/src/engine.rs
@@ -108,7 +108,6 @@ use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream};
use common_stat::get_total_memory_bytes;
use common_telemetry::{debug, info, tracing, warn};
use common_wal::options::WalOptions;
-use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool};
use futures::future::{join_all, try_join_all};
use futures::stream::{self, Stream, StreamExt};
use object_store::manager::ObjectStoreManagerRef;
@@ -138,7 +137,7 @@ use tokio::sync::{Semaphore, oneshot};
use crate::access_layer::RegionFilePathFactory;
use crate::cache::{CacheManagerRef, CacheStrategy};
-use crate::config::MitoConfig;
+use crate::config::{MitoConfig, ScanMemoryMode};
use crate::engine::puffin_index::{IndexEntryContext, collect_index_entries_from_puffin};
use crate::error::{
IncrementalQueryStaleSnafu, InvalidRequestSnafu, JoinSnafu, MitoManifestInfoSnafu, RecvSnafu,
@@ -244,26 +243,23 @@ impl<'a, S: LogStore> MitoEngineBuilder<'a, S> {
let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(self.log_store));
let total_memory = get_total_memory_bytes().max(0) as u64;
let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize;
- let scan_memory_pool = new_scan_memory_pool(scan_memory_limit);
- let scan_memory_tracker =
- QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted)
- .on_update(|usage| {
- SCAN_MEMORY_USAGE_BYTES.set(usage as i64);
- })
- .on_exhausted(|| {
- SCAN_MEMORY_EXHAUSTED_TOTAL.inc();
- })
- .on_reject(|| {
- SCAN_REQUESTS_REJECTED_TOTAL.inc();
- })
- .build();
+ let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit)
+ .on_update(|usage| {
+ SCAN_MEMORY_USAGE_BYTES.set(usage as i64);
+ })
+ .on_exhausted(|| {
+ SCAN_MEMORY_EXHAUSTED_TOTAL.inc();
+ })
+ .on_reject(|| {
+ SCAN_REQUESTS_REJECTED_TOTAL.inc();
+ })
+ .build();
let inner = EngineInner {
workers,
config,
wal_raw_entry_reader,
scan_memory_tracker,
- scan_memory_pool,
region_hook,
#[cfg(feature = "enterprise")]
extension_range_provider_factory: None,
@@ -749,8 +745,6 @@ struct EngineInner {
wal_raw_entry_reader: Arc,
/// Memory tracker for table scans.
scan_memory_tracker: QueryMemoryTracker,
- /// Memory pool shared by internal scan operators across all queries.
- scan_memory_pool: Arc,
/// The region hook (if any) registered via plugins; exposed for the GC worker
/// to fire [`RegionHook::on_region_gc`].
region_hook: Option,
@@ -1165,7 +1159,7 @@ impl EngineInner {
.with_ignore_range_index(!self.config.experimental_enable_range_index)
.with_query_stat_counters(region.region_stats.query_stat_counters())
.with_max_concurrent_scan_files(self.config.max_concurrent_scan_files)
- .with_scan_memory_pool(self.scan_memory_pool.clone())
+ .with_scan_memory_tracker(self.scan_memory_tracker.clone())
.with_experimental_series_scan_v2(self.config.experimental_series_scan_v2)
.with_ignore_inverted_index(self.config.inverted_index.apply_on_query.disabled())
.with_ignore_fulltext_index(self.config.fulltext_index.apply_on_query.disabled())
@@ -1457,7 +1451,8 @@ impl RegionEngine for MitoEngine {
}
fn query_memory_tracker(&self) -> Option {
- Some(self.inner.scan_memory_tracker.clone())
+ (self.inner.config.scan_memory_mode == ScanMemoryMode::Extended)
+ .then(|| self.inner.scan_memory_tracker.clone())
}
async fn get_committed_sequence(
@@ -1622,19 +1617,17 @@ impl MitoEngine {
let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(log_store.clone()));
let total_memory = get_total_memory_bytes().max(0) as u64;
let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize;
- let scan_memory_pool = new_scan_memory_pool(scan_memory_limit);
- let scan_memory_tracker =
- QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted)
- .on_update(|usage| {
- SCAN_MEMORY_USAGE_BYTES.set(usage as i64);
- })
- .on_exhausted(|| {
- SCAN_MEMORY_EXHAUSTED_TOTAL.inc();
- })
- .on_reject(|| {
- SCAN_REQUESTS_REJECTED_TOTAL.inc();
- })
- .build();
+ let scan_memory_tracker = QueryMemoryTracker::builder(scan_memory_limit)
+ .on_update(|usage| {
+ SCAN_MEMORY_USAGE_BYTES.set(usage as i64);
+ })
+ .on_exhausted(|| {
+ SCAN_MEMORY_EXHAUSTED_TOTAL.inc();
+ })
+ .on_reject(|| {
+ SCAN_REQUESTS_REJECTED_TOTAL.inc();
+ })
+ .build();
Ok(MitoEngine {
inner: Arc::new(EngineInner {
workers: WorkerGroup::start_for_test(
@@ -1653,7 +1646,6 @@ impl MitoEngine {
config,
wal_raw_entry_reader,
scan_memory_tracker,
- scan_memory_pool,
region_hook: None,
#[cfg(feature = "enterprise")]
extension_range_provider_factory: None,
@@ -1667,20 +1659,10 @@ impl MitoEngine {
}
}
-fn new_scan_memory_pool(scan_memory_limit: usize) -> Arc {
- if scan_memory_limit == 0 {
- Arc::new(UnboundedMemoryPool::default())
- } else {
- Arc::new(GreedyMemoryPool::new(scan_memory_limit))
- }
-}
-
#[cfg(test)]
mod tests {
use std::time::Duration;
- use datafusion::execution::memory_pool::MemoryConsumer;
-
use super::*;
use crate::sst::file::FileMeta;
@@ -1766,17 +1748,4 @@ mod tests {
};
assert!(!is_valid_region_edit(&edit));
}
-
- #[test]
- fn test_scan_memory_pool_is_shared_across_consumers() {
- let pool = new_scan_memory_pool(100);
- let cloned_pool = pool.clone();
- let first = MemoryConsumer::new("first-scan").register(&pool);
- let second = MemoryConsumer::new("second-scan").register(&cloned_pool);
-
- first.try_grow(60).unwrap();
- assert!(second.try_grow(50).is_err());
- second.try_grow(40).unwrap();
- assert_eq!(100, pool.reserved());
- }
}
diff --git a/src/mito2/src/engine/scan_test.rs b/src/mito2/src/engine/scan_test.rs
index f55f084ebd7..39613efb383 100644
--- a/src/mito2/src/engine/scan_test.rs
+++ b/src/mito2/src/engine/scan_test.rs
@@ -1275,15 +1275,26 @@ async fn test_scan_with_min_sst_sequence_with_format(flat_format: bool) {
#[tokio::test]
async fn test_max_concurrent_scan_files() {
- test_max_concurrent_scan_files_with_format(false).await;
- test_max_concurrent_scan_files_with_format(true).await;
+ test_max_concurrent_scan_files_with_format(false, false).await;
+ test_max_concurrent_scan_files_with_format(true, false).await;
}
-async fn test_max_concurrent_scan_files_with_format(flat_format: bool) {
+#[tokio::test]
+async fn test_scan_memory_limit_small_legacy_files() {
+ test_max_concurrent_scan_files_with_format(false, true).await;
+ test_max_concurrent_scan_files_with_format(true, true).await;
+}
+
+async fn test_max_concurrent_scan_files_with_format(flat_format: bool, estimated: bool) {
let mut env = TestEnv::with_prefix("test_max_concurrent_scan_files").await;
let config = MitoConfig {
default_flat_format: flat_format,
max_concurrent_scan_files: 2,
+ scan_memory_limit: if estimated {
+ common_base::memory_limit::MemoryLimit::Size(ReadableSize::mb(1))
+ } else {
+ common_base::memory_limit::MemoryLimit::Unlimited
+ },
..Default::default()
};
let engine = env.create_engine(config).await;
@@ -1311,13 +1322,45 @@ async fn test_max_concurrent_scan_files_with_format(flat_format: bool) {
put_and_flush(3, 7).await;
put_and_flush(6, 9).await;
+ if estimated {
+ // Simulate SSTs created before the manifest recorded row-group sizes.
+ let region = engine.get_region(region_id).unwrap();
+ let version = region.version_control.current().version;
+ let legacy_files = version
+ .ssts
+ .levels()
+ .iter()
+ .flat_map(|level| level.files())
+ .map(|file| {
+ let mut meta = file.meta_ref().clone();
+ meta.max_row_group_uncompressed_size = 0;
+ meta
+ })
+ .collect::>();
+ let version = crate::region::version::VersionBuilder::from_version(version)
+ .clear_files()
+ .add_files(
+ Arc::new(crate::sst::file_purger::NoopFilePurger),
+ legacy_files.into_iter(),
+ )
+ .build();
+ region.version_control.overwrite_current(Arc::new(version));
+ }
+
let request = ScanRequest::default();
let scanner = engine.scanner(region_id, request).await.unwrap();
let Scanner::Seq(scanner) = scanner else {
panic!("Scanner should be seq scan");
};
- let error = scanner.check_scan_limit().unwrap_err();
- assert_eq!(StatusCode::RateLimited, error.status_code());
+ if estimated {
+ scanner.check_scan_limit().unwrap();
+ RecordBatches::try_collect(scanner.build_stream().unwrap())
+ .await
+ .unwrap();
+ } else {
+ let error = scanner.check_scan_limit().unwrap_err();
+ assert_eq!(StatusCode::RateLimited, error.status_code());
+ }
let request = ScanRequest {
distribution: Some(TimeSeriesDistribution::PerSeries),
@@ -1327,8 +1370,15 @@ async fn test_max_concurrent_scan_files_with_format(flat_format: bool) {
let Scanner::Series(scanner) = scanner else {
panic!("Scanner should be series scan");
};
- let error = scanner.check_scan_limit().unwrap_err();
- assert_eq!(StatusCode::RateLimited, error.status_code());
+ if estimated {
+ scanner.check_scan_limit().unwrap();
+ RecordBatches::try_collect(scanner.build_stream().await.unwrap())
+ .await
+ .unwrap();
+ } else {
+ let error = scanner.check_scan_limit().unwrap_err();
+ assert_eq!(StatusCode::RateLimited, error.status_code());
+ }
}
#[tokio::test]
@@ -1443,9 +1493,15 @@ async fn test_series_scan_with_format(flat_format: bool) {
async fn test_two_phase_series_scan() {
for use_index in [false, true] {
for use_range_index in [false, true] {
- check_two_phase_series_scan(use_index, use_range_index, true).await;
+ for mode in [
+ crate::config::ScanMemoryMode::Basic,
+ crate::config::ScanMemoryMode::Extended,
+ ] {
+ check_two_phase_series_scan(use_index, use_range_index, true, mode).await;
+ }
}
- check_two_phase_series_scan(use_index, true, false).await;
+ check_two_phase_series_scan(use_index, true, false, crate::config::ScanMemoryMode::Basic)
+ .await;
}
}
@@ -1453,10 +1509,15 @@ async fn check_two_phase_series_scan(
use_index: bool,
use_range_index: bool,
enable_range_index: bool,
+ scan_memory_mode: crate::config::ScanMemoryMode,
) {
let mut env = TestEnv::with_prefix("test_two_phase_series_scan").await;
let engine = env
.create_engine(MitoConfig {
+ scan_memory_mode,
+ scan_memory_on_exhausted: common_memory_manager::OnExhaustedPolicy::Wait {
+ timeout: std::time::Duration::from_secs(30),
+ },
experimental_series_scan_v2: true,
experimental_enable_series_index: true,
experimental_enable_range_index: enable_range_index,
@@ -1685,7 +1746,12 @@ async fn check_two_phase_series_scan(
.unwrap();
let metrics_set = ExecutionPlanMetricsSet::default();
- let context = store_api::region_engine::QueryScanContext::default();
+ use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool};
+ let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024));
+ let context = store_api::region_engine::QueryScanContext {
+ memory_pool: pool.clone(),
+ ..Default::default()
+ };
let partition_batches = try_join_all((0..3).map(|partition| {
let stream = scanner
.scan_partition(&context, &metrics_set, partition)
@@ -1744,6 +1810,65 @@ async fn check_two_phase_series_scan(
assert_eq!(Some(&2), series_to_partition.get("b"));
assert_eq!(Some(&2), series_to_partition.get("d"));
+ assert_eq!(0, pool.reserved());
+ // Existing query operators consume the same pool in either scan mode.
+ scanner.reset_state();
+ let other = MemoryConsumer::new("other-query-operator").register(&pool);
+ other.try_grow(1024 * 1024).unwrap();
+ let failed = try_join_all((0..3).map(|partition| {
+ let stream = scanner
+ .scan_partition(&context, &metrics_set, partition)
+ .unwrap();
+ async move { stream.try_collect::>().await }
+ }))
+ .await;
+ assert!(failed.is_err());
+ other.free();
+ assert_eq!(0, pool.reserved());
+
+ // Exhausting the scan budget must fail even with an unbounded query pool.
+ scanner.reset_state();
+ let scan_tracker = &engine.inner.scan_memory_tracker;
+ // Cancelled background readers release their reservations asynchronously.
+ tokio::time::timeout(std::time::Duration::from_secs(5), async {
+ while scan_tracker.current() > 0 {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .unwrap();
+ let scan_guard = scan_tracker.reserve(scan_tracker.limit()).unwrap();
+ let fresh_context = store_api::region_engine::QueryScanContext::default();
+ let fresh_pool = Arc::downgrade(&fresh_context.memory_pool);
+ let failed = tokio::time::timeout(
+ std::time::Duration::from_secs(5),
+ try_join_all((0..3).map(|partition| {
+ let stream = scanner
+ .scan_partition(&fresh_context, &metrics_set, partition)
+ .unwrap();
+ async move { stream.try_collect::>().await }
+ })),
+ )
+ .await
+ .expect("legacy wait policy must fail immediately on scan memory exhaustion");
+ assert!(failed.is_err());
+ drop(fresh_context);
+ // Keep the external quota occupied while failure tears down the distributor
+ // and readers. Cleanup must not depend on quota becoming available.
+ tokio::time::timeout(std::time::Duration::from_secs(5), async {
+ while fresh_pool.upgrade().is_some() {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .unwrap();
+ assert_eq!(scan_tracker.limit(), scan_tracker.current());
+ drop(scan_guard);
+ assert_eq!(0, scan_tracker.current());
+
+ // Keep the previous execution's pool exhausted while replay uses a fresh pool.
+ other.try_grow(1024 * 1024).unwrap();
+ let context = store_api::region_engine::QueryScanContext::default();
scanner.reset_state();
assert_eq!("two_phase", scanner.mode());
assert_eq!(
@@ -1789,6 +1914,42 @@ async fn check_two_phase_series_scan(
replay_rows.sort();
assert_eq!(actual_rows, replay_rows);
+ // Dropping every partition must let the shared distributor release its pool.
+ scanner.reset_state();
+ let cancelled_context = store_api::region_engine::QueryScanContext::default();
+ let cancelled_pool = Arc::downgrade(&cancelled_context.memory_pool);
+ let streams = (0..3)
+ .map(|partition| {
+ scanner
+ .scan_partition(&cancelled_context, &metrics_set, partition)
+ .unwrap()
+ })
+ .collect::>();
+ drop(streams);
+ drop(cancelled_context);
+ tokio::time::timeout(std::time::Duration::from_secs(5), async {
+ while cancelled_pool.upgrade().is_some() {
+ tokio::task::yield_now().await;
+ }
+ })
+ .await
+ .unwrap();
+ scanner.reset_state();
+ let direct_batches = scanner
+ .build_stream()
+ .await
+ .unwrap()
+ .try_collect::>()
+ .await
+ .unwrap();
+ assert_eq!(
+ actual_rows.len(),
+ direct_batches
+ .iter()
+ .map(|batch| batch.num_rows())
+ .sum::()
+ );
+
// Exercise precise field/time filtering and candidate tag filtering on both paths.
let filtered = engine
.scanner(
diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs
index 0f9cea941b3..d2bc53167be 100644
--- a/src/mito2/src/read/scan_region.rs
+++ b/src/mito2/src/read/scan_region.rs
@@ -22,13 +22,12 @@ use std::time::Instant;
use api::v1::SemanticType;
use common_error::ext::BoxedError;
-use common_recordbatch::SendableRecordBatchStream;
use common_recordbatch::adapter::RegionQueryStatCounters;
use common_recordbatch::filter::SimpleFilterEvaluator;
+use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream};
use common_telemetry::tracing::Instrument;
use common_telemetry::{debug, error, tracing, warn};
use common_time::range::TimestampRange;
-use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool};
use datafusion::optimizer::simplify_expressions::ExprSimplifier;
use datafusion::physical_plan::expressions::DynamicFilterPhysicalExpr;
use datafusion_common::pruning::PruningStatistics;
@@ -244,8 +243,7 @@ pub(crate) struct ScanRegion {
cache_strategy: CacheStrategy,
/// Maximum number of SST files to scan concurrently.
max_concurrent_scan_files: usize,
- /// Memory pool shared by internal scan operators across all queries.
- scan_memory_pool: Arc,
+ scan_memory_tracker: Option,
/// Whether to enable the experimental two-phase metric series scan.
experimental_series_scan_v2: bool,
/// Whether to ignore range indexes during scans.
@@ -284,7 +282,7 @@ impl ScanRegion {
request,
cache_strategy,
max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES,
- scan_memory_pool: Arc::new(UnboundedMemoryPool::default()),
+ scan_memory_tracker: None,
experimental_series_scan_v2: false,
ignore_range_index: false,
ignore_inverted_index: false,
@@ -326,10 +324,8 @@ impl ScanRegion {
self
}
- /// Sets the memory pool shared by internal scan operators.
- #[must_use]
- pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self {
- self.scan_memory_pool = scan_memory_pool;
+ pub(crate) fn with_scan_memory_tracker(mut self, tracker: QueryMemoryTracker) -> Self {
+ self.scan_memory_tracker = Some(tracker);
self
}
@@ -583,7 +579,7 @@ impl ScanRegion {
.with_bloom_filter_index_appliers(bloom_filter_appliers)
.with_fulltext_index_appliers(fulltext_index_appliers)
.with_max_concurrent_scan_files(self.max_concurrent_scan_files)
- .with_scan_memory_pool(self.scan_memory_pool)
+ .with_scan_memory_tracker(self.scan_memory_tracker)
.with_start_time(self.start_time)
.with_append_mode(self.version.options.append_mode)
.with_filter_deleted(self.filter_deleted)
@@ -946,8 +942,8 @@ pub struct ScanInput {
ignore_file_not_found: bool,
/// Maximum number of SST files to scan concurrently.
pub(crate) max_concurrent_scan_files: usize,
- /// Memory pool shared by internal scan operators across all queries.
- pub(crate) scan_memory_pool: Arc,
+ /// Shared scan quota for projected SST row groups; absent during compaction.
+ scan_memory_tracker: Option,
/// Index appliers.
inverted_index_appliers: [Option; 2],
bloom_filter_index_appliers: [Option; 2],
@@ -1026,7 +1022,7 @@ impl ScanInput {
cache_strategy: CacheStrategy::Disabled,
ignore_file_not_found: false,
max_concurrent_scan_files: DEFAULT_MAX_CONCURRENT_SCAN_FILES,
- scan_memory_pool: Arc::new(UnboundedMemoryPool::default()),
+ scan_memory_tracker: None,
inverted_index_appliers: [None, None],
bloom_filter_index_appliers: [None, None],
fulltext_index_appliers: [None, None],
@@ -1119,6 +1115,9 @@ impl ScanInputBuilder {
/// Builds a finalized [ScanInput] and computes its scan analysis.
#[must_use]
pub(crate) fn build(mut self) -> ScanInput {
+ if self.input.compaction {
+ self.input.scan_memory_tracker = None;
+ }
let input = &self.input;
let eligible = !input.compaction
&& !input.files.is_empty()
@@ -1275,10 +1274,8 @@ impl ScanInputBuilder {
self
}
- /// Sets the memory pool shared by internal scan operators.
- #[must_use]
- pub(crate) fn with_scan_memory_pool(mut self, scan_memory_pool: Arc) -> Self {
- self.input.scan_memory_pool = scan_memory_pool;
+ pub(crate) fn with_scan_memory_tracker(mut self, tracker: Option) -> Self {
+ self.input.scan_memory_tracker = tracker;
self
}
@@ -1415,6 +1412,16 @@ impl ScanInputBuilder {
}
impl ScanInput {
+ pub(crate) fn uses_scan_memory_limit(&self) -> bool {
+ let covered = self
+ .scan_memory_tracker
+ .as_ref()
+ .is_some_and(|tracker| tracker.limit() > 0);
+ #[cfg(feature = "enterprise")]
+ let covered = covered && self.extension_ranges.is_empty();
+ covered
+ }
+
/// Builds memtable ranges to scan by `index`.
pub(crate) fn build_mem_ranges(&self, index: RowGroupIndex) -> SmallVec<[MemtableRange; 2]> {
let memtable = &self.memtables[index.index];
@@ -1547,6 +1554,7 @@ impl ScanInput {
.clone()
.filter(|_| !self.ignore_range_index),
)
+ .memory_tracker(self.scan_memory_tracker.clone())
.predicate(predicate)
.projection(Some(self.read_cols.clone()))
.json2_rewrite_targets(self.json2_rewrite_targets.clone())
@@ -1659,9 +1667,17 @@ impl ScanInput {
// We release the permit before sending result to avoid the task waiting on
// the channel with the permit held.
let maybe_batch = {
- // Safety: We never close the semaphore.
- let _permit = semaphore.acquire().await.unwrap();
- input.next().await
+ tokio::select! {
+ // Cancellation or early completion can drop the receiver while
+ // we wait for a permit or I/O. Exit and drop the
+ // input to release its reservations without waiting for a batch.
+ _ = sender.closed() => break,
+ batch = async {
+ // Safety: We never close the semaphore.
+ let _permit = semaphore.acquire().await.unwrap();
+ input.next().await
+ } => batch,
+ }
};
match maybe_batch {
Some(Ok(batch)) => {
diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs
index 24c35481669..ecfa41bcf2e 100644
--- a/src/mito2/src/read/seq_scan.rs
+++ b/src/mito2/src/read/seq_scan.rs
@@ -107,10 +107,9 @@ impl SeqScan {
)]
pub fn build_stream(&self) -> Result {
let metrics_set = ExecutionPlanMetricsSet::new();
+ let ctx = QueryScanContext::default();
let streams = (0..self.properties.partitions.len())
- .map(|partition: usize| {
- self.scan_partition(&QueryScanContext::default(), &metrics_set, partition)
- })
+ .map(|partition: usize| self.scan_partition(&ctx, &metrics_set, partition))
.collect::, _>>()?;
let aggr_stream = ChainedRecordBatchStream::new(streams).map_err(BoxedError::new)?;
@@ -506,6 +505,9 @@ impl SeqScan {
/// Checks resource limit for the scanner.
pub(crate) fn check_scan_limit(&self) -> Result<()> {
+ if self.stream_ctx.input.uses_scan_memory_limit() {
+ return Ok(());
+ }
// Check max file count limit for all partitions since we scan them in parallel.
let total_max_files: usize = self
.properties
diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs
index 4328a20f59c..a77cbebb557 100644
--- a/src/mito2/src/read/series_scan.rs
+++ b/src/mito2/src/read/series_scan.rs
@@ -25,6 +25,7 @@ use common_recordbatch::util::ChainedRecordBatchStream;
use common_recordbatch::{RecordBatchStreamWrapper, SendableRecordBatchStream};
use common_telemetry::tracing::{self, Instrument};
use common_telemetry::warn;
+use datafusion::execution::memory_pool::MemoryPool;
use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion::physical_plan::{DisplayAs, DisplayFormatType};
use datatypes::arrow::array::BinaryArray;
@@ -266,12 +267,9 @@ impl SeriesScan {
metrics_set,
ctx.explain_verbose,
),
- SeriesScanMode::TwoPhase => self.scan_two_phase_batch_in_partition(
- partition,
- part_metrics,
- metrics_set,
- ctx.explain_verbose,
- ),
+ SeriesScanMode::TwoPhase => {
+ self.scan_two_phase_batch_in_partition(ctx, partition, part_metrics, metrics_set)
+ }
}
}
@@ -311,12 +309,12 @@ impl SeriesScan {
fn scan_two_phase_batch_in_partition(
&self,
+ ctx: &QueryScanContext,
partition: usize,
part_metrics: PartitionMetrics,
metrics_set: &ExecutionPlanMetricsSet,
- explain_verbose: bool,
) -> Result {
- self.maybe_start_candidate_distributor(metrics_set, explain_verbose);
+ self.maybe_start_candidate_distributor(ctx, metrics_set);
// Safety: `scan_batch_in_partition` validates the index, and the receiver
// list always contains one entry for each output partition.
let mut receiver = self.candidate_receivers.lock().unwrap()[partition]
@@ -406,8 +404,8 @@ impl SeriesScan {
fn maybe_start_candidate_distributor(
&self,
+ ctx: &QueryScanContext,
metrics_set: &ExecutionPlanMetricsSet,
- explain_verbose: bool,
) {
let mut rx_list = self.candidate_receivers.lock().unwrap();
if !rx_list.is_empty() {
@@ -416,7 +414,10 @@ impl SeriesScan {
let (senders, receivers, active_receivers) =
new_candidate_channel_list(self.properties.num_partitions());
+ // Capture the first partition's runtime pool once for the shared distributor.
+ // reset_state clears the receivers so the next execution captures a fresh pool.
let mut distributor = SeriesCandidateDistributor {
+ memory_pool: ctx.memory_pool.clone(),
stream_ctx: self.stream_ctx.clone(),
range_semaphore: Arc::new(Semaphore::new(self.properties.num_partitions())),
partitions: self.properties.partitions.clone(),
@@ -425,7 +426,7 @@ impl SeriesScan {
active_receivers,
metrics_set: metrics_set.clone(),
metrics_list: self.metrics_list.clone(),
- explain_verbose,
+ explain_verbose: ctx.explain_verbose,
};
let region_id = distributor.stream_ctx.input.mapper.metadata().region_id;
let span = tracing::info_span!("SeriesScan::candidate_distributor", region_id = %region_id);
@@ -447,8 +448,9 @@ impl SeriesScan {
pub(crate) async fn build_stream(&self) -> Result {
let part_num = self.properties.num_partitions();
let metrics_set = ExecutionPlanMetricsSet::default();
+ let ctx = QueryScanContext::default();
let streams = (0..part_num)
- .map(|i| self.scan_partition(&QueryScanContext::default(), &metrics_set, i))
+ .map(|i| self.scan_partition(&ctx, &metrics_set, i))
.collect::, BoxedError>>()?;
let chained_stream = ChainedRecordBatchStream::new(streams).map_err(BoxedError::new)?;
Ok(Box::pin(chained_stream))
@@ -458,6 +460,7 @@ impl SeriesScan {
pub(crate) fn scan_all_partitions(&self) -> Result {
let metrics_set = ExecutionPlanMetricsSet::new();
+ let ctx = QueryScanContext::default();
let streams = (0..self.properties.partitions.len())
.map(|partition| {
let metrics = new_partition_metrics(
@@ -468,12 +471,7 @@ impl SeriesScan {
&self.metrics_list,
);
- self.scan_batch_in_partition(
- &QueryScanContext::default(),
- partition,
- metrics,
- &metrics_set,
- )
+ self.scan_batch_in_partition(&ctx, partition, metrics, &metrics_set)
})
.collect::>>()?;
@@ -482,6 +480,9 @@ impl SeriesScan {
/// Checks resource limit for the scanner.
pub(crate) fn check_scan_limit(&self) -> Result<()> {
+ if self.stream_ctx.input.uses_scan_memory_limit() {
+ return Ok(());
+ }
// Sum the total number of files across all partitions
let total_files: usize = self
.properties
@@ -842,6 +843,7 @@ impl SeriesDistributor {
/// Discovers metric series and sends one reader assignment to each partition.
struct SeriesCandidateDistributor {
+ memory_pool: Arc,
stream_ctx: Arc,
range_semaphore: Arc,
partitions: Vec>,
@@ -875,7 +877,7 @@ impl SeriesCandidateDistributor {
self.partitions.clone(),
self.pruner.clone(),
self.range_semaphore.clone(),
- self.stream_ctx.input.scan_memory_pool.clone(),
+ self.memory_pool.clone(),
self.metrics_set.clone(),
part_metrics.clone(),
)?;
diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs
index 8bc4f2e0038..9362898023f 100644
--- a/src/mito2/src/read/unordered_scan.rs
+++ b/src/mito2/src/read/unordered_scan.rs
@@ -89,8 +89,9 @@ impl UnorderedScan {
pub(crate) async fn build_stream(&self) -> Result {
let metrics_set = ExecutionPlanMetricsSet::new();
let part_num = self.properties.num_partitions();
+ let ctx = QueryScanContext::default();
let streams = (0..part_num)
- .map(|i| self.scan_partition(&QueryScanContext::default(), &metrics_set, i))
+ .map(|i| self.scan_partition(&ctx, &metrics_set, i))
.collect::, BoxedError>>()?;
let stream = stream! {
for mut stream in streams {
diff --git a/src/mito2/src/sst/parquet/push_decoder.rs b/src/mito2/src/sst/parquet/push_decoder.rs
index 4ae9057479a..1a1446c13c2 100644
--- a/src/mito2/src/sst/parquet/push_decoder.rs
+++ b/src/mito2/src/sst/parquet/push_decoder.rs
@@ -17,6 +17,7 @@
use std::ops::Range;
use bytes::{Bytes, BytesMut};
+use common_recordbatch::QueryMemoryTracker;
use datatypes::arrow::record_batch::RecordBatch;
use futures::StreamExt;
use futures::stream::BoxStream;
@@ -25,7 +26,7 @@ use parquet::DecodeResult;
use parquet::arrow::ProjectionMask;
use parquet::arrow::arrow_reader::{ArrowReaderMetadata, RowSelection};
use parquet::arrow::push_decoder::ParquetPushDecoderBuilder;
-use snafu::{ResultExt, ensure};
+use snafu::{OptionExt, ResultExt, ensure};
use crate::cache::file_cache::{FileType, IndexKey};
use crate::cache::{CacheStrategy, PageRangePart};
@@ -53,6 +54,8 @@ pub struct SstParquetRangeFetcher {
row_group_idx: usize,
/// Optional metrics for tracking fetch operations.
fetch_metrics: Option,
+ /// Shared quota for decoding this fetcher's row group.
+ memory_tracker: Option,
}
impl SstParquetRangeFetcher {
@@ -72,9 +75,16 @@ impl SstParquetRangeFetcher {
cache_strategy,
row_group_idx,
fetch_metrics,
+ memory_tracker: None,
}
}
+ /// Sets the shared quota for this decoder's projected row group.
+ pub(crate) fn with_memory_tracker(mut self, tracker: Option) -> Self {
+ self.memory_tracker = tracker;
+ self
+ }
+
/// Fetches byte ranges from page cache, write cache, or object store.
async fn fetch_bytes_with_cache(&self, ranges: Vec>) -> Result> {
let fetch_start = self
@@ -321,6 +331,24 @@ pub fn build_sst_parquet_record_batch_stream(
file_path: String,
batch_size: usize,
) -> Result>> {
+ let uncompressed_bytes = if fetcher.memory_tracker.is_some() {
+ let row_group = arrow_metadata.metadata().row_group(row_group_idx);
+ row_group
+ .columns()
+ .iter()
+ .enumerate()
+ .filter(|(index, _)| projection.leaf_included(*index))
+ .try_fold(0usize, |total, (_, column)| {
+ usize::try_from(column.uncompressed_size())
+ .ok()
+ .and_then(|size| total.checked_add(size))
+ })
+ .context(UnexpectedSnafu {
+ reason: "Invalid projected row-group uncompressed size",
+ })?
+ } else {
+ 0
+ };
let mut builder = ParquetPushDecoderBuilder::new_with_metadata(arrow_metadata)
.with_row_groups(vec![row_group_idx])
.with_projection(projection)
@@ -335,12 +363,17 @@ pub fn build_sst_parquet_record_batch_stream(
.context(ReadParquetSnafu { path: &file_path })?;
Ok(async_stream::try_stream! {
+ let _reservation = if let Some(tracker) = &fetcher.memory_tracker {
+ Some(tracker.reserve(uncompressed_bytes)
+ .context(crate::error::RecordBatchSnafu)?)
+ } else {
+ None
+ };
loop {
match decoder.try_decode().context(ReadParquetSnafu { path: &file_path })? {
DecodeResult::NeedsData(ranges) => {
let data = fetcher.fetch_bytes_with_cache(ranges.clone()).await?;
- decoder
- .push_ranges(ranges, data)
+ decoder.push_ranges(ranges, data)
.context(ReadParquetSnafu { path: &file_path })?;
}
DecodeResult::Data(batch) => yield batch,
@@ -355,6 +388,139 @@ pub fn build_sst_parquet_record_batch_stream(
mod tests {
use super::*;
+ #[tokio::test]
+ async fn test_decoder_memory_projection_lifetime_and_error() {
+ use std::sync::Arc;
+
+ use datatypes::arrow::array::{Int64Array, StringArray};
+ use datatypes::arrow::datatypes::{DataType, Field, Schema};
+ use object_store::services::Memory;
+ use parquet::arrow::ArrowWriter;
+ use parquet::arrow::arrow_reader::ArrowReaderOptions;
+ use parquet::file::properties::WriterProperties;
+
+ use crate::sst::parquet::metadata::MetadataLoader;
+ use crate::sst::parquet::reader::MetadataCacheMetrics;
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("value", DataType::Int64, false),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+ let batch = RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(Int64Array::from_iter_values(0..8)),
+ Arc::new(StringArray::from_iter_values(
+ (0..8).map(|i| format!("{i}:{}", "x".repeat(512))),
+ )),
+ ],
+ )
+ .unwrap();
+ let mut bytes = Vec::new();
+ let props = WriterProperties::builder()
+ .set_max_row_group_row_count(Some(4))
+ .build();
+ let mut writer = ArrowWriter::try_new(&mut bytes, schema, Some(props)).unwrap();
+ writer.write(&batch).unwrap();
+ writer.close().unwrap();
+ let store = ObjectStore::new(Memory::default()).unwrap();
+ let file_size = bytes.len() as u64;
+ store.write("memory.parquet", bytes).await.unwrap();
+ let metadata = MetadataLoader::new(store.clone(), "memory.parquet", file_size)
+ .load(&mut MetadataCacheMetrics::default())
+ .await
+ .unwrap();
+ let sizes: Vec<_> = metadata
+ .row_groups()
+ .iter()
+ .map(|rg| rg.column(0).uncompressed_size() as usize)
+ .collect();
+ let aligned = |bytes: usize| bytes.div_ceil(1024) * 1024;
+ let limit = aligned(*sizes.iter().max().unwrap());
+ assert!(metadata.row_group(0).total_byte_size() as usize > limit);
+ let projection = ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0]);
+ let metadata =
+ ArrowReaderMetadata::try_new(Arc::new(metadata), ArrowReaderOptions::new()).unwrap();
+ let tracker = QueryMemoryTracker::builder(limit).build();
+ let cache = CacheStrategy::EnableAll(Arc::new(
+ crate::cache::CacheManager::builder()
+ .page_cache_size(1024 * 1024)
+ .build(),
+ ));
+ let region_file_id = RegionFileId::new(
+ store_api::storage::RegionId::new(1, 1),
+ store_api::storage::FileId::random(),
+ );
+ let build = |row_group_idx, path: &str, selection| {
+ let fetcher = SstParquetRangeFetcher::new(
+ region_file_id,
+ path.into(),
+ store.clone(),
+ if path == "memory.parquet" {
+ cache.clone()
+ } else {
+ CacheStrategy::Disabled
+ },
+ row_group_idx,
+ None,
+ )
+ .with_memory_tracker(Some(tracker.clone()));
+ build_sst_parquet_record_batch_stream(
+ metadata.clone(),
+ row_group_idx,
+ selection,
+ projection.clone(),
+ fetcher,
+ path.into(),
+ 1,
+ )
+ .unwrap()
+ };
+ // Construction (including an unpolled stream's drop) must not reserve memory.
+ drop(build(0, "memory.parquet", None));
+ assert_eq!(0, tracker.current());
+ for (row_group_idx, size) in sizes.iter().enumerate() {
+ let mut first = build(row_group_idx, "memory.parquet", None);
+ assert_eq!(1, first.next().await.unwrap().unwrap().num_rows());
+ assert_eq!(aligned(*size), tracker.current());
+ let mut competing = build(row_group_idx, "memory.parquet", None);
+ assert!(competing.next().await.unwrap().is_err());
+ assert_eq!(aligned(*size), tracker.current());
+ // Reserve before decoding, even when the selection would require no page reads.
+ let mut skipped = build(
+ row_group_idx,
+ "memory.parquet",
+ Some(RowSelection::from(vec![
+ parquet::arrow::arrow_reader::RowSelector::skip(4),
+ ])),
+ );
+ assert!(skipped.next().await.unwrap().is_err());
+ assert_eq!(aligned(*size), tracker.current());
+ while first.next().await.transpose().unwrap().is_some() {}
+ // Completed streams may remain alive without holding their quota.
+ assert_eq!(0, tracker.current());
+ }
+ // The second read must acquire quota even when all pages are cached.
+ store.delete("memory.parquet").await.unwrap();
+ let mut cancelled = build(0, "memory.parquet", None);
+ cancelled.next().await.unwrap().unwrap();
+ drop(cancelled);
+ assert_eq!(0, tracker.current());
+ let mut failed = build(0, "missing.parquet", None);
+ assert!(failed.next().await.unwrap().is_err());
+ drop(failed);
+ assert_eq!(0, tracker.current());
+ let mut skipped = build(
+ 0,
+ "memory.parquet",
+ Some(RowSelection::from(vec![
+ parquet::arrow::arrow_reader::RowSelector::skip(4),
+ ])),
+ );
+ assert!(skipped.next().await.is_none());
+ assert_eq!(0, tracker.current());
+ }
+
#[test]
fn test_assemble_range_from_cached_subrange_and_fetched_tail() {
let cached_parts = vec![vec![PageRangePart {
diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs
index a3cbdef90f1..56488669f5a 100644
--- a/src/mito2/src/sst/parquet/reader.rs
+++ b/src/mito2/src/sst/parquet/reader.rs
@@ -22,6 +22,7 @@ use api::v1::SemanticType;
use arrow_schema::extension::{
EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, ExtensionType,
};
+use common_recordbatch::QueryMemoryTracker;
use common_recordbatch::filter::{SimpleFilterEvaluator, TimestampUnitCast};
use common_telemetry::{debug, error, tracing, warn};
use datafusion::physical_plan::PhysicalExpr;
@@ -226,9 +227,17 @@ pub struct ParquetReaderBuilder {
defer_optional_page_index: bool,
/// Scan-wide hint for rows in a decoded batch.
batch_size: usize,
+ /// Optional shared quota for projected row-group decoding.
+ memory_tracker: Option,
}
impl ParquetReaderBuilder {
+ /// Sets the quota used while decoding row groups.
+ pub(crate) fn memory_tracker(mut self, tracker: Option) -> Self {
+ self.memory_tracker = tracker;
+ self
+ }
+
/// Returns a new [ParquetReaderBuilder] to read specific SST.
pub fn new(
table_dir: String,
@@ -238,6 +247,7 @@ impl ParquetReaderBuilder {
) -> ParquetReaderBuilder {
ParquetReaderBuilder {
series_index: None,
+ memory_tracker: None,
table_dir,
path_type,
file_handle,
@@ -656,6 +666,7 @@ impl ParquetReaderBuilder {
.collect::>>()?;
let reader_builder = RowGroupReaderBuilder {
+ memory_tracker: self.memory_tracker.clone(),
file_handle: self.file_handle.clone(),
file_path,
parquet_meta,
@@ -1768,6 +1779,8 @@ pub(crate) struct RowGroupReaderBuilder {
prefilter_builder: Option,
/// Hint for rows in a decoded batch.
batch_size: usize,
+ /// Optional shared quota for projected row-group decoding.
+ memory_tracker: Option,
}
/// Context passed to [RowGroupReaderBuilder::build()] carrying all information
@@ -1956,7 +1969,8 @@ impl RowGroupReaderBuilder {
self.cache_strategy.clone(),
row_group_idx,
fetch_metrics.cloned(),
- );
+ )
+ .with_memory_tracker(self.memory_tracker.clone());
build_sst_parquet_record_batch_stream(
self.arrow_metadata.clone(),
@@ -2575,6 +2589,7 @@ mod tests {
ArrowReaderMetadata::try_new(parquet_meta.clone(), ArrowReaderOptions::new()).unwrap();
(
RowGroupReaderBuilder {
+ memory_tracker: None,
file_handle: file_handle.clone(),
file_path,
parquet_meta,
diff --git a/src/store-api/Cargo.toml b/src/store-api/Cargo.toml
index 4a3616afb8a..fefebeacaa0 100644
--- a/src/store-api/Cargo.toml
+++ b/src/store-api/Cargo.toml
@@ -22,6 +22,7 @@ common-sql.workspace = true
common-time.workspace = true
common-wal.workspace = true
datafusion-common.workspace = true
+datafusion-execution.workspace = true
datafusion-expr.workspace = true
datafusion-physical-plan.workspace = true
datatypes.workspace = true
diff --git a/src/store-api/src/region_engine.rs b/src/store-api/src/region_engine.rs
index 6577880b415..36b8a471394 100644
--- a/src/store-api/src/region_engine.rs
+++ b/src/store-api/src/region_engine.rs
@@ -26,6 +26,7 @@ use common_error::ext::BoxedError;
use common_recordbatch::adapter::RegionQueryStatCounters;
use common_recordbatch::{EmptyRecordBatchStream, QueryMemoryTracker, SendableRecordBatchStream};
use common_time::Timestamp;
+use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool};
use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr};
use datatypes::schema::SchemaRef;
@@ -461,10 +462,23 @@ impl PrepareRequest {
}
/// Necessary context of the query for the scanner.
-#[derive(Clone, Default)]
+/// Direct callers should create one context per scan and reuse it across partitions.
+/// The default pool is unbounded; cloning the context shares its pool.
+#[derive(Clone)]
pub struct QueryScanContext {
/// Whether the query is EXPLAIN ANALYZE VERBOSE.
pub explain_verbose: bool,
+ /// Executing node's query runtime pool, shared by all partitions of this scan.
+ pub memory_pool: Arc,
+}
+
+impl Default for QueryScanContext {
+ fn default() -> Self {
+ Self {
+ explain_verbose: false,
+ memory_pool: Arc::new(UnboundedMemoryPool::default()),
+ }
+ }
}
/// A scanner that provides a way to scan the region concurrently.
diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs
index fc7aaa24dc6..cae9f5f042c 100644
--- a/src/table/src/table/scan.rs
+++ b/src/table/src/table/scan.rs
@@ -456,6 +456,7 @@ impl ExecutionPlan for RegionScanExec {
let ctx = QueryScanContext {
explain_verbose: self.explain_verbose,
+ memory_pool: context.memory_pool().clone(),
};
let stream = self
.scanner
@@ -732,6 +733,7 @@ mod test {
properties: ScannerProperties,
metadata: store_api::metadata::RegionMetadataRef,
dynamic_filters: Arc>>>,
+ observed_pool: Arc>>>,
}
impl fmt::Display for RepeatableScanner {
@@ -770,10 +772,11 @@ mod test {
fn scan_partition(
&self,
- _ctx: &QueryScanContext,
+ ctx: &QueryScanContext,
_metrics_set: &ExecutionPlanMetricsSet,
_partition: usize,
) -> Result {
+ *self.observed_pool.lock().unwrap() = Some(ctx.memory_pool.clone());
Ok(self.batches.as_stream())
}
@@ -960,6 +963,7 @@ mod test {
batches,
properties: ScannerProperties::new(vec![vec![], vec![]], false, 0),
metadata: metadata.clone(),
+ observed_pool: Default::default(),
dynamic_filters: Arc::new(Mutex::new(Vec::new())),
});
let request = ScanRequest {
@@ -1069,6 +1073,7 @@ mod test {
properties,
metadata: metadata.clone(),
dynamic_filters: Arc::new(Mutex::new(Vec::new())),
+ observed_pool: Default::default(),
});
let plan = RegionScanExec::new(
scanner,
@@ -1166,6 +1171,7 @@ mod test {
properties,
metadata: metadata.clone(),
dynamic_filters: Arc::new(Mutex::new(Vec::new())),
+ observed_pool: Default::default(),
});
let plan = RegionScanExec::new(
scanner,
@@ -1219,6 +1225,31 @@ mod test {
}
}
+ #[test]
+ fn test_execute_forwards_runtime_memory_pool() {
+ let (schema, batch, metadata) = dynamic_filter_test_data(5686, false);
+ let observed_pool = Arc::new(Mutex::new(None));
+ let scanner = RepeatableScanner {
+ batches: RecordBatches::try_new(schema, vec![batch]).unwrap(),
+ properties: ScannerProperties::default(),
+ metadata,
+ dynamic_filters: Default::default(),
+ observed_pool: observed_pool.clone(),
+ };
+ let plan =
+ Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap());
+ for _ in 0..2 {
+ let context = SessionContext::new().task_ctx();
+ let pool = context.memory_pool().clone();
+ drop(plan.execute(0, context).unwrap());
+ assert!(Arc::ptr_eq(
+ observed_pool.lock().unwrap().as_ref().unwrap(),
+ &pool
+ ));
+ plan.scanner.lock().unwrap().reset_state();
+ }
+ }
+
#[test]
fn test_count_statistics_require_exact_source_rows() {
for (append_mode, exact, expected) in [
@@ -1236,6 +1267,7 @@ mod test {
.with_total_rows_is_exact(exact),
metadata,
dynamic_filters: Arc::new(Mutex::new(Vec::new())),
+ observed_pool: Default::default(),
});
let plan = RegionScanExec::new(scanner, ScanRequest::default(), None).unwrap();
let stats = StatisticsContext::new()
@@ -1271,6 +1303,7 @@ mod test {
properties: ScannerProperties::default(),
metadata,
dynamic_filters: scanner_dynamic_filters.clone(),
+ observed_pool: Default::default(),
};
let plan =
Arc::new(RegionScanExec::new(Box::new(scanner), ScanRequest::default(), None).unwrap());
diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs
index f522ad39ee5..0bf6228f4c9 100644
--- a/tests-integration/tests/http.rs
+++ b/tests-integration/tests/http.rs
@@ -3455,6 +3455,8 @@ manifest_cache_size = "256MiB"
sst_write_buffer_size = "8MiB"
max_concurrent_scan_files = 384
allow_stale_entries = false
+scan_memory_limit = "25%"
+scan_memory_mode = "basic"
scan_memory_on_exhausted = "fail"
min_compaction_interval = "0s"
schedule_compaction_after_edit = true
@@ -3548,7 +3550,6 @@ fn drop_lines_with_inconsistent_results(input: String) -> String {
"cache_capacity =",
"trace_aux_cache_size =",
"memory_pool_size =",
- "scan_memory_limit =",
"sas_token =",
"scope =",
"num_workers =",