Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore(python): Continue converting to PyO3 0.21 Bound<> APIs #16081

Merged
merged 5 commits into from
May 9, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions py-polars/src/arrow_interop/to_rust.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use pyo3::types::PyList;

use crate::error::PyPolarsErr;

pub fn field_to_rust(obj: &PyAny) -> PyResult<Field> {
pub fn field_to_rust(obj: Bound<'_, PyAny>) -> PyResult<Field> {
let schema = Box::new(ffi::ArrowSchema::empty());
let schema_ptr = &*schema as *const ffi::ArrowSchema;

Expand All @@ -20,7 +20,7 @@ pub fn field_to_rust(obj: &PyAny) -> PyResult<Field> {
}

// PyList<Field> which you get by calling `list(schema)`
pub fn pyarrow_schema_to_rust(obj: &PyList) -> PyResult<Schema> {
pub fn pyarrow_schema_to_rust(obj: &Bound<'_, PyList>) -> PyResult<Schema> {
obj.into_iter().map(field_to_rust).collect()
}

Expand Down
6 changes: 3 additions & 3 deletions py-polars/src/conversion/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,12 +61,12 @@ impl<T> From<T> for Wrap<T> {
}

// extract a Rust DataFrame from a python DataFrame, that is DataFrame<PyDataFrame<RustDataFrame>>
pub(crate) fn get_df(obj: &PyAny) -> PyResult<DataFrame> {
pub(crate) fn get_df(obj: &Bound<'_, PyAny>) -> PyResult<DataFrame> {
let pydf = obj.getattr(intern!(obj.py(), "_df"))?;
Ok(pydf.extract::<PyDataFrame>()?.df)
}

pub(crate) fn get_lf(obj: &PyAny) -> PyResult<LazyFrame> {
pub(crate) fn get_lf(obj: &Bound<'_, PyAny>) -> PyResult<LazyFrame> {
let pydf = obj.getattr(intern!(obj.py(), "_ldf"))?;
Ok(pydf.extract::<PyLazyFrame>()?.ldf)
}
Expand Down Expand Up @@ -455,7 +455,7 @@ impl FromPyObject<'_> for Wrap<Schema> {

impl IntoPy<PyObject> for Wrap<&Schema> {
fn into_py(self, py: Python<'_>) -> PyObject {
let dict = PyDict::new(py);
let dict = PyDict::new_bound(py);
for (k, v) in self.0.iter() {
dict.set_item(k.as_str(), Wrap(v.clone())).unwrap();
}
Expand Down
6 changes: 3 additions & 3 deletions py-polars/src/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -123,12 +123,12 @@ impl IntoPy<PyObject> for Wrap<PolarsWarning> {
fn into_py(self, py: Python<'_>) -> PyObject {
match self.0 {
PolarsWarning::CategoricalRemappingWarning => {
CategoricalRemappingWarning::type_object(py).to_object(py)
CategoricalRemappingWarning::type_object_bound(py).to_object(py)
},
PolarsWarning::MapWithoutReturnDtypeWarning => {
MapWithoutReturnDtypeWarning::type_object(py).to_object(py)
MapWithoutReturnDtypeWarning::type_object_bound(py).to_object(py)
},
PolarsWarning::UserWarning => PyUserWarning::type_object(py).to_object(py),
PolarsWarning::UserWarning => PyUserWarning::type_object_bound(py).to_object(py),
}
}
}
4 changes: 3 additions & 1 deletion py-polars/src/expr/array.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::borrow::Cow;

use polars::prelude::*;
use polars_ops::prelude::array::ArrToStructNameGenerator;
use pyo3::prelude::*;
Expand Down Expand Up @@ -112,7 +114,7 @@ impl PyExpr {
Arc::new(move |idx: usize| {
Python::with_gil(|py| {
let out = lambda.call1(py, (idx,)).unwrap();
let out: SmartString = out.extract::<&str>(py).unwrap().into();
let out: SmartString = out.extract::<Cow<str>>(py).unwrap().into();
out
})
}) as ArrToStructNameGenerator
Expand Down
9 changes: 6 additions & 3 deletions py-polars/src/expr/general.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::io::Cursor;
use std::ops::Neg;

use polars::lazy::dsl;
Expand All @@ -6,6 +7,7 @@ use polars::series::ops::NullBehavior;
use polars_core::series::IsSorted;
use pyo3::class::basic::CompareOp;
use pyo3::prelude::*;
use pyo3::pybacked::PyBackedBytes;
use pyo3::types::PyBytes;

use crate::conversion::{parse_fill_null_strategy, vec_extract_wrapped, Wrap};
Expand Down Expand Up @@ -83,14 +85,15 @@ impl PyExpr {
ciborium::ser::into_writer(&self.inner, &mut writer)
.map_err(|e| PyPolarsErr::Other(format!("{}", e)))?;

Ok(PyBytes::new(py, &writer).to_object(py))
Ok(PyBytes::new_bound(py, &writer).to_object(py))
}

fn __setstate__(&mut self, py: Python, state: PyObject) -> PyResult<()> {
// Used in pickle/pickling
match state.extract::<&PyBytes>(py) {
match state.extract::<PyBackedBytes>(py) {
Ok(s) => {
self.inner = ciborium::de::from_reader(s.as_bytes())
let cursor = Cursor::new(&*s);
self.inner = ciborium::de::from_reader(cursor)
.map_err(|e| PyPolarsErr::Other(format!("{}", e)))?;
Ok(())
},
Expand Down
4 changes: 3 additions & 1 deletion py-polars/src/expr/list.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::borrow::Cow;

use polars::prelude::*;
use polars::series::ops::NullBehavior;
use pyo3::prelude::*;
Expand Down Expand Up @@ -212,7 +214,7 @@ impl PyExpr {
Arc::new(move |idx: usize| {
Python::with_gil(|py| {
let out = lambda.call1(py, (idx,)).unwrap();
let out: SmartString = out.extract::<&str>(py).unwrap().into();
let out: SmartString = out.extract::<Cow<str>>(py).unwrap().into();
out
})
}) as NameGenerator
Expand Down
4 changes: 3 additions & 1 deletion py-polars/src/expr/name.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::borrow::Cow;

use polars::prelude::*;
use pyo3::prelude::*;
use smartstring::alias::String as SmartString;
Expand Down Expand Up @@ -46,7 +48,7 @@ impl PyExpr {
let name_mapper = Arc::new(move |name: &str| {
Python::with_gil(|py| {
let out = name_mapper.call1(py, (name,)).unwrap();
let out: SmartString = out.extract::<&str>(py).unwrap().into();
let out: SmartString = out.extract::<Cow<str>>(py).unwrap().into();
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do we use Cow in favor of pybackedstr here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Less typing to convert to SmartString. Could do it the other way too, though.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok, looking at the FromPyObject implementations in Pyo3, I now see what the "fastest" way is to get a rust &str.

And that is indeed going via Cow. As they write obj.downcast::<PyString>()?.to_cow().map(Cow::into_owned) for String, where to_cow says:

"""
Converts the PyString into a Rust string, avoiding copying when possible.
"""

So I do think this is the fastest way to get a borrowing string out of python. I think we should favor this than over pybytes in other places.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. Is that specific to this PR too? If not, can it be merged?
  2. I am not sure that Cow<str> is always the best option. In particular, it has a lifetime that is tied to the original object. PyBackedStr in contrast is an owned object, so it's a lot easier to work with some contexts. And it has some overhead but I don't think its copying overhead, it's more "having a Py<>" overhead, so like... a Python incref basically?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I mean in the borrowing &str case. Not related to this PR indeed. Thanks!

out
})
}) as FieldsNameMapper;
Expand Down
2 changes: 1 addition & 1 deletion py-polars/src/expr/rolling.rs
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,7 @@ impl PyExpr {
},
Err(_) => {
let obj = out;
let is_float = obj.as_ref(py).is_instance_of::<PyFloat>();
let is_float = obj.bind(py).is_instance_of::<PyFloat>();

let dtype = s.dtype();

Expand Down
14 changes: 7 additions & 7 deletions py-polars/src/functions/eager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,20 +7,20 @@ use crate::error::PyPolarsErr;
use crate::{PyDataFrame, PySeries};

#[pyfunction]
pub fn concat_df(dfs: &PyAny, py: Python) -> PyResult<PyDataFrame> {
pub fn concat_df(dfs: &Bound<'_, PyAny>, py: Python) -> PyResult<PyDataFrame> {
use polars_core::error::PolarsResult;
use polars_core::utils::rayon::prelude::*;

let mut iter = dfs.iter()?;
let first = iter.next().unwrap()?;

let first_rdf = get_df(first)?;
let first_rdf = get_df(&first)?;
let identity_df = first_rdf.clear();

let mut rdfs: Vec<PolarsResult<DataFrame>> = vec![Ok(first_rdf)];

for item in iter {
let rdf = get_df(item?)?;
let rdf = get_df(&item?)?;
rdfs.push(Ok(rdf));
}

Expand Down Expand Up @@ -63,13 +63,13 @@ pub fn concat_series(series: &Bound<'_, PyAny>) -> PyResult<PySeries> {
}

#[pyfunction]
pub fn concat_df_diagonal(dfs: &PyAny) -> PyResult<PyDataFrame> {
pub fn concat_df_diagonal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
let iter = dfs.iter()?;

let dfs = iter
.map(|item| {
let item = item?;
get_df(item)
get_df(&item)
})
.collect::<PyResult<Vec<_>>>()?;

Expand All @@ -78,13 +78,13 @@ pub fn concat_df_diagonal(dfs: &PyAny) -> PyResult<PyDataFrame> {
}

#[pyfunction]
pub fn concat_df_horizontal(dfs: &PyAny) -> PyResult<PyDataFrame> {
pub fn concat_df_horizontal(dfs: &Bound<'_, PyAny>) -> PyResult<PyDataFrame> {
let iter = dfs.iter()?;

let dfs = iter
.map(|item| {
let item = item?;
get_df(item)
get_df(&item)
})
.collect::<PyResult<Vec<_>>>()?;

Expand Down
21 changes: 8 additions & 13 deletions py-polars/src/functions/lazy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ pub fn cols(names: Vec<String>) -> PyExpr {

#[pyfunction]
pub fn concat_lf(
seq: &PyAny,
seq: &Bound<'_, PyAny>,
rechunk: bool,
parallel: bool,
to_supertypes: bool,
Expand All @@ -163,7 +163,7 @@ pub fn concat_lf(

for res in seq.iter()? {
let item = res?;
let lf = get_lf(item)?;
let lf = get_lf(&item)?;
lfs.push(lf);
}

Expand Down Expand Up @@ -270,7 +270,7 @@ pub fn datetime(

#[pyfunction]
pub fn concat_lf_diagonal(
lfs: &PyAny,
lfs: &Bound<'_, PyAny>,
rechunk: bool,
parallel: bool,
to_supertypes: bool,
Expand All @@ -280,7 +280,7 @@ pub fn concat_lf_diagonal(
let lfs = iter
.map(|item| {
let item = item?;
get_lf(item)
get_lf(&item)
})
.collect::<PyResult<Vec<_>>>()?;

Expand All @@ -298,13 +298,13 @@ pub fn concat_lf_diagonal(
}

#[pyfunction]
pub fn concat_lf_horizontal(lfs: &PyAny, parallel: bool) -> PyResult<PyLazyFrame> {
pub fn concat_lf_horizontal(lfs: &Bound<'_, PyAny>, parallel: bool) -> PyResult<PyLazyFrame> {
let iter = lfs.iter()?;

let lfs = iter
.map(|item| {
let item = item?;
get_lf(item)
get_lf(&item)
})
.collect::<PyResult<Vec<_>>>()?;

Expand Down Expand Up @@ -387,7 +387,7 @@ pub fn last() -> PyExpr {
}

#[pyfunction]
pub fn lit(value: &PyAny, allow_object: bool) -> PyResult<PyExpr> {
pub fn lit(value: &Bound<'_, PyAny>, allow_object: bool) -> PyResult<PyExpr> {
if value.is_instance_of::<PyBool>() {
let val = value.extract::<bool>().unwrap();
Ok(dsl::lit(val).into())
Expand All @@ -401,12 +401,7 @@ pub fn lit(value: &PyAny, allow_object: bool) -> PyResult<PyExpr> {
let val = float.extract::<f64>().unwrap();
Ok(Expr::Literal(LiteralValue::Float(val)).into())
} else if let Ok(pystr) = value.downcast::<PyString>() {
Ok(dsl::lit(
pystr
.to_str()
.expect("could not transform Python string to Rust Unicode"),
)
.into())
Ok(dsl::lit(pystr.to_string()).into())
} else if let Ok(series) = value.extract::<PySeries>() {
Ok(dsl::lit(series.series).into())
} else if value.is_none() {
Expand Down
6 changes: 3 additions & 3 deletions py-polars/src/functions/range.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,9 @@ pub fn int_range(start: PyExpr, end: PyExpr, step: i64, dtype: Wrap<DataType>) -
/// Eager version of `int_range` to avoid overhead from the expression engine.
#[pyfunction]
pub fn eager_int_range(
lower: &PyAny,
upper: &PyAny,
step: &PyAny,
lower: &Bound<'_, PyAny>,
upper: &Bound<'_, PyAny>,
step: &Bound<'_, PyAny>,
dtype: Wrap<DataType>,
) -> PyResult<PySeries> {
let dtype = dtype.0;
Expand Down
Loading
Loading