Skip to content

Commit

Permalink
Store static module imports in GILOnceCell
Browse files Browse the repository at this point in the history
  • Loading branch information
@bschoenmaeckers
bschoenmaeckers committed Dec 4, 2024
1 parent b42e04d commit 8aa59a8
Show file tree
Hide file tree
Showing 12 changed files with 72 additions and 60 deletions.
8 changes: 4 additions & 4 deletions crates/polars-python/src/conversion/any_value.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use super::datetime::{
};
use super::{decimal_to_digits, struct_dict, ObjectValue, Wrap};
use crate::error::PyPolarsErr;
use crate::py_modules::{SERIES, UTILS};
use crate::py_modules::{pl_series, pl_utils};
use crate::series::PySeries;

impl IntoPy<PyObject> for Wrap<AnyValue<'_>> {
Expand All @@ -43,7 +43,7 @@ impl<'py> FromPyObject<'py> for Wrap<AnyValue<'py>> {
}

pub(crate) fn any_value_into_py_object(av: AnyValue, py: Python) -> PyObject {
let utils = UTILS.bind(py);
let utils = pl_utils(py).bind(py);
match av {
AnyValue::UInt8(v) => v.into_py(py),
AnyValue::UInt16(v) => v.into_py(py),
Expand Down Expand Up @@ -266,7 +266,7 @@ pub(crate) fn py_object_to_any_value<'py>(
// Probably needs to wait for
// https://github.com/pola-rs/polars/issues/16199 to do it a faster way.
Python::with_gil(|py| {
let date = UTILS
let date = pl_utils(py)
.bind(py)
.getattr(intern!(py, "datetime_to_int"))
.unwrap()
Expand Down Expand Up @@ -353,7 +353,7 @@ pub(crate) fn py_object_to_any_value<'py>(
let py = ob.py();
let kwargs = PyDict::new(py);
kwargs.set_item("strict", strict)?;
let s = SERIES.call_bound(py, (ob,), Some(&kwargs))?;
let s = pl_series(py).call_bound(py, (ob,), Some(&kwargs))?;
get_list_from_series(s.bind(py), strict)
}

Expand Down
6 changes: 3 additions & 3 deletions crates/polars-python/src/conversion/chunked_array.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use super::datetime::{
};
use super::{decimal_to_digits, struct_dict};
use crate::prelude::*;
use crate::py_modules::UTILS;
use crate::py_modules::pl_utils;

impl ToPyObject for Wrap<&StringChunked> {
fn to_object(&self, py: Python) -> PyObject {
Expand Down Expand Up @@ -61,7 +61,7 @@ impl ToPyObject for Wrap<&DatetimeChunked> {
if time_zone.is_some() {
// Switch to more efficient code path in
// https://github.com/pola-rs/polars/issues/16199
let utils = UTILS.bind(py);
let utils = pl_utils(py).bind(py);
let convert = utils.getattr(intern!(py, "to_py_datetime")).unwrap();
let time_unit = self.0.time_unit().to_ascii();
let time_zone = time_zone.as_deref().to_object(py);
Expand Down Expand Up @@ -113,7 +113,7 @@ pub(crate) fn decimal_to_pyobject_iter<'a>(
py: Python<'a>,
ca: &'a DecimalChunked,
) -> impl ExactSizeIterator<Item = Option<Bound<'a, PyAny>>> {
let utils = UTILS.bind(py);
let utils = pl_utils(py).bind(py);
let convert = utils.getattr(intern!(py, "to_py_decimal")).unwrap();
let py_scale = (-(ca.scale() as i32)).to_object(py);
// if we don't know precision, the only safe bet is to set it to 39
Expand Down
6 changes: 3 additions & 3 deletions crates/polars-python/src/conversion/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ use crate::file::{get_python_scan_source_input, PythonScanSourceInput};
#[cfg(feature = "object")]
use crate::object::OBJECT_NAME;
use crate::prelude::*;
use crate::py_modules::{POLARS, SERIES};
use crate::py_modules::{pl_series, polars};
use crate::series::PySeries;
use crate::{PyDataFrame, PyLazyFrame};

Expand Down Expand Up @@ -109,7 +109,7 @@ pub(crate) fn get_series(obj: &Bound<'_, PyAny>) -> PyResult<Series> {
}

pub(crate) fn to_series(py: Python, s: PySeries) -> PyObject {
let series = SERIES.bind(py);
let series = pl_series(py).bind(py);
let constructor = series
.getattr(intern!(series.py(), "_from_pyseries"))
.unwrap();
Expand Down Expand Up @@ -184,7 +184,7 @@ fn decimal_to_digits(v: i128, buf: &mut [u128; 3]) -> usize {

impl ToPyObject for Wrap<DataType> {
fn to_object(&self, py: Python) -> PyObject {
let pl = POLARS.bind(py);
let pl = polars(py).bind(py);

match &self.0 {
DataType::Int8 => {
Expand Down
3 changes: 2 additions & 1 deletion crates/polars-python/src/dataframe/general.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ use crate::map::dataframe::{
apply_lambda_with_string_out_type,
};
use crate::prelude::strings_to_pl_smallstr;
use crate::py_modules::polars;
use crate::series::{PySeries, ToPySeries, ToSeries};
use crate::{PyExpr, PyLazyFrame};

Expand Down Expand Up @@ -402,7 +403,7 @@ impl PyDataFrame {

let function = move |df: DataFrame| {
Python::with_gil(|py| {
let pypolars = PyModule::import(py, "polars").unwrap();
let pypolars = polars(py).bind(py);
let pydf = PyDataFrame::new(df);
let python_df_wrapper =
pypolars.getattr("wrap_df").unwrap().call1((pydf,)).unwrap();
Expand Down
3 changes: 2 additions & 1 deletion crates/polars-python/src/lazygroupby.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ use pyo3::prelude::*;
use crate::conversion::Wrap;
use crate::error::PyPolarsErr;
use crate::expr::ToExprs;
use crate::py_modules::polars;
use crate::{PyDataFrame, PyExpr, PyLazyFrame};

#[pyclass]
Expand Down Expand Up @@ -51,7 +52,7 @@ impl PyLazyGroupBy {
let function = move |df: DataFrame| {
Python::with_gil(|py| {
// get the pypolars module
let pypolars = PyModule::import(py, "polars").unwrap();
let pypolars = polars(py).bind(py);

// create a PyDataFrame struct/object for Python
let pydf = PyDataFrame::new(df);
Expand Down
15 changes: 8 additions & 7 deletions crates/polars-python/src/map/lazy.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use pyo3::ffi::Py_uintptr_t;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};

use crate::py_modules::POLARS;
use crate::py_modules::polars;
use crate::series::PySeries;
use crate::{PyExpr, Wrap};

Expand Down Expand Up @@ -90,7 +90,7 @@ pub(crate) fn call_lambda_with_series(
s: Series,
lambda: &PyObject,
) -> PyResult<PyObject> {
let pypolars = POLARS.bind(py);
let pypolars = polars(py).bind(py);

// create a PySeries struct/object for Python
let pyseries = PySeries::new(s);
Expand All @@ -112,7 +112,7 @@ pub(crate) fn binary_lambda(
) -> PolarsResult<Option<Series>> {
Python::with_gil(|py| {
// get the pypolars module
let pypolars = PyModule::import(py, "polars").unwrap();
let pypolars = polars(py).bind(py);
// create a PySeries struct/object for Python
let pyseries_a = PySeries::new(a);
let pyseries_b = PySeries::new(b);
Expand Down Expand Up @@ -151,7 +151,8 @@ pub(crate) fn binary_lambda(
let s = out.select_at_idx(0).unwrap().clone();
PySeries::new(s.take_materialized_series())
} else {
return Some(result_series_wrapper.to_series(py, &pypolars.unbind(), "")).transpose();
return Some(result_series_wrapper.to_series(py, pypolars.as_unbound(), ""))
.transpose();
};

// Finally get the actual Series
Expand All @@ -178,9 +179,9 @@ pub(crate) fn call_lambda_with_columns_slice(
py: Python,
s: &[Column],
lambda: &PyObject,
polars_module: &Py<PyModule>,
pypolars: &Py<PyModule>,
) -> PyObject {
let pypolars = polars_module.bind(py);
let pypolars = pypolars.bind(py);

// create a PySeries struct/object for Python
let iter = s.iter().map(|s| {
Expand Down Expand Up @@ -210,7 +211,7 @@ pub fn map_mul(
) -> PyExpr {
// get the pypolars module
// do the import outside of the function to prevent import side effects in a hot loop.
let pypolars = PyModule::import(py, "polars").unwrap().unbind();
let pypolars = polars(py).clone_ref(py);

let function = move |s: &mut [Column]| {
Python::with_gil(|py| {
Expand Down
47 changes: 23 additions & 24 deletions crates/polars-python/src/map/series.rs
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,9 @@
use polars::prelude::*;
use pyo3::prelude::*;
use pyo3::pybacked::PyBackedStr;
use pyo3::types::{PyBool, PyCFunction, PyFloat, PyList, PyString, PyTuple};

use super::*;
use crate::py_modules::SERIES;
use crate::py_modules::{pl_series, polars};

/// Find the output type and dispatch to that implementation.
fn infer_and_finish<'a, A: ApplyLambda<'a>>(
Expand Down Expand Up @@ -48,7 +47,7 @@ fn infer_and_finish<'a, A: ApplyLambda<'a>>(
)
.map(|ca| ca.into_series().into())
} else if out.is_instance_of::<PyList>() || out.is_instance_of::<PyTuple>() {
let series = SERIES.call1(py, (out,))?;
let series = pl_series(py).call1(py, (out,))?;
let py_pyseries = series.getattr(py, "_s").unwrap();
let series = py_pyseries.extract::<PySeries>(py).unwrap().series;

Expand All @@ -70,7 +69,7 @@ fn infer_and_finish<'a, A: ApplyLambda<'a>>(
Python::with_gil(|py| {
let out = lambda_owned.call1(py, args)?;
// check if Series, if not, call series constructor on it
SERIES.call1(py, (out,))
pl_series(py).call1(py, (out,))
})
})?
.into_any()
Expand Down Expand Up @@ -1178,7 +1177,7 @@ fn call_series_lambda(

impl<'a> ApplyLambda<'a> for ListChunked {
fn apply_lambda_unknown(&'a self, py: Python, lambda: &Bound<'a, PyAny>) -> PyResult<PySeries> {
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
let mut null_count = 0;
for opt_v in self.into_iter() {
if let Some(v) = opt_v {
Expand Down Expand Up @@ -1215,7 +1214,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
) -> PyResult<PySeries> {
let skip = 1;
// get the pypolars module
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if !self.has_nulls() {
let it = self
.into_no_null_iter()
Expand Down Expand Up @@ -1279,7 +1278,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
D::Native: ToPyObject + FromPyObject<'a>,
{
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
Expand Down Expand Up @@ -1339,7 +1338,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
first_value: Option<bool>,
) -> PyResult<BooleanChunked> {
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
Expand Down Expand Up @@ -1400,7 +1399,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
) -> PyResult<StringChunked> {
let skip = usize::from(first_value.is_some());
// get the pypolars module
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);

if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
Expand Down Expand Up @@ -1462,15 +1461,15 @@ impl<'a> ApplyLambda<'a> for ListChunked {
dt: &DataType,
) -> PyResult<ListChunked> {
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
let lambda = lambda.bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
let it = self
.into_no_null_iter()
.skip(init_null_count + skip)
.map(|val| call_series_lambda(&pypolars, lambda, val));
.map(|val| call_series_lambda(pypolars, lambda, val));

iterator_to_list(
dt,
Expand All @@ -1484,7 +1483,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
let it = self
.into_iter()
.skip(init_null_count + skip)
.map(|opt_val| opt_val.and_then(|val| call_series_lambda(&pypolars, lambda, val)));
.map(|opt_val| opt_val.and_then(|val| call_series_lambda(pypolars, lambda, val)));
iterator_to_list(
dt,
it,
Expand All @@ -1503,7 +1502,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
init_null_count: usize,
first_value: AnyValue<'a>,
) -> PyResult<Series> {
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
let mut avs = Vec::with_capacity(self.len());
avs.extend(std::iter::repeat(AnyValue::Null).take(init_null_count));
avs.push(first_value);
Expand Down Expand Up @@ -1550,7 +1549,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
first_value: Option<ObjectValue>,
) -> PyResult<ObjectChunked<ObjectValue>> {
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
Expand Down Expand Up @@ -1607,7 +1606,7 @@ impl<'a> ApplyLambda<'a> for ListChunked {
#[cfg(feature = "dtype-array")]
impl<'a> ApplyLambda<'a> for ArrayChunked {
fn apply_lambda_unknown(&'a self, py: Python, lambda: &Bound<'a, PyAny>) -> PyResult<PySeries> {
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
let mut null_count = 0;
for opt_v in self.into_iter() {
if let Some(v) = opt_v {
Expand Down Expand Up @@ -1644,7 +1643,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
) -> PyResult<PySeries> {
let skip = 1;
// get the pypolars module
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if !self.has_nulls() {
let it = self
.into_no_null_iter()
Expand Down Expand Up @@ -1708,7 +1707,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
D::Native: ToPyObject + FromPyObject<'a>,
{
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
Expand Down Expand Up @@ -1768,7 +1767,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
first_value: Option<bool>,
) -> PyResult<BooleanChunked> {
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
Expand Down Expand Up @@ -1829,7 +1828,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
) -> PyResult<StringChunked> {
let skip = usize::from(first_value.is_some());
// get the pypolars module
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);

if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
Expand Down Expand Up @@ -1891,15 +1890,15 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
dt: &DataType,
) -> PyResult<ListChunked> {
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
let lambda = lambda.bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
let it = self
.into_no_null_iter()
.skip(init_null_count + skip)
.map(|val| call_series_lambda(&pypolars, lambda, val));
.map(|val| call_series_lambda(pypolars, lambda, val));

iterator_to_list(
dt,
Expand All @@ -1913,7 +1912,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
let it = self
.into_iter()
.skip(init_null_count + skip)
.map(|opt_val| opt_val.and_then(|val| call_series_lambda(&pypolars, lambda, val)));
.map(|opt_val| opt_val.and_then(|val| call_series_lambda(pypolars, lambda, val)));
iterator_to_list(
dt,
it,
Expand All @@ -1932,7 +1931,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
init_null_count: usize,
first_value: AnyValue<'a>,
) -> PyResult<Series> {
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
let mut avs = Vec::with_capacity(self.len());
avs.extend(std::iter::repeat(AnyValue::Null).take(init_null_count));
avs.push(first_value);
Expand Down Expand Up @@ -1979,7 +1978,7 @@ impl<'a> ApplyLambda<'a> for ArrayChunked {
first_value: Option<ObjectValue>,
) -> PyResult<ObjectChunked<ObjectValue>> {
let skip = usize::from(first_value.is_some());
let pypolars = PyModule::import(py, "polars")?;
let pypolars = polars(py).bind(py);
if init_null_count == self.len() {
Ok(ChunkedArray::full_null(self.name().clone(), self.len()))
} else if !self.has_nulls() {
Expand Down
Loading

0 comments on commit 8aa59a8

Please sign in to comment.