diff --git a/Cargo.lock b/Cargo.lock index 8fe6d74dd2..94428d8b13 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11796,6 +11796,7 @@ dependencies = [ "table", "tokio", "tokio-stream", + "toml 0.8.23", "tracing", "uuid", ] diff --git a/config/config.md b/config/config.md index e8a89e94e1..52362b55c4 100644 --- a/config/config.md +++ b/config/config.md @@ -125,6 +125,11 @@ | `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 | `50%` | 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 (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by 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". | +| `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". | +| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | | `storage.copy_root` | String | `./greptimedb_data/copy` | Root directory for standalone SQL access to local files.
Relative SQL paths are resolved below this directory. Absolute paths are accepted only when
they are inside this directory. Defaults to `/copy`.
Distributed deployments always reject SQL access to local files.
Upgrade note: COPY commands and existing external tables that reference paths outside this
directory will fail. Move those files below the copy root, set this option to a dedicated
directory containing them, or migrate the files to object storage before upgrading. | @@ -345,6 +350,11 @@ | `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 | `50%` | 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 (unbounded, default behavior).
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.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". | +| `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". | +| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `datanode` | -- | -- | Datanode options. | | `datanode.client` | -- | -- | Datanode client options. | | `datanode.client.connect_timeout` | String | `10s` | -- | @@ -548,6 +558,11 @@ | `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 | `50%` | 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 (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by 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". | +| `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". | +| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `storage` | -- | -- | The data storage options. | | `storage.data_home` | String | `./greptimedb_data` | The working home directory. | | `storage.type` | String | `File` | The storage type used to store the data.
- `File`: the data is stored in the local file system.
- `S3`: the data is stored in the S3 object storage.
- `Gcs`: the data is stored in the Google Cloud Storage.
- `Azblob`: the data is stored in the Azure Blob Storage.
- `Oss`: the data is stored in the Aliyun OSS. | @@ -721,5 +736,10 @@ | `query` | -- | -- | -- | | `query.parallelism` | Integer | `1` | Parallelism of the query engine for query sent by flownode.
Default to 1, so it won't use too much cpu or memory | | `query.memory_pool_size` | String | `50%` | Memory pool size for query execution operators (aggregation, sorting, join).
Supports absolute size (e.g., "1GB", "2GB") or percentage of system memory (e.g., "20%").
Setting it to 0 disables the limit (unbounded, default behavior).
When this limit is reached, queries will fail with ResourceExhausted error.
NOTE: This does NOT limit memory used by 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". | +| `query.experimental_spill_max_temp_directory_size` | String | `1GiB` | Maximum total size of spill directory (default: "1GiB").
Ignored unless mode is "custom". | +| `query.experimental_spill_compression` | String | `uncompressed` | Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd".
Ignored unless mode is "custom". | | `memory` | -- | -- | The memory options. | | `memory.enable_heap_profiling` | Bool | `true` | Whether to enable heap profiling activation during startup.
When enabled, heap profiling will be activated if the `MALLOC_CONF` environment variable
is set to "prof:true,prof_active:false". The official image adds this env variable.
Default is true. | diff --git a/config/datanode.example.toml b/config/datanode.example.toml index 20c84fd9ee..4846ebcea1 100644 --- a/config/datanode.example.toml +++ b/config/datanode.example.toml @@ -278,6 +278,30 @@ parallelism = 0 ## NOTE: This does NOT limit memory used by table scans. memory_pool_size = "50%" +## 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). +#+ experimental_memory_pool_policy = "greedy" + +# --- Experimental: DataFusion spill-to-disk controls --- +## 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. +#+ experimental_spill_mode = "default" +## Spill directory path. Ignored unless mode is "custom". +## @toml2docs:none-default +#+ experimental_spill_path = "/path/to/spill" +## Maximum total size of spill directory (default: "1GiB"). +## Ignored unless mode is "custom". +#+ experimental_spill_max_temp_directory_size = "1GiB" +## Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd". +## Ignored unless mode is "custom". +#+ experimental_spill_compression = "uncompressed" +# --- End spill-to-disk --- + ## The data storage options. [storage] ## The working home directory. diff --git a/config/flownode.example.toml b/config/flownode.example.toml index 41212d0ab0..66caf2430a 100644 --- a/config/flownode.example.toml +++ b/config/flownode.example.toml @@ -157,6 +157,30 @@ parallelism = 1 ## NOTE: This does NOT limit memory used by table scans. memory_pool_size = "50%" +## 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). +#+ experimental_memory_pool_policy = "greedy" + +# --- Experimental: DataFusion spill-to-disk controls --- +## 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. +#+ experimental_spill_mode = "default" +## Spill directory path. Ignored unless mode is "custom". +## @toml2docs:none-default +#+ experimental_spill_path = "/path/to/spill" +## Maximum total size of spill directory (default: "1GiB"). +## Ignored unless mode is "custom". +#+ experimental_spill_max_temp_directory_size = "1GiB" +## Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd". +## Ignored unless mode is "custom". +#+ experimental_spill_compression = "uncompressed" +# --- End spill-to-disk --- + ## The memory options. [memory] ## Whether to enable heap profiling activation during startup. diff --git a/config/frontend.example.toml b/config/frontend.example.toml index 8053485572..048a736055 100644 --- a/config/frontend.example.toml +++ b/config/frontend.example.toml @@ -315,6 +315,30 @@ allow_query_fallback = false ## NOTE: This does NOT limit memory used by table scans (only applies to datanodes). memory_pool_size = "50%" +## 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). +#+ experimental_memory_pool_policy = "greedy" + +# --- Experimental: DataFusion spill-to-disk controls --- +## 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. +#+ experimental_spill_mode = "default" +## Spill directory path. Ignored unless mode is "custom". +## @toml2docs:none-default +#+ experimental_spill_path = "/path/to/spill" +## Maximum total size of spill directory (default: "1GiB"). +## Ignored unless mode is "custom". +#+ experimental_spill_max_temp_directory_size = "1GiB" +## Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd". +## Ignored unless mode is "custom". +#+ experimental_spill_compression = "uncompressed" +# --- End spill-to-disk --- + ## Datanode options. [datanode] ## Datanode client options. diff --git a/config/standalone.example.toml b/config/standalone.example.toml index 1fd8fe5a52..b08251b259 100644 --- a/config/standalone.example.toml +++ b/config/standalone.example.toml @@ -445,6 +445,30 @@ parallelism = 0 ## NOTE: This does NOT limit memory used by table scans. memory_pool_size = "50%" +## 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). +#+ experimental_memory_pool_policy = "greedy" + +# --- Experimental: DataFusion spill-to-disk controls --- +## 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. +#+ experimental_spill_mode = "default" +## Spill directory path. Ignored unless mode is "custom". +## @toml2docs:none-default +#+ experimental_spill_path = "/path/to/spill" +## Maximum total size of spill directory (default: "1GiB"). +## Ignored unless mode is "custom". +#+ experimental_spill_max_temp_directory_size = "1GiB" +## Compression for spilled data files: "uncompressed" (default), "lz4_frame", "zstd". +## Ignored unless mode is "custom". +#+ experimental_spill_compression = "uncompressed" +# --- End spill-to-disk --- + ## The data storage options. [storage] ## The working home directory. diff --git a/src/cmd/tests/load_config_test.rs b/src/cmd/tests/load_config_test.rs index 8af216256f..8049e8b8ab 100644 --- a/src/cmd/tests/load_config_test.rs +++ b/src/cmd/tests/load_config_test.rs @@ -307,6 +307,7 @@ fn test_load_flownode_example_config() { allow_query_fallback: false, memory_pool_size: MemoryLimit::Percentage(50), enable_per_region_metrics: false, + ..Default::default() }, meta_client: Some(MetaClientOptions { metasrv_addrs: vec!["127.0.0.1:3002".to_string()], diff --git a/src/flow/src/adapter.rs b/src/flow/src/adapter.rs index e8dd10ad7d..c2f95e1258 100644 --- a/src/flow/src/adapter.rs +++ b/src/flow/src/adapter.rs @@ -147,6 +147,7 @@ impl Default for FlownodeOptions { allow_query_fallback: false, memory_pool_size: MemoryLimit::default(), enable_per_region_metrics: false, + ..Default::default() }, memory: MemoryOptions::default(), } diff --git a/src/query/Cargo.toml b/src/query/Cargo.toml index 3ce1b3df39..9e43bb53e4 100644 --- a/src/query/Cargo.toml +++ b/src/query/Cargo.toml @@ -98,3 +98,4 @@ session = { workspace = true, features = ["testing"] } store-api.workspace = true table = { workspace = true, features = ["testing"] } tokio-stream.workspace = true +toml.workspace = true diff --git a/src/query/src/options.rs b/src/query/src/options.rs index c998839bf2..24cd3853f0 100644 --- a/src/query/src/options.rs +++ b/src/query/src/options.rs @@ -13,9 +13,11 @@ // limitations under the License. use std::collections::HashMap; +use std::path::PathBuf; use chrono::{DateTime, Utc}; use common_base::memory_limit::MemoryLimit; +use common_base::readable_size::ReadableSize; use datafusion::config::{ConfigEntry, ConfigExtension, ExtensionOptions}; use serde::{Deserialize, Serialize}; use session::context::QueryContextRef; @@ -38,6 +40,41 @@ pub const QUERY_ENABLE_REMOTE_DYNAMIC_FILTER_PUSHDOWN: &str = pub const FLOW_INCREMENTAL_MODE_MEMTABLE_ONLY: &str = "memtable_only"; +/// Query spill mode controlling disk manager behavior. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum QuerySpillMode { + /// Preserve DataFusion default disk manager behavior (OS temp directory). + Default, + /// Explicitly configure spill path, quota, and compression. + Custom, + /// Explicitly disable disk spilling; temporary file creation will error. + Disabled, +} + +/// Compression for spilled data files. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum QuerySpillCompression { + /// No compression (default, matches DataFusion default). + Uncompressed, + /// LZ4 frame compression. + Lz4Frame, + /// Zstandard compression. + Zstd, +} + +/// Memory pool allocation policy. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum QueryMemoryPoolPolicy { + /// Greedy first-come-first-served (default). + Greedy, + /// Fair divides memory available after unspillable reservations evenly among + /// spillable reservations and may trigger earlier spills. + Fair, +} + /// Query engine config #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(default)] @@ -53,6 +90,25 @@ pub struct QueryOptions { /// Whether to expose per-region query load metrics. #[serde(skip)] pub enable_per_region_metrics: bool, + /// Experimental: spill-to-disk mode. + /// - `default`: preserve DataFusion built-in OS temp directory behavior. + /// - `custom`: explicitly configure spill path, max directory size, and compression. + /// - `disabled`: explicitly disable disk spilling. + pub experimental_spill_mode: QuerySpillMode, + /// Experimental: spill directory path. Ignored unless `experimental_spill_mode` is + /// `"custom"`. When set, spill files are written into this directory. + pub experimental_spill_path: Option, + /// Experimental: maximum total size of the spill directory (data written to spill files). + /// Ignored unless `experimental_spill_mode` is `"custom"`. Default: `1GiB`. + pub experimental_spill_max_temp_directory_size: ReadableSize, + /// Experimental: compression algorithm applied to spilled data. + /// Ignored unless `experimental_spill_mode` is `"custom"`. Default: `uncompressed`. + pub experimental_spill_compression: QuerySpillCompression, + /// Experimental: memory pool allocation policy. + /// - `greedy`: Greedy first-come-first-served (default). + /// - `fair`: Fair divides memory available after unspillable reservations + /// evenly among spillable reservations and may trigger earlier spills. + pub experimental_memory_pool_policy: QueryMemoryPoolPolicy, } #[allow(clippy::derivable_impls)] @@ -63,6 +119,11 @@ impl Default for QueryOptions { allow_query_fallback: false, memory_pool_size: MemoryLimit::default(), enable_per_region_metrics: false, + experimental_spill_mode: QuerySpillMode::Default, + experimental_spill_path: None, + experimental_spill_max_temp_directory_size: ReadableSize::gb(1), + experimental_spill_compression: QuerySpillCompression::Uncompressed, + experimental_memory_pool_policy: QueryMemoryPoolPolicy::Greedy, } } } @@ -700,3 +761,51 @@ mod flow_extension_tests { assert!(format!("{err}").contains("Out-of-range")); } } + +#[cfg(test)] +mod query_options_tests { + use super::*; + + #[test] + fn test_parse_spill_options_from_toml() { + let toml_str = r#" +experimental_spill_mode = "custom" +experimental_spill_path = "/tmp/spill" +experimental_spill_max_temp_directory_size = "50GiB" +experimental_spill_compression = "zstd" +experimental_memory_pool_policy = "fair" +"#; + let opts: QueryOptions = toml::from_str(toml_str).unwrap(); + assert_eq!(opts.experimental_spill_mode, QuerySpillMode::Custom); + assert_eq!( + opts.experimental_spill_path, + Some(PathBuf::from("/tmp/spill")) + ); + assert_eq!( + opts.experimental_spill_max_temp_directory_size, + ReadableSize::gb(50) + ); + assert_eq!( + opts.experimental_spill_compression, + QuerySpillCompression::Zstd + ); + assert_eq!( + opts.experimental_memory_pool_policy, + QueryMemoryPoolPolicy::Fair + ); + } + + #[test] + fn test_parse_invalid_spill_option_values() { + for toml_str in [ + r#"experimental_spill_mode = "invalid""#, + r#"experimental_spill_compression = "gzip""#, + r#"experimental_memory_pool_policy = "none""#, + ] { + assert!( + toml::from_str::(toml_str).is_err(), + "{toml_str}" + ); + } + } +} diff --git a/src/query/src/query_engine/runtime.rs b/src/query/src/query_engine/runtime.rs index 3169db4e83..4dbf50babe 100644 --- a/src/query/src/query_engine/runtime.rs +++ b/src/query/src/query_engine/runtime.rs @@ -16,9 +16,11 @@ use std::sync::Arc; use datafusion::error::Result as DfResult; use datafusion::execution::context::SessionConfig; +use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; +use datafusion_common::config::SpillCompression; -use crate::options::QueryOptions; +use crate::options::{QueryOptions, QuerySpillCompression, QuerySpillMode}; use crate::query_engine::state::MetricsMemoryPool; /// Reference-counted query runtime provider. @@ -67,14 +69,72 @@ pub struct DefaultQueryRuntimeProvider; impl DefaultQueryRuntimeProvider { /// Creates a default DataFusion runtime environment builder. pub fn runtime_env_builder(ctx: QueryRuntimeContext<'_>) -> RuntimeEnvBuilder { + let mut builder = RuntimeEnvBuilder::new(); + + // Attach the bounded metrics memory pool only when a limit is set + // (>0). When unbounded (0), keep the DataFusion default + // (UnboundedMemoryPool). if ctx.resolved_memory_pool_size > 0 { - RuntimeEnvBuilder::new().with_memory_pool(Arc::new(MetricsMemoryPool::new( + builder = builder.with_memory_pool(Arc::new(MetricsMemoryPool::new( ctx.resolved_memory_pool_size, - ))) - } else { - RuntimeEnvBuilder::new() + ctx.query_options.experimental_memory_pool_policy, + ))); + } + + match ctx.query_options.experimental_spill_mode { + QuerySpillMode::Default => { + // No custom disk manager; preserve DataFusion default OS temp directory. + } + QuerySpillMode::Custom => { + let mut dm_builder = DiskManagerBuilder::default(); + if let Some(ref path) = ctx.query_options.experimental_spill_path { + dm_builder = + dm_builder.with_mode(DiskManagerMode::Directories(vec![path.clone()])); + } + let max_temp_directory_size = + ctx.query_options.experimental_spill_max_temp_directory_size; + dm_builder = + dm_builder.with_max_temp_directory_size(max_temp_directory_size.as_bytes()); + common_telemetry::info!( + "Configured custom query spill: path={:?}, max_temp_directory_size={}, compression={:?}", + ctx.query_options.experimental_spill_path, + max_temp_directory_size, + ctx.query_options.experimental_spill_compression, + ); + builder = builder.with_disk_manager_builder(dm_builder); + } + QuerySpillMode::Disabled => { + let dm_builder = DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled); + builder = builder.with_disk_manager_builder(dm_builder); + } + } + + builder + } +} + +impl QueryRuntimeProvider for DefaultQueryRuntimeProvider { + fn configure_session_config(&self, ctx: QueryRuntimeContext<'_>, config: &mut SessionConfig) { + // Set spill compression on the session config only when spill mode is + // Custom. In Default/Disabled modes, DataFusion's own default + // (Uncompressed) is preserved—setting compression when spill is not + // explicitly configured would be misleading. + if ctx.query_options.experimental_spill_mode == QuerySpillMode::Custom { + config.options_mut().execution.spill_compression = + spill_compression_from_options(ctx.query_options.experimental_spill_compression); } } } -impl QueryRuntimeProvider for DefaultQueryRuntimeProvider {} +/// Map [`QuerySpillCompression`] to DataFusion's [`SpillCompression`]. +/// +/// This conversion is intentionally not a `From` impl because the +/// semantics depend on the spill mode; callers should only invoke +/// this when `experimental_spill_mode == Custom`. +fn spill_compression_from_options(comp: QuerySpillCompression) -> SpillCompression { + match comp { + QuerySpillCompression::Uncompressed => SpillCompression::Uncompressed, + QuerySpillCompression::Lz4Frame => SpillCompression::Lz4Frame, + QuerySpillCompression::Zstd => SpillCompression::Zstd, + } +} diff --git a/src/query/src/query_engine/state.rs b/src/query/src/query_engine/state.rs index 50898aaa58..ecb1a5e33f 100644 --- a/src/query/src/query_engine/state.rs +++ b/src/query/src/query_engine/state.rs @@ -35,7 +35,7 @@ use datafusion::error::Result as DfResult; use datafusion::execution::SessionStateBuilder; use datafusion::execution::context::{QueryPlanner, SessionConfig, SessionContext, SessionState}; use datafusion::execution::memory_pool::{ - GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, + FairSpillPool, GreedyMemoryPool, MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, TrackConsumersPool, }; use datafusion::physical_optimizer::PhysicalOptimizerRule; @@ -76,11 +76,11 @@ use crate::optimizer::string_normalization::StringNormalizationRule; use crate::optimizer::transcribe_atat::TranscribeAtatRule; use crate::optimizer::type_conversion::TypeConversionRule; use crate::optimizer::windowed_sort::WindowedSortPhysicalRule; -use crate::options::QueryOptions as QueryOptionsNew; +use crate::options::{QueryMemoryPoolPolicy, QueryOptions as QueryOptionsNew}; use crate::query_engine::DefaultSerializer; use crate::query_engine::options::QueryOptions; use crate::query_engine::runtime::{ - DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProviderRef, + DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProvider, QueryRuntimeProviderRef, }; use crate::range_select::planner::RangeSelectPlanner; use crate::region_query::RegionQueryHandlerRef; @@ -148,9 +148,7 @@ impl QueryEngineState { ) -> DfResult { let total_memory = get_total_memory_bytes().max(0) as u64; let memory_pool_size = options.memory_pool_size.resolve(total_memory) as usize; - let runtime_provider = plugins - .get::() - .unwrap_or_else(|| Arc::new(DefaultQueryRuntimeProvider)); + let runtime_provider = plugins.get::(); let mut session_config = SessionConfig::new().with_create_default_catalog_and_schema(false); if options.parallelism > 0 { session_config = session_config.with_target_partitions(options.parallelism); @@ -172,9 +170,16 @@ impl QueryEngineState { .skip_physical_aggregate_schema_check = true; let runtime_context = QueryRuntimeContext::new(&options, memory_pool_size); - runtime_provider.configure_session_config(runtime_context, &mut session_config); + let default_runtime_provider = DefaultQueryRuntimeProvider; + default_runtime_provider.configure_session_config(runtime_context, &mut session_config); + if let Some(provider) = runtime_provider.as_ref() { + provider.configure_session_config(runtime_context, &mut session_config); + } let runtime_builder = DefaultQueryRuntimeProvider::runtime_env_builder(runtime_context); - let runtime_env = runtime_provider.build_runtime_env(runtime_context, runtime_builder)?; + let runtime_env = match runtime_provider { + Some(provider) => provider.build_runtime_env(runtime_context, runtime_builder)?, + None => default_runtime_provider.build_runtime_env(runtime_context, runtime_builder)?, + }; // Apply extension rules let mut extension_rules = Vec::new(); @@ -566,26 +571,35 @@ impl DfQueryPlanner { } } -/// A wrapper around TrackConsumersPool that records metrics. +/// A wrapper around a memory pool that records metrics. /// /// This wrapper intercepts all memory pool operations and updates /// Prometheus metrics for monitoring query memory usage and rejections. +/// +/// The inner pool is wrapped with `TrackConsumersPool` to preserve +/// top-consumer error context on rejection. #[derive(Debug)] pub(super) struct MetricsMemoryPool { - inner: Arc>, + inner: Arc, } impl MetricsMemoryPool { // Number of top memory consumers to report in OOM error messages const TOP_CONSUMERS_TO_REPORT: usize = 5; - pub(super) fn new(limit: usize) -> Self { - Self { - inner: Arc::new(TrackConsumersPool::new( - GreedyMemoryPool::new(limit), - NonZeroUsize::new(Self::TOP_CONSUMERS_TO_REPORT).unwrap(), - )), - } + /// Create a new metrics-wrapped memory pool with the given size limit and + /// allocation policy. + pub(super) fn new(limit: usize, policy: QueryMemoryPoolPolicy) -> Self { + let top_n = NonZeroUsize::new(Self::TOP_CONSUMERS_TO_REPORT).unwrap(); + let inner: Arc = match policy { + QueryMemoryPoolPolicy::Greedy => { + Arc::new(TrackConsumersPool::new(GreedyMemoryPool::new(limit), top_n)) + } + QueryMemoryPoolPolicy::Fair => { + Arc::new(TrackConsumersPool::new(FairSpillPool::new(limit), top_n)) + } + }; + Self { inner } } #[inline] @@ -645,11 +659,15 @@ mod tests { use datafusion::error::DataFusionError; use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryLimit as DfMemoryLimit}; use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + use datafusion_common::config::SpillCompression; use session::context::QueryContext; use super::*; - use crate::options::QueryOptions; - use crate::query_engine::runtime::{QueryRuntimeProvider, QueryRuntimeProviderRef}; + use crate::options::{QueryOptions, QuerySpillCompression, QuerySpillMode}; + use crate::query_engine::runtime::{ + DefaultQueryRuntimeProvider, QueryRuntimeContext, QueryRuntimeProvider, + QueryRuntimeProviderRef, + }; fn new_query_engine_state() -> QueryEngineState { new_query_engine_state_with(Plugins::default(), QueryOptions::default()) @@ -775,12 +793,23 @@ mod tests { plugins, QueryOptions { memory_pool_size: MemoryLimit::Size(ReadableSize(1024)), + experimental_spill_mode: QuerySpillMode::Custom, + experimental_spill_compression: QuerySpillCompression::Zstd, ..Default::default() }, ); assert!(provider.configure_called.load(Ordering::SeqCst)); assert_eq!(7, state.session_state().config().target_partitions()); + assert_eq!( + SpillCompression::Zstd, + state + .session_state() + .config() + .options() + .execution + .spill_compression + ); } #[test] @@ -868,4 +897,259 @@ mod tests { second_query_ctx.remote_query_id_value().unwrap() ); } + + /// Builds a runtime env through the default provider seam, mirroring what + /// [`QueryEngineState::try_new`] does. + fn build_runtime_env(options: &QueryOptions, memory_pool_size: usize) -> Arc { + let ctx = QueryRuntimeContext::new(options, memory_pool_size); + let builder = DefaultQueryRuntimeProvider::runtime_env_builder(ctx); + DefaultQueryRuntimeProvider + .build_runtime_env(ctx, builder) + .expect("Failed to build RuntimeEnv") + } + + #[test] + fn test_build_runtime_env_custom_mode_with_path() { + // Use a temp directory managed manually so we don't need the `tempfile` crate. + let spill_dir = std::env::temp_dir().join(format!("df_spill_test_{}", std::process::id())); + let _ = std::fs::remove_dir_all(&spill_dir); + let opts = QueryOptions { + experimental_spill_mode: QuerySpillMode::Custom, + experimental_spill_path: Some(spill_dir.clone()), + experimental_spill_max_temp_directory_size: ReadableSize::gb(1), + ..Default::default() + }; + let env = build_runtime_env(&opts, 0); + + assert!(env.disk_manager.tmp_files_enabled()); + let tmp_file = env.disk_manager.create_tmp_file("test spill"); + assert!(tmp_file.is_ok()); + assert!(spill_dir.exists()); + + let _ = std::fs::remove_dir_all(&spill_dir); + } + + #[test] + fn test_build_runtime_env_disabled_mode() { + let opts = QueryOptions { + experimental_spill_mode: QuerySpillMode::Disabled, + ..Default::default() + }; + let env = build_runtime_env(&opts, 0); + + assert!(!env.disk_manager.tmp_files_enabled()); + let result = env.disk_manager.create_tmp_file("test spill"); + assert!(result.is_err()); + assert!(format!("{}", result.unwrap_err()).contains("DiskManager is disabled")); + } + + #[test] + fn test_metrics_memory_pool_policy_differs_with_multiple_spillable_consumers() { + for (policy, greedy) in [ + (QueryMemoryPoolPolicy::Greedy, true), + (QueryMemoryPoolPolicy::Fair, false), + ] { + let pool: Arc = Arc::new(MetricsMemoryPool::new(100, policy)); + let first = MemoryConsumer::new("first") + .with_can_spill(true) + .register(&pool); + let second = MemoryConsumer::new("second") + .with_can_spill(true) + .register(&pool); + + let first_result = first.try_grow(75); + assert_eq!(first_result.is_ok(), greedy); + if greedy { + assert!(second.try_grow(26).is_err()); + } else { + first.try_grow(50).unwrap(); + second.try_grow(50).unwrap(); + } + } + } + + /// Builds a runtime with custom spill configuration and a bounded memory + /// pool, then runs sort queries that probe spill-to-disk behaviour: + /// + /// - A pool large enough for merge chunks but still much smaller + /// than the total data. Executes the physical plan directly so we can + /// walk the plan tree afterwards and sum `spill_count` / `spilled_rows` / + /// `spilled_bytes` on `SortExec` nodes. All three must be > 0. + #[tokio::test] + async fn test_sort_spill_smoke_with_custom_runtime() { + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::collect as df_collect; + + let spill_dir = + std::env::temp_dir().join(format!("greptime_spill_smoke_{}", std::process::id())); + let _ = std::fs::remove_dir_all(&spill_dir); + + let opts = QueryOptions { + experimental_spill_mode: QuerySpillMode::Custom, + experimental_spill_path: Some(spill_dir.clone()), + experimental_spill_max_temp_directory_size: ReadableSize::gb(1), + experimental_memory_pool_policy: QueryMemoryPoolPolicy::Greedy, + ..Default::default() + }; + + let session_config = SessionConfig::new() + .with_target_partitions(1) + .with_sort_in_place_threshold_bytes(0) + .with_sort_spill_reservation_bytes(64 * 1024); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("val", DataType::Int32, false), + ])); + let n_rows: i32 = 200_000; + let batch_size: i32 = 5_000; + let partitions = 1usize; + + let mut table_partitions: Vec> = Vec::with_capacity(partitions); + for _ in 0..partitions { + let mut batches = Vec::new(); + let mut row_offset: i32 = 0; + while row_offset < n_rows { + let chunk_end = (row_offset + batch_size).min(n_rows); + let chunk_len = (chunk_end - row_offset) as usize; + let mut id_builder = Int32Array::builder(chunk_len); + let mut val_builder = Int32Array::builder(chunk_len); + for i in row_offset..chunk_end { + id_builder.append_value(i); + val_builder.append_value(n_rows - 1 - i); + } + batches.push( + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(id_builder.finish()), + Arc::new(val_builder.finish()), + ], + ) + .unwrap(), + ); + row_offset = chunk_end; + } + table_partitions.push(batches); + } + + fn build_ctx( + session_config: &SessionConfig, + runtime: &Arc, + schema: &Arc, + partitions: &[Vec], + ) -> SessionContext { + let session_state = SessionStateBuilder::new() + .with_config(session_config.clone()) + .with_runtime_env(Arc::clone(runtime)) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(session_state); + let table = MemTable::try_new(Arc::clone(schema), partitions.to_vec()).unwrap(); + ctx.register_table("t", Arc::new(table)).unwrap(); + ctx + } + + { + let mem_limit: usize = 512 * 1024; // 512 KB pool, 64 KB reserved for merge + + let runtime = build_runtime_env(&opts, mem_limit); + assert_eq!(runtime.memory_pool.reserved(), 0); + let ctx = build_ctx(&session_config, &runtime, &schema, &table_partitions); + + let df = ctx + .sql("SELECT val FROM t ORDER BY val ASC") + .await + .expect("planning ORDER BY"); + let plan = df + .create_physical_plan() + .await + .expect("creating physical plan"); + let task_ctx = ctx.task_ctx(); + + let batches = df_collect(Arc::clone(&plan), task_ctx) + .await + .expect("executing ORDER BY"); + + let (spill_count, spilled_rows, spilled_bytes) = sum_sort_spill_metrics(&plan); + assert!( + spill_count > 0, + "expected SortExec spill_count > 0 (pool={} KB, reservation=64 KB), got 0", + mem_limit / 1024, + ); + assert!( + spilled_rows > 0, + "expected SortExec spilled_rows > 0, got 0 \ + (spill_count={spill_count}, spilled_bytes={spilled_bytes})", + ); + assert!( + spilled_bytes > 0, + "expected SortExec spilled_bytes > 0, got 0 \ + (spill_count={spill_count}, spilled_rows={spilled_rows})", + ); + + let vals: Vec = batches + .iter() + .flat_map(|b| { + let col = b.column(0).as_any().downcast_ref::().unwrap(); + (0..b.num_rows()).map(move |i| col.value(i)) + }) + .collect(); + assert_eq!(vals.len(), n_rows as usize); + for w in vals.windows(2) { + assert!(w[0] <= w[1], "sort order violation: {} > {}", w[0], w[1]); + } + + assert!( + std::fs::read_dir(&spill_dir) + .ok() + .map(|mut entries| entries.any(|e| { + e.as_ref() + .map(|de| de.file_name().to_string_lossy().starts_with("datafusion-")) + .unwrap_or(false) + })) + .unwrap_or(false), + "Expected 'datafusion-*' directory in spill path {:?} \ + (DiskManager should create it on build).", + spill_dir, + ); + + assert_eq!(runtime.memory_pool.reserved(), 0); + } + + let _ = std::fs::remove_dir_all(&spill_dir); + } + + /// Walk a physical plan tree, summing spill metrics for every node whose + /// name starts with `"SortExec"`. + fn sum_sort_spill_metrics(plan: &Arc) -> (usize, usize, usize) { + let mut spill_count = 0usize; + let mut spilled_rows = 0usize; + let mut spilled_bytes = 0usize; + + fn walk(plan: &Arc, sc: &mut usize, sr: &mut usize, sb: &mut usize) { + let name = plan.name(); + if name.starts_with("SortExec") + && let Some(m) = plan.metrics() + { + *sc += m.spill_count().unwrap_or(0); + *sr += m.spilled_rows().unwrap_or(0); + *sb += m.spilled_bytes().unwrap_or(0); + } + for child in plan.children() { + walk(child, sc, sr, sb); + } + } + + walk( + plan, + &mut spill_count, + &mut spilled_rows, + &mut spilled_bytes, + ); + (spill_count, spilled_rows, spilled_bytes) + } } diff --git a/tests-integration/tests/http.rs b/tests-integration/tests/http.rs index f4cfda24dc..5e9012514e 100644 --- a/tests-integration/tests/http.rs +++ b/tests-integration/tests/http.rs @@ -2345,6 +2345,10 @@ ttl = "2months 29days 2h 52m 48s" [query] parallelism = 0 allow_query_fallback = false +experimental_spill_mode = "default" +experimental_spill_max_temp_directory_size = "1GiB" +experimental_spill_compression = "uncompressed" +experimental_memory_pool_policy = "greedy" [memory] enable_heap_profiling = true