polars_plan/dsl/function_expr/range/
mod.rs

1#[cfg(feature = "dtype-date")]
2mod date_range;
3#[cfg(feature = "dtype-datetime")]
4mod datetime_range;
5mod int_range;
6mod linear_space;
7#[cfg(feature = "dtype-time")]
8mod time_range;
9mod utils;
10
11use std::fmt::{Display, Formatter};
12
13use polars_core::prelude::*;
14use polars_ops::series::ClosedInterval;
15#[cfg(feature = "temporal")]
16use polars_time::{ClosedWindow, Duration};
17#[cfg(feature = "serde")]
18use serde::{Deserialize, Serialize};
19
20use super::{FunctionExpr, FunctionOptions};
21use crate::dsl::SpecialEq;
22use crate::dsl::function_expr::FieldsMapper;
23use crate::map_as_slice;
24use crate::prelude::{ColumnsUdf, FunctionFlags};
25
26#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
27#[derive(Clone, PartialEq, Debug, Eq, Hash)]
28pub enum RangeFunction {
29    IntRange {
30        step: i64,
31        dtype: DataType,
32    },
33    IntRanges,
34    LinearSpace {
35        closed: ClosedInterval,
36    },
37    LinearSpaces {
38        closed: ClosedInterval,
39        array_width: Option<usize>,
40    },
41    #[cfg(feature = "dtype-date")]
42    DateRange {
43        interval: Duration,
44        closed: ClosedWindow,
45    },
46    #[cfg(feature = "dtype-date")]
47    DateRanges {
48        interval: Duration,
49        closed: ClosedWindow,
50    },
51    #[cfg(feature = "dtype-datetime")]
52    DatetimeRange {
53        interval: Duration,
54        closed: ClosedWindow,
55        time_unit: Option<TimeUnit>,
56        time_zone: Option<TimeZone>,
57    },
58    #[cfg(feature = "dtype-datetime")]
59    DatetimeRanges {
60        interval: Duration,
61        closed: ClosedWindow,
62        time_unit: Option<TimeUnit>,
63        time_zone: Option<TimeZone>,
64    },
65    #[cfg(feature = "dtype-time")]
66    TimeRange {
67        interval: Duration,
68        closed: ClosedWindow,
69    },
70    #[cfg(feature = "dtype-time")]
71    TimeRanges {
72        interval: Duration,
73        closed: ClosedWindow,
74    },
75}
76
77fn map_linspace_dtype(mapper: &FieldsMapper) -> PolarsResult<DataType> {
78    let fields = mapper.args();
79    let start_dtype = fields[0].dtype();
80    let end_dtype = fields[1].dtype();
81    Ok(match (start_dtype, end_dtype) {
82        (&DataType::Float32, &DataType::Float32) => DataType::Float32,
83        // A linear space of a Date produces a sequence of Datetimes
84        (dt1, dt2) if dt1.is_temporal() && dt1 == dt2 => {
85            if dt1 == &DataType::Date {
86                DataType::Datetime(TimeUnit::Milliseconds, None)
87            } else {
88                dt1.clone()
89            }
90        },
91        (dt1, dt2) if !dt1.is_primitive_numeric() || !dt2.is_primitive_numeric() => {
92            polars_bail!(ComputeError:
93                "'start' and 'end' have incompatible dtypes, got {:?} and {:?}",
94                dt1, dt2
95            )
96        },
97        _ => DataType::Float64,
98    })
99}
100
101impl RangeFunction {
102    pub(super) fn get_field(&self, mapper: FieldsMapper) -> PolarsResult<Field> {
103        use RangeFunction::*;
104        match self {
105            IntRange { dtype, .. } => mapper.with_dtype(dtype.clone()),
106            IntRanges => mapper.with_dtype(DataType::List(Box::new(DataType::Int64))),
107            LinearSpace { .. } => mapper.with_dtype(map_linspace_dtype(&mapper)?),
108            LinearSpaces {
109                closed: _,
110                array_width,
111            } => {
112                let inner = Box::new(map_linspace_dtype(&mapper)?);
113                let dt = match array_width {
114                    Some(width) => DataType::Array(inner, *width),
115                    None => DataType::List(inner),
116                };
117                mapper.with_dtype(dt)
118            },
119            #[cfg(feature = "dtype-date")]
120            DateRange { .. } => mapper.with_dtype(DataType::Date),
121            #[cfg(feature = "dtype-date")]
122            DateRanges { .. } => mapper.with_dtype(DataType::List(Box::new(DataType::Date))),
123            #[cfg(feature = "dtype-datetime")]
124            DatetimeRange {
125                interval: _,
126                closed: _,
127                time_unit,
128                time_zone,
129            } => {
130                // output dtype may change based on `interval`, `time_unit`, and `time_zone`
131                let dtype =
132                    mapper.map_to_datetime_range_dtype(time_unit.as_ref(), time_zone.as_ref())?;
133                mapper.with_dtype(dtype)
134            },
135            #[cfg(feature = "dtype-datetime")]
136            DatetimeRanges {
137                interval: _,
138                closed: _,
139                time_unit,
140                time_zone,
141            } => {
142                // output dtype may change based on `interval`, `time_unit`, and `time_zone`
143                let inner_dtype =
144                    mapper.map_to_datetime_range_dtype(time_unit.as_ref(), time_zone.as_ref())?;
145                mapper.with_dtype(DataType::List(Box::new(inner_dtype)))
146            },
147            #[cfg(feature = "dtype-time")]
148            TimeRange { .. } => mapper.with_dtype(DataType::Time),
149            #[cfg(feature = "dtype-time")]
150            TimeRanges { .. } => mapper.with_dtype(DataType::List(Box::new(DataType::Time))),
151        }
152    }
153
154    pub fn function_options(&self) -> FunctionOptions {
155        use RangeFunction as R;
156        match self {
157            R::IntRange { .. } => {
158                FunctionOptions::row_separable().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
159            },
160            R::LinearSpace { .. } => {
161                FunctionOptions::row_separable().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
162            },
163            #[cfg(feature = "dtype-date")]
164            R::DateRange { .. } => {
165                FunctionOptions::row_separable().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
166            },
167            #[cfg(feature = "dtype-datetime")]
168            R::DatetimeRange { .. } => FunctionOptions::row_separable()
169                .with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
170                .with_supertyping(Default::default()),
171            #[cfg(feature = "dtype-time")]
172            R::TimeRange { .. } => {
173                FunctionOptions::row_separable().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
174            },
175            R::IntRanges => {
176                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
177            },
178            R::LinearSpaces { .. } => {
179                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
180            },
181            #[cfg(feature = "dtype-date")]
182            R::DateRanges { .. } => {
183                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
184            },
185            #[cfg(feature = "dtype-datetime")]
186            R::DatetimeRanges { .. } => FunctionOptions::elementwise()
187                .with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
188                .with_supertyping(Default::default()),
189            #[cfg(feature = "dtype-time")]
190            R::TimeRanges { .. } => {
191                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
192            },
193        }
194    }
195}
196
197impl Display for RangeFunction {
198    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
199        use RangeFunction::*;
200        let s = match self {
201            IntRange { .. } => "int_range",
202            IntRanges => "int_ranges",
203            LinearSpace { .. } => "linear_space",
204            LinearSpaces { .. } => "linear_spaces",
205            #[cfg(feature = "dtype-date")]
206            DateRange { .. } => "date_range",
207            #[cfg(feature = "temporal")]
208            DateRanges { .. } => "date_ranges",
209            #[cfg(feature = "dtype-datetime")]
210            DatetimeRange { .. } => "datetime_range",
211            #[cfg(feature = "dtype-datetime")]
212            DatetimeRanges { .. } => "datetime_ranges",
213            #[cfg(feature = "dtype-time")]
214            TimeRange { .. } => "time_range",
215            #[cfg(feature = "dtype-time")]
216            TimeRanges { .. } => "time_ranges",
217        };
218        write!(f, "{s}")
219    }
220}
221
222impl From<RangeFunction> for SpecialEq<Arc<dyn ColumnsUdf>> {
223    fn from(func: RangeFunction) -> Self {
224        use RangeFunction::*;
225        match func {
226            IntRange { step, dtype } => {
227                map_as_slice!(int_range::int_range, step, dtype.clone())
228            },
229            IntRanges => {
230                map_as_slice!(int_range::int_ranges)
231            },
232            LinearSpace { closed } => {
233                map_as_slice!(linear_space::linear_space, closed)
234            },
235            LinearSpaces {
236                closed,
237                array_width,
238            } => {
239                map_as_slice!(linear_space::linear_spaces, closed, array_width)
240            },
241            #[cfg(feature = "dtype-date")]
242            DateRange { interval, closed } => {
243                map_as_slice!(date_range::date_range, interval, closed)
244            },
245            #[cfg(feature = "dtype-date")]
246            DateRanges { interval, closed } => {
247                map_as_slice!(date_range::date_ranges, interval, closed)
248            },
249            #[cfg(feature = "dtype-datetime")]
250            DatetimeRange {
251                interval,
252                closed,
253                time_unit,
254                time_zone,
255            } => {
256                map_as_slice!(
257                    datetime_range::datetime_range,
258                    interval,
259                    closed,
260                    time_unit,
261                    time_zone.clone()
262                )
263            },
264            #[cfg(feature = "dtype-datetime")]
265            DatetimeRanges {
266                interval,
267                closed,
268                time_unit,
269                time_zone,
270            } => {
271                map_as_slice!(
272                    datetime_range::datetime_ranges,
273                    interval,
274                    closed,
275                    time_unit,
276                    time_zone.clone()
277                )
278            },
279            #[cfg(feature = "dtype-time")]
280            TimeRange { interval, closed } => {
281                map_as_slice!(time_range::time_range, interval, closed)
282            },
283            #[cfg(feature = "dtype-time")]
284            TimeRanges { interval, closed } => {
285                map_as_slice!(time_range::time_ranges, interval, closed)
286            },
287        }
288    }
289}
290
291impl From<RangeFunction> for FunctionExpr {
292    fn from(value: RangeFunction) -> Self {
293        Self::Range(value)
294    }
295}
polars_plan/dsl/function_expr/range/mod.rs

polars_plan/dsl/function_expr/range/
mod.rs
