polars_plan/dsl/functions/
range.rs

1use polars_ops::series::ClosedInterval;
2use polars_time::ClosedWindow;
3
4use super::*;
5
6/// Generate a range of integers.
7///
8/// Alias for `int_range`.
9pub fn arange(start: Expr, end: Expr, step: i64, dtype: DataType) -> Expr {
10    int_range(start, end, step, dtype)
11}
12
13/// Generate a range of integers.
14pub fn int_range(start: Expr, end: Expr, step: i64, dtype: DataType) -> Expr {
15    let input = vec![start, end];
16
17    Expr::Function {
18        input,
19        function: FunctionExpr::Range(RangeFunction::IntRange { step, dtype }),
20        options: FunctionOptions {
21            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
22            ..Default::default()
23        },
24    }
25}
26
27/// Generate a range of integers for each row of the input columns.
28pub fn int_ranges(start: Expr, end: Expr, step: Expr) -> Expr {
29    let input = vec![start, end, step];
30
31    Expr::Function {
32        input,
33        function: FunctionExpr::Range(RangeFunction::IntRanges),
34        options: FunctionOptions {
35            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
36            ..Default::default()
37        },
38    }
39}
40
41/// Create a date range from a `start` and `stop` expression.
42#[cfg(feature = "temporal")]
43pub fn date_range(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
44    let input = vec![start, end];
45
46    Expr::Function {
47        input,
48        function: FunctionExpr::Range(RangeFunction::DateRange { interval, closed }),
49        options: FunctionOptions {
50            collect_groups: ApplyOptions::GroupWise,
51            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
52            ..Default::default()
53        },
54    }
55}
56
57/// Create a column of date ranges from a `start` and `stop` expression.
58#[cfg(feature = "temporal")]
59pub fn date_ranges(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
60    let input = vec![start, end];
61
62    Expr::Function {
63        input,
64        function: FunctionExpr::Range(RangeFunction::DateRanges { interval, closed }),
65        options: FunctionOptions {
66            collect_groups: ApplyOptions::GroupWise,
67            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
68            ..Default::default()
69        },
70    }
71}
72
73/// Create a datetime range from a `start` and `stop` expression.
74#[cfg(feature = "dtype-datetime")]
75pub fn datetime_range(
76    start: Expr,
77    end: Expr,
78    interval: Duration,
79    closed: ClosedWindow,
80    time_unit: Option<TimeUnit>,
81    time_zone: Option<TimeZone>,
82) -> Expr {
83    let input = vec![start, end];
84
85    Expr::Function {
86        input,
87        function: FunctionExpr::Range(RangeFunction::DatetimeRange {
88            interval,
89            closed,
90            time_unit,
91            time_zone,
92        }),
93        options: FunctionOptions {
94            collect_groups: ApplyOptions::GroupWise,
95            cast_options: Some(CastingRules::cast_to_supertypes()),
96            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
97            ..Default::default()
98        },
99    }
100}
101
102/// Create a column of datetime ranges from a `start` and `stop` expression.
103#[cfg(feature = "dtype-datetime")]
104pub fn datetime_ranges(
105    start: Expr,
106    end: Expr,
107    interval: Duration,
108    closed: ClosedWindow,
109    time_unit: Option<TimeUnit>,
110    time_zone: Option<TimeZone>,
111) -> Expr {
112    let input = vec![start, end];
113
114    Expr::Function {
115        input,
116        function: FunctionExpr::Range(RangeFunction::DatetimeRanges {
117            interval,
118            closed,
119            time_unit,
120            time_zone,
121        }),
122        options: FunctionOptions {
123            collect_groups: ApplyOptions::GroupWise,
124            cast_options: Some(CastingRules::cast_to_supertypes()),
125            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
126            ..Default::default()
127        },
128    }
129}
130
131/// Generate a time range.
132#[cfg(feature = "dtype-time")]
133pub fn time_range(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
134    let input = vec![start, end];
135
136    Expr::Function {
137        input,
138        function: FunctionExpr::Range(RangeFunction::TimeRange { interval, closed }),
139        options: FunctionOptions {
140            collect_groups: ApplyOptions::GroupWise,
141            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
142            ..Default::default()
143        },
144    }
145}
146
147/// Create a column of time ranges from a `start` and `stop` expression.
148#[cfg(feature = "dtype-time")]
149pub fn time_ranges(start: Expr, end: Expr, interval: Duration, closed: ClosedWindow) -> Expr {
150    let input = vec![start, end];
151
152    Expr::Function {
153        input,
154        function: FunctionExpr::Range(RangeFunction::TimeRanges { interval, closed }),
155        options: FunctionOptions {
156            collect_groups: ApplyOptions::GroupWise,
157            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
158            ..Default::default()
159        },
160    }
161}
162
163/// Generate a series of equally-spaced points.
164pub fn linear_space(start: Expr, end: Expr, num_samples: Expr, closed: ClosedInterval) -> Expr {
165    let input = vec![start, end, num_samples];
166
167    Expr::Function {
168        input,
169        function: FunctionExpr::Range(RangeFunction::LinearSpace { closed }),
170        options: FunctionOptions {
171            collect_groups: ApplyOptions::GroupWise,
172            flags: FunctionFlags::default() | FunctionFlags::ALLOW_RENAME,
173            ..Default::default()
174        },
175    }
176}
polars_plan/dsl/functions/range.rs

polars_plan/dsl/functions/
range.rs
