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

Code before:
    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.iter().any(|r| *r == PortState::Ready);
            send[0] = if any_ready {
                PortState::Ready
            } else {
                PortState::Blocked
            };
        }

        if send[0] == PortState::Blocked {
            for r in recv.iter_mut() {
                if *r == PortState::Ready {
                    *r = 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_schema = self.output_schema.clone();

        let output_senders = send_ports[0].take().unwrap().parallel();
        let n = output_senders.len();
        assert_eq!(n, _state.num_pipelines);

Review comment (bug):
This is incorrect and overcomplicated. I think you can just do `recv.fill(send[0])`. From the input's perspective it's directly connected to our send side.

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();

        for recv_port in recv_ports {