polars_python/functions/
range.rs

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

polars_python/functions/
range.rs
