Repository: pola-rs/polars
File: crates/polars-stream/src/nodes/unordered_union.rs (Rust)

Code before:
        }
    }
}

impl ComputeNode for UnorderedUnionNode {
    fn name(&self) -> &str {
        "unordered-union"
    }

    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        _state: &StreamingExecutionState,
    ) -> PolarsResult<()> {
        assert_eq!(send.len(), 1);

        let all_done = recv.iter().all(|r| *r == PortState::Done);

        if all_done {
            send[0] = PortState::Done;
        } else {
            let any_ready = recv.contains(&PortState::Ready);
            send[0] = if any_ready {
                PortState::Ready
            } else {
                PortState::Blocked
            };
        }

        recv.fill(send[0]);

        Ok(())
    }

    fn spawn<'env, 's>(
        &'env mut self,
        scope: &'s TaskScope<'s, 'env>,
        recv_ports: &mut [Option<RecvPort<'_>>],
        send_ports: &mut [Option<SendPort<'_>>],
        state: &'s StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert_eq!(send_ports.len(), 1);
        let output_senders = send_ports[0].take().unwrap().parallel();
        let n = output_senders.len();
        assert_eq!(n, state.num_pipelines);

        let (mpsc_senders, mpsc_receivers): (Vec<_>, Vec<_>) =
            (0..n).map(|_| mpsc::channel::<Morsel>(1)).unzip();


Review comment (bug):
Actually this is still incorrect (should've been more clear, sorry). What we need to send here is the *original* `send` state, not the one we just calculated (that's ours).

So I'd suggest

```rust
let done = send[0] == PortState::Done ||
    recv.iter().all(|r| *r == PortState::Done);
if done {
    send[0] = PortState::Done;
    recv.fill(PortStone::Done);
    return Ok(());
}


let any_ready = recv.contains(&PortState::Ready);
recv.fill(send[0]);
send[0] = if any_ready { PortState::Ready } else { PortState::Blocked };
Ok(())
```

Code after:
        }
    }
}

impl ComputeNode for UnorderedUnionNode {
    fn name(&self) -> &str {
        "unordered-union"
    }

    fn update_state(
        &mut self,
        recv: &mut [PortState],
        send: &mut [PortState],
        _state: &StreamingExecutionState,
    ) -> PolarsResult<()> {
        assert_eq!(send.len(), 1);

        let done = send[0] == PortState::Done || recv.iter().all(|r| *r == PortState::Done);
        if done {
            send[0] = PortState::Done;
            recv.fill(PortState::Done);
            return Ok(());
        }

        let any_ready = recv.contains(&PortState::Ready);
        recv.fill(send[0]);
        send[0] = if any_ready {
            PortState::Ready
        } else {
            PortState::Blocked
        };
        Ok(())
    }

    fn spawn<'env, 's>(
        &'env mut self,
        scope: &'s TaskScope<'s, 'env>,
        recv_ports: &mut [Option<RecvPort<'_>>],
        send_ports: &mut [Option<SendPort<'_>>],
        state: &'s StreamingExecutionState,
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert_eq!(send_ports.len(), 1);
        let output_senders = send_ports[0].take().unwrap().parallel();
        let num_pipelines = output_senders.len();
        assert_eq!(num_pipelines, state.num_pipelines);

        let (mpsc_senders, mpsc_receivers): (Vec<_>, Vec<_>) = (0..num_pipelines)
            .map(|_| mpsc::channel::<Morsel>(1))
            .unzip();