use change buffer for compactor

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

Author: BugenZhao

src/common/src/array/data_chunk_iter.rs

@@ -125,7 +125,7 @@ unsafe impl TrustedLen for DataChunkRefIterWithHoles<'_> {}
 mod row_ref {
     use super::*;
 
-    #[derive(Clone, Copy)]
+    #[derive(Clone, Copy, Default)]
     pub struct RowRef<'a> {
         columns: &'a [ArrayRef],
 

src/common/src/row/project.rs

@@ -18,7 +18,7 @@ use super::Row;
 use crate::types::DatumRef;
 
 /// Row for the [`project`](super::RowExt::project) method.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Default)]
 pub struct Project<'i, R> {
     row: R,
     indices: &'i [usize],

src/stream/src/common/change_buffer.rs

@@ -131,13 +131,36 @@ where
             },
         }
     }
+
+    pub fn apply_record(&mut self, record: Record<R>, key_fn: impl Fn(&R) -> K) {
+        match record {
+            Record::Insert { new_row } => self.insert(key_fn(&new_row), new_row),
+            Record::Delete { old_row } => self.delete(key_fn(&old_row), old_row),
+            Record::Update { old_row, new_row } => {
+                let old_key = key_fn(&old_row);
+                let new_key = key_fn(&new_row);
+                if old_key != new_key {
+                    self.ib
+                        .report("inconsistent changes: mismatched key in update");
+                }
+                self.update(old_key, old_row, new_row);
+            }
+        }
+    }
 }
 
 impl<K, R> ChangeBuffer<K, R> {
     pub fn new() -> Self {
         Self::default()
     }
 
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            buffer: IndexMap::with_capacity(capacity),
+            ..Default::default()
+        }
+    }
+
     pub fn with_inconsistency_behavior(mut self, ib: InconsistencyBehavior) -> Self {
         self.ib = ib;
         self
@@ -165,4 +188,16 @@ impl<K, R: Row> ChangeBuffer<K, R> {
         }
         builder.take()
     }
+
+    pub fn into_chunks(self, data_types: Vec<DataType>, chunk_size: usize) -> Vec<StreamChunk> {
+        let mut res = Vec::new();
+        let mut builder = StreamChunkBuilder::new(chunk_size, data_types);
+        for record in self.into_records() {
+            if let Some(chunk) = builder.append_record_eliminate_noop_update(record) {
+                res.push(chunk);
+            }
+        }
+        res.extend(builder.take());
+        res
+    }
 }

src/stream/src/common/compact_chunk.rs

@@ -20,8 +20,6 @@ use std::mem;
 use itertools::Itertools;
 use prehash::{Passthru, Prehashed, new_prehashed_map_with_capacity};
 use risingwave_common::array::stream_chunk::{OpRowMutRef, StreamChunkMut};
-use risingwave_common::array::stream_chunk_builder::StreamChunkBuilder;
-use risingwave_common::array::stream_record::Record;
 use risingwave_common::array::{Op, RowRef, StreamChunk};
 use risingwave_common::row::{Project, RowExt};
 use risingwave_common::types::DataType;
@@ -39,6 +37,7 @@ use risingwave_common::util::hash_util::Crc32FastBuilder;
 //   of the same key, construct new chunks. A combination of `into_compacted_chunks`, `compact`,
 //   and `StreamChunkBuilder`.
 pub use super::change_buffer::InconsistencyBehavior;
