polars_utils/
python_convert_registry.rs

1use std::any::Any;
2use std::ops::Deref;
3use std::sync::{Arc, LazyLock, RwLock};
4
5use pyo3::sync::PyOnceLock;
6use pyo3::types::PyAnyMethods;
7use pyo3::{Py, PyAny, PyResult, Python};
8
9pub type FromPython = Arc<dyn Fn(Py<PyAny>) -> PyResult<Box<dyn Any>> + Send + Sync>;
10pub type ToPython = Arc<dyn for<'a> Fn(&'a dyn Any) -> PyResult<Py<PyAny>> + Send + Sync>;
11
12#[derive(Clone)]
13pub struct FromPythonConvertRegistry {
14    pub file_provider_result: FromPython,
15    pub series: FromPython,
16    pub df: FromPython,
17    pub dsl_plan: FromPython,
18    pub schema: FromPython,
19}
20
21#[derive(Clone)]
22pub struct ToPythonConvertRegistry {
23    pub df: ToPython,
24    pub series: ToPython,
25    pub dsl_plan: ToPython,
26    pub schema: ToPython,
27}
28
29impl ToPythonConvertRegistry {
30    /// Convert a Rust `DataFrame` to a Python `pl.DataFrame` object.
31    pub fn df_to_wrapped_pydf(&self, df: &dyn Any) -> PyResult<Py<PyAny>> {
32        static WRAP_DF: LazyLock<Py<PyAny>> = LazyLock::new(|| {
33            Python::attach(|py| {
34                py.import("polars._utils.wrap")
35                    .unwrap()
36                    .getattr("wrap_df")
37                    .unwrap()
38                    .unbind()
39            })
40        });
41
42        let pydf = (self.df)(df)?;
43
44        Python::attach(|py| WRAP_DF.call1(py, (pydf,)))
45    }
46}
47
48#[derive(Clone)]
49pub struct PythonConvertRegistry {
50    pub from_py: FromPythonConvertRegistry,
51    pub to_py: ToPythonConvertRegistry,
52}
53
54impl PythonConvertRegistry {
55    pub fn py_file_provider_args_dataclass(&self, py: Python<'_>) -> &'static Py<PyAny> {
56        static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
57
58        CLS.get_or_init(py, || {
59            py.import("polars.io.partition")
60                .unwrap()
61                .getattr("FileProviderArgs")
62                .unwrap()
63                .unbind()
64        })
65    }
66
67    pub fn py_sinked_paths_callback_args_dataclass(&self, py: Python<'_>) -> &'static Py<PyAny> {
68        static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
69
70        CLS.get_or_init(py, || {
71            py.import("polars.io.partition")
72                .unwrap()
73                .getattr("SinkedPathsCallbackArgs")
74                .unwrap()
75                .unbind()
76        })
77    }
78
79    pub fn py_sinked_path_dataclass(&self, py: Python<'_>) -> &'static Py<PyAny> {
80        static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
81
82        CLS.get_or_init(py, || {
83            py.import("polars.io.partition")
84                .unwrap()
85                .getattr("SinkedPath")
86                .unwrap()
87                .unbind()
88        })
89    }
90
91    pub fn py_iceberg_sink_state_class(&self, py: Python<'_>) -> &'static Py<PyAny> {
92        static CLS: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
93
94        CLS.get_or_init(py, || {
95            py.import("polars.io.iceberg._sink")
96                .unwrap()
97                .getattr("IcebergSinkState")
98                .unwrap()
99                .unbind()
100        })
101    }
102}
103
104static PYTHON_CONVERT_REGISTRY: LazyLock<RwLock<Option<PythonConvertRegistry>>> =
105    LazyLock::new(Default::default);
106
107pub fn get_python_convert_registry() -> PythonConvertRegistry {
108    PYTHON_CONVERT_REGISTRY
109        .deref()
110        .read()
111        .unwrap()
112        .as_ref()
113        .unwrap()
114        .clone()
115}
116
117pub fn register_converters(registry: PythonConvertRegistry) {
118    *PYTHON_CONVERT_REGISTRY.deref().write().unwrap() = Some(registry);
119}
polars_utils/python_convert_registry.rs

polars_utils/
python_convert_registry.rs
