pyo3_polars/
types.rs

1use std::convert::Infallible;
2
3use arrow;
4use polars_core::datatypes::{CompatLevel, DataType};
5use polars_core::prelude::*;
6use polars_core::utils::materialize_dyn_int;
7#[cfg(feature = "lazy")]
8use polars_lazy::frame::LazyFrame;
9#[cfg(feature = "lazy")]
10use polars_plan::dsl::DslPlan;
11#[cfg(feature = "lazy")]
12use polars_plan::dsl::Expr;
13#[cfg(feature = "lazy")]
14use polars_utils::pl_serialize;
15use pyo3::exceptions::{PyTypeError, PyValueError};
16use pyo3::ffi::Py_uintptr_t;
17use pyo3::intern;
18use pyo3::prelude::*;
19use pyo3::pybacked::PyBackedStr;
20#[cfg(feature = "lazy")]
21use pyo3::types::PyBytes;
22#[cfg(feature = "dtype-struct")]
23use pyo3::types::PyList;
24use pyo3::types::{PyDict, PyString};
25
26use super::*;
27use crate::error::PyPolarsErr;
28use crate::ffi::to_py::to_py_array;
29
30#[cfg(feature = "dtype-categorical")]
31pub(crate) fn get_series(obj: &Bound<'_, PyAny>) -> PyResult<Series> {
32    let s = obj.getattr(intern!(obj.py(), "_s"))?;
33    Ok(s.extract::<PySeries>()?.0)
34}
35
36#[repr(transparent)]
37#[derive(Debug, Clone)]
38/// A wrapper around a [`Series`] that can be converted to and from python with `pyo3`.
39pub struct PySeries(pub Series);
40
41#[repr(transparent)]
42#[derive(Debug, Clone)]
43/// A wrapper around a [`DataFrame`] that can be converted to and from python with `pyo3`.
44pub struct PyDataFrame(pub DataFrame);
45
46#[cfg(feature = "lazy")]
47#[repr(transparent)]
48#[derive(Clone)]
49/// A wrapper around a [`DataFrame`] that can be converted to and from python with `pyo3`.
50/// # Warning
51/// If the [`LazyFrame`] contains in memory data,
52/// such as a [`DataFrame`] this will be serialized/deserialized.
53///
54/// It is recommended to only have `LazyFrame`s that scan data
55/// from disk
56pub struct PyLazyFrame(pub LazyFrame);
57
58#[cfg(feature = "lazy")]
59#[repr(transparent)]
60#[derive(Clone)]
61pub struct PyExpr(pub Expr);
62
63#[repr(transparent)]
64#[derive(Clone)]
65pub struct PySchema(pub SchemaRef);
66
67#[repr(transparent)]
68#[derive(Clone)]
69pub struct PyDataType(pub DataType);
70
71#[repr(transparent)]
72#[derive(Clone, Copy)]
73pub struct PyTimeUnit(TimeUnit);
74
75#[repr(transparent)]
76#[derive(Clone)]
77pub struct PyField(Field);
78
79impl<'py> FromPyObject<'py> for PyField {
80    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
81        let py = ob.py();
82        let name = ob
83            .getattr(intern!(py, "name"))?
84            .str()?
85            .extract::<PyBackedStr>()?;
86        let dtype = ob.getattr(intern!(py, "dtype"))?.extract::<PyDataType>()?;
87        let name: &str = name.as_ref();
88        Ok(PyField(Field::new(name.into(), dtype.0)))
89    }
90}
91
92impl<'py> FromPyObject<'py> for PyTimeUnit {
93    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
94        let parsed = match &*ob.extract::<PyBackedStr>()? {
95            "ns" => TimeUnit::Nanoseconds,
96            "us" => TimeUnit::Microseconds,
97            "ms" => TimeUnit::Milliseconds,
98            v => {
99                return Err(PyValueError::new_err(format!(
100                    "`time_unit` must be one of {{'ns', 'us', 'ms'}}, got {v}",
101                )));
102            },
103        };
104        Ok(PyTimeUnit(parsed))
105    }
106}
107
108impl<'py> IntoPyObject<'py> for PyTimeUnit {
109    type Target = PyString;
110    type Output = Bound<'py, Self::Target>;
111    type Error = Infallible;
112
113    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
114        let time_unit = match self.0 {
115            TimeUnit::Nanoseconds => "ns",
116            TimeUnit::Microseconds => "us",
117            TimeUnit::Milliseconds => "ms",
118        };
119        time_unit.into_pyobject(py)
120    }
121}
122
123impl From<PyDataFrame> for DataFrame {
124    fn from(value: PyDataFrame) -> Self {
125        value.0
126    }
127}
128
129impl From<PySeries> for Series {
130    fn from(value: PySeries) -> Self {
131        value.0
132    }
133}
134
135#[cfg(feature = "lazy")]
136impl From<PyLazyFrame> for LazyFrame {
137    fn from(value: PyLazyFrame) -> Self {
138        value.0
139    }
140}
141
142impl From<PySchema> for SchemaRef {
143    fn from(value: PySchema) -> Self {
144        value.0
145    }
146}
147
148impl AsRef<Series> for PySeries {
149    fn as_ref(&self) -> &Series {
150        &self.0
151    }
152}
153
154impl AsRef<DataFrame> for PyDataFrame {
155    fn as_ref(&self) -> &DataFrame {
156        &self.0
157    }
158}
159
160#[cfg(feature = "lazy")]
161impl AsRef<LazyFrame> for PyLazyFrame {
162    fn as_ref(&self) -> &LazyFrame {
163        &self.0
164    }
165}
166
167impl AsRef<Schema> for PySchema {
168    fn as_ref(&self) -> &Schema {
169        self.0.as_ref()
170    }
171}
172
173impl<'a> FromPyObject<'a> for PySeries {
174    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
175        let ob = ob.call_method0("rechunk")?;
176
177        let name = ob.getattr("name")?;
178        let py_name = name.str()?;
179        let name = py_name.to_cow()?;
180
181        let kwargs = PyDict::new(ob.py());
182        if let Ok(compat_level) = ob.call_method0("_newest_compat_level") {
183            // Choose the maximum supported between both us and Python's compatibility level.
184            let compat_level = compat_level.extract().unwrap();
185            let compat_level =
186                CompatLevel::with_level(compat_level).unwrap_or(CompatLevel::newest());
187            let compat_level_type = POLARS_INTERCHANGE
188                .bind(ob.py())
189                .getattr("CompatLevel")
190                .unwrap();
191            let py_compat_level =
192                compat_level_type.call_method1("_with_version", (compat_level.get_level(),))?;
193            kwargs.set_item("compat_level", py_compat_level)?;
194        }
195        let arr = ob.call_method("to_arrow", (), Some(&kwargs))?;
196        let arr = ffi::to_rust::array_to_rust(&arr)?;
197        let name = name.as_ref();
198        Ok(PySeries(
199            Series::try_from((PlSmallStr::from(name), arr)).map_err(PyPolarsErr::from)?,
200        ))
201    }
202}
203
204impl<'a> FromPyObject<'a> for PyDataFrame {
205    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
206        let series = ob.call_method0("get_columns")?;
207        let n = ob.getattr("width")?.extract::<usize>()?;
208        let mut columns = Vec::with_capacity(n);
209        for pyseries in series.try_iter()? {
210            let pyseries = pyseries?;
211            let s = pyseries.extract::<PySeries>()?.0;
212            columns.push(s.into_column());
213        }
214        unsafe {
215            Ok(PyDataFrame(DataFrame::new_no_checks_height_from_first(
216                columns,
217            )))
218        }
219    }
220}
221
222#[cfg(feature = "lazy")]
223impl<'a> FromPyObject<'a> for PyLazyFrame {
224    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
225        let s = ob.call_method0("__getstate__")?;
226        let b = s.extract::<Bound<'_, PyBytes>>()?;
227        let b = b.as_bytes();
228
229        let lp = DslPlan::deserialize_versioned(b).map_err(
230            |e| PyPolarsErr::Other(
231                format!("Error when deserializing LazyFrame. This may be due to mismatched polars versions. {e}")
232            ))
233            ?;
234
235        Ok(PyLazyFrame(LazyFrame::from(lp)))
236    }
237}
238
239#[cfg(feature = "lazy")]
240impl<'a> FromPyObject<'a> for PyExpr {
241    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
242        let s = ob.call_method0("__getstate__")?.extract::<Vec<u8>>()?;
243
244        let e: Expr = pl_serialize::SerializeOptions::default()
245            .deserialize_from_reader::<Expr, &[u8], false>(&*s)
246            .map_err(
247            |e| PyPolarsErr::Other(
248                format!("Error when deserializing 'Expr'. This may be due to mismatched polars versions. {e}")
249            )
250        )?;
251
252        Ok(PyExpr(e))
253    }
254}
255
256impl<'py> IntoPyObject<'py> for PySeries {
257    type Target = PyAny;
258    type Output = Bound<'py, Self::Target>;
259    type Error = PyErr;
260
261    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
262        let polars = POLARS.bind(py);
263        let s = SERIES.bind(py);
264        match s
265            .getattr("_import_arrow_from_c")
266            .or_else(|_| s.getattr("_import_from_c"))
267        {
268            // Go via polars
269            Ok(import_arrow_from_c) => {
270                // Get supported compatibility level
271                let compat_level = CompatLevel::with_level(
272                    s.getattr("_newest_compat_level")
273                        .map_or(1, |newest_compat_level| {
274                            newest_compat_level.call0().unwrap().extract().unwrap()
275                        }),
276                )
277                .unwrap_or(CompatLevel::newest());
278                // Prepare pointers on the heap.
279                let mut chunk_ptrs = Vec::with_capacity(self.0.n_chunks());
280                for i in 0..self.0.n_chunks() {
281                    let array = self.0.to_arrow(i, compat_level);
282                    let schema = Box::new(arrow::ffi::export_field_to_c(&ArrowField::new(
283                        "".into(),
284                        array.dtype().clone(),
285                        true,
286                    )));
287                    let array = Box::new(arrow::ffi::export_array_to_c(array.clone()));
288
289                    let schema_ptr: *const arrow::ffi::ArrowSchema = Box::leak(schema);
290                    let array_ptr: *const arrow::ffi::ArrowArray = Box::leak(array);
291
292                    chunk_ptrs.push((schema_ptr as Py_uintptr_t, array_ptr as Py_uintptr_t))
293                }
294
295                // Somehow we need to clone the Vec, because pyo3 doesn't accept a slice here.
296                let pyseries = import_arrow_from_c
297                    .call1((self.0.name().as_str(), chunk_ptrs.clone()))
298                    .unwrap();
299                // Deallocate boxes
300                for (schema_ptr, array_ptr) in chunk_ptrs {
301                    let schema_ptr = schema_ptr as *mut arrow::ffi::ArrowSchema;
302                    let array_ptr = array_ptr as *mut arrow::ffi::ArrowArray;
303                    unsafe {
304                        // We can drop both because the `schema` isn't read in an owned matter on the other side.
305                        let _ = Box::from_raw(schema_ptr);
306
307                        // The array is `ptr::read_unaligned` so there are two owners.
308                        // We drop the box, and forget the content so the other process is the owner.
309                        let array = Box::from_raw(array_ptr);
310                        // We must forget because the other process will call the release callback.
311                        // Read *array as Box::into_inner
312                        let array = *array;
313                        std::mem::forget(array);
314                    }
315                }
316
317                Ok(pyseries)
318            },
319            // Go via pyarrow
320            Err(_) => {
321                let s = self.0.rechunk();
322                let name = s.name().as_str();
323                let arr = s.to_arrow(0, CompatLevel::oldest());
324                let pyarrow = py.import("pyarrow").expect("pyarrow not installed");
325
326                let arg = to_py_array(arr, pyarrow).unwrap();
327                let s = polars.call_method1("from_arrow", (arg,)).unwrap();
328                let s = s.call_method1("rename", (name,)).unwrap();
329                Ok(s)
330            },
331        }
332    }
333}
334
335impl<'py> IntoPyObject<'py> for PyDataFrame {
336    type Target = PyAny;
337    type Output = Bound<'py, Self::Target>;
338    type Error = PyErr;
339
340    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
341        let pyseries = self
342            .0
343            .get_columns()
344            .iter()
345            .map(|s| PySeries(s.as_materialized_series().clone()).into_pyobject(py))
346            .collect::<PyResult<Vec<_>>>()?;
347
348        let polars = POLARS.bind(py);
349        polars.call_method1("DataFrame", (pyseries,))
350    }
351}
352
353#[cfg(feature = "lazy")]
354impl<'py> IntoPyObject<'py> for PyLazyFrame {
355    type Target = PyAny;
356    type Output = Bound<'py, Self::Target>;
357    type Error = PyErr;
358
359    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
360        use polars::prelude::PlanSerializationContext;
361
362        let polars = POLARS.bind(py);
363        let cls = polars.getattr("LazyFrame")?;
364        let instance = cls.call_method1(intern!(py, "__new__"), (&cls,)).unwrap();
365
366        let mut v = vec![];
367        self.0
368            .logical_plan
369            .serialize_versioned(&mut v, PlanSerializationContext::default())
370            .unwrap();
371        instance.call_method1("__setstate__", (&v,))?;
372        Ok(instance)
373    }
374}
375
376#[cfg(feature = "lazy")]
377impl<'py> IntoPyObject<'py> for PyExpr {
378    type Target = PyAny;
379    type Output = Bound<'py, Self::Target>;
380    type Error = PyErr;
381
382    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
383        let polars = POLARS.bind(py);
384        let cls = polars.getattr("Expr")?;
385        let instance = cls.call_method1(intern!(py, "__new__"), (&cls,))?;
386
387        let buf = pl_serialize::SerializeOptions::default()
388            .serialize_to_bytes::<Expr, false>(&self.0)
389            .unwrap();
390
391        instance
392            .call_method1("__setstate__", (&buf,))
393            .map_err(|err| {
394                let msg = format!("deserialization failed: {err}");
395                PyValueError::new_err(msg)
396            })
397    }
398}
399
400#[cfg(feature = "dtype-categorical")]
401pub(crate) fn to_series(py: Python, s: PySeries) -> Py<PyAny> {
402    let series = SERIES.bind(py);
403    let constructor = series
404        .getattr(intern!(series.py(), "_from_pyseries"))
405        .unwrap();
406    constructor
407        .call1((s,))
408        .unwrap()
409        .into_pyobject(py)
410        .unwrap()
411        .into()
412}
413
414impl<'py> IntoPyObject<'py> for PyDataType {
415    type Target = PyAny;
416    type Output = Bound<'py, Self::Target>;
417    type Error = PyErr;
418
419    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
420        let pl = POLARS.bind(py);
421
422        match &self.0 {
423            DataType::Int8 => {
424                let class = pl.getattr(intern!(py, "Int8")).unwrap();
425                class.call0()
426            },
427            DataType::Int16 => {
428                let class = pl.getattr(intern!(py, "Int16")).unwrap();
429                class.call0()
430            },
431            DataType::Int32 => {
432                let class = pl.getattr(intern!(py, "Int32")).unwrap();
433                class.call0()
434            },
435            DataType::Int64 => {
436                let class = pl.getattr(intern!(py, "Int64")).unwrap();
437                class.call0()
438            },
439            DataType::Int128 => {
440                let class = pl.getattr(intern!(py, "Int128")).unwrap();
441                class.call0()
442            },
443            DataType::UInt8 => {
444                let class = pl.getattr(intern!(py, "UInt8")).unwrap();
445                class.call0()
446            },
447            DataType::UInt16 => {
448                let class = pl.getattr(intern!(py, "UInt16")).unwrap();
449                class.call0()
450            },
451            DataType::UInt32 => {
452                let class = pl.getattr(intern!(py, "UInt32")).unwrap();
453                class.call0()
454            },
455            DataType::UInt64 => {
456                let class = pl.getattr(intern!(py, "UInt64")).unwrap();
457                class.call0()
458            },
459            DataType::UInt128 => {
460                let class = pl.getattr(intern!(py, "UInt128")).unwrap();
461                class.call0()
462            },
463            DataType::Float32 => {
464                let class = pl.getattr(intern!(py, "Float32")).unwrap();
465                class.call0()
466            },
467            DataType::Float64 | DataType::Unknown(UnknownKind::Float) => {
468                let class = pl.getattr(intern!(py, "Float64")).unwrap();
469                class.call0()
470            },
471            #[cfg(feature = "dtype-decimal")]
472            DataType::Decimal(precision, scale) => {
473                let class = pl.getattr(intern!(py, "Decimal")).unwrap();
474                let args = (*precision, *scale);
475                class.call1(args)
476            },
477            DataType::Boolean => {
478                let class = pl.getattr(intern!(py, "Boolean")).unwrap();
479                class.call0()
480            },
481            DataType::String | DataType::Unknown(UnknownKind::Str) => {
482                let class = pl.getattr(intern!(py, "String")).unwrap();
483                class.call0()
484            },
485            DataType::Binary => {
486                let class = pl.getattr(intern!(py, "Binary")).unwrap();
487                class.call0()
488            },
489            #[cfg(feature = "dtype-array")]
490            DataType::Array(inner, size) => {
491                let class = pl.getattr(intern!(py, "Array")).unwrap();
492                let inner = PyDataType(*inner.clone()).into_pyobject(py)?;
493                let args = (inner, *size);
494                class.call1(args)
495            },
496            DataType::List(inner) => {
497                let class = pl.getattr(intern!(py, "List")).unwrap();
498                let inner = PyDataType(*inner.clone()).into_pyobject(py)?;
499                class.call1((inner,))
500            },
501            DataType::Date => {
502                let class = pl.getattr(intern!(py, "Date")).unwrap();
503                class.call0()
504            },
505            DataType::Datetime(tu, tz) => {
506                let datetime_class = pl.getattr(intern!(py, "Datetime")).unwrap();
507                datetime_class.call1((tu.to_ascii(), tz.as_ref().map(|s| s.as_str())))
508            },
509            DataType::Duration(tu) => {
510                let duration_class = pl.getattr(intern!(py, "Duration")).unwrap();
511                duration_class.call1((tu.to_ascii(),))
512            },
513            #[cfg(feature = "object")]
514            DataType::Object(_) => {
515                let class = pl.getattr(intern!(py, "Object")).unwrap();
516                class.call0()
517            },
518            #[cfg(feature = "dtype-categorical")]
519            DataType::Categorical(_, _) => {
520                let class = pl.getattr(intern!(py, "Categorical")).unwrap();
521                class.call1(())
522            },
523            #[cfg(feature = "dtype-categorical")]
524            DataType::Enum(categories, _) => {
525                // we should always have an initialized rev_map coming from rust
526                let class = pl.getattr(intern!(py, "Enum")).unwrap();
527                let s =
528                    Series::from_arrow("category".into(), categories.categories().clone().boxed())
529                        .unwrap();
530                let series = to_series(py, PySeries(s));
531                class.call1((series,))
532            },
533            DataType::Time => pl.getattr(intern!(py, "Time")),
534            #[cfg(feature = "dtype-struct")]
535            DataType::Struct(fields) => {
536                let field_class = pl.getattr(intern!(py, "Field")).unwrap();
537                let iter = fields
538                    .iter()
539                    .map(|fld| {
540                        let name = fld.name().as_str();
541                        let dtype = PyDataType(fld.dtype().clone()).into_pyobject(py)?;
542                        field_class.call1((name, dtype))
543                    })
544                    .collect::<PyResult<Vec<_>>>()?;
545                let fields = PyList::new(py, iter)?;
546                let struct_class = pl.getattr(intern!(py, "Struct")).unwrap();
547                struct_class.call1((fields,))
548            },
549            DataType::Null => {
550                let class = pl.getattr(intern!(py, "Null")).unwrap();
551                class.call0()
552            },
553            DataType::Unknown(UnknownKind::Int(v)) => {
554                PyDataType(materialize_dyn_int(*v).dtype()).into_pyobject(py)
555            },
556            DataType::Unknown(_) => {
557                let class = pl.getattr(intern!(py, "Unknown")).unwrap();
558                class.call0()
559            },
560            DataType::BinaryOffset => {
561                panic!("this type isn't exposed to python")
562            },
563            #[allow(unreachable_patterns)]
564            _ => panic!("activate dtype"),
565        }
566    }
567}
568
569impl<'py> IntoPyObject<'py> for PySchema {
570    type Target = PyDict;
571    type Output = Bound<'py, Self::Target>;
572    type Error = PyErr;
573
574    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
575        let dict = PyDict::new(py);
576        for (k, v) in self.0.iter() {
577            dict.set_item(k.as_str(), PyDataType(v.clone()).into_pyobject(py)?)?;
578        }
579        Ok(dict)
580    }
581}
582
583impl<'py> FromPyObject<'py> for PyDataType {
584    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
585        let py = ob.py();
586        let type_name = ob.get_type().qualname()?.to_string();
587
588        let dtype = match type_name.as_ref() {
589            "DataTypeClass" => {
590                // just the class, not an object
591                let name = ob
592                    .getattr(intern!(py, "__name__"))?
593                    .str()?
594                    .extract::<PyBackedStr>()?;
595                match &*name {
596                    "Int8" => DataType::Int8,
597                    "Int16" => DataType::Int16,
598                    "Int32" => DataType::Int32,
599                    "Int64" => DataType::Int64,
600                    "Int128" => DataType::Int128,
601                    "UInt8" => DataType::UInt8,
602                    "UInt16" => DataType::UInt16,
603                    "UInt32" => DataType::UInt32,
604                    "UInt64" => DataType::UInt64,
605                    "UInt128" => DataType::UInt128,
606                    "Float32" => DataType::Float32,
607                    "Float64" => DataType::Float64,
608                    "Boolean" => DataType::Boolean,
609                    "String" => DataType::String,
610                    "Binary" => DataType::Binary,
611                    #[cfg(feature = "dtype-categorical")]
612                    "Categorical" => {
613                        DataType::Categorical(Categories::global(), Categories::global().mapping())
614                    },
615                    #[cfg(feature = "dtype-categorical")]
616                    "Enum" => {
617                        let categories = FrozenCategories::new([]).unwrap();
618                        let mapping = categories.mapping().clone();
619                        DataType::Enum(categories, mapping)
620                    },
621                    "Date" => DataType::Date,
622                    "Time" => DataType::Time,
623                    "Datetime" => DataType::Datetime(TimeUnit::Microseconds, None),
624                    "Duration" => DataType::Duration(TimeUnit::Microseconds),
625                    #[cfg(feature = "dtype-decimal")]
626                    "Decimal" => {
627                        return Err(PyTypeError::new_err("Decimal without specifying precision and scale is not a valid Polars data type".to_string()));
628                    },
629                    "List" => DataType::List(Box::new(DataType::Null)),
630                    #[cfg(feature = "dtype-array")]
631                    "Array" => DataType::Array(Box::new(DataType::Null), 0),
632                    #[cfg(feature = "dtype-struct")]
633                    "Struct" => DataType::Struct(vec![]),
634                    "Null" => DataType::Null,
635                    #[cfg(feature = "object")]
636                    "Object" => todo!(),
637                    "Unknown" => DataType::Unknown(Default::default()),
638                    dt => {
639                        return Err(PyTypeError::new_err(format!(
640                            "'{dt}' is not a Polars data type, or the plugin isn't compiled with the right features",
641                        )));
642                    },
643                }
644            },
645            "Int8" => DataType::Int8,
646            "Int16" => DataType::Int16,
647            "Int32" => DataType::Int32,
648            "Int64" => DataType::Int64,
649            "Int128" => DataType::Int128,
650            "UInt8" => DataType::UInt8,
651            "UInt16" => DataType::UInt16,
652            "UInt32" => DataType::UInt32,
653            "UInt64" => DataType::UInt64,
654            "UInt128" => DataType::UInt128,
655            "Float32" => DataType::Float32,
656            "Float64" => DataType::Float64,
657            "Boolean" => DataType::Boolean,
658            "String" => DataType::String,
659            "Binary" => DataType::Binary,
660            #[cfg(feature = "dtype-categorical")]
661            "Categorical" => {
662                DataType::Categorical(Categories::global(), Categories::global().mapping())
663            },
664            #[cfg(feature = "dtype-categorical")]
665            "Enum" => {
666                let categories = ob.getattr(intern!(py, "categories")).unwrap();
667                let s = get_series(&categories.as_borrowed())?;
668                let ca = s.str().map_err(PyPolarsErr::from)?;
669                let categories = ca.iter();
670                let categories = FrozenCategories::new(categories.map(|v| v.unwrap())).unwrap();
671                let mapping = categories.mapping().clone();
672                DataType::Enum(categories, mapping)
673            },
674            "Date" => DataType::Date,
675            "Time" => DataType::Time,
676            "Datetime" => {
677                let time_unit = ob.getattr(intern!(py, "time_unit")).unwrap();
678                let time_unit = time_unit.extract::<PyTimeUnit>()?.0;
679                let time_zone = ob.getattr(intern!(py, "time_zone")).unwrap();
680                let time_zone: Option<String> = time_zone.extract()?;
681                DataType::Datetime(time_unit, TimeZone::opt_try_new(time_zone).unwrap())
682            },
683            "Duration" => {
684                let time_unit = ob.getattr(intern!(py, "time_unit")).unwrap();
685                let time_unit = time_unit.extract::<PyTimeUnit>()?.0;
686                DataType::Duration(time_unit)
687            },
688            #[cfg(feature = "dtype-decimal")]
689            "Decimal" => {
690                let precision = ob.getattr(intern!(py, "precision"))?.extract()?;
691                let scale = ob.getattr(intern!(py, "scale"))?.extract()?;
692                DataType::Decimal(precision, scale)
693            },
694            "List" => {
695                let inner = ob.getattr(intern!(py, "inner")).unwrap();
696                let inner = inner.extract::<PyDataType>()?;
697                DataType::List(Box::new(inner.0))
698            },
699            #[cfg(feature = "dtype-array")]
700            "Array" => {
701                let inner = ob.getattr(intern!(py, "inner")).unwrap();
702                let size = ob.getattr(intern!(py, "size")).unwrap();
703                let inner = inner.extract::<PyDataType>()?;
704                let size = size.extract::<usize>()?;
705                DataType::Array(Box::new(inner.0), size)
706            },
707            #[cfg(feature = "dtype-struct")]
708            "Struct" => {
709                let fields = ob.getattr(intern!(py, "fields"))?;
710                let fields = fields
711                    .extract::<Vec<PyField>>()?
712                    .into_iter()
713                    .map(|f| f.0)
714                    .collect::<Vec<Field>>();
715                DataType::Struct(fields)
716            },
717            "Null" => DataType::Null,
718            #[cfg(feature = "object")]
719            "Object" => panic!("object not supported"),
720            "Unknown" => DataType::Unknown(Default::default()),
721            dt => {
722                return Err(PyTypeError::new_err(format!(
723                    "'{dt}' is not a Polars data type, or the plugin isn't compiled with the right features",
724                )));
725            },
726        };
727        Ok(PyDataType(dtype))
728    }
729}
pyo3_polars/types.rs

pyo3_polars/
types.rs
