partial restore

Author: zhyass

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

@@ -211,7 +211,8 @@ impl Processor for TransformHilbertCollect {
             }
             State::Flush => {
                 if let Some((partition_id, data_block)) = self.immediate_output_blocks.pop() {
-                    let mut restored_data_blocks = self.buffer.restore_by_id(partition_id).await?;
+                    let mut restored_data_blocks =
+                        self.buffer.restore_by_id(partition_id, true).await?;
                     restored_data_blocks.push(data_block);
                     self.state = State::Concat(restored_data_blocks);
                 }

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

@@ -145,15 +145,19 @@ impl WindowPartitionBuffer {
         while self.next_to_restore_partition_id + 1 < self.num_partitions as isize {
             self.next_to_restore_partition_id += 1;
             let partition_id = self.next_to_restore_partition_id as usize;
-            let result = self.restore_by_id(partition_id).await?;
+            let result = self.restore_by_id(partition_id, false).await?;
             if !result.is_empty() {
                 return Ok(result);
             }
         }
         Ok(vec![])
     }
 
-    pub async fn restore_by_id(&mut self, partition_id: usize) -> Result<Vec<DataBlock>> {
+    pub async fn restore_by_id(
+        &mut self,
+        partition_id: usize,
+        partial_restore: bool,
+    ) -> Result<Vec<DataBlock>> {
         // Restore large partitions from spilled files.
         let mut result = self.spiller.take_spilled_partition(&partition_id).await?;
 
@@ -171,7 +175,7 @@ impl WindowPartitionBuffer {
                 location,
                 partitions,
             } = merged_partitions;
-            if out_of_memory_limit || *partial_restored {
+            if out_of_memory_limit || *partial_restored || partial_restore {
                 if let Some(pos) = partitions.iter().position(|(id, _)| *id == partition_id) {
                     let data_block = self
                         .spiller

src/query/storages/fuse/src/io/write/stream/block_builder.rs

@@ -242,7 +242,7 @@ impl StreamBlockBuilder {
         self.row_count >= self.properties.block_thresholds.min_rows_per_block
             || self.block_size >= self.properties.block_thresholds.max_bytes_per_block
             || (file_size >= self.properties.block_thresholds.min_compressed_per_block
-            && self.block_size >= self.properties.block_thresholds.min_bytes_per_block)
+                && self.block_size >= self.properties.block_thresholds.min_bytes_per_block)
     }
 
     pub fn write(&mut self, block: DataBlock) -> Result<()> {