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

polars_python/functions/
range.rs
