polars_python/interop/arrow/
mod.rs1use 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}