feat(stream): compact stream chunk by dropping useless rows in sink

ci/scripts/e2e-kafka-sink-test.sh

@@ -56,15 +56,15 @@ fi
 # test upsert kafka sink after update
 echo "testing upsert kafka sink after updating data"
 diff ./e2e_test/sink/kafka/upsert2.result \
-<((./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-upsert --from-beginning --property print.key=true --max-messages 11 | sort) 2> /dev/null)
+<((./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-upsert --from-beginning --property print.key=true --max-messages 12 | sort) 2> /dev/null)
 if [ $? -ne 0 ]; then
   echo "The output for upsert sink after update is not as expected."
   exit 1
 fi
 
 # test debezium kafka sink after update
 echo "testing debezium kafka sink after updating data"
-(./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-debezium --property print.key=true --from-beginning --max-messages 11 | sort) > ./e2e_test/sink/kafka/debezium2.tmp.result 2> /dev/null
+(./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-debezium --property print.key=true --from-beginning --max-messages 13 | sort) > ./e2e_test/sink/kafka/debezium2.tmp.result 2> /dev/null
 python3 e2e_test/sink/kafka/debezium.py e2e_test/sink/kafka/debezium2.result e2e_test/sink/kafka/debezium2.tmp.result
 if [ $? -ne 0 ]; then
   echo "The output for debezium sink after update is not as expected."
@@ -81,15 +81,15 @@ psql -h localhost -p 4566 -d dev -U root -c "delete from t_kafka where id = 1;"
 # test upsert kafka sink after delete
 echo "testing upsert kafka sink after deleting data"
 diff ./e2e_test/sink/kafka/upsert3.result \
-<((./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-upsert --from-beginning --property print.key=true --max-messages 12 | sort) 2> /dev/null)
+<((./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-upsert --from-beginning --property print.key=true --max-messages 13 | sort) 2> /dev/null)
 if [ $? -ne 0 ]; then
   echo "The output for upsert sink after update is not as expected."
   exit 1
 fi
 
 # test debezium kafka sink after delete
 echo "testing debezium kafka sink after deleting data"
-(./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-debezium --property print.key=true --from-beginning --max-messages 13 | sort) > ./e2e_test/sink/kafka/debezium3.tmp.result 2> /dev/null
+(./.risingwave/bin/kafka/bin/kafka-console-consumer.sh --bootstrap-server 127.0.0.1:29092 --topic test-rw-sink-debezium --property print.key=true --from-beginning --max-messages 15 | sort) > ./e2e_test/sink/kafka/debezium3.tmp.result 2> /dev/null
 python3 e2e_test/sink/kafka/debezium.py e2e_test/sink/kafka/debezium3.result e2e_test/sink/kafka/debezium3.tmp.result
 if [ $? -ne 0 ]; then
   echo "The output for debezium sink after delete is not as expected."

e2e_test/sink/kafka/upsert2.result

@@ -1,4 +1,5 @@
 {"id":10}	{"id":10,"v_bigint":20674,"v_double":9042.404483827513,"v_float":19387.23828125,"v_integer":20674,"v_smallint":26951,"v_timestamp":1681404058888,"v_varchar":"0oVqRIHqkb"}
+{"id":1}	null
 {"id":1}	{"id":1,"v_bigint":0,"v_double":0.0,"v_float":0.0,"v_integer":0,"v_smallint":0,"v_timestamp":0,"v_varchar":""}
 {"id":1}	{"id":1,"v_bigint":1872,"v_double":23956.39329760601,"v_float":26261.416015625,"v_integer":1872,"v_smallint":31031,"v_timestamp":1681453634104,"v_varchar":"8DfUFencLe"}
 {"id":2}	{"id":2,"v_bigint":4598,"v_double":31923.077305746086,"v_float":27031.224609375,"v_integer":4598,"v_smallint":22690,"v_timestamp":1681429444869,"v_varchar":"sIo1XXVeHZ"}

e2e_test/sink/kafka/upsert3.result

