refactor(spill): extract MemoryPool from SpillsBufferPool and make pool size per-session configurable (#19805)

* refactor(spill): extract MemoryPool from SpillsBufferPool and make pool size per-session configurable

- Extract MemoryPool struct from SpillsBufferPool with chunk-based allocation and blocking wait
- Add session-level setting 'spill_writer_memory_pool_size_mb' (default 20MB) to control per-writer pool size
- Remove global 'buffer_pool_memory' from SpillConfig, making memory management more flexible
- MemoryPool uses lazy allocation up to max_chunks, then blocks waiting for buffer release

* z

* z

Author: zhang2014

src/query/config/src/config.rs

@@ -3503,10 +3503,6 @@ pub struct SpillConfig {
     #[clap(long, value_name = "PERCENT", default_value = "0")]
     pub result_set_spilling_disk_quota_ratio: u64,
 
-    /// Total memory for the spill buffer pool in bytes.
-    #[clap(long, value_name = "VALUE", default_value = "209715200")]
-    pub spill_buffer_pool_memory: u64,
-
     /// Number of worker tasks in the spill buffer pool.
     #[clap(long, value_name = "VALUE", default_value = "2")]
     pub spill_buffer_pool_workers: usize,
@@ -3623,7 +3619,6 @@ mod config_converters {
             window_partition_spilling_disk_quota_ratio: spill
                 .window_partition_spilling_disk_quota_ratio,
             result_set_spilling_disk_quota_ratio: spill.result_set_spilling_disk_quota_ratio,
-            buffer_pool_memory: spill.spill_buffer_pool_memory,
             buffer_pool_workers: spill.spill_buffer_pool_workers,
         })
     }
@@ -3647,7 +3642,6 @@ mod config_converters {
                 window_partition_spilling_disk_quota_ratio: value
                     .window_partition_spilling_disk_quota_ratio,
                 result_set_spilling_disk_quota_ratio: value.result_set_spilling_disk_quota_ratio,
-                spill_buffer_pool_memory: value.buffer_pool_memory,
                 spill_buffer_pool_workers: value.buffer_pool_workers,
             }
         }

src/query/config/src/inner.rs

@@ -434,9 +434,6 @@ pub struct SpillConfig {
     /// result-set spilling may use for one query.
     pub result_set_spilling_disk_quota_ratio: u64,
 
-    /// Total memory for the spill buffer pool in bytes.
-    pub buffer_pool_memory: u64,
-
     /// Number of worker tasks in the spill buffer pool.
     pub buffer_pool_workers: usize,
 }
@@ -509,7 +506,6 @@ impl SpillConfig {
             window_partition_spilling_disk_quota_ratio: 60,
             // TODO: keep 0 to avoid deleting local result-set spill dir before HTTP pagination finishes.
             result_set_spilling_disk_quota_ratio: 0,
-            buffer_pool_memory: 200 * 1024 * 1024,
             buffer_pool_workers: 2,
         }
     }
@@ -527,7 +523,6 @@ impl Default for SpillConfig {
             window_partition_spilling_disk_quota_ratio: 60,
             // TODO: keep 0 to avoid deleting local result-set spill dir before HTTP pagination finishes.
             result_set_spilling_disk_quota_ratio: 0,
-            buffer_pool_memory: 200 * 1024 * 1024,
             buffer_pool_workers: 2,
         }
     }

src/query/config/src/mask.rs

@@ -218,7 +218,6 @@ impl SpillConfig {
             sort_spilling_disk_quota_ratio,
             window_partition_spilling_disk_quota_ratio,
             result_set_spilling_disk_quota_ratio,
-            spill_buffer_pool_memory,
             spill_buffer_pool_workers,
         } = *self;
 
@@ -230,7 +229,6 @@ impl SpillConfig {
             sort_spilling_disk_quota_ratio,
             window_partition_spilling_disk_quota_ratio,
             result_set_spilling_disk_quota_ratio,
-            spill_buffer_pool_memory,
             spill_buffer_pool_workers,
         }
     }
