revert for materialize

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Author: BugenZhao

src/stream/src/executor/mview/materialize.rs

@@ -13,7 +13,8 @@
 // limitations under the License.
 
 use std::assert_matches::assert_matches;
-use std::collections::HashSet;
+use std::collections::hash_map::Entry;
+use std::collections::{HashMap, HashSet};
 use std::marker::PhantomData;
 use std::ops::{Bound, Deref, Index};
 
@@ -30,6 +31,7 @@ use risingwave_common::catalog::{
 use risingwave_common::hash::{VirtualNode, VnodeBitmapExt};
 use risingwave_common::row::{CompactedRow, OwnedRow, RowExt};
 use risingwave_common::types::{DEBEZIUM_UNAVAILABLE_VALUE, DataType, ScalarImpl};
+use risingwave_common::util::chunk_coalesce::DataChunkBuilder;
 use risingwave_common::util::iter_util::{ZipEqDebug, ZipEqFast};
 use risingwave_common::util::sort_util::{ColumnOrder, OrderType, cmp_datum};
 use risingwave_common::util::value_encoding::{BasicSerde, ValueRowSerializer};
@@ -505,7 +507,7 @@ impl<S: StateStore, SD: ValueRowSerde> MaterializeExecutor<S, SD> {
                                             )
                                             .await?;
 
-                                        match change_buffer.into_chunk(data_types.clone()) {
+                                        match generate_output(change_buffer, data_types.clone())? {
                                             Some(output_chunk) => {
                                                 self.state_table.write_chunk(output_chunk.clone());
                                                 self.state_table.try_flush().await?;
@@ -1193,8 +1195,134 @@ fn handle_toast_columns_for_postgres_cdc(
     OwnedRow::new(fixed_row_data)
 }
 
-type ChangeBuffer = crate::common::change_buffer::ChangeBuffer<Vec<u8>, OwnedRow>;
+/// Construct output `StreamChunk` from given buffer.
+fn generate_output(
+    change_buffer: ChangeBuffer,
+    data_types: Vec<DataType>,
+) -> StreamExecutorResult<Option<StreamChunk>> {
+    // construct output chunk
+    // TODO(st1page): when materialize partial columns(), we should construct some columns in the pk
+    let mut new_ops: Vec<Op> = vec![];
+    let mut new_rows: Vec<OwnedRow> = vec![];
+    for (_, row_op) in change_buffer.into_parts() {
+        match row_op {
+            ChangeBufferKeyOp::Insert(value) => {
+                new_ops.push(Op::Insert);
+                new_rows.push(value);
+            }
+            ChangeBufferKeyOp::Delete(old_value) => {
+                new_ops.push(Op::Delete);
+                new_rows.push(old_value);
+            }
+            ChangeBufferKeyOp::Update((old_value, new_value)) => {
+                // if old_value == new_value, we don't need to emit updates to downstream.
+                if old_value != new_value {
+                    new_ops.push(Op::UpdateDelete);
+                    new_ops.push(Op::UpdateInsert);
+                    new_rows.push(old_value);
+                    new_rows.push(new_value);
+                }
+            }
+        }
+    }
+    let mut data_chunk_builder = DataChunkBuilder::new(data_types, new_rows.len() + 1);
+
+    for row in new_rows {
+        let res = data_chunk_builder.append_one_row(row);
+        debug_assert!(res.is_none());
+    }
+
+    if let Some(new_data_chunk) = data_chunk_builder.consume_all() {
+        let new_stream_chunk = StreamChunk::new(new_ops, new_data_chunk.columns().to_vec());
+        Ok(Some(new_stream_chunk))
+    } else {
+        Ok(None)
+    }
+}
+
+/// `ChangeBuffer` is a buffer to handle chunk into `KeyOp`.
+/// TODO(rc): merge with `TopNStaging`.
+pub struct ChangeBuffer {
+    buffer: HashMap<Vec<u8>, ChangeBufferKeyOp>,
+}
 
+/// `KeyOp` variant for `ChangeBuffer` that stores `OwnedRow` instead of Bytes
+enum ChangeBufferKeyOp {
+    Insert(OwnedRow),
+    Delete(OwnedRow),
+    /// (`old_value`, `new_value`)
+    Update((OwnedRow, OwnedRow)),
+}
+
+impl ChangeBuffer {
+    fn new() -> Self {
+        Self {
+            buffer: HashMap::new(),
+        }
+    }
+
+    fn insert(&mut self, pk: Vec<u8>, value: OwnedRow) {
+        let entry = self.buffer.entry(pk);
+        match entry {
+            Entry::Vacant(e) => {
+                e.insert(ChangeBufferKeyOp::Insert(value));
+            }
+            Entry::Occupied(mut e) => {
+                if let ChangeBufferKeyOp::Delete(old_value) = e.get_mut() {
+                    let old_val = std::mem::take(old_value);
+                    e.insert(ChangeBufferKeyOp::Update((old_val, value)));
+                } else {
+                    unreachable!();
+                }
+            }
+        }
+    }
+
+    fn delete(&mut self, pk: Vec<u8>, old_value: OwnedRow) {
+        let entry: Entry<'_, Vec<u8>, ChangeBufferKeyOp> = self.buffer.entry(pk);
+        match entry {
+            Entry::Vacant(e) => {
+                e.insert(ChangeBufferKeyOp::Delete(old_value));
+            }
+            Entry::Occupied(mut e) => match e.get_mut() {
+                ChangeBufferKeyOp::Insert(_) => {
+                    e.remove();
+                }
+                ChangeBufferKeyOp::Update((prev, _curr)) => {
+                    let prev = std::mem::take(prev);
+                    e.insert(ChangeBufferKeyOp::Delete(prev));
+                }
+                ChangeBufferKeyOp::Delete(_) => {
+                    unreachable!();
+                }
+            },
+        }
+    }
+
+    fn update(&mut self, pk: Vec<u8>, old_value: OwnedRow, new_value: OwnedRow) {
+        let entry = self.buffer.entry(pk);
+        match entry {
+            Entry::Vacant(e) => {
+                e.insert(ChangeBufferKeyOp::Update((old_value, new_value)));
+            }
+            Entry::Occupied(mut e) => match e.get_mut() {
+                ChangeBufferKeyOp::Insert(_) => {
+                    e.insert(ChangeBufferKeyOp::Insert(new_value));
+                }
+                ChangeBufferKeyOp::Update((_prev, curr)) => {
+                    *curr = new_value;
+                }
+                ChangeBufferKeyOp::Delete(_) => {
+                    unreachable!()
+                }
+            },
+        }
+    }
+
+    fn into_parts(self) -> HashMap<Vec<u8>, ChangeBufferKeyOp> {
+        self.buffer
+    }
+}
 impl<S: StateStore, SD: ValueRowSerde> Execute for MaterializeExecutor<S, SD> {
     fn execute(self: Box<Self>) -> BoxedMessageStream {
         self.execute_inner().boxed()