1use core::fmt::Formatter;
16use std::fmt::Debug;
17use std::sync::Arc;
18
19use futures::prelude::stream::StreamExt;
20use futures_async_stream::try_stream;
21use itertools::Itertools;
22use risingwave_common::array::stream_chunk_builder::StreamChunkBuilder;
23use risingwave_common::array::{ArrayImpl, I16Array, Op, SerialArray, StreamChunk};
24use risingwave_common::bitmap::Bitmap;
25use risingwave_common::hash::VirtualNode;
26use risingwave_common::row::{OwnedRow, RowExt};
27use risingwave_common::types::{DataType, ScalarImpl, Serial};
28use risingwave_common::util::row_id::ChangelogRowIdGenerator;
29use risingwave_pb::stream_plan::stream_node::StreamKind;
30
31use super::{ActorContextRef, BoxedMessageStream, Execute, Executor, Message, StreamExecutorError};
32use crate::common::change_buffer::InconsistencyBehavior;
33use crate::common::change_buffer::output_kind::RETRACT;
34use crate::common::compact_chunk::StreamChunkCompactor;
35
36pub struct ChangeLogExecutor {
37 ctx: ActorContextRef,
38 input: Executor,
39 need_op: bool,
40 all_vnode_count: usize,
41 mode: ChangeLogMode,
42 distribution_keys: Vec<usize>,
43 stream_keys: Vec<usize>,
44 changelog_row_id_generator: ChangelogRowIdGenerator,
45}
46
47impl Execute for ChangeLogExecutor {
48 fn execute(self: Box<Self>) -> BoxedMessageStream {
49 if matches!(self.mode, ChangeLogMode::Keyed { .. }) {
50 self.execute_keyed().boxed()
51 } else {
52 self.execute_normal().boxed()
53 }
54 }
55}
56
57impl ChangeLogExecutor {
58 #[allow(clippy::too_many_arguments)]
59 pub fn new(
60 ctx: ActorContextRef,
61 input: Executor,
62 need_op: bool,
63 all_vnode_count: usize,
64 vnodes: Bitmap,
65 mode: ChangeLogMode,
66 distribution_keys: Vec<usize>,
67 stream_keys: Vec<usize>,
68 ) -> Self {
69 let changelog_row_id_generator = ChangelogRowIdGenerator::new(vnodes, all_vnode_count);
70 Self {
71 ctx,
72 input,
73 need_op,
74 all_vnode_count,
75 mode,
76 distribution_keys,
77 stream_keys,
78 changelog_row_id_generator,
79 }
80 }
81
82 #[try_stream(ok = Message, error = StreamExecutorError)]
83 async fn execute_normal(mut self) {
84 let input = self.input.execute();
85 #[for_await]
86 for msg in input {
87 let msg = msg?;
88 match msg {
89 Message::Chunk(chunk) => {
90 let data_chunk = chunk.data_chunk();
91 let vnodes = VirtualNode::compute_chunk(
92 data_chunk,
93 &self.distribution_keys,
94 self.all_vnode_count,
95 );
96
97 let (ops, mut columns, bitmap) = chunk.into_inner();
98
99 if self.need_op {
100 let ops = ops.iter().map(|op| Some(op.to_i16())).collect_vec();
101 let ops_array = Arc::new(ArrayImpl::Int16(I16Array::from_iter(ops)));
102
103 columns.push(ops_array);
104 }
105
106 let changelog_row_ids = vnodes
107 .iter()
108 .map(|vnode| self.changelog_row_id_generator.next(vnode));
109 let changelog_row_id_array = Arc::new(ArrayImpl::Serial(
110 SerialArray::from_iter(changelog_row_ids.map(Serial::from)),
111 ));
112
113 columns.push(changelog_row_id_array);
114
115 let new_ops = vec![Op::Insert; ops.len()];
116 let new_chunk = StreamChunk::with_visibility(new_ops, columns, bitmap);
117
118 yield Message::Chunk(new_chunk);
119 }
120 Message::Watermark(_w) => {}
121 Message::Barrier(barrier) => {
122 if let Some(vnodes) = barrier.as_update_vnode_bitmap(self.ctx.id) {
123 self.changelog_row_id_generator = ChangelogRowIdGenerator::new(
124 vnodes.as_ref().clone(),
125 self.all_vnode_count,
126 );
127 }
128 yield Message::Barrier(barrier);
129 }
130 }
131 }
132 }
133
134 #[try_stream(ok = Message, error = StreamExecutorError)]
135 async fn execute_keyed(mut self) {
136 let ChangeLogMode::Keyed {
137 additional_key_columns,
138 } = &self.mode
139 else {
140 unreachable!("keyed execution requires keyed mode");
141 };
142 let input_kind = self.input.stream_kind();
143 let compaction_keys = if input_kind == StreamKind::Upsert {
144 self.stream_keys.clone()
147 } else {
148 self.stream_keys
149 .iter()
150 .chain(&self.distribution_keys)
151 .chain(additional_key_columns)
152 .copied()
153 .unique()
154 .collect_vec()
155 };
156 let mut builder = StreamChunkBuilder::new(
157 self.ctx.config.developer.chunk_size,
158 self.output_data_types(),
159 );
160 let mut chunk_buffer: Vec<StreamChunk> = vec![];
161 let input = self.input.execute();
162
163 #[for_await]
164 for msg in input {
165 match msg? {
166 Message::Chunk(chunk) => chunk_buffer.push(chunk),
167 Message::Watermark(_w) => {}
168 Message::Barrier(barrier) => {
169 let chunks = StreamChunkCompactor::new(
170 compaction_keys.clone(),
171 std::mem::take(&mut chunk_buffer),
172 input_kind,
173 )
174 .into_compacted_chunks_inline::<RETRACT>(InconsistencyBehavior::Warn);
175
176 let mut deletes = Vec::with_capacity(chunks.len());
177 let mut inserts = Vec::with_capacity(chunks.len());
178
179 for chunk in chunks {
180 let vnodes: Arc<[VirtualNode]> = VirtualNode::compute_chunk(
181 chunk.data_chunk(),
182 &self.distribution_keys,
183 self.all_vnode_count,
184 )
185 .into();
186
187 let delete_view = chunk.clone().retain_ops(&[Op::Delete, Op::UpdateDelete]);
188 if delete_view.has_visible_rows() {
189 deletes.push((delete_view, vnodes.clone()));
190 }
191
192 let insert_view = chunk.retain_ops(&[Op::Insert, Op::UpdateInsert]);
193 if insert_view.has_visible_rows() {
194 inserts.push((insert_view, vnodes));
195 }
196 }
197
198 for (chunk, vnodes) in deletes.into_iter().chain(inserts) {
199 for (event, vnode) in chunk.rows_with_holes().zip_eq(vnodes.iter()) {
200 let Some((op, row)) = event else { continue };
201
202 let id = Serial::from(self.changelog_row_id_generator.next(vnode));
203
204 let mut suffix = Vec::with_capacity(2);
205
206 if self.need_op {
207 suffix.push(Some(ScalarImpl::Int16(op.to_i16())));
208 }
209
210 suffix.push(Some(ScalarImpl::Serial(id)));
211
212 if let Some(chunk) =
213 builder.append_row(Op::Insert, row.chain(OwnedRow::new(suffix)))
214 {
215 yield Message::Chunk(chunk);
216 }
217 }
218 }
219
220 if let Some(chunk) = builder.take() {
221 yield Message::Chunk(chunk);
222 }
223
224 if let Some(vnodes) = barrier.as_update_vnode_bitmap(self.ctx.id) {
225 self.changelog_row_id_generator = ChangelogRowIdGenerator::new(
226 vnodes.as_ref().clone(),
227 self.all_vnode_count,
228 );
229 }
230
231 yield Message::Barrier(barrier);
232 }
233 }
234 }
235 }
236
237 fn output_data_types(&self) -> Vec<DataType> {
238 let mut data_types = self.input.schema().data_types();
239
240 if self.need_op {
241 data_types.push(DataType::Int16);
242 }
243
244 data_types.push(DataType::Serial);
245 data_types
246 }
247}
248
249#[derive(Debug)]
250pub enum ChangeLogMode {
251 Normal,
252 Keyed { additional_key_columns: Vec<usize> },
253}
254
255impl Debug for ChangeLogExecutor {
256 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
257 f.debug_struct("ChangeLogExecutor").finish()
258 }
259}
260
261#[cfg(test)]
262mod tests {
263 use itertools::Itertools;
264 use risingwave_common::catalog::{Field, Schema};
265 use risingwave_common::row::{Row, RowExt};
266 use risingwave_common::test_prelude::StreamChunkTestExt;
267 use risingwave_common::types::{DataType, ScalarRefImpl};
268 use risingwave_common::util::epoch::test_epoch;
269
270 use super::*;
271 use crate::executor::test_utils::MockSource;
272 use crate::executor::{ActorContext, Barrier};
273
274 #[tokio::test]
275 async fn test_keyed_changelog_upsert_retains_delete() {
276 use crate::executor::ExecutorInfo;
277
278 let source = Executor::new(
279 ExecutorInfo {
280 stream_kind: StreamKind::Upsert,
281 ..ExecutorInfo::for_test(
282 Schema::new(vec![Field::unnamed(DataType::Int64); 2]),
283 vec![0],
284 "MockSource".to_owned(),
285 0,
286 )
287 },
288 MockSource::with_messages(vec![
289 Message::Barrier(Barrier::new_test_barrier(test_epoch(1))),
290 Message::Chunk(StreamChunk::from_pretty("I I\n+ 1 10")),
291 Message::Barrier(Barrier::new_test_barrier(test_epoch(2))),
292 Message::Chunk(StreamChunk::from_pretty("I I\n+ 1 20")),
293 Message::Chunk(StreamChunk::from_pretty("I I\n- 1 .")),
294 Message::Barrier(Barrier::new_test_barrier(test_epoch(3))),
295 ])
296 .stop_on_finish(false)
297 .boxed(),
298 );
299 let mut output = ChangeLogExecutor::new(
300 ActorContext::for_test(1),
301 source,
302 true,
303 VirtualNode::COUNT_FOR_TEST,
304 Bitmap::ones(VirtualNode::COUNT_FOR_TEST),
305 ChangeLogMode::Keyed {
306 additional_key_columns: vec![],
307 },
308 vec![0],
309 vec![0],
310 )
311 .boxed()
312 .execute();
313
314 let mut epochs = Vec::new();
315 let mut events = Vec::new();
316 while let Some(message) = output.next().await {
317 match message.unwrap() {
318 Message::Chunk(chunk) => {
319 for (op, row) in chunk.rows() {
320 assert_eq!(op, Op::Insert);
321 events.push(row.project(&[0, 1, 2]).into_owned_row());
322 }
323 }
324 Message::Barrier(_) => epochs.push(std::mem::take(&mut events)),
325 Message::Watermark(_) => panic!("unexpected watermark"),
326 }
327 }
328 assert_eq!(
329 epochs,
330 vec![
331 vec![],
332 vec![OwnedRow::new(vec![
333 Some(1_i64.into()),
334 Some(10_i64.into()),
335 Some(Op::Insert.to_i16().into())
336 ])],
337 vec![OwnedRow::new(vec![
338 Some(1_i64.into()),
339 None,
340 Some(Op::Delete.to_i16().into())
341 ])],
342 ]
343 );
344 }
345
346 #[tokio::test]
347 async fn test_keyed_changelog_upsert_subset_key_preserves_distinct_changes() {
348 use crate::executor::ExecutorInfo;
349
350 let source = Executor::new(
352 ExecutorInfo {
353 stream_kind: StreamKind::Upsert,
354 ..ExecutorInfo::for_test(
355 Schema::new(vec![Field::unnamed(DataType::Int64); 3]),
356 vec![0, 1],
357 "MockSource".to_owned(),
358 0,
359 )
360 },
361 MockSource::with_messages(vec![
362 Message::Barrier(Barrier::new_test_barrier(test_epoch(1))),
363 Message::Chunk(StreamChunk::from_pretty("I I I\n+ 1 10 100")),
364 Message::Barrier(Barrier::new_test_barrier(test_epoch(2))),
365 Message::Chunk(StreamChunk::from_pretty("I I I\n+ 1 20 200\n+ 1 10 101")),
366 Message::Chunk(StreamChunk::from_pretty("I I I\n- 1 10 .")),
367 Message::Barrier(Barrier::new_test_barrier(test_epoch(3))),
368 Message::Chunk(StreamChunk::from_pretty("I I I\n- 1 20 .")),
369 Message::Barrier(Barrier::new_test_barrier(test_epoch(4))),
370 ])
371 .stop_on_finish(false)
372 .boxed(),
373 );
374 let mut output = ChangeLogExecutor::new(
375 ActorContext::for_test(1),
376 source,
377 true,
378 VirtualNode::COUNT_FOR_TEST,
379 Bitmap::ones(VirtualNode::COUNT_FOR_TEST),
380 ChangeLogMode::Keyed {
381 additional_key_columns: vec![],
382 },
383 vec![0],
384 vec![0, 1],
385 )
386 .boxed()
387 .execute();
388
389 let mut epochs = Vec::new();
390 let mut events = Vec::new();
391 let mut previous_id = None;
392 while let Some(message) = output.next().await {
393 match message.unwrap() {
394 Message::Chunk(chunk) => {
395 for (op, row) in chunk.rows() {
396 assert_eq!(op, Op::Insert);
397 let Some(ScalarRefImpl::Serial(id)) = row.datum_at(4) else {
398 panic!("expected a non-null changelog event ID");
399 };
400 if let Some(previous_id) = previous_id {
401 assert!(previous_id < id);
402 }
403 previous_id = Some(id);
404 events.push(row.project(&[0, 1, 2, 3]).into_owned_row());
405 }
406 }
407 Message::Barrier(_) => epochs.push(std::mem::take(&mut events)),
408 Message::Watermark(_) => panic!("unexpected watermark"),
409 }
410 }
411
412 let expected = [
415 "I I I",
416 "I I I\n+ 1 10 100",
417 "I I I\n- 1 10 .\n+ 1 20 200",
418 "I I I\n- 1 20 .",
419 ]
420 .map(|chunk| {
421 StreamChunk::from_pretty(chunk)
422 .rows()
423 .map(|(op, row)| {
424 row.chain(OwnedRow::new(vec![Some(op.to_i16().into())]))
425 .into_owned_row()
426 })
427 .collect_vec()
428 });
429 assert_eq!(epochs, expected);
430 }
431
432 #[tokio::test]
433 async fn test_keyed_changelog_compacts_with_additional_key_columns() {
434 for (additional_key_columns, expected_ops) in [
435 (vec![], [Op::Insert, Op::UpdateDelete, Op::UpdateInsert]),
436 (vec![1], [Op::Insert, Op::Delete, Op::Insert]),
437 ] {
438 let source = MockSource::with_messages(vec![
439 Message::Barrier(Barrier::new_test_barrier(test_epoch(1))),
440 Message::Chunk(StreamChunk::from_pretty("I I I\n+ 10 100 7")),
441 Message::Barrier(Barrier::new_test_barrier(test_epoch(2))),
442 Message::Chunk(StreamChunk::from_pretty("I I I\nU- 10 100 7\nU+ 10 200 7")),
443 Message::Barrier(Barrier::new_test_barrier(test_epoch(3))),
444 ])
445 .stop_on_finish(false)
446 .into_executor(
447 Schema::new(vec![
448 Field::unnamed(DataType::Int64),
449 Field::unnamed(DataType::Int64),
450 Field::unnamed(DataType::Int64),
451 ]),
452 vec![2],
453 );
454
455 let mut output = ChangeLogExecutor::new(
456 ActorContext::for_test(1),
457 source,
458 true,
459 VirtualNode::COUNT_FOR_TEST,
460 Bitmap::ones(VirtualNode::COUNT_FOR_TEST),
461 ChangeLogMode::Keyed {
462 additional_key_columns,
463 },
464 vec![0],
465 vec![2],
466 )
467 .boxed()
468 .execute();
469
470 let mut events = Vec::new();
471 while let Some(message) = output.next().await {
472 match message.unwrap() {
473 Message::Chunk(chunk) => {
474 for (op, row) in chunk.rows() {
475 assert_eq!(op, Op::Insert, "changelog output must be append-only");
476 let Some(ScalarRefImpl::Int64(key)) = row.datum_at(1) else {
477 panic!("expected a non-null changelog key");
478 };
479 let Some(ScalarRefImpl::Int16(change_op)) = row.datum_at(3) else {
480 panic!("expected a non-null changelog operation");
481 };
482 events.push((key, change_op));
483 }
484 }
485 Message::Barrier(_) => {}
486 Message::Watermark(_) => panic!("unexpected watermark"),
487 }
488 }
489
490 assert_eq!(
491 events,
492 [100, 100, 200]
493 .into_iter()
494 .zip_eq(expected_ops.map(Op::to_i16))
495 .collect_vec()
496 );
497 }
498 }
499
500 #[tokio::test]
501 async fn test_keyed_changelog_reorders_replacement() {
502 let source = MockSource::with_messages(vec![
503 Message::Barrier(Barrier::new_test_barrier(test_epoch(1))),
504 Message::Chunk(StreamChunk::from_pretty("I T B\n+ 42 Beginner false")),
505 Message::Barrier(Barrier::new_test_barrier(test_epoch(2))),
506 Message::Chunk(StreamChunk::from_pretty("I T B\n+ 42 Advanced true")),
507 Message::Chunk(StreamChunk::from_pretty("I T B\n- 42 Beginner false")),
508 Message::Barrier(Barrier::new_test_barrier(test_epoch(3))),
509 ])
510 .stop_on_finish(false)
511 .into_executor(
512 Schema::new(vec![
513 Field::unnamed(DataType::Int64),
514 Field::unnamed(DataType::Varchar),
515 Field::unnamed(DataType::Boolean),
516 ]),
517 vec![0, 1, 2],
518 );
519
520 let mut output = ChangeLogExecutor::new(
521 ActorContext::for_test(1),
522 source,
523 true,
524 VirtualNode::COUNT_FOR_TEST,
525 Bitmap::ones(VirtualNode::COUNT_FOR_TEST),
526 ChangeLogMode::Keyed {
527 additional_key_columns: vec![],
528 },
529 vec![0],
530 vec![0, 1, 2],
531 )
532 .boxed()
533 .execute();
534
535 let mut events = Vec::new();
536
537 while let Some(message) = output.next().await {
538 match message.unwrap() {
539 Message::Chunk(chunk) => {
540 for (op, row) in chunk.rows() {
541 assert_eq!(op, Op::Insert, "changelog output must be append-only");
542
543 let Some(ScalarRefImpl::Serial(id)) = row.datum_at(4) else {
544 panic!("expected a non-null changelog event ID");
545 };
546
547 let Some(ScalarRefImpl::Int16(change_op)) = row.datum_at(3) else {
548 panic!("expected a non-null changelog operation");
549 };
550
551 let row = row.project(&[0, 1, 2]).into_owned_row();
552
553 events.push((row, change_op, id));
554 }
555 }
556 Message::Barrier(_) => {}
557 Message::Watermark(_) => panic!("unexpected watermark"),
558 }
559 }
560
561 let expected = StreamChunk::from_pretty(
562 "I T B\n+ 42 Beginner false\n- 42 Beginner false\n+ 42 Advanced true",
563 );
564 let ids = expected
565 .rows()
566 .map(|(op, row)| {
567 let row = row.to_owned_row();
568
569 events
570 .iter()
571 .find(|(event, change_op, _)| *event == row && *change_op == op.to_i16())
572 .unwrap_or_else(|| panic!("missing changelog event: {row:?}"))
573 .2
574 })
575 .collect_vec();
576
577 let [old_insert_id, old_delete_id, new_insert_id] = ids[..] else {
578 panic!("expected exactly three changelog events");
579 };
580
581 assert!(
582 old_insert_id < old_delete_id,
583 "original insert must precede deletion"
584 );
585
586 assert!(
587 old_delete_id < new_insert_id,
588 "old deletion must precede replacement insert"
589 );
590 }
591
592 #[tokio::test]
593 async fn test_keyed_changelog_does_not_pair_different_keys() {
594 let source = MockSource::with_messages(vec![
595 Message::Barrier(Barrier::new_test_barrier(test_epoch(1))),
596 Message::Chunk(StreamChunk::from_pretty("I T I\n+ 1 Alice 7")),
597 Message::Barrier(Barrier::new_test_barrier(test_epoch(2))),
598 Message::Chunk(StreamChunk::from_pretty("I T I\n- 1 Alice 7\n+ 2 Alice 7")),
599 Message::Barrier(Barrier::new_test_barrier(test_epoch(3))),
600 ])
601 .stop_on_finish(false)
602 .into_executor(
603 Schema::new(vec![
604 Field::unnamed(DataType::Int64),
605 Field::unnamed(DataType::Varchar),
606 Field::unnamed(DataType::Int64),
607 ]),
608 vec![2],
609 );
610
611 let mut output = ChangeLogExecutor::new(
612 ActorContext::for_test(1),
613 source,
614 true,
615 VirtualNode::COUNT_FOR_TEST,
616 Bitmap::ones(VirtualNode::COUNT_FOR_TEST),
617 ChangeLogMode::Keyed {
618 additional_key_columns: vec![],
619 },
620 vec![0],
621 vec![2],
622 )
623 .boxed()
624 .execute();
625
626 let mut events = Vec::new();
627 while let Some(message) = output.next().await {
628 match message.unwrap() {
629 Message::Chunk(chunk) => {
630 for (op, row) in chunk.rows() {
631 assert_eq!(op, Op::Insert, "changelog output must be append-only");
632 let Some(ScalarRefImpl::Int64(key)) = row.datum_at(0) else {
633 panic!("expected a non-null changelog key");
634 };
635 let Some(ScalarRefImpl::Int16(change_op)) = row.datum_at(3) else {
636 panic!("expected a non-null changelog operation");
637 };
638 events.push((key, change_op));
639 }
640 }
641 Message::Barrier(_) => {}
642 Message::Watermark(_) => panic!("unexpected watermark"),
643 }
644 }
645
646 assert_eq!(
647 events,
648 vec![
649 (1, Op::Insert.to_i16()),
650 (1, Op::Delete.to_i16()),
651 (2, Op::Insert.to_i16())
652 ]
653 );
654 }
655}