pyo3_polars/
types.rs

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

pyo3_polars/
types.rs
