refactor(executor): use event cause to refactor shuffle processor

databendlabs · Sep 12, 2024 · 1d21645 · 1d21645
1 parent 5d756fc
commit 1d21645
Showing 1 changed file with 27 additions and 3 deletions.
diff --git a/src/query/pipeline/core/src/processors/shuffle_processor.rs b/src/query/pipeline/core/src/processors/shuffle_processor.rs
@@ -74,9 +74,7 @@ impl Processor for ShuffleProcessor {
 
     fn event_with_cause(&mut self, cause: EventCause) -> Result<Event> {
         let (input, output) = match cause {
-            EventCause::Other => {
-                return Ok(Event::NeedConsume);
-            }
+            EventCause::Other => unreachable!(),
             EventCause::Input(index) => {
                 (&self.inputs[index], &self.outputs[self.input2output[index]])
             }
@@ -87,6 +85,19 @@ impl Processor for ShuffleProcessor {
 
         if output.is_finished() {
             input.finish();
+
+            for input in &self.inputs {
+                if !input.is_finished() {
+                    return Ok(Event::NeedConsume);
+                }
+            }
+
+            for output in &self.outputs {
+                if !output.is_finished() {
+                    return Ok(Event::NeedConsume);
+                }
+            }
+
             return Ok(Event::Finished);
         }
 
@@ -102,6 +113,19 @@ impl Processor for ShuffleProcessor {
 
         if input.is_finished() {
             output.finish();
+
+            for input in &self.inputs {
+                if !input.is_finished() {
+                    return Ok(Event::NeedConsume);
+                }
+            }
+
+            for output in &self.outputs {
+                if !output.is_finished() {
+                    return Ok(Event::NeedConsume);
+                }
+            }
+
             return Ok(Event::Finished);
         }