polars_compute/rolling/
mod.rs

1mod min_max;
2pub mod moment;
3pub mod no_nulls;
4pub mod nulls;
5pub mod quantile_filter;
6pub(super) mod window;
7use std::hash::Hash;
8use std::ops::{Add, AddAssign, Div, Mul, Sub, SubAssign};
9
10use arrow::array::{ArrayRef, PrimitiveArray};
11use arrow::bitmap::{Bitmap, MutableBitmap};
12use arrow::types::NativeType;
13use num_traits::{Bounded, Float, NumCast, One, Zero};
14use polars_utils::float::IsFloat;
15#[cfg(feature = "serde")]
16use serde::{Deserialize, Serialize};
17use strum_macros::IntoStaticStr;
18use window::*;
19
20type Start = usize;
21type End = usize;
22type Idx = usize;
23type WindowSize = usize;
24type Len = usize;
25
26#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash, IntoStaticStr)]
27#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
28#[strum(serialize_all = "snake_case")]
29pub enum QuantileMethod {
30    #[default]
31    Nearest,
32    Lower,
33    Higher,
34    Midpoint,
35    Linear,
36    Equiprobable,
37}
38
39#[deprecated(note = "use QuantileMethod instead")]
40pub type QuantileInterpolOptions = QuantileMethod;
41
42#[derive(Clone, Copy, Debug, PartialEq, Hash)]
43#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
44pub enum RollingFnParams {
45    Quantile(RollingQuantileParams),
46    Var(RollingVarParams),
47    Skew { bias: bool },
48    Kurtosis { fisher: bool, bias: bool },
49}
50
51fn det_offsets(i: Idx, window_size: WindowSize, _len: Len) -> (usize, usize) {
52    (i.saturating_sub(window_size - 1), i + 1)
53}
54fn det_offsets_center(i: Idx, window_size: WindowSize, len: Len) -> (usize, usize) {
55    let right_window = window_size.div_ceil(2);
56    (
57        i.saturating_sub(window_size - right_window),
58        std::cmp::min(len, i + right_window),
59    )
60}
61
62fn create_validity<Fo>(
63    min_periods: usize,
64    len: usize,
65    window_size: usize,
66    det_offsets_fn: Fo,
67) -> Option<MutableBitmap>
68where
69    Fo: Fn(Idx, WindowSize, Len) -> (Start, End),
70{
71    if min_periods > 1 {
72        let mut validity = MutableBitmap::with_capacity(len);
73        validity.extend_constant(len, true);
74
75        // Set the null values at the boundaries
76
77        // Head.
78        for i in 0..len {
79            let (start, end) = det_offsets_fn(i, window_size, len);
80            if (end - start) < min_periods {
81                validity.set(i, false)
82            } else {
83                break;
84            }
85        }
86        // Tail.
87        for i in (0..len).rev() {
88            let (start, end) = det_offsets_fn(i, window_size, len);
89            if (end - start) < min_periods {
90                validity.set(i, false)
91            } else {
92                break;
93            }
94        }
95
96        Some(validity)
97    } else {
98        None
99    }
100}
101
102// Parameters allowed for rolling operations.
103#[derive(Clone, Copy, Debug, PartialEq, Hash)]
104#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
105pub struct RollingVarParams {
106    pub ddof: u8,
107}
108
109#[derive(Clone, Copy, Debug, PartialEq)]
110#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
111pub struct RollingQuantileParams {
112    pub prob: f64,
113    pub method: QuantileMethod,
114}
115
116impl Hash for RollingQuantileParams {
117    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
118        // Will not be NaN, so hash + eq symmetry will hold.
119        self.prob.to_bits().hash(state);
120        self.method.hash(state);
121    }
122}
polars_compute/rolling/mod.rs

polars_compute/rolling/
mod.rs
