Skip to content

Commit

Permalink
Switch to newer PyO3 0.21 APIs (Bound<> and friends) (#82)
Browse files Browse the repository at this point in the history
Co-authored-by: Itamar Turner-Trauring <[email protected]>
  • Loading branch information
itamarst and pythonspeed authored May 21, 2024
1 parent 19146ab commit 73e9a30
Show file tree
Hide file tree
Showing 4 changed files with 13 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ fn lazy_parallel_jaccard(pydf: PyLazyFrame, col_a: &str, col_b: &str) -> PyResul

/// A Python module implemented in Rust.
#[pymodule]
fn extend_polars(_py: Python, m: &PyModule) -> PyResult<()> {
fn extend_polars(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(parallel_jaccard, m)?)?;
m.add_function(wrap_pyfunction!(lazy_parallel_jaccard, m)?)?;
m.add_function(wrap_pyfunction!(debug, m)?)?;
Expand Down
2 changes: 1 addition & 1 deletion pyo3-polars/src/ffi/to_rust.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use polars::prelude::*;
use pyo3::ffi::Py_uintptr_t;
use pyo3::prelude::*;

pub fn array_to_rust(obj: &PyAny) -> PyResult<ArrayRef> {
pub fn array_to_rust(obj: &Bound<PyAny>) -> PyResult<ArrayRef> {
// prepare a pointer to receive the Array struct
let array = Box::new(ffi::ArrowArray::empty());
let schema = Box::new(ffi::ArrowSchema::empty());
Expand Down
19 changes: 10 additions & 9 deletions pyo3-polars/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
//!
//! /// A Python module implemented in Rust.
//! #[pymodule]
//! fn expression_lib(_py: Python, m: &PyModule) -> PyResult<()> {
//! fn expression_lib(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
//! m.add_function(wrap_pyfunction!(my_cool_function, m)?)?;
//! Ok(())
//! }
Expand Down Expand Up @@ -119,22 +119,23 @@ impl AsRef<LazyFrame> for PyLazyFrame {
}

impl<'a> FromPyObject<'a> for PySeries {
fn extract(ob: &'a PyAny) -> PyResult<Self> {
fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
let ob = ob.call_method0("rechunk")?;

let name = ob.getattr("name")?;
let name = name.str()?.to_str()?;
let py_name = name.str()?;
let name = py_name.to_cow()?;

let arr = ob.call_method0("to_arrow")?;
let arr = ffi::to_rust::array_to_rust(arr)?;
let arr = ffi::to_rust::array_to_rust(&arr)?;
Ok(PySeries(
Series::try_from((name, arr)).map_err(PyPolarsErr::from)?,
Series::try_from((&*name, arr)).map_err(PyPolarsErr::from)?,
))
}
}

impl<'a> FromPyObject<'a> for PyDataFrame {
fn extract(ob: &'a PyAny) -> PyResult<Self> {
fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
let series = ob.call_method0("get_columns")?;
let n = ob.getattr("width")?.extract::<usize>()?;
let mut columns = Vec::with_capacity(n);
Expand All @@ -149,7 +150,7 @@ impl<'a> FromPyObject<'a> for PyDataFrame {

#[cfg(feature = "lazy")]
impl<'a> FromPyObject<'a> for PyLazyFrame {
fn extract(ob: &'a PyAny) -> PyResult<Self> {
fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
let s = ob.call_method0("__getstate__")?.extract::<Vec<u8>>()?;
let lp: LogicalPlan = ciborium::de::from_reader(&*s).map_err(
|e| PyPolarsErr::Other(
Expand Down Expand Up @@ -237,9 +238,9 @@ impl IntoPy<PyObject> for PyDataFrame {
#[cfg(feature = "lazy")]
impl IntoPy<PyObject> for PyLazyFrame {
fn into_py(self, py: Python<'_>) -> PyObject {
let polars = py.import("polars").expect("polars not installed");
let polars = py.import_bound("polars").expect("polars not installed");
let cls = polars.getattr("LazyFrame").unwrap();
let instance = cls.call_method1("__new__", (cls,)).unwrap();
let instance = cls.call_method1("__new__", (&cls,)).unwrap();
let mut writer: Vec<u8> = vec![];
ciborium::ser::into_writer(&self.0.logical_plan, &mut writer).unwrap();

Expand Down
2 changes: 1 addition & 1 deletion rust-toolchain.toml
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
[toolchain]
channel = "nightly-2024-03-28"
channel = "nightly-2024-05-14"

0 comments on commit 73e9a30

Please sign in to comment.