refactor(window): simplify TransformWindowPartitionCollect event loop

Reduce from 5-state machine (Collect/Spill/Process/Restore/Finish) to
3-state (Collect/Output/Finish). Since spill and restore are now both
synchronous, the indirection through separate Spill/Restore steps with
async fallback is unnecessary.

- Collect: pull data, buffer, spill inline when memory pressure is high
- Output: restore next partition, process (sort), push downstream
- Remove next_step(), collect(), output() helper methods
- Remove is_collect_finished, restored_data_blocks fields

Author: dqhl76

src/query/service/src/pipelines/processors/transforms/materialized_cte.rs

@@ -69,14 +69,8 @@ impl MaterializedCteSpilledPayload {
         let operator = data_operator.spill_operator();
         let buffer_pool = SpillsBufferPool::instance();
         let settings = ReadSettings::default();
-        let mut reader = buffer_pool.reader(
-            operator,
-            self.path,
-            self.schema,
-            selected,
-            target,
-            settings,
-        )?;
+        let mut reader =
+            buffer_pool.reader(operator, self.path, self.schema, selected, target, settings)?;
         match reader.read()? {
             Some(block) => Ok(block),
             None => Err(ErrorCode::Internal(

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

@@ -15,7 +15,7 @@
 mod data_processor_strategy;
 mod hilbert_partition_exchange;
 mod transform_window_partition_collect;
-mod window_partition_buffer_v2;
+mod window_partition_buffer;
 mod window_partition_exchange;
 mod window_partition_meta;
 mod window_partition_partial_top_n_exchange;

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

@@ -31,36 +31,28 @@ use databend_common_pipeline_transforms::MemorySettings;
 use databend_common_settings::Settings;
 
 use super::WindowPartitionMeta;
-use super::window_partition_buffer_v2::WindowPartitionBufferV2;
+use super::window_partition_buffer::WindowPartitionBuffer;
 use crate::pipelines::processors::transforms::DataProcessorStrategy;
 use crate::sessions::QueryContext;
 
+/// Two-phase state machine:
+/// - Collect: pull input data into buffer, spill when memory pressure is high.
+/// - Output: restore partitions one by one, process (sort), and push downstream.
 #[derive(Debug, Clone, Copy)]
 enum Step {
     Collect,
-    Process,
-    Spill,
-    Restore,
-    Finish,
+    Output,
 }
 
 pub struct TransformWindowPartitionCollect<S: DataProcessorStrategy> {
     input: Arc<InputPort>,
     output: Arc<OutputPort>,
-
-    restored_data_blocks: Vec<DataBlock>,
     output_data_blocks: VecDeque<DataBlock>,
 
-    // The partition id is used to map the partition id to the new partition id.
     index_map: Vec<usize>,
-    // The buffer is used to control the memory usage of the window operator.
-    buffer: WindowPartitionBufferV2,
-
+    buffer: WindowPartitionBuffer,
     strategy: S,
-
-    // Event variables.
     step: Step,
-    is_collect_finished: bool,
 }
 
 impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
@@ -75,12 +67,10 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
         memory_settings: MemorySettings,
         strategy: S,
     ) -> Result<Self> {
-        // Calculate the partition ids collected by the processor.
         let partitions: Vec<usize> = (0..num_partitions)
             .filter(|&partition| partition % num_processors == processor_id)
             .collect();
 
-        // Map each partition id to new partition index.
         let mut index_map = vec![0; num_partitions];
         for (index, partition) in partitions.iter().enumerate() {
             index_map[*partition] = index;
@@ -91,9 +81,8 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
             .get_spill_writer_memory_pool_size_mb()?
             .saturating_mul(1024 * 1024);
 
-        // Create the window partition buffer.
         let sort_block_size = settings.get_window_partition_sort_block_size()? as usize;
-        let buffer = WindowPartitionBufferV2::new(
+        let buffer = WindowPartitionBuffer::new(
             prefix,
             writer_pool_bytes,
             partitions.len(),
@@ -104,83 +93,22 @@ impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
         Ok(Self {
             input,
             output,
+            output_data_blocks: VecDeque::new(),
             index_map,
             buffer,
             strategy,
-            is_collect_finished: false,
-            output_data_blocks: VecDeque::new(),
-            restored_data_blocks: Vec::new(),
             step: Step::Collect,
         })
     }
 
-    fn next_step(&mut self, step: Step) -> Result<Event> {
-        let event = match step {
-            Step::Finish => {
-                self.input.finish();
-                self.output.finish();
-                Event::Finished
+    fn collect_data_block(&mut self, data_block: DataBlock) {
+        if let Some(meta) = data_block
+            .get_owned_meta()
+            .and_then(WindowPartitionMeta::downcast_from)
+        {
+            for (id, block) in meta.partitioned_data {
+                self.buffer.add_data_block(self.index_map[id], block);
             }
-            _ => Event::Sync,
-        };
-        self.step = step;
-        Ok(event)
-    }
-
-    fn collect(&mut self) -> Result<Event> {
-        if self.output.is_finished() {
-            self.input.finish();
-            return self.next_step(Step::Finish);
-        }
-
-        // First check. flush memory data to external storage if need
-        if self.need_spill() {
-            return self.next_step(Step::Spill);
-        }
-
-        if self.input.has_data() {
-            Self::collect_data_block(
-                self.input.pull_data().unwrap()?,
-                &self.index_map,
-                &mut self.buffer,
-            );
-        }
-
-        // Check again. flush memory data to external storage if need
-        if self.need_spill() {
-            return self.next_step(Step::Spill);
-        }
-
-        if self.input.is_finished() {
-            self.is_collect_finished = true;
-            return self.next_step(Step::Restore);
-        }
-
-        self.input.set_need_data();
-        Ok(Event::NeedData)
-    }
-
-    fn output(&mut self) -> Result<Event> {
-        if self.output.is_finished() {
-            return self.next_step(Step::Finish);
-        }
-
-        if !self.output.can_push() {
-            return Ok(Event::NeedConsume);
-        }
-
-        if self.need_spill() {
-            return self.next_step(Step::Spill);
-        }
-
-        if let Some(data_block) = self.output_data_blocks.pop_front() {
-            self.output.push_data(Ok(data_block));
-            return Ok(Event::NeedConsume);
-        }
-
-        match self.buffer.is_empty() {
-            true => self.next_step(Step::Finish),
-            false => self.next_step(Step::Restore),
         }
     }
 }
@@ -195,66 +123,66 @@ impl<S: DataProcessorStrategy> Processor for TransformWindowPartitionCollect<S>
     }
 
     fn event(&mut self) -> Result<Event> {
-        // (collect <--> spill) -> (process <--> restore) -> finish
+        if self.output.is_finished() {
+            self.input.finish();
+            return Ok(Event::Finished);
+        }
+
         match self.step {
-            Step::Collect => self.collect(),
-            Step::Spill => {
-                if self.is_collect_finished {
-                    self.step = Step::Process;
-                    self.output()
-                } else {
-                    // collect data again.
-                    self.step = Step::Collect;
-                    self.collect()
+            Step::Collect => {
+                if self.input.has_data() {
+                    return Ok(Event::Sync);
                 }
+
+                if self.input.is_finished() {
+                    self.step = Step::Output;
+                    return Ok(Event::Sync);
+                }
+
+                self.input.set_need_data();
+                Ok(Event::NeedData)
             }
-            Step::Process => self.output(),
-            Step::Restore => {
-                if self.restored_data_blocks.is_empty() {
-                    self.next_step(Step::Finish)
-                } else {
-                    self.next_step(Step::Process)
+            Step::Output => {
+                if !self.output.can_push() {
+                    return Ok(Event::NeedConsume);
+                }
+
+                if let Some(block) = self.output_data_blocks.pop_front() {
+                    self.output.push_data(Ok(block));
+                    return Ok(Event::NeedConsume);
                 }
+
+                if self.buffer.is_empty() {
+                    self.output.finish();
+                    return Ok(Event::Finished);
+                }
+
+                Ok(Event::Sync)
             }
-            Step::Finish => Ok(Event::Finished),
         }
     }
 
     fn process(&mut self) -> Result<()> {
         match self.step {
-            Step::Process => {
-                let restored_data_blocks = std::mem::take(&mut self.restored_data_blocks);
-                let processed_blocks = self.strategy.process_data_blocks(restored_data_blocks)?;
-                self.output_data_blocks.extend(processed_blocks);
+            Step::Collect => {
+                let data_block = self.input.pull_data().unwrap()?;
+                self.collect_data_block(data_block);
+
+                // Spill if memory pressure is high.
+                while self.buffer.need_spill() {
+                    self.buffer.spill()?;
+                }
                 Ok(())
             }
-            Step::Spill => self.buffer.spill(),
-            Step::Restore => {
-                self.restored_data_blocks = self.buffer.restore()?;
+            Step::Output => {
+                // Restore next non-empty partition and process it.
+                let restored = self.buffer.restore()?;
+                if !restored.is_empty() {
+                    let processed = self.strategy.process_data_blocks(restored)?;
+                    self.output_data_blocks.extend(processed);
+                }
                 Ok(())
             }
-            _ => unreachable!(),
         }
     }
 }
-
-impl<S: DataProcessorStrategy> TransformWindowPartitionCollect<S> {
-    fn collect_data_block(
-        data_block: DataBlock,
-        index_map: &[usize],
-        buffer: &mut WindowPartitionBufferV2,
-    ) {
-        if let Some(meta) = data_block
-            .get_owned_meta()
-            .and_then(WindowPartitionMeta::downcast_from)
-        {
-            for (id, data_block) in meta.partitioned_data {
-                buffer.add_data_block(index_map[id], data_block);
-            }
-        }
-    }
-
-    fn need_spill(&mut self) -> bool {
-        self.buffer.need_spill()
-    }
-}

…/partition/window_partition_buffer_v2.rs …dow/partition/window_partition_buffer.rssrc/query/service/src/pipelines/processors/transforms/window/partition/window_partition_buffer_v2.rs renamed to src/query/service/src/pipelines/processors/transforms/window/partition/window_partition_buffer.rs src/query/service/src/pipelines/processors/transforms/window/partition/window_partition_buffer_v2.rs renamed to src/query/service/src/pipelines/processors/transforms/window/partition/window_partition_buffer.rs

@@ -23,7 +23,6 @@ use databend_common_storage::DataOperator;
 use databend_common_storages_parquet::ReadSettings;
 use parquet::file::metadata::RowGroupMetaData;
 
-
 fn concat_data_blocks(data_blocks: Vec<DataBlock>, target_size: usize) -> Result<Vec<DataBlock>> {
     let mut num_rows = 0;
     let mut result = Vec::new();
@@ -107,7 +106,10 @@ struct PartitionFileReader {
 
 impl PartitionFileReader {
     fn restore(&self, indices: Vec<usize>) -> Result<Vec<DataBlock>> {
-        let selected: Vec<_> = indices.iter().map(|&i| self.row_groups[i].clone()).collect();
+        let selected: Vec<_> = indices
+            .iter()
+            .map(|&i| self.row_groups[i].clone())
+            .collect();
         if selected.is_empty() {
             return Ok(Vec::new());
         }
@@ -233,10 +235,6 @@ impl PartitionSlot {
     }
 }
 
-// --- Public PartitionBuffer (WindowPartitionBufferV2) ---
-
-pub(super) type WindowPartitionBufferV2 = WindowPartitionBuffer;
-
 pub(super) struct WindowPartitionBuffer {
     prefix: String,
     writer_pool_bytes: usize,