polars_python/interop/arrow/
mod.rs

1use polars::prelude::{ArrowDataType, DataType};
2use polars_error::polars_err;
3use pyo3::exceptions::PyValueError;
4use pyo3::types::{PyAnyMethods, PyTuple};
5use pyo3::{Bound, IntoPyObject, PyAny, PyResult, intern, pyfunction};
6
7use crate::prelude::Wrap;
8use crate::series::import_schema_pycapsule;
9use crate::utils::to_py_err;
10
11pub mod to_py;
12pub mod to_rust;
13
14#[pyfunction]
15pub fn init_polars_schema_from_arrow_c_schema(
16    polars_schema: Bound<PyAny>,
17    schema_object: Bound<PyAny>,
18) -> PyResult<()> {
19    let py = polars_schema.py();
20
21    let schema_capsule = schema_object
22        .getattr(intern!(py, "__arrow_c_schema__"))?
23        .call0()?;
24
25    let field = import_schema_pycapsule(&schema_capsule.extract()?)?;
26
27    let ArrowDataType::Struct(fields) = field.dtype else {
28        return Err(PyValueError::new_err(format!(
29            "__arrow_c_schema__ of object passed to pl.Schema did not return struct dtype: \
30            object: {}, dtype: {:?}",
31            schema_object, &field.dtype
32        )));
33    };
34
35    for field in fields {
36        let dtype = DataType::from_arrow_field(&field);
37
38        let name = field.name.into_pyobject(py)?;
39        let dtype = Wrap(dtype).into_pyobject(py)?;
40
41        if polars_schema.contains(&name)? {
42            return Err(to_py_err(polars_err!(
43                Duplicate:
44                "arrow schema contained duplicate name: {}",
45                name
46            )));
47        }
48
49        polars_schema.set_item(name, dtype)?;
50    }
51
52    Ok(())
53}
54
55#[pyfunction]
56pub fn polars_schema_field_from_arrow_c_schema(
57    schema_object: Bound<PyAny>,
58) -> PyResult<Bound<PyTuple>> {
59    let py = schema_object.py();
60
61    let schema_capsule = schema_object
62        .getattr(intern!(py, "__arrow_c_schema__"))?
63        .call0()?;
64
65    let field = import_schema_pycapsule(&schema_capsule.extract()?)?;
66    let dtype = DataType::from_arrow_field(&field);
67
68    let name = field.name.into_pyobject(py)?.into_any();
69    let dtype = Wrap(dtype).into_pyobject(py)?.into_any();
70
71    PyTuple::new(py, [name, dtype])
72}
polars_python/interop/arrow/mod.rs

polars_python/interop/arrow/
mod.rs
