polars_plan/dsl/function_expr/
range.rs

1use std::fmt;
2
3use polars_core::prelude::*;
4use polars_ops::series::ClosedInterval;
5#[cfg(feature = "temporal")]
6use polars_time::{ClosedWindow, Duration};
7#[cfg(feature = "serde")]
8use serde::{Deserialize, Serialize};
9
10use super::{DataTypeExpr, FunctionExpr};
11
12#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
13#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
14#[derive(Clone, PartialEq, Debug, Hash)]
15pub enum RangeFunction {
16    IntRange {
17        step: i64,
18        dtype: DataTypeExpr,
19    },
20    IntRanges {
21        dtype: DataTypeExpr,
22    },
23    LinearSpace {
24        closed: ClosedInterval,
25    },
26    LinearSpaces {
27        closed: ClosedInterval,
28        array_width: Option<usize>,
29    },
30    #[cfg(feature = "dtype-date")]
31    DateRange {
32        interval: Duration,
33        closed: ClosedWindow,
34    },
35    #[cfg(feature = "dtype-date")]
36    DateRanges {
37        interval: Duration,
38        closed: ClosedWindow,
39    },
40    #[cfg(feature = "dtype-datetime")]
41    DatetimeRange {
42        interval: Duration,
43        closed: ClosedWindow,
44        time_unit: Option<TimeUnit>,
45        time_zone: Option<TimeZone>,
46    },
47    #[cfg(feature = "dtype-datetime")]
48    DatetimeRanges {
49        interval: Duration,
50        closed: ClosedWindow,
51        time_unit: Option<TimeUnit>,
52        time_zone: Option<TimeZone>,
53    },
54    #[cfg(feature = "dtype-time")]
55    TimeRange {
56        interval: Duration,
57        closed: ClosedWindow,
58    },
59    #[cfg(feature = "dtype-time")]
60    TimeRanges {
61        interval: Duration,
62        closed: ClosedWindow,
63    },
64}
65
66impl From<RangeFunction> for FunctionExpr {
67    fn from(value: RangeFunction) -> Self {
68        Self::Range(value)
69    }
70}
71
72impl fmt::Display for RangeFunction {
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::fmt::Result {
74        use RangeFunction::*;
75        let s = match self {
76            IntRange { .. } => "int_range",
77            IntRanges { .. } => "int_ranges",
78            LinearSpace { .. } => "linear_space",
79            LinearSpaces { .. } => "linear_spaces",
80            #[cfg(feature = "dtype-date")]
81            DateRange { .. } => "date_range",
82            #[cfg(feature = "temporal")]
83            DateRanges { .. } => "date_ranges",
84            #[cfg(feature = "dtype-datetime")]
85            DatetimeRange { .. } => "datetime_range",
86            #[cfg(feature = "dtype-datetime")]
87            DatetimeRanges { .. } => "datetime_ranges",
88            #[cfg(feature = "dtype-time")]
89            TimeRange { .. } => "time_range",
90            #[cfg(feature = "dtype-time")]
91            TimeRanges { .. } => "time_ranges",
92        };
93        write!(f, "{s}")
94    }
95}