Skip to content

Commit

Permalink
Merge #2954
Browse files Browse the repository at this point in the history
2954: optimize mapping conversion for dict r=davidhewitt a=davidhewitt

Equivalent of #2944 for dicts -> mapping.

The benchmark diff is not as significant as in #2944. Still something like 80ns -> 2ns on my machine.

Co-authored-by: David Hewitt <1939362+davidhewitt@users.noreply.github.com>
  • Loading branch information
bors[bot] and davidhewitt authored Feb 15, 2023
2 parents 03ecb80 + 59630be commit 18e7094
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 21 deletions.
17 changes: 16 additions & 1 deletion benches/bench_dict.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::IntoPyDict;
use pyo3::{prelude::*, types::PyMapping};
use std::collections::{BTreeMap, HashMap};

fn iter_dict(b: &mut Bencher<'_>) {
Expand Down Expand Up @@ -63,6 +63,19 @@ fn extract_hashbrown_map(b: &mut Bencher<'_>) {
});
}

fn mapping_from_dict(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 100_000;
let dict = (0..LEN as u64)
.map(|i| (i, i * 2))
.into_py_dict(py)
.to_object(py);
b.iter(|| {
let _: &PyMapping = dict.extract(py).unwrap();
});
});
}

fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("iter_dict", iter_dict);
c.bench_function("dict_new", dict_new);
Expand All @@ -72,6 +85,8 @@ fn criterion_benchmark(c: &mut Criterion) {

#[cfg(feature = "hashbrown")]
c.bench_function("extract_hashbrown_map", extract_hashbrown_map);

c.bench_function("mapping_from_dict", mapping_from_dict);
}

criterion_group!(benches, criterion_benchmark);
Expand Down
1 change: 1 addition & 0 deletions newsfragments/2954.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Optimize `PyMapping` conversion for `dict` inputs.
37 changes: 17 additions & 20 deletions src/types/mapping.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,8 @@
use crate::err::{PyDowncastError, PyErr, PyResult};
use crate::once_cell::GILOnceCell;
use crate::type_object::PyTypeInfo;
use crate::types::{PyAny, PySequence, PyType};
use crate::{
ffi, AsPyPointer, IntoPy, IntoPyPointer, Py, PyNativeType, PyTryFrom, Python, ToPyObject,
};

static MAPPING_ABC: GILOnceCell<PyResult<Py<PyType>>> = GILOnceCell::new();
use crate::types::{PyAny, PyDict, PySequence, PyType};
use crate::{ffi, AsPyPointer, IntoPyPointer, Py, PyNativeType, PyTryFrom, Python, ToPyObject};

/// Represents a reference to a Python object supporting the mapping protocol.
#[repr(transparent)]
Expand Down Expand Up @@ -119,17 +115,14 @@ impl PyMapping {
}
}

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

fn get_mapping_abc(py: Python<'_>) -> PyResult<&PyType> {
MAPPING_ABC
.get_or_init(py, || {
Ok(py
.import("collections.abc")?
.getattr("Mapping")?
.downcast::<PyType>()?
.into_py(py))
.get_or_try_init(py, || {
py.import("collections.abc")?.getattr("Mapping")?.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 PyMapping {
Expand All @@ -139,11 +132,15 @@ impl<'v> PyTryFrom<'v> for PyMapping {
fn try_from<V: Into<&'v PyAny>>(value: V) -> Result<&'v PyMapping, PyDowncastError<'v>> {
let value = value.into();

// TODO: surface specific errors in this chain to the user
if let Ok(abc) = get_mapping_abc(value.py()) {
if value.is_instance(abc).unwrap_or(false) {
unsafe { return Ok(value.downcast_unchecked()) }
}
// Using `is_instance` for `collections.abc.Mapping` is slow, so provide
// optimized case dict as a well-known mapping
if PyDict::is_type_of(value)
|| get_mapping_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()) }
}

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

0 comments on commit 18e7094

Please sign in to comment.