polars_python/
datatypes.rs

1use polars::prelude::*;
2use polars_core::utils::arrow::array::Utf8ViewArray;
3use polars_lazy::dsl;
4use pyo3::prelude::*;
5
6use crate::error::PyPolarsErr;
7#[cfg(feature = "object")]
8use crate::object::OBJECT_NAME;
9use crate::{PyExpr, Wrap};
10
11// Don't change the order of these!
12#[repr(u8)]
13#[derive(Clone)]
14pub(crate) enum PyDataType {
15    Int8,
16    Int16,
17    Int32,
18    Int64,
19    UInt8,
20    UInt16,
21    UInt32,
22    UInt64,
23    Float32,
24    Float64,
25    Bool,
26    String,
27    List,
28    Date,
29    Datetime(TimeUnit, Option<TimeZone>),
30    Duration(TimeUnit),
31    Time,
32    #[cfg(feature = "object")]
33    Object,
34    Categorical,
35    Struct,
36    Binary,
37    Decimal(Option<usize>, usize),
38    Array(usize),
39    Enum(Utf8ViewArray),
40    Int128,
41}
42
43impl From<&DataType> for PyDataType {
44    fn from(dt: &DataType) -> Self {
45        use PyDataType::*;
46        match dt {
47            DataType::Int8 => Int8,
48            DataType::Int16 => Int16,
49            DataType::Int32 => Int32,
50            DataType::Int64 => Int64,
51            DataType::Int128 => Int128,
52            DataType::UInt8 => UInt8,
53            DataType::UInt16 => UInt16,
54            DataType::UInt32 => UInt32,
55            DataType::UInt64 => UInt64,
56            DataType::Float32 => Float32,
57            DataType::Float64 => Float64,
58            DataType::Decimal(p, s) => Decimal(*p, s.expect("unexpected null decimal scale")),
59            DataType::Boolean => Bool,
60            DataType::String => String,
61            DataType::Binary => Binary,
62            DataType::Array(_, width) => Array(*width),
63            DataType::List(_) => List,
64            DataType::Date => Date,
65            DataType::Datetime(tu, tz) => Datetime(*tu, tz.clone()),
66            DataType::Duration(tu) => Duration(*tu),
67            DataType::Time => Time,
68            #[cfg(feature = "object")]
69            DataType::Object(_) => Object,
70            DataType::Categorical(_, _) => Categorical,
71            DataType::Enum(rev_map, _) => Enum(rev_map.as_ref().unwrap().get_categories().clone()),
72            DataType::Struct(_) => Struct,
73            DataType::Null | DataType::Unknown(_) | DataType::BinaryOffset => {
74                panic!("null or unknown not expected here")
75            },
76        }
77    }
78}
79
80impl From<DataType> for PyDataType {
81    fn from(dt: DataType) -> Self {
82        (&dt).into()
83    }
84}
85
86impl From<PyDataType> for DataType {
87    fn from(pdt: PyDataType) -> DataType {
88        use DataType::*;
89        match pdt {
90            PyDataType::Int8 => Int8,
91            PyDataType::Int16 => Int16,
92            PyDataType::Int32 => Int32,
93            PyDataType::Int64 => Int64,
94            PyDataType::UInt8 => UInt8,
95            PyDataType::UInt16 => UInt16,
96            PyDataType::UInt32 => UInt32,
97            PyDataType::UInt64 => UInt64,
98            PyDataType::Float32 => Float32,
99            PyDataType::Float64 => Float64,
100            PyDataType::Bool => Boolean,
101            PyDataType::String => String,
102            PyDataType::Binary => Binary,
103            PyDataType::List => List(DataType::Null.into()),
104            PyDataType::Date => Date,
105            PyDataType::Datetime(tu, tz) => Datetime(tu, tz),
106            PyDataType::Duration(tu) => Duration(tu),
107            PyDataType::Time => Time,
108            #[cfg(feature = "object")]
109            PyDataType::Object => Object(OBJECT_NAME),
110            PyDataType::Categorical => Categorical(None, Default::default()),
111            PyDataType::Enum(categories) => create_enum_dtype(categories),
112            PyDataType::Struct => Struct(vec![]),
113            PyDataType::Decimal(p, s) => Decimal(p, Some(s)),
114            PyDataType::Array(width) => Array(DataType::Null.into(), width),
115            PyDataType::Int128 => Int128,
116        }
117    }
118}
119
120impl<'py> FromPyObject<'py> for PyDataType {
121    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
122        let dt = ob.extract::<Wrap<DataType>>()?;
123        Ok(dt.0.into())
124    }
125}
126
127#[pyfunction]
128pub fn _get_dtype_max(dt: Wrap<DataType>) -> PyResult<PyExpr> {
129    let v = dt.0.max().map_err(PyPolarsErr::from)?;
130    Ok(dsl::lit(v).into())
131}
132
133#[pyfunction]
134pub fn _get_dtype_min(dt: Wrap<DataType>) -> PyResult<PyExpr> {
135    let v = dt.0.min().map_err(PyPolarsErr::from)?;
136    Ok(dsl::lit(v).into())
137}
138
139#[pyfunction]
140pub fn _known_timezones() -> PyResult<Vec<String>> {
141    use polars_time::prelude::known_timezones;
142    Ok(known_timezones()
143        .iter()
144        .map(|tz| tz.to_string())
145        .collect::<Vec<_>>())
146}
polars_python/datatypes.rs

polars_python/
datatypes.rs
