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

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

        let mut mpsc_senders = Vec::new();
        let mut mpsc_receivers = Vec::new();

        for _ in 0..n {
            let (tx, rx) = mpsc::channel::<Morsel>(1000); // what should be? 1000 is just random
            mpsc_senders.push(tx);
            mpsc_receivers.push(rx);
        }

        for recv_port in recv_ports {
            let receivers = recv_port.take().unwrap().parallel();
            let output_schema = output_schema.clone();

            let mpsc_senders_clone: Vec<_> = mpsc_senders.iter().map(|s| s.clone()).collect();

            join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                let mut sub_tasks = Vec::new();

                for (lane_idx, mut receiver) in receivers.into_iter().enumerate() {
                    let mpsc_senders = mpsc_senders_clone.clone();
                    let output_schema = output_schema.clone();

                    sub_tasks.push(scope.spawn_task(TaskPriority::High, async move {
                        while let Ok(mut morsel) = receiver.recv().await {
                            morsel.df_mut().ensure_matches_schema(&output_schema)?;

                            if mpsc_senders[lane_idx].send(morsel).await.is_err() {
                                break;
                            }
                        }
                        PolarsResult::Ok(())
                    }));
                }

                for sub_task in sub_tasks {
                    sub_task.await?;

Review comment (suggestion):
You shouldn't `unwrap` here, it's possible some inputs are already done, or blocked. You should skip the receive port then.

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

        drop(mpsc_senders);

        // Each pipeline relabels morsel sequences independently of the others.
        // We first compute the `morsel_offset` as (max morsel sequence sent so far + 1), so this
        // phase never reuses sequence numbers from earlier phases.
        //
        // Then, each pipeline assigns sequences by:
        // - starting at `morsel_offset + pipeline_idx` (so pipelines start at different values),
        // - advancing by `num_pipelines` each time it emits a morsel.
        //
        // Example with 2 pipelines (num_pipelines = 2) and morsel_offset = 1000:
        // pipeline 0: 1000, 1002, 1004, ...
        // pipeline 1: 1001, 1003, 1005, ...
        //
        // This guarantees: