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

optimize sequence conversion for list and tuple #2944

Merged
merged 2 commits into from
Feb 16, 2023
Merged
Show file tree
Hide file tree
Changes from all 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
13 changes: 12 additions & 1 deletion benches/bench_list.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use criterion::{criterion_group, criterion_main, Bencher, Criterion};

use pyo3::prelude::*;
use pyo3::types::PyList;
use pyo3::types::{PyList, PySequence};

fn iter_list(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
Expand Down Expand Up @@ -53,12 +53,23 @@ fn list_get_item_unchecked(b: &mut Bencher<'_>) {
});
}

fn sequence_from_list(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 50_000;
let list = PyList::new(py, 0..LEN).to_object(py);
b.iter(|| {
let _: &PySequence = list.extract(py).unwrap();
});
});
}

fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("iter_list", iter_list);
c.bench_function("list_new", list_new);
c.bench_function("list_get_item", list_get_item);
#[cfg(not(Py_LIMITED_API))]
c.bench_function("list_get_item_unchecked", list_get_item_unchecked);
c.bench_function("sequence_from_list", sequence_from_list);
}

criterion_group!(benches, criterion_benchmark);
Expand Down
11 changes: 10 additions & 1 deletion benches/bench_tuple.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use criterion::{criterion_group, criterion_main, Bencher, Criterion};

use pyo3::prelude::*;
use pyo3::types::PyTuple;
use pyo3::types::{PySequence, PyTuple};

fn iter_tuple(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
Expand Down Expand Up @@ -53,12 +53,21 @@ fn tuple_get_item_unchecked(b: &mut Bencher<'_>) {
});
}

fn sequence_from_tuple(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 50_000;
let tuple = PyTuple::new(py, 0..LEN).to_object(py);
b.iter(|| tuple.extract::<&PySequence>(py).unwrap());
});
}

fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("iter_tuple", iter_tuple);
c.bench_function("tuple_new", tuple_new);
c.bench_function("tuple_get_item", tuple_get_item);
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
c.bench_function("tuple_get_item_unchecked", tuple_get_item_unchecked);
c.bench_function("sequence_from_tuple", sequence_from_tuple);
}

criterion_group!(benches, criterion_benchmark);
Expand Down
1 change: 1 addition & 0 deletions newsfragments/2944.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Optimize `PySequence` conversion for `list` and `tuple` inputs.
34 changes: 17 additions & 17 deletions src/types/sequence.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,9 @@ use crate::once_cell::GILOnceCell;
use crate::type_object::PyTypeInfo;
use crate::types::{PyAny, PyList, PyString, PyTuple, PyType};
use crate::{ffi, PyNativeType, ToPyObject};
use crate::{AsPyPointer, IntoPy, IntoPyPointer, Py, Python};
use crate::{AsPyPointer, IntoPyPointer, Py, Python};
use crate::{FromPyObject, PyTryFrom};

static SEQUENCE_ABC: GILOnceCell<PyResult<Py<PyType>>> = GILOnceCell::new();

/// Represents a reference to a Python object supporting the sequence protocol.
#[repr(transparent)]
pub struct PySequence(PyAny);
Expand Down Expand Up @@ -318,17 +316,14 @@ where
Ok(v)
}

fn get_sequence_abc(py: Python<'_>) -> Result<&PyType, PyErr> {
static SEQUENCE_ABC: GILOnceCell<Py<PyType>> = GILOnceCell::new();

fn get_sequence_abc(py: Python<'_>) -> PyResult<&PyType> {
SEQUENCE_ABC
.get_or_init(py, || {
Ok(py
.import("collections.abc")?
.getattr("Sequence")?
.downcast::<PyType>()?
.into_py(py))
.get_or_try_init(py, || {
py.import("collections.abc")?.getattr("Sequence")?.extract()
})
.as_ref()
.map_or_else(|e| Err(e.clone_ref(py)), |t| Ok(t.as_ref(py)))
.map(|ty| ty.as_ref(py))
}

impl<'v> PyTryFrom<'v> for PySequence {
Expand All @@ -338,11 +333,16 @@ impl<'v> PyTryFrom<'v> for PySequence {
fn try_from<V: Into<&'v PyAny>>(value: V) -> Result<&'v PySequence, PyDowncastError<'v>> {
let value = value.into();
adamreichold marked this conversation as resolved.
Show resolved Hide resolved

// TODO: surface specific errors in this chain to the user
if let Ok(abc) = get_sequence_abc(value.py()) {
if value.is_instance(abc).unwrap_or(false) {
unsafe { return Ok(value.downcast_unchecked::<PySequence>()) }
}
// Using `is_instance` for `collections.abc.Sequence` is slow, so provide
// optimized cases for list and tuples as common well-known sequences
if PyList::is_type_of(value)
|| PyTuple::is_type_of(value)
|| get_sequence_abc(value.py())
.and_then(|abc| value.is_instance(abc))
// TODO: surface errors in this chain to the user
.unwrap_or(false)
{
unsafe { return Ok(value.downcast_unchecked::<PySequence>()) }
}

Err(PyDowncastError::new(value, "Sequence"))
Expand Down