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

Code before:
            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();

        for recv_port in recv_ports {
            if let Some(recv) = recv_port.take() {
                let receivers = recv.parallel();
                let mpsc_senders_clone = mpsc_senders.clone();

                for (mut receiver, sender) in receivers.into_iter().zip(mpsc_senders_clone) {
                    let output_schema = self.output_schema.clone();
                    join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                        while let Ok(mut morsel) = receiver.recv().await {
                            // Ensure the morsel matches the expected output schema,
                            // casting nulls to the appropriate output type.
                            morsel.df_mut().ensure_matches_schema(&output_schema)?;

                            if sender.send(morsel).await.is_err() {
                                break;
                            }
                        }
                        PolarsResult::Ok(())
                    }));
                }

Review comment (suggestion):
I think `n` is too short a variable name to use for this. Just `num_pipelines` is fine I think.

Code after:
            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 {
            if let Some(recv) = recv_port.take() {
                let receivers = recv.parallel();
                let mpsc_senders_clone = mpsc_senders.clone();

                for (mut receiver, sender) in receivers.into_iter().zip(mpsc_senders_clone) {
                    let output_schema = self.output_schema.clone();
                    join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                        while let Ok(mut morsel) = receiver.recv().await {
                            // Ensure the morsel matches the expected output schema,
                            // casting nulls to the appropriate output type.
                            morsel.df_mut().ensure_matches_schema(&output_schema)?;

                            if sender.send(morsel).await.is_err() {
                                break;
                            }
                        }
                        PolarsResult::Ok(())
                    }));
                }