Repository: pola-rs/polars
File: crates/polars-compute/src/rolling/no_nulls/mod.rs (Rust)

Code before:
        .map(|idx| {
            let (start, end) = det_offsets_fn(idx, window_size, len);
            let vals = unsafe { values.get_unchecked(start..end) };
            let win_len = end - start;
            let weights_slice = if start == 0 {
                // Truncated at the start: take the last win_len weights
                &weights[weights.len() - win_len..]
            } else if end == len {
                // Truncated at the end: take the first win_len weights
                &weights[..win_len]
            } else {
                // Full window
                weights
            };

            if normalize && win_len != window_size {
                // Renormalize weights so they sum to 1
                let wsum = weights_slice
                    .iter()
                    .copied()
                    .fold(T::zero(), |acc, x| acc + x);
                if wsum == T::zero() {
                    panic!("Weighted mean is undefined if weights sum to 0");
                }
                let normed_weights: Vec<T> = weights_slice.iter().map(|&w| w / wsum).collect();
                aggregator(vals, &normed_weights)
            } else {
                aggregator(vals, weights_slice)
            }
        })
        .collect_trusted::<Vec<T>>();

    let validity = create_validity(min_periods, len, window_size, det_offsets_fn);
    Ok(Box::new(PrimitiveArray::new(
        ArrowDataType::from(T::PRIMITIVE),
        out.into(),
        validity.map(|b| b.into()),
    )))
}

fn compute_var_weights<T>(vals: &[T], weights: &[T]) -> T
where
    T: Float + std::ops::AddAssign,
{
    // Assumes the weights have already been standardized to 1
    debug_assert!(
        weights.iter().fold(T::zero(), |acc, x| acc + *x) == T::one(),
        "Rolling weighted variance Weights don't sum to 1"
    );
    let (wssq, wmean) = vals
        .iter()

Review comment (suggestion):
I'm *really* not a fan of normalizing the weights for each call. Instead the inner aggregator if necessary (`mean` / `var`) should not assume the weights have been pre-normalized and calculate the total weight themselves and divide by it.

Code after:
            let (start, end) = det_offsets_fn(idx, window_size, len);
            let vals = unsafe { values.get_unchecked(start..end) };
            let win_len = end - start;
            let weights_slice = if start == 0 {
                // Truncated at the start: take the last win_len weights
                &weights[weights.len() - win_len..]
            } else if end == len {
                // Truncated at the end: take the first win_len weights
                &weights[..win_len]
            } else {
                // Full window
                weights
            };
            aggregator(vals, weights_slice)
        })
        .collect_trusted::<Vec<T>>();

    let validity = create_validity(min_periods, len, window_size, det_offsets_fn);
    Ok(Box::new(PrimitiveArray::new(
        ArrowDataType::from(T::PRIMITIVE),
        out.into(),
        validity.map(|b| b.into()),
    )))
}

fn compute_var_weights<T>(vals: &[T], weights: &[T]) -> T
where
    T: Float + std::ops::AddAssign,
{
    // Compute weighted mean and weighted sum of squares in a single pass
    let (wssq, wmean, total_weight) = vals.iter().zip(weights).fold(
        (T::zero(), T::zero(), T::zero()),
        |(wssq, wsum, wtot), (&v, &w)| (wssq + v * v * w, wsum + v * w, wtot + w),
    );
    if total_weight.is_zero() {
        panic!("Weighted variance is undefined if weights sum to 0");
    }
    let mean = wmean / total_weight;
    (wssq / total_weight) - (mean * mean)
}

pub(crate) fn compute_sum_weights<T>(values: &[T], weights: &[T]) -> T
where
    T: std::iter::Sum<T> + Copy + std::ops::Mul<Output = T>,
{
    values.iter().zip(weights).map(|(v, w)| *v * *w).sum()
}

/// Compute the weighted mean of values, given weights (not necessarily normalized).
/// Returns sum_i(values[i] * weights[i]) / sum_i(weights[i])
pub(crate) fn compute_mean_weights<T>(values: &[T], weights: &[T]) -> T