feat(query): add experimental DataFusion spill-to-disk controls (#8884)

* feat(query): add experimental DataFusion spill-to-disk controls Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * docs(config): regenerate configuration reference Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * test: update config API for spill defaults Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): address spill configuration review Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> * fix(query): preserve spill settings with runtime plugins Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --------- Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
2026-09-05 21:18:57 +00:00 · 2026-09-01 07:40:55 +00:00
parent 27165c2fdc
commit 00d43b29ad
13 changed files with 602 additions and 25 deletions
@@ -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()],
@@ -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(),
        }
@@ -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
@@ -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<PathBuf>,
+    /// 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::<QueryOptions>(toml_str).is_err(),
+                "{toml_str}"
+            );
+        }
+    }
+}
@@ -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,
+    }
+}
@@ -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<Self> {
        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::<QueryRuntimeProviderRef>()
-            .unwrap_or_else(|| Arc::new(DefaultQueryRuntimeProvider));
+        let runtime_provider = plugins.get::<QueryRuntimeProviderRef>();
        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<TrackConsumersPool<GreedyMemoryPool>>,
+    inner: Arc<dyn MemoryPool>,
 }

 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<dyn MemoryPool> = 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<RuntimeEnv> {
+        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<dyn MemoryPool> = 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<RecordBatch>> = 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<RuntimeEnv>,
+            schema: &Arc<Schema>,
+            partitions: &[Vec<RecordBatch>],
+        ) -> 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<i32> = batches
+                .iter()
+                .flat_map(|b| {
+                    let col = b.column(0).as_any().downcast_ref::<Int32Array>().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<dyn ExecutionPlan>) -> (usize, usize, usize) {
+        let mut spill_count = 0usize;
+        let mut spilled_rows = 0usize;
+        let mut spilled_bytes = 0usize;
+
+        fn walk(plan: &Arc<dyn ExecutionPlan>, 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)
+    }
 }