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

polars_python/functions/
range.rs