+use crate::common::change_buffer::ChangeBuffer;
 
 /// A helper to compact the stream chunks by modifying the `Ops` and visibility of the chunk.
 pub struct StreamChunkCompactor {
@@ -110,107 +109,6 @@ impl<'a> OpRowMutRefTuple<'a> {
 type OpRowMap<'a, 'b> =
     HashMap<Prehashed<Project<'b, RowRef<'a>>>, OpRowMutRefTuple<'a>, BuildHasherDefault<Passthru>>;
 
-#[derive(Clone, Debug)]
-pub enum RowOp<'a> {
-    Insert(RowRef<'a>),
-    Delete(RowRef<'a>),
-    /// (`old_value`, `new_value`)
-    Update((RowRef<'a>, RowRef<'a>)),
-}
-
-pub struct RowOpMap<'a, 'b> {
-    map: HashMap<Prehashed<Project<'b, RowRef<'a>>>, RowOp<'a>, BuildHasherDefault<Passthru>>,
-    ib: InconsistencyBehavior,
-}
-
-impl<'a, 'b> RowOpMap<'a, 'b> {
-    fn with_capacity(estimate_size: usize, ib: InconsistencyBehavior) -> Self {
-        Self {
-            map: new_prehashed_map_with_capacity(estimate_size),
-            ib,
-        }
-    }
-
-    pub fn insert(&mut self, k: Prehashed<Project<'b, RowRef<'a>>>, v: RowRef<'a>) {
-        let entry = self.map.entry(k);
-        match entry {
-            Entry::Vacant(e) => {
-                e.insert(RowOp::Insert(v));
-            }
-            Entry::Occupied(mut e) => match e.get() {
-                RowOp::Delete(old_v) => {
-                    e.insert(RowOp::Update((*old_v, v)));
-                }
-                RowOp::Insert(_) => {
-                    self.ib
-                        .report("double insert for the same pk, breaking the sink's pk constraint");
-                    e.insert(RowOp::Insert(v));
-                }
-                RowOp::Update((old_v, _)) => {
-                    self.ib
-                        .report("double insert for the same pk, breaking the sink's pk constraint");
-                    e.insert(RowOp::Update((*old_v, v)));
-                }
-            },
-        }
-    }
-
-    pub fn delete(&mut self, k: Prehashed<Project<'b, RowRef<'a>>>, v: RowRef<'a>) {
-        let entry = self.map.entry(k);
-        match entry {
-            Entry::Vacant(e) => {
-                e.insert(RowOp::Delete(v));
-            }
-            Entry::Occupied(mut e) => match e.get() {
-                RowOp::Insert(_) => {
-                    e.remove();
-                }
-                RowOp::Update((prev, _)) => {
-                    e.insert(RowOp::Delete(*prev));
-                }
-                RowOp::Delete(_) => {
-                    self.ib.report("double delete for the same pk");
-                    e.insert(RowOp::Delete(v));
-                }
-            },
-        }
-    }
-
-    pub fn into_chunks(self, chunk_size: usize, data_types: Vec<DataType>) -> Vec<StreamChunk> {
-        let mut ret = vec![];
-        let mut builder = StreamChunkBuilder::new(chunk_size, data_types);
-        for (_, row_op) in self.map {
-            match row_op {
-                RowOp::Insert(row) => {
-                    if let Some(c) = builder.append_record(Record::Insert { new_row: row }) {
-                        ret.push(c)
-                    }
-                }
-                RowOp::Delete(row) => {
-                    if let Some(c) = builder.append_record(Record::Delete { old_row: row }) {
-                        ret.push(c)
-                    }
-                }
-                RowOp::Update((old, new)) => {
-                    if old == new {
-                        continue;
-                    }
-                    if let Some(c) = builder.append_record(Record::Update {
-                        old_row: old,
-                        new_row: new,
-                    }) {
-                        ret.push(c)
-                    }
-                }
-            }
-        }
-        if let Some(c) = builder.take() {
-            ret.push(c);
-        }
-        ret
-    }
-}
-
 impl StreamChunkCompactor {
     pub fn new(key: Vec<usize>, chunks: Vec<StreamChunk>) -> Self {
         Self { chunks, key }
@@ -292,32 +190,15 @@ impl StreamChunkCompactor {
         let (chunks, key_indices) = self.into_inner();
 
         let estimate_size = chunks.iter().map(|c| c.cardinality()).sum();
-        let chunks: Vec<(_, _, _)> = chunks
-            .into_iter()
-            .map(|c| {
-                let (c, ops) = c.into_parts();
-                let hash_values = c
-                    .get_hash_values(&key_indices, Crc32FastBuilder)
-                    .into_iter()
-                    .map(|hash| hash.value())
-                    .collect_vec();
-                (hash_values, ops, c)
-            })
-            .collect_vec();
-        let mut map = RowOpMap::with_capacity(estimate_size, ib);
-        for (hash_values, ops, c) in &chunks {
-            for row in c.rows() {
-                let hash = hash_values[row.index()];
-                let op = ops[row.index()];
-                let key = row.project(&key_indices);
-                let k = Prehashed::new(key, hash);
-                match op {
-                    Op::Insert | Op::UpdateInsert => map.insert(k, row),
-                    Op::Delete | Op::UpdateDelete => map.delete(k, row),
-                }
+        let mut cb = ChangeBuffer::with_capacity(estimate_size).with_inconsistency_behavior(ib);
+
+        for chunk in chunks.iter() {
+            for record in chunk.records() {
+                cb.apply_record(record, |&row| row.project(&key_indices));
             }
         }
-        map.into_chunks(chunk_size, data_types)
+
+        cb.into_chunks(data_types, chunk_size)
     }
 }