Repository: pola-rs/polars
File: crates/polars-arrow/src/legacy/kernels/ewm/average.rs (Rust)

Code before:
    T: NativeType + num_traits::Float + std::ops::MulAssign,
{
    pub fn update(&mut self, values: &PrimitiveArray<T>) -> PrimitiveArray<T> {
        values.iter().map(|x| x.copied()).collect()
    }

    pub fn update_iter<I>(&mut self, values: I) -> impl Iterator<Item = Option<T>>
    where
        I: IntoIterator<Item = Option<T>>,
    {
        let new_value_weight = if self.adjust { T::one() } else { self.alpha };

        values.into_iter().map(move |opt_v| {
            if self.non_null_count == 0
                && let Some(v) = opt_v
            {
                // Initialize. This is done here to skip leading NULLs.
                self.non_null_count = 1;
                self.weighted_mean = v;
                self.weight = T::one();
            } else {
                if opt_v.is_some() || !self.ignore_nulls {
                    self.weight *= T::one() - self.alpha;
                }

                if let Some(new_v) = opt_v {
                    let new_weight = self.weight + new_value_weight;

                    self.weighted_mean = self.weighted_mean
                        + (new_v - self.weighted_mean) * (new_value_weight / new_weight);

                    self.weight = if self.adjust {
                        self.weight + T::one()
                    } else {
                        T::one()
                    };

                    self.non_null_count += 1;
                }
            }

            (opt_v.is_some() && self.non_null_count >= self.min_periods)
                .then_some(self.weighted_mean)
        })
    }
}
#[cfg(test)]
mod test {
    use super::super::assert_allclose;
    use super::*;
    const ALPHA: f64 = 0.5;

Review comment (suggestion):
Note, weighted mean update rearranged to match what we do in `polars-compute/moment.rs`. The updated code is mathematically identical - see derivation below

```
// (old_wt * w_avg + new_wt * x) / (old_wt + new_wt)
// = (old_wt * w_avg + new_wt * x) / (old_wt + new_wt) - w_avg + w_avg
// = (new_wt * x - new_wt * w_avg) / (old_wt + new_wt) + w_avg
// = new_wt * (x - w_avg) / (old_wt + new_wt) + w_avg
// = w_avg + (x - w_avg) * (new_wt / (old_wt + new_wt))
// = w_avg + (x - w_avg) * (new_wt / new_weight)
// (then rename variables)
```

Relevant code in `polars-compute/moment.rs` - 

https://github.com/pola-rs/polars/blob/e91e1c51147951df97e6efc87b3ec89f97700de5/crates/polars-compute/src/moment.rs#L116-L119


Code after:
    T: NativeType + num_traits::Float + std::ops::MulAssign,
{
    pub fn update_iter<I>(&mut self, values: I) -> impl Iterator<Item = Option<T>>
    where
        I: IntoIterator<Item = Option<T>>,
    {
        let new_value_weight = if self.adjust { T::one() } else { self.alpha };

        values.into_iter().map(move |opt_v| {
            if self.non_null_count == 0
                && let Some(v) = opt_v
            {
                // Initialize
                self.non_null_count = 1;
                self.weighted_mean = v;
                self.weight = T::one();
            } else {
                if opt_v.is_some() || !self.ignore_nulls {
                    self.weight *= T::one() - self.alpha;
                }

                if let Some(new_v) = opt_v {
                    let new_weight = self.weight + new_value_weight;

                    self.weighted_mean = self.weighted_mean
                        + (new_v - self.weighted_mean) * (new_value_weight / new_weight);

                    self.weight = if self.adjust {
                        self.weight + T::one()
                    } else {
                        T::one()
                    };

                    self.non_null_count += 1;
                }
            }

            (opt_v.is_some() && self.non_null_count >= self.min_periods)
                .then_some(self.weighted_mean)
        })
    }
}
#[cfg(test)]
mod test {
    use super::super::assert_allclose;
    use super::*;
    const ALPHA: f64 = 0.5;
    const EPS: f64 = 1e-15;

    #[test]
    fn test_ewm_mean_without_null() {