polars-python 0.52.0

Enable running Polars workloads in Python
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

use polars::prelude::{ArrowDataType, DataType};
use polars_error::polars_err;
use pyo3::exceptions::PyValueError;
use pyo3::types::{PyAnyMethods, PyTuple};
use pyo3::{Bound, IntoPyObject, PyAny, PyResult, intern, pyfunction};

use crate::interop::arrow::to_rust::normalize_arrow_fields;
use crate::prelude::Wrap;
use crate::series::import_schema_pycapsule;
use crate::utils::to_py_err;

pub mod to_py;
pub mod to_rust;

#[pyfunction]
pub fn init_polars_schema_from_arrow_c_schema(
    polars_schema: Bound<PyAny>,
    schema_object: Bound<PyAny>,
) -> PyResult<()> {
    let py = polars_schema.py();

    let schema_capsule = schema_object
        .getattr(intern!(py, "__arrow_c_schema__"))?
        .call0()?;

    let field = import_schema_pycapsule(&schema_capsule.extract()?)?;
    let field = normalize_arrow_fields(&field);

    let ArrowDataType::Struct(fields) = field.dtype else {
        return Err(PyValueError::new_err(format!(
            "__arrow_c_schema__ of object passed to pl.Schema did not return struct dtype: \
            object: {}, dtype: {:?}",
            schema_object, &field.dtype
        )));
    };

    for field in fields {
        let dtype = DataType::from_arrow_field(&field);

        let name = field.name.into_pyobject(py)?;
        let dtype = Wrap(dtype).into_pyobject(py)?;

        if polars_schema.contains(&name)? {
            return Err(to_py_err(polars_err!(
                Duplicate:
                "arrow schema contained duplicate name: {}",
                name
            )));
        }

        polars_schema.set_item(name, dtype)?;
    }

    Ok(())
}

#[pyfunction]
pub fn polars_schema_field_from_arrow_c_schema(
    schema_object: Bound<PyAny>,
) -> PyResult<Bound<PyTuple>> {
    let py = schema_object.py();

    let schema_capsule = schema_object
        .getattr(intern!(py, "__arrow_c_schema__"))?
        .call0()?;

    let field = import_schema_pycapsule(&schema_capsule.extract()?)?;
    let field = normalize_arrow_fields(&field);
    let dtype = DataType::from_arrow_field(&field);

    let name = field.name.into_pyobject(py)?.into_any();
    let dtype = Wrap(dtype).into_pyobject(py)?.into_any();

    PyTuple::new(py, [name, dtype])
}