polars_plan/dsl/functions/
range.rs

1use polars_ops::series::ClosedInterval;
2#[cfg(feature = "temporal")]
3use polars_time::ClosedWindow;
4
5use super::*;
6
7/// Generate a range of integers.
8///
9/// Alias for `int_range`.
10pub fn arange(start: Expr, end: Expr, step: i64, dtype: DataType) -> Expr {
11    int_range(start, end, step, dtype)
12}
13
14/// Generate a range of integers.
15pub fn int_range(start: Expr, end: Expr, step: i64, dtype: DataType) -> Expr {
16    Expr::n_ary(RangeFunction::IntRange { step, dtype }, vec![start, end])
17}
18
19/// Generate a range of integers for each row of the input columns.
20pub fn int_ranges(start: Expr, end: Expr, step: Expr) -> Expr {
21    Expr::n_ary(RangeFunction::IntRanges, vec![start, end, step])
22}
23
24/// Create a date range from a `start` and `stop` expression.
25#[cfg(feature = "temporal")]
26pub fn date_range(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
27    Expr::n_ary(
28        RangeFunction::DateRange { interval, closed },
29        vec![start, end],
30    )
31}
32
33/// Create a column of date ranges from a `start` and `stop` expression.
34#[cfg(feature = "temporal")]
35pub fn date_ranges(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
36    Expr::n_ary(
37        RangeFunction::DateRanges { interval, closed },
38        vec![start, end],
39    )
40}
41
42/// Create a datetime range from a `start` and `stop` expression.
43#[cfg(feature = "dtype-datetime")]
44pub fn datetime_range(
45    start: Expr,
46    end: Expr,
47    interval: Duration,
48    closed: ClosedWindow,
49    time_unit: Option<TimeUnit>,
50    time_zone: Option<TimeZone>,
51) -> Expr {
52    Expr::n_ary(
53        RangeFunction::DatetimeRange {
54            interval,
55            closed,
56            time_unit,
57            time_zone,
58        },
59        vec![start, end],
60    )
61}
62
63/// Create a column of datetime ranges from a `start` and `stop` expression.
64#[cfg(feature = "dtype-datetime")]
65pub fn datetime_ranges(
66    start: Expr,
67    end: Expr,
68    interval: Duration,
69    closed: ClosedWindow,
70    time_unit: Option<TimeUnit>,
71    time_zone: Option<TimeZone>,
72) -> Expr {
73    Expr::n_ary(
74        RangeFunction::DatetimeRanges {
75            interval,
76            closed,
77            time_unit,
78            time_zone,
79        },
80        vec![start, end],
81    )
82}
83
84/// Generate a time range.
85#[cfg(feature = "dtype-time")]
86pub fn time_range(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
87    Expr::n_ary(
88        RangeFunction::TimeRange { interval, closed },
89        vec![start, end],
90    )
91}
92
93/// Create a column of time ranges from a `start` and `stop` expression.
94#[cfg(feature = "dtype-time")]
95pub fn time_ranges(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
96    Expr::n_ary(
97        RangeFunction::TimeRanges { interval, closed },
98        vec![start, end],
99    )
100}
101
102/// Generate a series of equally-spaced points.
103pub fn linear_space(start: Expr, end: Expr, num_samples: Expr, closed: ClosedInterval) -> Expr {
104    Expr::n_ary(
105        RangeFunction::LinearSpace { closed },
106        vec![start, end, num_samples],
107    )
108}
109
110/// Create a column of linearly-spaced sequences from 'start', 'end', and 'num_samples' expressions.
111pub fn linear_spaces(
112    start: Expr,
113    end: Expr,
114    num_samples: Expr,
115    closed: ClosedInterval,
116    as_array: bool,
117) -> PolarsResult<Expr> {
118    let mut input = Vec::<Expr>::with_capacity(3);
119    input.push(start);
120    input.push(end);
121    let array_width = if as_array {
122        Some(num_samples.extract_usize().map_err(|_| {
123            polars_err!(InvalidOperation: "'as_array' is only valid when 'num_samples' is a constant integer")
124        })?)
125    } else {
126        input.push(num_samples);
127        None
128    };
129
130    Ok(Expr::n_ary(
131        RangeFunction::LinearSpaces {
132            closed,
133            array_width,
134        },
135        input,
136    ))
137}
polars_plan/dsl/functions/range.rs

polars_plan/dsl/functions/
range.rs
