x

Author: forsaken628

src/query/service/src/physical_plans/physical_recluster.rs

@@ -350,6 +350,7 @@ impl IPhysicalPlan for HilbertPartition {
                     self.num_partitions,
                     window_spill_settings.clone(),
                     disk_spill.clone(),
+                    false,
                     CompactStrategy::new(self.rows_per_block, max_bytes_per_block),
                 )?,
             )))

src/query/service/src/physical_plans/physical_window_partition.rs

@@ -181,6 +181,7 @@ impl IPhysicalPlan for WindowPartition {
                     num_partitions,
                     window_spill_settings.clone(),
                     disk_spill.clone(),
+                    true,
                     strategy,
                 )?,
             )))

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

@@ -20,6 +20,7 @@ use std::any::Any;
 use std::collections::VecDeque;
 use std::sync::Arc;
 
+use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_exception::Result;
 use databend_common_expression::BlockMetaInfoDowncast;
 use databend_common_expression::DataBlock;
@@ -30,17 +31,19 @@ use databend_common_pipeline_core::processors::Processor;
 use databend_common_pipeline_transforms::MemorySettings;
 use databend_common_settings::Settings;
 use databend_common_storage::DataOperator;
+use either::Either;
 
 use super::window_partition_buffer_v2::WindowPartitionBufferV2;
 use super::WindowPartitionBuffer;
 use super::WindowPartitionMeta;
 use crate::pipelines::processors::transforms::DataProcessorStrategy;
 use crate::sessions::QueryContext;
+use crate::spillers::BackpressureSpiller;
+use crate::spillers::BufferPool;
 use crate::spillers::Spiller;
 use crate::spillers::SpillerConfig;
 use crate::spillers::SpillerDiskConfig;
 use crate::spillers::SpillerType;
