polars_python/functions/
range.rs

1use polars::lazy::dsl;
2use polars_core::with_match_physical_integer_polars_type;
3use polars_ops::series::ClosedInterval;
4use polars_plan::constants::get_literal_name;
5use pyo3::prelude::*;
6
7use crate::error::PyPolarsErr;
8use crate::expr::datatype::PyDataTypeExpr;
9use crate::prelude::*;
10use crate::utils::EnterPolarsExt;
11use crate::{PyExpr, PySeries};
12
13#[pyfunction]
14pub fn int_range(start: PyExpr, end: PyExpr, step: i64, dtype: PyDataTypeExpr) -> PyExpr {
15    let start = start.inner;
16    let end = end.inner;
17    let dtype = dtype.inner;
18    dsl::int_range(start, end, step, dtype).into()
19}
20
21/// Eager version of `int_range` to avoid overhead from the expression engine.
22#[pyfunction]
23pub fn eager_int_range(
24    py: Python<'_>,
25    lower: &Bound<'_, PyAny>,
26    upper: &Bound<'_, PyAny>,
27    step: &Bound<'_, PyAny>,
28    dtype: PyDataTypeExpr,
29) -> PyResult<PySeries> {
30    let dtype = dtype.inner;
31    let Some(dtype) = dtype.into_literal() else {
32        return Err(PyPolarsErr::from(
33            polars_err!(ComputeError: "eager `int_range` cannot be given lazy datatype expression"),
34        )
35        .into());
36    };
37
38    if !dtype.is_integer() {
39        return Err(PyPolarsErr::from(
40            polars_err!(SchemaMismatch: "non-integer `dtype` passed to `int_range`: '{}'", dtype),
41        )
42        .into());
43    }
44
45    with_match_physical_integer_polars_type!(dtype, |$T| {
46        let start_v: <$T as PolarsNumericType>::Native = lower.extract()?;
47        let end_v: <$T as PolarsNumericType>::Native = upper.extract()?;
48        let step: i64 = step.extract()?;
49        py.enter_polars_series(|| new_int_range::<$T>(start_v, end_v, step, get_literal_name()))
50    })
51}
52
53#[pyfunction]
54pub fn int_ranges(
55    start: PyExpr,
56    end: PyExpr,
57    step: PyExpr,
58    dtype: PyDataTypeExpr,
59) -> PyResult<PyExpr> {
60    let dtype = dtype.inner;
61    Ok(dsl::int_ranges(start.inner, end.inner, step.inner, dtype).into())
62}
63
64#[pyfunction]
65pub fn date_range(
66    start: PyExpr,
67    end: PyExpr,
68    interval: &str,
69    closed: Wrap<ClosedWindow>,
70) -> PyResult<PyExpr> {
71    let start = start.inner;
72    let end = end.inner;
73    let interval = Duration::try_parse(interval).map_err(PyPolarsErr::from)?;
74    let closed = closed.0;
75    let out = dsl::date_range(
76        Some(start),
77        Some(end),
78        Some(interval),
79        None, // TODO: num_samples
80        closed,
81    )
82    .map_err(PyPolarsErr::from)?;
83    Ok(out.into())
84}
85
86#[pyfunction]
87pub fn date_ranges(
88    start: PyExpr,
89    end: PyExpr,
90    interval: &str,
91    closed: Wrap<ClosedWindow>,
92) -> PyResult<PyExpr> {
93    let start = start.inner;
94    let end = end.inner;
95    let interval = Duration::try_parse(interval).map_err(PyPolarsErr::from)?;
96    let closed = closed.0;
97    let out = dsl::date_ranges(
98        Some(start),
99        Some(end),
100        Some(interval),
101        None, // TODO: num_samples
102        closed,
103    )
104    .map_err(PyPolarsErr::from)?;
105    Ok(out.into())
106}
107
108#[pyfunction]
109#[pyo3(signature = (start, end, interval, closed, time_unit, time_zone))]
110pub fn datetime_range(
111    start: PyExpr,
112    end: PyExpr,
113    interval: &str,
114    closed: Wrap<ClosedWindow>,
115    time_unit: Option<Wrap<TimeUnit>>,
116    time_zone: Wrap<Option<TimeZone>>,
117) -> PyResult<PyExpr> {
118    let start = start.inner;
119    let end = end.inner;
120    let interval = Duration::try_parse(interval).map_err(PyPolarsErr::from)?;
121    let closed = closed.0;
122    let time_unit = time_unit.map(|x| x.0);
123    let time_zone = time_zone.0;
124    let out = dsl::datetime_range(
125        Some(start),
126        Some(end),
127        Some(interval),
128        None, // TODO: num_samples
129        closed,
130        time_unit,
131        time_zone,
132    )
133    .map_err(PyPolarsErr::from)?;
134    Ok(out.into())
135}
136
137#[pyfunction]
138#[pyo3(signature = (start, end, interval, closed, time_unit, time_zone))]
139pub fn datetime_ranges(
140    start: PyExpr,
141    end: PyExpr,
142    interval: &str,
143    closed: Wrap<ClosedWindow>,
144    time_unit: Option<Wrap<TimeUnit>>,
145    time_zone: Wrap<Option<TimeZone>>,
146) -> PyResult<PyExpr> {
147    let start = start.inner;
148    let end = end.inner;
149    let interval = Duration::try_parse(interval).map_err(PyPolarsErr::from)?;
150    let closed = closed.0;
151    let time_unit = time_unit.map(|x| x.0);
152    let time_zone = time_zone.0;
153    let out = dsl::datetime_ranges(
154        Some(start),
155        Some(end),
156        Some(interval),
157        None, // TODO: num_samples
158        closed,
159        time_unit,
160        time_zone,
161    )
162    .map_err(PyPolarsErr::from)?;
163    Ok(out.into())
164}
165
166#[pyfunction]
167pub fn time_range(
168    start: PyExpr,
169    end: PyExpr,
170    every: &str,
171    closed: Wrap<ClosedWindow>,
172) -> PyResult<PyExpr> {
173    let start = start.inner;
174    let end = end.inner;
175    let every = Duration::try_parse(every).map_err(PyPolarsErr::from)?;
176    let closed = closed.0;
177    Ok(dsl::time_range(start, end, every, closed).into())
178}
179
180#[pyfunction]
181pub fn time_ranges(
182    start: PyExpr,
183    end: PyExpr,
184    every: &str,
185    closed: Wrap<ClosedWindow>,
186) -> PyResult<PyExpr> {
187    let start = start.inner;
188    let end = end.inner;
189    let every = Duration::try_parse(every).map_err(PyPolarsErr::from)?;
190    let closed = closed.0;
191    Ok(dsl::time_ranges(start, end, every, closed).into())
192}
193
194#[pyfunction]
195pub fn linear_space(
196    start: PyExpr,
197    end: PyExpr,
198    num_samples: PyExpr,
199    closed: Wrap<ClosedInterval>,
200) -> PyResult<PyExpr> {
201    let start = start.inner;
202    let end = end.inner;
203    let num_samples = num_samples.inner;
204    let closed = closed.0;
205    Ok(dsl::linear_space(start, end, num_samples, closed).into())
206}
207
208#[pyfunction]
209pub fn linear_spaces(
210    start: PyExpr,
211    end: PyExpr,
212    num_samples: PyExpr,
213    closed: Wrap<ClosedInterval>,
214    as_array: bool,
215) -> PyResult<PyExpr> {
216    let start = start.inner;
217    let end = end.inner;
218    let num_samples = num_samples.inner;
219    let closed = closed.0;
220    let out =
221        dsl::linear_spaces(start, end, num_samples, closed, as_array).map_err(PyPolarsErr::from)?;
222    Ok(out.into())
223}
polars_python/functions/range.rs

polars_python/functions/
range.rs