@@ -1,5 +1,6 @@
 {"id":10}	{"id":10,"v_bigint":20674,"v_double":9042.404483827513,"v_float":19387.23828125,"v_integer":20674,"v_smallint":26951,"v_timestamp":1681404058888,"v_varchar":"0oVqRIHqkb"}
 {"id":1}	null
+{"id":1}	null
 {"id":1}	{"id":1,"v_bigint":0,"v_double":0.0,"v_float":0.0,"v_integer":0,"v_smallint":0,"v_timestamp":0,"v_varchar":""}
 {"id":1}	{"id":1,"v_bigint":1872,"v_double":23956.39329760601,"v_float":26261.416015625,"v_integer":1872,"v_smallint":31031,"v_timestamp":1681453634104,"v_varchar":"8DfUFencLe"}
 {"id":2}	{"id":2,"v_bigint":4598,"v_double":31923.077305746086,"v_float":27031.224609375,"v_integer":4598,"v_smallint":22690,"v_timestamp":1681429444869,"v_varchar":"sIo1XXVeHZ"}

src/common/Cargo.toml

@@ -62,6 +62,7 @@ postgres-types = { version = "0.2.6", features = [
     "with-chrono-0_4",
     "with-serde_json-1",
 ] }
+prehash = "1"
 prometheus = { version = "0.13" }
 prost = "0.11"
 rand = "0.8"

src/common/src/array/compact_chunk.rs

