fix(streaming): map watermark in dispatcher with output indices (risi…

…ngwavelabs#8506)

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

src/stream/src/executor/backfill.rs

@@ -398,10 +398,7 @@ where
     }
 
     fn mapping_watermark(watermark: Watermark, upstream_indices: &[usize]) -> Option<Watermark> {
-        upstream_indices
-            .iter()
-            .position(|&idx| idx == watermark.col_idx)
-            .map(|idx| watermark.with_idx(idx))
+        watermark.transform_with_indices(upstream_indices)
     }
 
     fn mapping_message(msg: Message, upstream_indices: &[usize]) -> Option<Message> {

src/stream/src/executor/chain.rs

@@ -53,10 +53,7 @@ fn mapping_chunk(chunk: StreamChunk, upstream_indices: &[usize]) -> StreamChunk
 }
 
 fn mapping_watermark(watermark: Watermark, upstream_indices: &[usize]) -> Option<Watermark> {
-    upstream_indices
-        .iter()
-        .position(|&idx| idx == watermark.col_idx)
-        .map(|idx| watermark.with_idx(idx))
+    watermark.transform_with_indices(upstream_indices)
 }
 
 impl ChainExecutor {

src/stream/src/executor/dispatch.rs

@@ -446,14 +446,20 @@ pub trait Dispatcher: Debug + 'static {
 #[derive(Debug)]
 pub struct RoundRobinDataDispatcher {
     outputs: Vec<BoxedOutput>,
+    output_indices: Vec<usize>,
     cur: usize,
     dispatcher_id: DispatcherId,
 }
 
 impl RoundRobinDataDispatcher {
-    pub fn new(outputs: Vec<BoxedOutput>, dispatcher_id: DispatcherId) -> Self {
+    pub fn new(
+        outputs: Vec<BoxedOutput>,
+        output_indices: Vec<usize>,
+        dispatcher_id: DispatcherId,
+    ) -> Self {
         Self {
             outputs,
+            output_indices,
             cur: 0,
             dispatcher_id,
         }
@@ -465,6 +471,7 @@ impl Dispatcher for RoundRobinDataDispatcher {
 
     fn dispatch_data(&mut self, chunk: StreamChunk) -> Self::DataFuture<'_> {
         async move {
+            let chunk = chunk.reorder_columns(&self.output_indices);
             self.outputs[self.cur].send(Message::Chunk(chunk)).await?;
             self.cur += 1;
             self.cur %= self.outputs.len();
@@ -484,9 +491,11 @@ impl Dispatcher for RoundRobinDataDispatcher {
 
     fn dispatch_watermark(&mut self, watermark: Watermark) -> Self::WatermarkFuture<'_> {
         async move {
-            // always broadcast watermark
-            for output in &mut self.outputs {
-                output.send(Message::Watermark(watermark.clone())).await?;
+            if let Some(watermark) = watermark.transform_with_indices(&self.output_indices) {
+                // always broadcast watermark
+                for output in &mut self.outputs {
+                    output.send(Message::Watermark(watermark.clone())).await?;
+                }
             }
             Ok(())
         }
@@ -569,9 +578,11 @@ impl Dispatcher for HashDataDispatcher {
 
     fn dispatch_watermark(&mut self, watermark: Watermark) -> Self::WatermarkFuture<'_> {
         async move {
-            // always broadcast watermark
-            for output in &mut self.outputs {
-                output.send(Message::Watermark(watermark.clone())).await?;
+            if let Some(watermark) = watermark.transform_with_indices(&self.output_indices) {
+                // always broadcast watermark
+                for output in &mut self.outputs {
+                    output.send(Message::Watermark(watermark.clone())).await?;
+                }
             }
             Ok(())
         }
@@ -751,8 +762,11 @@ impl Dispatcher for BroadcastDispatcher {
 
     fn dispatch_watermark(&mut self, watermark: Watermark) -> Self::WatermarkFuture<'_> {
         async move {
-            for output in self.outputs.values_mut() {
-                output.send(Message::Watermark(watermark.clone())).await?;
+            if let Some(watermark) = watermark.transform_with_indices(&self.output_indices) {
+                // always broadcast watermark
+                for output in self.outputs.values_mut() {
+                    output.send(Message::Watermark(watermark.clone())).await?;
+                }
             }
             Ok(())
         }
@@ -851,7 +865,10 @@ impl Dispatcher for SimpleDispatcher {
                 .exactly_one()
                 .expect("expect exactly one output");
 
-            output.send(Message::Watermark(watermark)).await
+            if let Some(watermark) = watermark.transform_with_indices(&self.output_indices) {
+                output.send(Message::Watermark(watermark)).await?;
+            }
+            Ok(())
         }
     }
 

src/stream/src/executor/integration_tests.rs

@@ -124,7 +124,9 @@ async fn test_merger_sum_aggr() {
     let dispatcher = DispatchExecutor::new(
         receiver_op,
         vec![DispatcherImpl::RoundRobin(RoundRobinDataDispatcher::new(
-            inputs, 0,
+            inputs,
+            vec![0],
+            0,
         ))],
         0,
         ctx,

src/stream/src/executor/mod.rs

@@ -615,6 +615,15 @@ impl Watermark {
         })
     }
 
+    /// Transform the watermark with the given output indices. If this watermark is not in the
+    /// output, return `None`.
+    pub fn transform_with_indices(self, output_indices: &[usize]) -> Option<Self> {
+        output_indices
+            .iter()
+            .position(|p| *p == self.col_idx)
+            .map(|new_col_idx| self.with_idx(new_col_idx))
+    }
+
     pub fn to_protobuf(&self) -> ProstWatermark {
         ProstWatermark {
             column: Some(ProstInputRef {

src/stream/src/executor/rearranged_chain.rs

@@ -71,10 +71,7 @@ fn mapping(upstream_indices: &[usize], msg: Message) -> Option<Message> {
 }
 
 fn mapping_watermark(watermark: Watermark, upstream_indices: &[usize]) -> Option<Watermark> {
-    upstream_indices
-        .iter()
-        .position(|&idx| idx == watermark.col_idx)
-        .map(|idx| watermark.with_idx(idx))
+    watermark.transform_with_indices(upstream_indices)
 }
 
 #[derive(Debug)]