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

Code before:

        join_handles.push(scope.spawn_task(TaskPriority::High, async move {
            let mut lengths = Vec::new();
            while let Ok(mut m) = recv.recv().await {
                let df = m.df_mut();
                if df.height() == 0 {
                    continue;
                }

                let column = if df.width() == 1 {
                    Cow::Borrowed(&df[0])
                } else {
                    Cow::Owned(
                        std::mem::take(df)
                            .into_struct(PlSmallStr::EMPTY)
                            .into_column(),
                    )
                };
                let column = column.as_ref();

                lengths.clear();
                polars_ops::series::rle_lengths(column, &mut lengths)?;

                // If the last value seen is different from this first value here, bump the index
                // by 1.
                self.index += u32::from(self.last.take().is_some_and(|last| {
                    let fst = Scalar::new(self.dtype.clone(), column.get(0).unwrap().into_static());
                    let last = Scalar::new(self.dtype.clone(), last);
                    fst != last
                }));
                self.last = Some(column.get(column.len() - 1).unwrap().into_static());

                let column = if lengths.len() == 1 {
                    // If we only have one unique value, just give a Scalar column.
                    Column::new_scalar(
                        self.name.clone(),
                        Scalar::from(self.index),
                        lengths[0] as usize,
                    )
                } else {
                    let mut values = Vec::with_capacity(column.len());
                    values.extend(std::iter::repeat_n(self.index, lengths[0] as usize));
                    for length in lengths.iter().skip(1) {
                        self.index += 1;
                        values.extend(std::iter::repeat_n(self.index, *length as usize));
                    }
                    let mut column = Column::new(self.name.clone(), values);
                    column.set_sorted_flag(polars_core::series::IsSorted::Ascending);
                    column
                };


Review comment (suggestion):
I think this would be a lot cleaner with an `if let Some(last) = last.take()` instead of a `is_some_and`.

Code after:
            while let Ok(mut m) = recv.recv().await {
                let df = m.df_mut();
                if df.height() == 0 {
                    continue;
                }

                assert_eq!(df.width(), 1);
                let column = &df[0];

                lengths.clear();
                polars_ops::series::rle_lengths(column, &mut lengths)?;

                // If the last value seen is different from this first value here, bump the index
                // by 1.
                if let Some(last) = self.last.take() {
                    let fst = Scalar::new(self.dtype.clone(), column.get(0).unwrap().into_static());
                    let last = Scalar::new(self.dtype.clone(), last);
                    self.index += IdxSize::from(fst != last);
                }
                self.last = Some(column.get(column.len() - 1).unwrap().into_static());

                let column = if lengths.len() == 1 {
                    // If we only have one unique value, just give a Scalar column.
                    Column::new_scalar(
                        self.name.clone(),
                        Scalar::from(self.index),
                        lengths[0] as usize,
                    )
                } else {
                    let mut values = Vec::with_capacity(column.len());
                    values.extend(std::iter::repeat_n(self.index, lengths[0] as usize));
                    for length in lengths.iter().skip(1) {
                        self.index += 1;
                        values.extend(std::iter::repeat_n(self.index, *length as usize));
                    }
                    let mut column = Column::new(self.name.clone(), values);
                    column.set_sorted_flag(polars_core::series::IsSorted::Ascending);
                    column
                };

                *df = unsafe { DataFrame::new_no_checks(column.len(), vec![column]) };

                if send.send(m).await.is_err() {
                    break;
                }
            }

            Ok(())
        }));
    }
}