@@ -0,0 +1,209 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
+use std::hash::BuildHasherDefault;
+use std::mem;
+
+use itertools::Itertools;
+use prehash::{new_prehashed_map_with_capacity, Passthru, Prehashed};
+
+use super::stream_chunk::{OpRowMutRef, StreamChunkMut};
+use crate::array::{Op, RowRef, StreamChunk};
+use crate::row::{Project, RowExt};
+use crate::util::hash_util::Crc32FastBuilder;
+pub struct StreamChunkCompactor {
+    chunks: Vec<StreamChunk>,
+    stream_key: Vec<usize>,
+}
+
+struct OpRowMutRefTuple<'a> {
+    previous: Option<OpRowMutRef<'a>>,
+    latest: OpRowMutRef<'a>,
+}
+
+impl<'a> OpRowMutRefTuple<'a> {
+    /// return true if no row left
+    fn push(&mut self, mut op_row: OpRowMutRef<'a>) -> bool {
+        debug_assert!(self.latest.vis());
+        match (self.latest.op(), op_row.op()) {
+            (Op::Insert, Op::Insert) => panic!("receive duplicated insert on the stream"),
+            (Op::Delete, Op::Delete) => panic!("receive duplicated delete on the stream"),
+            (Op::Insert, Op::Delete) => {
+                self.latest.set_vis(false);
+                op_row.set_vis(false);
+                self.latest = if let Some(prev) = self.previous.take() {
+                    prev
+                } else {
+                    return true;
+                }
+            }
+            (Op::Delete, Op::Insert) => {
+                // The operation for the key must be (+, -, +) or (-, +). And the (+, -) must has
+                // been filtered.
+                debug_assert!(self.previous.is_none());
+                self.previous = Some(mem::replace(&mut self.latest, op_row));
+            }
+            // `all the updateDelete` and `updateInsert` should be normalized to `delete`
+            // and`insert`
+            _ => unreachable!(),
+        };
+        false
+    }
+
+    fn as_update_op(&mut self) -> Option<(&mut OpRowMutRef<'a>, &mut OpRowMutRef<'a>)> {
+        self.previous.as_mut().map(|prev| {
+            debug_assert_eq!(prev.op(), Op::Delete);
+            debug_assert_eq!(self.latest.op(), Op::Insert);
+            (prev, &mut self.latest)
+        })
+    }
+}
+
+type OpRowMap<'a, 'b> =
+    HashMap<Prehashed<Project<'b, RowRef<'a>>>, OpRowMutRefTuple<'a>, BuildHasherDefault<Passthru>>;
+
+impl StreamChunkCompactor {
+    pub fn new(stream_key: Vec<usize>) -> Self {
+        Self {
+            stream_key,
+            chunks: vec![],
+        }
+    }
+
+    pub fn into_inner(self) -> (Vec<StreamChunk>, Vec<usize>) {
+        (self.chunks, self.stream_key)
+    }
+
+    pub fn push_chunk(&mut self, c: StreamChunk) {
+        self.chunks.push(c);
+    }
+
+    /// Compact a chunk by modifying the ops and the visibility of a stream chunk. All UPDATE INSERT
+    /// and UPDATE DELETE will be converted to INSERT and DELETE, and dropped according to
+    /// certain rules (see `merge_insert` and `merge_delete` for more details).
+
+    pub fn into_compacted_chunks(self) -> Vec<StreamChunk> {
+        let (chunks, key_indices) = self.into_inner();
+
+        let estimate_size = chunks.iter().map(|c| c.cardinality()).sum();
+        let mut chunks: Vec<(Vec<u64>, StreamChunkMut)> = chunks
+            .into_iter()
+            .map(|c| {
+                let hash_values = c
+                    .data_chunk()
+                    .get_hash_values(&key_indices, Crc32FastBuilder)
+                    .into_iter()
+                    .map(|hash| hash.value())
+                    .collect_vec();
+                (hash_values, StreamChunkMut::from(c))
+            })
+            .collect_vec();
+
+        let mut op_row_map: OpRowMap<'_, '_> = new_prehashed_map_with_capacity(estimate_size);
+        for (hash_values, c) in &mut chunks {
+            for (row, mut op_row) in c.to_rows_mut() {
+                if !op_row.vis() {
+                    continue;
+                }
+                op_row.set_op(op_row.op().normalize_update());
+                let hash = hash_values[row.index()];
+                let stream_key = row.project(&key_indices);
+                match op_row_map.entry(Prehashed::new(stream_key, hash)) {
+                    Entry::Vacant(v) => {
+                        v.insert(OpRowMutRefTuple {
+                            previous: None,
+                            latest: op_row,
+                        });
+                    }
+                    Entry::Occupied(mut o) => {
+                        if o.get_mut().push(op_row) {
+                            o.remove_entry();
+                        }
+                    }
+                }
+            }
+        }
+        for tuple in op_row_map.values_mut() {
+            if let Some((prev, latest)) = tuple.as_update_op() {
+                if prev.row_ref() == latest.row_ref() {
+                    prev.set_vis(false);
+                    latest.set_vis(false);
+                }
+            }
+        }
+        chunks.into_iter().map(|(_, c)| c.into()).collect_vec()
+    }
+}
+
+pub fn merge_chunk_row(stream_chunk: StreamChunk, pk_indices: &[usize]) -> StreamChunk {
+    let mut compactor = StreamChunkCompactor::new(pk_indices.to_vec());
+    compactor.push_chunk(stream_chunk);
+    compactor
+        .into_compacted_chunks()
+        .into_iter()
+        .next()
+        .unwrap()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::array::StreamChunk;
+    use crate::test_prelude::StreamChunkTestExt;
+
+    #[test]
+    fn test_merge_chunk_row() {
+        let pk_indices = [0, 1];
+        let mut compactor = StreamChunkCompactor::new(pk_indices.to_vec());
+        compactor.push_chunk(StreamChunk::from_pretty(
+            " I I I
+            + 2 5 7
+            + 4 9 2
+            - 2 5 7
+            + 2 5 5
+            - 6 6 9
+            + 6 6 9
+            - 9 9 1",
+        ));
+        compactor.push_chunk(StreamChunk::from_pretty(
+            " I I I
+            - 6 6 9
+            + 9 9 9
+            - 9 9 4
+            + 2 2 2
+            + 9 9 1",
+        ));
+        let mut iter = compactor.into_compacted_chunks().into_iter();
+        assert_eq!(
+            iter.next().unwrap().compact(),
+            StreamChunk::from_pretty(
+                " I I I
+                + 4 9 2
+                + 2 5 5
+                - 6 6 9",
+            )
+        );
+        assert_eq!(
+            iter.next().unwrap().compact(),
+            StreamChunk::from_pretty(
+                " I I I
+                + 2 2 2",
+            )
+        );
+
+        assert_eq!(iter.next(), None);
+    }
+}