@@ -388,7 +386,6 @@ mod tests {
             sort_spilling_disk_quota_ratio: 60,
             window_partition_spilling_disk_quota_ratio: 30,
             result_set_spilling_disk_quota_ratio: 0,
-            spill_buffer_pool_memory: 209715200,
             spill_buffer_pool_workers: 2,
             storage: Some(StorageConfig {
                 typ: "s3".to_string(),

src/query/service/src/pipelines/builders/builder_sort.rs

@@ -138,6 +138,9 @@ impl SortPipelineBuilder {
                 location_prefix,
                 disk_spill,
                 use_parquet: settings.get_spilling_file_format()?.is_parquet(),
+                writer_pool_bytes: settings
+                    .get_spill_writer_memory_pool_size_mb()?
+                    .saturating_mul(1024 * 1024),
             };
             let op = DataOperator::instance().spill_operator();
             SortSpillerImpl::new(self.ctx.clone(), op, config)?
@@ -224,6 +227,9 @@ impl SortPipelineBuilder {
                 location_prefix,
                 disk_spill: None,
                 use_parquet: settings.get_spilling_file_format()?.is_parquet(),
+                writer_pool_bytes: settings
+                    .get_spill_writer_memory_pool_size_mb()?
+                    .saturating_mul(1024 * 1024),
             };
             let op = DataOperator::instance().spill_operator();
             SortSpillerImpl::new(self.ctx.clone(), op, config)?

src/query/service/src/pipelines/processors/transforms/aggregator/new_aggregate/new_aggregate_spiller.rs

@@ -45,18 +45,21 @@ use crate::spillers::SpillTarget;
 use crate::spillers::SpillsBufferPool;
 use crate::spillers::SpillsDataWriter;
 
+const BYTES_PER_MIB: usize = 1024 * 1024;
+
 struct PayloadWriter {
     path: String,
     writer: SpillsDataWriter,
 }
 
 impl PayloadWriter {
-    fn try_create(prefix: &str) -> Result<Self> {
+    fn try_create(prefix: &str, writer_pool_bytes: usize) -> Result<Self> {
         let data_operator = DataOperator::instance();
         let operator = data_operator.spill_operator();
         let buffer_pool = SpillsBufferPool::instance();
         let file_path = format!("{}/{}", prefix, GlobalUniq::unique());
-        let spills_data_writer = buffer_pool.writer(operator, file_path.clone())?;
+        let spills_data_writer =
+            buffer_pool.writer(operator, file_path.clone(), writer_pool_bytes)?;
 
         Ok(PayloadWriter {
             path: file_path,
@@ -109,16 +112,23 @@ impl WriteStats {
 struct AggregatePayloadWriters {
     spill_prefix: String,
     partition_count: usize,
+    writer_pool_bytes: usize,
     writers: Vec<Option<PayloadWriter>>,
     write_stats: WriteStats,
     ctx: Arc<QueryContext>,
 }
 
 impl AggregatePayloadWriters {
-    pub fn create(prefix: &str, partition_count: usize, ctx: Arc<QueryContext>) -> Self {
+    pub fn create(
+        prefix: &str,
+        partition_count: usize,
+        writer_pool_bytes: usize,
+        ctx: Arc<QueryContext>,
+    ) -> Self {
         AggregatePayloadWriters {
             spill_prefix: prefix.to_string(),
             partition_count,
+            writer_pool_bytes,
             writers: Self::empty_writers(partition_count),
             write_stats: WriteStats::default(),
             ctx,
@@ -133,7 +143,10 @@ impl AggregatePayloadWriters {
 
     fn ensure_writer(&mut self, bucket: usize) -> Result<&mut PayloadWriter> {
         if self.writers[bucket].is_none() {
-            self.writers[bucket] = Some(PayloadWriter::try_create(&self.spill_prefix)?);
+            self.writers[bucket] = Some(PayloadWriter::try_create(
+                &self.spill_prefix,
+                self.writer_pool_bytes,
+            )?);
         }
 
         Ok(self.writers[bucket].as_mut().unwrap())
@@ -335,8 +348,13 @@ impl<P: PartitionStream> NewAggregateSpiller<P> {
         let table_ctx: Arc<dyn TableContext> = ctx.clone();
         let read_setting = ReadSettings::from_settings(&table_ctx.get_settings())?;
         let spill_prefix = ctx.query_id_spill_prefix();
+        let writer_pool_bytes = ctx
+            .get_settings()
+            .get_spill_writer_memory_pool_size_mb()?
+            .saturating_mul(BYTES_PER_MIB);
 
-        let payload_writers = AggregatePayloadWriters::create(&spill_prefix, partition_count, ctx);
+        let payload_writers =
+            AggregatePayloadWriters::create(&spill_prefix, partition_count, writer_pool_bytes, ctx);
 
         Ok(Self {
             memory_settings,

src/query/service/src/pipelines/processors/transforms/aggregator/serde/transform_aggregate_spill_writer.rs

@@ -68,6 +68,10 @@ impl TransformAggregateSpillWriter {
             location_prefix,
             disk_spill: None,
             use_parquet: ctx.get_settings().get_spilling_file_format()?.is_parquet(),
+            writer_pool_bytes: ctx
+                .get_settings()
+                .get_spill_writer_memory_pool_size_mb()?
+                .saturating_mul(1024 * 1024),
         };
 
         let spiller = Spiller::create(ctx.clone(), operator, config.clone())?;

src/query/service/src/pipelines/processors/transforms/aggregator/serde/transform_exchange_aggregate_serializer.rs

@@ -97,6 +97,10 @@ impl TransformExchangeAggregateSerializer {
             location_prefix,
             disk_spill: None,
             use_parquet: ctx.get_settings().get_spilling_file_format()?.is_parquet(),
+            writer_pool_bytes: ctx
+                .get_settings()
+                .get_spill_writer_memory_pool_size_mb()?
+                .saturating_mul(1024 * 1024),
         };
 
         let spiller = Spiller::create(ctx.clone(), operator, config.clone())?;

src/query/service/src/pipelines/processors/transforms/hash_join/hash_join_spiller.rs

@@ -81,6 +81,10 @@ impl HashJoinSpiller {
             location_prefix,
             disk_spill: None,
             use_parquet: ctx.get_settings().get_spilling_file_format()?.is_parquet(),
+            writer_pool_bytes: ctx
+                .get_settings()
+                .get_spill_writer_memory_pool_size_mb()?
+                .saturating_mul(1024 * 1024),
         };
         let operator = DataOperator::instance().spill_operator();
         let spiller = Spiller::create(ctx.clone(), operator, spill_config)?;

src/query/service/src/pipelines/processors/transforms/new_hash_join/grace/grace_join.rs

@@ -46,12 +46,15 @@ use crate::spillers::SpillsBufferPool;
 use crate::spillers::SpillsDataReader;
 use crate::spillers::SpillsDataWriter;
 
+const BYTES_PER_MIB: usize = 1024 * 1024;
+
 pub struct GraceHashJoin<T: GraceMemoryJoin> {
     pub(crate) desc: Arc<HashJoinDesc>,
     pub(crate) hash_method_kind: HashMethodKind,
     pub(crate) function_context: FunctionContext,
 
     pub(crate) location_prefix: String,
+    pub(crate) writer_pool_bytes: usize,
     pub(crate) shift_bits: usize,
 
     pub(crate) stage: RestoreStage,
@@ -93,7 +96,8 @@ impl<T: GraceMemoryJoin> Join for GraceHashJoin<T> {
 
         let mut ready_partitions = Vec::with_capacity(self.partitions.len());
         for id in 0..self.partitions.len() {
-            let mut partition = GraceJoinPartition::create(&self.location_prefix)?;
+            let mut partition =
+                GraceJoinPartition::create(&self.location_prefix, self.writer_pool_bytes)?;
 
             std::mem::swap(&mut self.partitions[id], &mut partition);
             ready_partitions.push(partition);
@@ -206,18 +210,25 @@ impl<T: GraceMemoryJoin> GraceHashJoin<T> {
         let rows = settings.get_max_block_size()? as usize;
         let bytes = settings.get_max_block_bytes()? as usize;
         let location_prefix = ctx.query_id_spill_prefix();
+        let writer_pool_bytes = settings
+            .get_spill_writer_memory_pool_size_mb()?
+            .saturating_mul(BYTES_PER_MIB);
 
         let mut partitions = Vec::with_capacity(16);
 
         for _ in 0..16 {
-            partitions.push(GraceJoinPartition::create(&location_prefix)?);
+            partitions.push(GraceJoinPartition::create(
+                &location_prefix,
+                writer_pool_bytes,
+            )?);
         }
 
         Ok(GraceHashJoin {
             desc,
             state,
             shift_bits,
             location_prefix,
+            writer_pool_bytes,
             hash_method_kind,
             memory_hash_join,
             function_context: function_ctx,
@@ -416,13 +427,14 @@ pub struct GraceJoinPartition {
 }
 
 impl GraceJoinPartition {
-    pub fn create(prefix: &str) -> Result<GraceJoinPartition> {
+    pub fn create(prefix: &str, writer_pool_bytes: usize) -> Result<GraceJoinPartition> {
         let data_operator = DataOperator::instance();
 
         let operator = data_operator.spill_operator();
         let buffer_pool = SpillsBufferPool::instance();
         let file_path = format!("{}/{}", prefix, GlobalUniq::unique());
-        let spills_data_writer = buffer_pool.writer(operator, file_path.clone())?;
+        let spills_data_writer =
+            buffer_pool.writer(operator, file_path.clone(), writer_pool_bytes)?;
 
         Ok(GraceJoinPartition {
             path: file_path,

src/query/service/src/pipelines/processors/transforms/window/partition/transform_window_partition_collect.rs

@@ -193,6 +193,9 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
             location_prefix,
             disk_spill,
             use_parquet: settings.get_spilling_file_format()?.is_parquet(),
+            writer_pool_bytes: settings
+                .get_spill_writer_memory_pool_size_mb()?
+                .saturating_mul(1024 * 1024),
         };
 
         // Create spillers for window operator.