-use crate::spillers::WindowSpiller;
 
 enum WindowBuffer {
     V1(WindowPartitionBuffer),
@@ -49,29 +52,30 @@ enum WindowBuffer {
 
 impl WindowBuffer {
     fn new(
-        is_v2: bool,
-        partition_spiller: Spiller,
-        writer_spiller: WindowSpiller,
+        spiller: Either<Spiller, BackpressureSpiller>,
         num_partitions: usize,
         sort_block_size: usize,
         memory_settings: MemorySettings,
     ) -> Result<Self> {
-        if is_v2 {
-            let inner = WindowPartitionBufferV2::new(
-                writer_spiller,
-                num_partitions,
-                sort_block_size,
-                memory_settings,
-            )?;
-            Ok(Self::V2(inner))
-        } else {
-            let inner = WindowPartitionBuffer::new(
-                partition_spiller,
-                num_partitions,
-                sort_block_size,
-                memory_settings,
-            )?;
-            Ok(Self::V1(inner))
+        match spiller {
+            Either::Left(spiller) => {
+                let inner = WindowPartitionBuffer::new(
+                    spiller,
+                    num_partitions,
+                    sort_block_size,
+                    memory_settings,
+                )?;
+                Ok(Self::V1(inner))
+            }
+            Either::Right(spiller) => {
+                let inner = WindowPartitionBufferV2::new(
+                    spiller,
+                    num_partitions,
+                    sort_block_size,
+                    memory_settings,
+                )?;
+                Ok(Self::V2(inner))
+            }
         }
     }
 
@@ -150,6 +154,7 @@ pub struct TransformWindowPartitionCollect<S: DataProcessorStrategy> {
 }
 
 impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
+    #[expect(clippy::too_many_arguments)]
     pub fn new(
         ctx: Arc<QueryContext>,
         input: Arc<InputPort>,
@@ -160,6 +165,7 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
         num_partitions: usize,
         memory_settings: MemorySettings,
         disk_spill: Option<SpillerDiskConfig>,
+        enable_backpressure_spiller: bool,
         strategy: S,
     ) -> Result<Self> {
         // Calculate the partition ids collected by the processor.
@@ -183,20 +189,24 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
 
         // Create spillers for window operator.
         let operator = DataOperator::instance().spill_operator();
-        let partition_spiller =
-            Spiller::create(ctx.clone(), operator.clone(), spill_config.clone())?;
-        let window_spiller = WindowSpiller::create(ctx, operator, spill_config)?;
+        let spiller = if !enable_backpressure_spiller {
+            Either::Left(Spiller::create(ctx, operator, spill_config)?)
+        } else {
+            let runtime = GlobalIORuntime::instance();
+            let buffer_pool = BufferPool::create(runtime, 128 * 1024 * 1024, 3);
+            Either::Right(BackpressureSpiller::create(
+                ctx,
+                operator,
+                spill_config,
+                buffer_pool,
+                8 * 1024 * 1024,
+            )?)
+        };
 
         // Create the window partition buffer.
         let sort_block_size = settings.get_window_partition_sort_block_size()? as usize;
-        let buffer = WindowBuffer::new(
-            true,
-            partition_spiller,
-            window_spiller,
-            partitions.len(),
-            sort_block_size,
-            memory_settings,
-        )?;
+        let buffer =
+            WindowBuffer::new(spiller, partitions.len(), sort_block_size, memory_settings)?;
 
         Ok(Self {
             input,

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

@@ -20,9 +20,9 @@ use databend_common_expression::DataSchema;
 use databend_common_pipeline_transforms::MemorySettings;
 
 use super::concat_data_blocks;
+use crate::spillers::BackpressureSpiller;
 use crate::spillers::SpillReader;
 use crate::spillers::SpillWriter;
-use crate::spillers::WindowSpiller;
 
 #[async_trait::async_trait]
 pub trait Reader: Send {
@@ -46,7 +46,7 @@ pub trait Builder: Send + Sync {
 }
 
 #[async_trait::async_trait]
-impl Builder for WindowSpiller {
+impl Builder for BackpressureSpiller {
     type Writer = SpillWriter;
 
     async fn create(&self, schema: Arc<DataSchema>) -> Result<SpillWriter> {
@@ -181,7 +181,7 @@ where
     }
 }
 
-pub(super) type WindowPartitionBufferV2 = PartitionBuffer<WindowSpiller>;
+pub(super) type WindowPartitionBufferV2 = PartitionBuffer<BackpressureSpiller>;
 
 pub(super) struct PartitionBuffer<B>
 where B: Builder

src/query/service/src/spillers/adapter.rs

@@ -24,7 +24,6 @@ use std::time::Instant;
 use databend_common_base::base::dma_buffer_to_bytes;
 use databend_common_base::base::dma_read_file_range;
 use databend_common_base::base::ProgressValues;
-use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
@@ -357,13 +356,13 @@ impl Spiller {
 }
 
 #[derive(Clone)]
-pub struct WindowWriterAdapter {
+pub struct BackpressureAdapter {
     ctx: Arc<QueryContext>,
     buffer_pool: Arc<BufferPool>,
     chunk_size: usize,
 }
 
-impl WindowWriterAdapter {
+impl BackpressureAdapter {
     fn add_spill_file(&self, location: Location, layout: Layout, size: usize) {
         if location.is_remote() {
             self.ctx.as_ref().incr_spill_progress(1, size);
@@ -374,25 +373,21 @@ impl WindowWriterAdapter {
     }
 }
 
-pub type WindowSpiller = SpillerInner<WindowWriterAdapter>;
+pub type BackpressureSpiller = SpillerInner<BackpressureAdapter>;
 
-impl WindowSpiller {
+impl BackpressureSpiller {
     pub fn create(
         ctx: Arc<QueryContext>,
         operator: Operator,
         config: SpillerConfig,
+        buffer_pool: Arc<BufferPool>,
+        chunk_size: usize,
     ) -> Result<Self> {
-        let runtime = GlobalIORuntime::instance();
-        let buffer_pool = BufferPool::create(
-            runtime,
-            WINDOW_SPILL_BUFFER_MEMORY_BYTES,
-            WINDOW_SPILL_BUFFER_WORKERS,
-        );
         Self::new(
-            WindowWriterAdapter {
+            BackpressureAdapter {
                 ctx,
                 buffer_pool,
-                chunk_size: WINDOW_SPILL_CHUNK_SIZE,
+                chunk_size,
             },
             operator,
             config,
@@ -419,12 +414,8 @@ pub struct Chunk {
     pub layout: Layout,
 }
 
-const WINDOW_SPILL_BUFFER_MEMORY_BYTES: usize = 64 * 1024 * 1024;
-const WINDOW_SPILL_BUFFER_WORKERS: usize = 2;
-const WINDOW_SPILL_CHUNK_SIZE: usize = 8 * 1024 * 1024;
-
 pub struct SpillWriter {
-    spiller: WindowSpiller,
+    spiller: BackpressureSpiller,
     chunk_size: usize,
     schema: Arc<DataSchema>,
     file_writer: Option<FileWriter<UnionFileWriter>>,
@@ -517,7 +508,7 @@ impl SpillWriter {
 }
 
 pub struct SpillReader {
-    spiller: WindowSpiller,
+    spiller: BackpressureSpiller,
     schema: Arc<DataSchema>,
     parquet_metadata: Arc<ParquetMetaData>,
     union_file: UnionFile,

src/query/service/src/spillers/mod.rs

@@ -24,6 +24,7 @@ mod test_memory;
 mod union_file;
 
 pub use adapter::*;
+pub use async_buffer::BufferPool;
 pub use block_writer::*;
 pub use databend_common_pipeline_transforms::traits::Location;
 pub use inner::*;