Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
 Path: blob/main/crates/polars-expr/src/dispatch/rolling_by.rs
7884 views
1
use arrow::temporal_conversions::MICROSECONDS_IN_DAY as US_IN_DAY;

2
use polars_core::error::PolarsResult;

3
use polars_core::prelude::{Column, DataType, IntoColumn, TimeUnit};

4
use polars_core::series::Series;

5
use polars_time::prelude::{RollingOptionsDynamicWindow, SeriesOpsTime};

6


7
fn roll_by_with_temporal_conversion<F: FnOnce(&Series, &Series) -> PolarsResult<Series>>(

8
    s: &[Column],

9
    op: F,

10
) -> PolarsResult<Column> {

11
    let by = &s[1];

12
    let dt = &s[0].dtype();

13
    let s = if dt.is_temporal() {

14
        &s[0].to_physical_repr()

15
    } else {

16
        &s[0]

17
    };

18


19
    // @scalar-opt

20
    let out = op(s.as_materialized_series(), by.as_materialized_series())?;

21


22
    Ok(match dt {

23
        DataType::Date => (out * US_IN_DAY as f64)

24
            .cast(&DataType::Int64)?

25
            .into_datetime(TimeUnit::Microseconds, None),

26
        DataType::Datetime(tu, tz) => out.cast(&DataType::Int64)?.into_datetime(*tu, tz.clone()),

27
        DataType::Duration(tu) => out.cast(&DataType::Int64)?.into_duration(*tu),

28
        DataType::Time => out.cast(&DataType::Int64)?.into_time(),

29
        _ => out,

30
    }

31
    .into_column())

32
}

33


34
pub(super) fn rolling_min_by(

35
    s: &[Column],

36
    options: RollingOptionsDynamicWindow,

37
) -> PolarsResult<Column> {

38
    // @scalar-opt

39
    s[0].as_materialized_series()

40
        .rolling_min_by(s[1].as_materialized_series(), options)

41
        .map(Column::from)

42
}

43


44
pub(super) fn rolling_max_by(

45
    s: &[Column],

46
    options: RollingOptionsDynamicWindow,

47
) -> PolarsResult<Column> {

48
    // @scalar-opt

49
    s[0].as_materialized_series()

50
        .rolling_max_by(s[1].as_materialized_series(), options)

51
        .map(Column::from)

52
}

53


54
pub(super) fn rolling_mean_by(

55
    s: &[Column],

56
    options: RollingOptionsDynamicWindow,

57
) -> PolarsResult<Column> {

58
    roll_by_with_temporal_conversion(s, |s, by| s.rolling_mean_by(by, options))

59
}

60


61
pub(super) fn rolling_sum_by(

62
    s: &[Column],

63
    options: RollingOptionsDynamicWindow,

64
) -> PolarsResult<Column> {

65
    // @scalar-opt

66
    s[0].as_materialized_series()

67
        .rolling_sum_by(s[1].as_materialized_series(), options)

68
        .map(Column::from)

69
}

70


71
pub(super) fn rolling_quantile_by(

72
    s: &[Column],

73
    options: RollingOptionsDynamicWindow,

74
) -> PolarsResult<Column> {

75
    roll_by_with_temporal_conversion(s, |s, by| s.rolling_quantile_by(by, options))

76
}

77


78
pub(super) fn rolling_var_by(

79
    s: &[Column],

80
    options: RollingOptionsDynamicWindow,

81
) -> PolarsResult<Column> {

82
    // @scalar-opt

83
    s[0].as_materialized_series()

84
        .rolling_var_by(s[1].as_materialized_series(), options)

85
        .map(Column::from)

86
}

87


88
pub(super) fn rolling_std_by(

89
    s: &[Column],

90
    options: RollingOptionsDynamicWindow,

91
) -> PolarsResult<Column> {

92
    // @scalar-opt

93
    s[0].as_materialized_series()

94
        .rolling_std_by(s[1].as_materialized_series(), options)

95
        .map(Column::from)

96
}

97


98
pub(super) fn rolling_rank_by(

99
    s: &[Column],

100
    options: RollingOptionsDynamicWindow,

101
) -> PolarsResult<Column> {

102
    // @scalar-opt

103
    s[0].as_materialized_series()

104
        .rolling_rank_by(s[1].as_materialized_series(), options)

105
        .map(Column::from)

106
}

107


108