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

polars_utils/
python_convert_registry.rs
