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

polars_python/interop/arrow/
mod.rs
