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

Code before:
    min_periods: usize,
    det_offsets_fn: Fo,
    aggregator: Fa,
    weights: &[T],
    centered: bool,
) -> PolarsResult<ArrayRef>
where
    T: NativeType + num_traits::Zero + std::ops::Div<Output = T> + Copy,
    Fo: Fn(Idx, WindowSize, Len) -> (Start, End),
    Fa: Fn(&[T], &[T]) -> T,
{
    assert_eq!(weights.len(), window_size);
    let len = values.len();
    let out = (0..len)
        .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 && end == len && centered {
                // Edge case: There are fewer elements than the window size
                // in the array. We'll get the same array of values at each idx
                // so we need to use the idx to know how to slide the weights
                // along when centered.
                let slice_start = weights.len() - win_len - idx;
                let slice_end = slice_start + win_len;
                &weights[slice_start..slice_end]
            } else 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,

Review comment (suggestion):
I'm not 100% sure this formula is correct. There should be something to determine the center which I think is lacking? I think it now takes from the end.

```python
pl.Series([1]).rolling_sum(6, center=True, min_samples=0, weights=[1, 10, 100, 1000, 10_000, 100_000])
```

If I'm not mistaken this should return 1000.


Code after:
    min_periods: usize,
    det_offsets_fn: Fo,
    aggregator: Fa,
    weights: &[T],
    centered: bool,
) -> PolarsResult<ArrayRef>
where
    T: NativeType + num_traits::Zero + std::ops::Div<Output = T> + Copy,
    Fo: Fn(Idx, WindowSize, Len) -> (Start, End),
    Fa: Fn(&[T], &[T]) -> T,
{
    assert_eq!(weights.len(), window_size);
    let len = values.len();
    let out = (0..len)
        .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_start = if centered {
                // When using centered weights, we need to find the right location
                // in the weights array specifically by aligning the center of the
                // window with idx, to handle cases where the window is smaller than
                // weights array.
                let center = (window_size / 2) as isize;
                let offset = center - (idx as isize - start as isize);
                offset.max(0) as usize
            } else if start == 0 {
                // When start is 0, we need to work backwards from the end of the
                // weights array to ensure we are lined up correctly (since the
                // start of the values array is implicitly cut off)
                weights.len() - win_len
            } else {
                0
            };
            let weights_slice = &weights[weights_start..weights_start + win_len];
            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,
{