Optimize handling of tail fragments and block serialization

Author: zhyass

src/query/storages/fuse/src/operations/common/processors/transform_block_writer.rs

@@ -43,7 +43,9 @@ use crate::FUSE_OPT_KEY_ROW_PER_BLOCK;
 #[allow(clippy::large_enum_variant)]
 enum State {
     Consume,
-    Serialize(DataBlock),
+    Collect(DataBlock),
+    Serialize,
+    Finalize,
     Flush,
     Write(BlockSerialization),
 }
@@ -56,6 +58,10 @@ pub struct TransformBlockWriter {
     properties: Arc<StreamBlockProperties>,
 
     builder: Option<StreamBlockBuilder>,
+    need_flush: bool,
+    input_data_size: usize,
+    input_num_rows: usize,
+
     dal: Operator,
     // Only used in multi table insert
     table_id: Option<u64>,
@@ -86,8 +92,11 @@ impl TransformBlockWriter {
             properties,
             builder: None,
             dal: table.get_operator(),
+            need_flush: false,
             table_id: if with_tid { Some(table.get_id()) } else { None },
             input_data: VecDeque::new(),
+            input_data_size: 0,
+            input_num_rows: 0,
             output_data: None,
             max_block_size,
         })))
@@ -126,12 +135,12 @@ impl Processor for TransformBlockWriter {
     }
 
     fn event(&mut self) -> Result<Event> {
-        if matches!(self.state, State::Serialize { .. } | State::Flush) {
-            return Ok(Event::Sync);
-        }
-
-        if matches!(self.state, State::Write { .. }) {
-            return Ok(Event::Async);
+        match &self.state {
+            State::Collect(_) | State::Serialize | State::Flush | State::Finalize => {
+                return Ok(Event::Sync)
+            }
+            State::Write(_) => return Ok(Event::Async),
+            _ => {}
         }
 
         if self.output.is_finished() {
@@ -147,59 +156,80 @@ impl Processor for TransformBlockWriter {
             return Ok(Event::NeedConsume);
         }
 
-        if let Some(block) = self.input_data.pop_front() {
-            self.state = State::Serialize(block);
+        // To avoid tail fragments, flush only when the input is large enough.
+        if self.need_flush
+            && self
+                .properties
+                .block_thresholds
+                .check_large_enough(self.input_num_rows, self.input_data_size)
+        {
+            self.state = State::Flush;
+            return Ok(Event::Sync);
+        }
+
+        if !self.need_flush && !self.input_data.is_empty() {
+            self.state = State::Serialize;
+            return Ok(Event::Sync);
+        }
+
+        if self.input.has_data() {
+            let input_data = self.input.pull_data().unwrap()?;
+            self.state = State::Collect(input_data);
             return Ok(Event::Sync);
         }
 
         if self.input.is_finished() {
-            if self.builder.is_some() {
-                self.state = State::Flush;
+            if !self.input_data.is_empty() || self.builder.is_some() {
+                self.state = State::Finalize;
                 return Ok(Event::Sync);
             }
             self.output.finish();
             return Ok(Event::Finished);
         }
 
-        if !self.input.has_data() {
-            self.input.set_need_data();
-            return Ok(Event::NeedData);
-        }
-
-        let input_data = self.input.pull_data().unwrap()?;
-        self.state = State::Serialize(input_data);
-        Ok(Event::Sync)
+        self.input.set_need_data();
+        Ok(Event::NeedData)
     }
 
     fn process(&mut self) -> Result<()> {
         match std::mem::replace(&mut self.state, State::Consume) {
-            State::Serialize(block) => {
+            State::Collect(block) => {
                 // Check if the datablock is valid, this is needed to ensure data is correct
                 block.check_valid()?;
+                self.input_data_size += block.estimate_block_size();
+                self.input_num_rows += block.num_rows();
                 let max_rows_per_block = self.calc_max_block_size(&block);
                 let blocks = block.split_by_rows_if_needed_no_tail(max_rows_per_block);
-                let mut blocks = VecDeque::from(blocks);
+                self.input_data.extend(blocks);
+            }
+            State::Serialize => {
+                while let Some(b) = self.input_data.pop_front() {
+                    self.input_data_size -= b.estimate_block_size();
+                    self.input_num_rows -= b.num_rows();
 
-                let builder = self.get_or_create_builder()?;
-                while let Some(b) = blocks.pop_front() {
+                    let builder = self.get_or_create_builder()?;
                     builder.write(b)?;
 
                     if builder.need_flush() {
-                        self.state = State::Flush;
-
-                        for left in blocks {
-                            self.input_data.push_back(left);
-                        }
+                        self.need_flush = true;
                         return Ok(());
                     }
                 }
             }
+            State::Finalize => {
+                while let Some(b) = self.input_data.pop_front() {
+                    let builder = self.get_or_create_builder()?;
+                    builder.write(b)?;
+                }
+                self.state = State::Flush;
+            }
             State::Flush => {
                 let builder = self.builder.take().unwrap();
                 if !builder.is_empty() {
                     let serialized = builder.finish()?;
                     self.state = State::Write(serialized);
                 }
+                self.need_flush = false;
             }
             _ => return Err(ErrorCode::Internal("It's a bug.")),
         }