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

Code before:
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert!(recv_ports.len() == 1 && send_ports.len() == 1);
        let mut receiver = recv_ports[0].take().unwrap().serial();
        let senders = send_ports[0].take().unwrap().parallel();

        let (mut distributor, distr_receivers) =
            distributor_channel(senders.len(), *DEFAULT_DISTRIBUTOR_BUFFER_SIZE);

        let n = self.n;

        // To figure out the correct offsets we need to be serial.
        join_handles.push(scope.spawn_task(TaskPriority::High, async move {
            while let Ok(morsel) = receiver.recv().await {
                let height = morsel.df().height();
                if self.offset > height {
                    self.offset -= height;
                    continue;
                }

                if distributor.send((morsel, self.offset)).await.is_err() {
                    break;
                }

                self.offset += height.next_multiple_of(self.n) - height;
                self.offset %= self.n;
            }

            Ok(())
        }));

        // But gathering the column can be done in parallel.
        for (mut send, mut recv) in senders.into_iter().zip(distr_receivers) {
            join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                let wait_group = WaitGroup::default();
                while let Ok((morsel, offset)) = recv.recv().await {
                    let mut morsel = morsel.try_map(|mut df| {
                        let column = &df.get_columns()[0];
                        let out = column
                            .gather_every(n, offset)?
                            .with_name(column.name().clone());
                        unsafe { df.get_columns_mut()[0] = out };
                        PolarsResult::Ok(df)
                    })?;
                    morsel.set_consume_token(wait_group.token());
                    if send.send(morsel).await.is_err() {
                        break;
                    }
                    wait_group.wait().await;
                }


Review comment (bug):
Maybe add a comment that this is simply `(self.offset - height) mod n` without overflow.

Code after:
        join_handles: &mut Vec<JoinHandle<PolarsResult<()>>>,
    ) {
        assert!(recv_ports.len() == 1 && send_ports.len() == 1);
        let mut receiver = recv_ports[0].take().unwrap().serial();
        let senders = send_ports[0].take().unwrap().parallel();

        let (mut distributor, distr_receivers) =
            distributor_channel(senders.len(), *DEFAULT_DISTRIBUTOR_BUFFER_SIZE);

        let n = self.n;

        // To figure out the correct offsets we need to be serial.
        join_handles.push(scope.spawn_task(TaskPriority::High, async move {
            while let Ok(morsel) = receiver.recv().await {
                let height = morsel.df().height();
                if self.offset >= height {
                    self.offset -= height;
                    continue;
                }

                if distributor.send((morsel, self.offset)).await.is_err() {
                    break;
                }

                // Calculates `offset = (offset - height) mod n` without under- and overflow.
                self.offset += height.next_multiple_of(self.n) - height;
                self.offset %= self.n;
            }

            Ok(())
        }));

        // But gathering the column can be done in parallel.
        for (mut send, mut recv) in senders.into_iter().zip(distr_receivers) {
            join_handles.push(scope.spawn_task(TaskPriority::High, async move {
                let wait_group = WaitGroup::default();
                while let Ok((morsel, offset)) = recv.recv().await {
                    let mut morsel = morsel.try_map(|mut df| {
                        let column = &df.get_columns()[0];
                        let out = column
                            .gather_every(n, offset)?
                            .with_name(column.name().clone());
                        unsafe { df.get_columns_mut()[0] = out };
                        PolarsResult::Ok(df)
                    })?;
                    morsel.set_consume_token(wait_group.token());
                    if send.send(morsel).await.is_err() {
                        break;
                    }
                    wait_group.wait().await;
                }