polars_compute/rolling/
mod.rs

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

polars_compute/rolling/
mod.rs
