add assert on stream key consistency

Signed-off-by: Richard Chien <[email protected]>
risingwavelabs · Sep 27, 2023 · 007548f · 007548f
1 parent 0726b59
commit 007548f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 22 deletions.
diff --git a/src/common/src/array/stream_chunk_iter.rs b/src/common/src/array/stream_chunk_iter.rs
@@ -84,17 +84,26 @@ impl<'a> Iterator for StreamChunkRefIter<'a> {
             Op::Insert => Some(Record::Insert { new_row: row }),
             Op::Delete => Some(Record::Delete { old_row: row }),
             Op::UpdateDelete => {
-                let insert_row = self.inner.next().expect("expect a row after U-");
+                let next_row = self
+                    .inner
+                    .next()
+                    .unwrap_or_else(|| panic!("expect a row after U-, U- row: {:?}", row));
                 // SAFETY: index is checked since `insert_row` is `Some`.
-                let op = unsafe { *self.chunk.ops().get_unchecked(insert_row.index()) };
-                debug_assert_eq!(op, Op::UpdateInsert, "expect a U+ after U-");
+                let op = unsafe { *self.chunk.ops().get_unchecked(next_row.index()) };
+                debug_assert_eq!(
+                    op,
+                    Op::UpdateInsert,
+                    "expect a U+ after U-, U- row: {:?}, row after U-: {:?}",
+                    row,
+                    next_row
+                );
 
                 Some(Record::Update {
                     old_row: row,
-                    new_row: insert_row,
+                    new_row: next_row,
                 })
             }
-            Op::UpdateInsert => panic!("expect a U- before U+"),
+            Op::UpdateInsert => panic!("expect a U- before U+, U+ row: {:?}", row),
         }
     }
 

diff --git a/src/stream/src/executor/wrapper/update_check.rs b/src/stream/src/executor/wrapper/update_check.rs
@@ -12,12 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::iter::once;
 use std::sync::Arc;
 
 use futures_async_stream::try_stream;
-use itertools::Itertools;
-use risingwave_common::array::Op;
+use risingwave_common::array::stream_record::Record;
+use risingwave_common::row::RowExt;
 
 use crate::executor::error::StreamExecutorError;
 use crate::executor::{ExecutorInfo, Message, MessageStream};
@@ -31,20 +30,13 @@ pub async fn update_check(info: Arc<ExecutorInfo>, input: impl MessageStream) {
         let message = message?;
 
         if let Message::Chunk(chunk) = &message {
-            for ((op1, row1), (op2, row2)) in once(None)
-                .chain(chunk.rows().map(Some))
-                .chain(once(None))
-                .map(|r| (r.unzip()))
-                .tuple_windows()
-            {
-                if (op1.is_none() && op2 == Some(Op::UpdateInsert)) // the first row is U+
-                    || (op1 == Some(Op::UpdateDelete) && op2 != Some(Op::UpdateInsert))
-                {
-                    panic!(
-                        "update check failed on `{}`: expect U+ after  U-:\n first row: {:?}\nsecond row: {:?}",
-                        info.identity,
-                        row1,
-                        row2,
+            for record in chunk.records() {
+                // `chunk.records()` will check U-/U+ pairing
+                if let Record::Update { old_row, new_row } = record {
+                    debug_assert!(
+                        old_row.project(&info.pk_indices) == new_row.project(&info.pk_indices),
+                        "U- and U+ should have same stream key, U- row: {:?}, U+ row: {:?}, stream key: {:?}, executor id: {}",
+                        old_row, new_row, info.pk_indices, info.identity
                     )
                 }
             }