polars_python/
on_startup.rs

1#![allow(unsafe_op_in_unsafe_fn)]
2use std::any::Any;
3use std::sync::OnceLock;
4
5use polars::prelude::*;
6use polars_core::chunked_array::object::builder::ObjectChunkedBuilder;
7use polars_core::chunked_array::object::registry::AnonymousObjectBuilder;
8use polars_core::chunked_array::object::{registry, set_polars_allow_extension};
9use polars_error::PolarsWarning;
10use polars_error::signals::register_polars_keyboard_interrupt_hook;
11use polars_ffi::version_0::SeriesExport;
12use polars_plan::plans::python_df_to_rust;
13use polars_utils::python_convert_registry::{FromPythonConvertRegistry, PythonConvertRegistry};
14use pyo3::prelude::*;
15use pyo3::{IntoPyObjectExt, intern};
16
17use crate::Wrap;
18use crate::dataframe::PyDataFrame;
19use crate::map::lazy::{ToSeries, call_lambda_with_series};
20use crate::prelude::ObjectValue;
21use crate::py_modules::{pl_df, pl_utils, polars, polars_rs};
22use crate::series::PySeries;
23
24fn python_function_caller_series(
25    s: Column,
26    output_dtype: Option<DataType>,
27    lambda: &PyObject,
28) -> PolarsResult<Column> {
29    Python::with_gil(|py| {
30        let object =
31            call_lambda_with_series(py, s.as_materialized_series(), Some(output_dtype), lambda)?;
32        object.to_series(py, polars(py), s.name()).map(Column::from)
33    })
34}
35
36fn python_function_caller_df(df: DataFrame, lambda: &PyObject) -> PolarsResult<DataFrame> {
37    Python::with_gil(|py| {
38        let pypolars = polars(py).bind(py);
39
40        // create a PySeries struct/object for Python
41        let mut pydf = PyDataFrame::new(df);
42        // Wrap this PySeries object in the python side Series wrapper
43        let mut python_df_wrapper = pypolars
44            .getattr("wrap_df")
45            .unwrap()
46            .call1((pydf.clone(),))
47            .unwrap();
48
49        if !python_df_wrapper
50            .getattr("_df")
51            .unwrap()
52            .is_instance(polars_rs(py).getattr(py, "PyDataFrame").unwrap().bind(py))
53            .unwrap()
54        {
55            let pldf = pl_df(py).bind(py);
56            let width = pydf.width();
57            // Don't resize the Vec to avoid calling SeriesExport's Drop impl
58            // The import takes ownership and is responsible for dropping
59            let mut columns: Vec<SeriesExport> = Vec::with_capacity(width);
60            unsafe {
61                pydf._export_columns(columns.as_mut_ptr() as usize);
62            }
63            // Wrap this PyDataFrame object in the python side DataFrame wrapper
64            python_df_wrapper = pldf
65                .getattr("_import_columns")
66                .unwrap()
67                .call1((columns.as_mut_ptr() as usize, width))
68                .unwrap();
69        }
70        // call the lambda and get a python side df wrapper
71        let result_df_wrapper = lambda.call1(py, (python_df_wrapper,))?;
72
73        // unpack the wrapper in a PyDataFrame
74        let py_pydf = result_df_wrapper.getattr(py, "_df").map_err(|_| {
75            let pytype = result_df_wrapper.bind(py).get_type();
76            PolarsError::ComputeError(
77                format!("Expected 'LazyFrame.map' to return a 'DataFrame', got a '{pytype}'",)
78                    .into(),
79            )
80        })?;
81        // Downcast to Rust
82        match py_pydf.extract::<PyDataFrame>(py) {
83            Ok(pydf) => Ok(pydf.df),
84            Err(_) => python_df_to_rust(py, result_df_wrapper.into_bound(py)),
85        }
86    })
87}
88
89fn warning_function(msg: &str, warning: PolarsWarning) {
90    Python::with_gil(|py| {
91        let warn_fn = pl_utils(py)
92            .bind(py)
93            .getattr(intern!(py, "_polars_warn"))
94            .unwrap();
95
96        if let Err(e) = warn_fn.call1((msg, Wrap(warning).into_pyobject(py).unwrap())) {
97            eprintln!("{e}")
98        }
99    });
100}
101
102static POLARS_REGISTRY_INIT_LOCK: OnceLock<()> = OnceLock::new();
103
104/// # Safety
105/// Caller must ensure that no other threads read the objects set by this registration.
106pub unsafe fn register_startup_deps(catch_keyboard_interrupt: bool) {
107    // TODO: should we throw an error if we try to initialize while already initialized?
108    POLARS_REGISTRY_INIT_LOCK.get_or_init(|| {
109        set_polars_allow_extension(true);
110
111        // Stack frames can get really large in debug mode.
112        #[cfg(debug_assertions)]
113        {
114            recursive::set_minimum_stack_size(1024 * 1024);
115            recursive::set_stack_allocation_size(1024 * 1024 * 16);
116        }
117
118        // Register object type builder.
119        let object_builder = Box::new(|name: PlSmallStr, capacity: usize| {
120            Box::new(ObjectChunkedBuilder::<ObjectValue>::new(name, capacity))
121                as Box<dyn AnonymousObjectBuilder>
122        });
123
124        let object_converter = Arc::new(|av: AnyValue| {
125            let object = Python::with_gil(|py| ObjectValue {
126                inner: Wrap(av).into_py_any(py).unwrap(),
127            });
128            Box::new(object) as Box<dyn Any>
129        });
130        let pyobject_converter = Arc::new(|av: AnyValue| {
131            let object = Python::with_gil(|py| Wrap(av).into_py_any(py).unwrap());
132            Box::new(object) as Box<dyn Any>
133        });
134
135        polars_utils::python_convert_registry::register_converters(PythonConvertRegistry {
136            from_py: FromPythonConvertRegistry {
137                partition_target_cb_result: Arc::new(|py_f| {
138                    Python::with_gil(|py| {
139                        Ok(Box::new(
140                            py_f.extract::<Wrap<polars_plan::dsl::PartitionTargetCallbackResult>>(
141                                py,
142                            )?
143                            .0,
144                        ) as _)
145                    })
146                }),
147                series: Arc::new(|py_f| {
148                    Python::with_gil(|py| Ok(Box::new(py_f.extract::<PySeries>(py)?.series) as _))
149                }),
150            },
151            to_py: polars_utils::python_convert_registry::ToPythonConvertRegistry {
152                df: Arc::new(|df| {
153                    Python::with_gil(|py| PyDataFrame::new(*df.downcast().unwrap()).into_py_any(py))
154                }),
155                series: Arc::new(|series| {
156                    Python::with_gil(|py| {
157                        PySeries::new(*series.downcast().unwrap()).into_py_any(py)
158                    })
159                }),
160            },
161        });
162
163        let object_size = size_of::<ObjectValue>();
164        let physical_dtype = ArrowDataType::FixedSizeBinary(object_size);
165        registry::register_object_builder(
166            object_builder,
167            object_converter,
168            pyobject_converter,
169            physical_dtype,
170        );
171        // Register SERIES UDF.
172        python_dsl::CALL_COLUMNS_UDF_PYTHON = Some(python_function_caller_series);
173        // Register DATAFRAME UDF.
174        python_dsl::CALL_DF_UDF_PYTHON = Some(python_function_caller_df);
175        // Register warning function for `polars_warn!`.
176        polars_error::set_warning_function(warning_function);
177
178        if catch_keyboard_interrupt {
179            register_polars_keyboard_interrupt_hook();
180        }
181    });
182}
polars_python/on_startup.rs

polars_python/
on_startup.rs
