feat(sink): merge small chunks for sink executor (#17825)

risingwavelabs · Jul 29, 2024 · f8adc74 · f8adc74
1 parent 9f3af5a
commit f8adc74
Showing 1 changed file with 19 additions and 7 deletions.
diff --git a/src/stream/src/executor/sink.rs b/src/stream/src/executor/sink.rs
@@ -339,20 +339,32 @@ impl<F: LogStoreFactory> SinkExecutor<F> {
                         } else {
                             chunks
                         };
-                        let chunks = if re_construct_with_sink_pk {
-                            StreamChunkCompactor::new(down_stream_pk.clone(), chunks)
+                        if re_construct_with_sink_pk {
+                            let chunks = StreamChunkCompactor::new(down_stream_pk.clone(), chunks)
                                 .reconstructed_compacted_chunks(
                                     chunk_size,
                                     input_data_types.clone(),
                                     sink_type != SinkType::ForceAppendOnly,
-                                )
+                                );
+                            for c in chunks {
+                                yield Message::Chunk(c);
+                            }
                         } else {
-                            chunks
+                            let mut chunk_builder =
+                                StreamChunkBuilder::new(chunk_size, input_data_types.clone());
+                            for chunk in chunks {
+                                for (op, row) in chunk.rows() {
+                                    if let Some(c) = chunk_builder.append_row(op, row) {
+                                        yield Message::Chunk(c);
+                                    }
+                                }
+                            }
+
+                            if let Some(c) = chunk_builder.take() {
+                                yield Message::Chunk(c);
+                            }
                         };
 
-                        for c in chunks {
-                            yield Message::Chunk(c);
-                        }
                         if let Some(w) = mem::take(&mut watermark) {
                             yield Message::Watermark(w)
                         }