perf(convert): intern decode-path dict keys

CHANGELOG.md

@@ -69,6 +69,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
   `accumulate_arrays_are_independent_of_thread_count` checks every output
   array and a `slice_by` re-accumulation bitwise across 1 to 16 threads on a
   dataset with tied scores across images.
+- **Decoding a Python dict into a dataset interns its field-name keys instead of
+  allocating a new string per key per record.** `COCO(dict)`, `loadRes(list of
+  dicts)`, and the RLE/segmentation decoders looked up each field with a bare
+  string literal, and `PyDict.get_item` allocates a fresh `PyString` for that on
+  every call; decoding an RF-DETR-shaped annotation list calls this once per
+  field per detection — millions of throwaway strings for a fixed set of ~10
+  field names. Those lookups now go through `pyo3::intern!`, which builds each
+  literal's `PyString` once per process and reuses it. `COCO(dict)` construction
+  is 18–20% faster across three shapes (a small dict, a 1.5M-annotation
+  segmentation dict, and a bounding-box-only dict), with no numeric or
+  structural change to the decoded dataset. A new test,
+  `TestKnownKeysRoundTrip`, round-trips every known field of an annotation,
+  image, and category through `COCO(dict)` and fails if any interned key
+  literal drifts from the field it names, whether the field is required
+  (raises instead of decoding) or optional (silently dropped instead of
+  decoding) — both failure shapes are pinned by injection. A separate
+  reprofile found that `extra`-field extraction, not covered by this change,
+  now costs roughly half of what remains of dict decoding on the same
+  workloads; that cost is unaddressed here.
 
 ### Fixed
 
@@ -316,7 +335,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
   back to. Previously the marker existed only on `EvalReport`, which nothing that
   writes a file uses, so comparability died with the process.
 
-
 - **`hotcoco.metrics` and `hotcoco.primitives` — the functional layer.** Metric
   functions you can call on plain arrays, with no evaluator, no dataset, and no COCO
   JSON:

crates/hotcoco-pyo3/src/convert.rs

@@ -4,17 +4,27 @@ use pyo3::prelude::*;
 use pyo3::types::{PyBytes, PyDict, PyList};
 
 /// Extract an optional field from a Python dict.
+///
+/// `$key` is interned (`pyo3::intern!`) rather than passed as a bare `&str`:
+/// `get_item` takes anything `IntoPyObject`, and for a plain `&str` that means
+/// allocating a fresh `PyString` on every call. Decoding an RF-DETR-shaped
+/// annotation list calls this once per key per annotation — millions of
+/// throwaway strings for a fixed set of ~10 field names. Interning builds each
+/// literal's `PyString` once per process and reuses it from then on.
 macro_rules! opt {
-    ($dict:expr, $key:expr) => {
-        $dict.get_item($key)?.map(|v| v.extract()).transpose()?
+    ($dict:expr, $key:literal) => {
+        $dict
+            .get_item(pyo3::intern!($dict.py(), $key))?
+            .map(|v| v.extract())
+            .transpose()?
     };
 }
 
 /// Extract a required field from a Python dict, raising `PyValueError` if missing.
 macro_rules! req {
-    ($dict:expr, $key:expr) => {
+    ($dict:expr, $key:literal) => {
         $dict
-            .get_item($key)?
+            .get_item(pyo3::intern!($dict.py(), $key))?
             .ok_or_else(|| {
                 pyo3::exceptions::PyValueError::new_err(concat!("dict missing '", $key, "'"))
             })?
@@ -116,19 +126,28 @@ fn counts_as_str(counts: &Bound<'_, PyAny>) -> PyResult<Option<String>> {
     Ok(None)
 }
 
-/// `req!` with an explicit reader — [`extract_int`] or [`extract_flag`].
+/// `req!` with an explicit reader — [`extract_int`] or [`extract_flag`]. Same
+/// interning rationale as `opt!` above.
 macro_rules! req_with {
-    ($dict:expr, $key:expr, $read:expr) => {
-        $read(&$dict.get_item($key)?.ok_or_else(|| {
-            pyo3::exceptions::PyValueError::new_err(concat!("dict missing '", $key, "'"))
-        })?)?
+    ($dict:expr, $key:literal, $read:expr) => {
+        $read(
+            &$dict
+                .get_item(pyo3::intern!($dict.py(), $key))?
+                .ok_or_else(|| {
+                    pyo3::exceptions::PyValueError::new_err(concat!("dict missing '", $key, "'"))
+                })?,
+        )?
     };
 }
 
-/// `opt!` with an explicit reader — [`extract_int`] or [`extract_flag`].
+/// `opt!` with an explicit reader — [`extract_int`] or [`extract_flag`]. Same
+/// interning rationale as `opt!` above.
 macro_rules! opt_with {
-    ($dict:expr, $key:expr, $read:expr) => {
-        $dict.get_item($key)?.map(|v| $read(&v)).transpose()?
+    ($dict:expr, $key:literal, $read:expr) => {
+        $dict
+            .get_item(pyo3::intern!($dict.py(), $key))?
+            .map(|v| $read(&v))
+            .transpose()?
     };
 }
 
@@ -290,12 +309,12 @@ pub fn py_to_annotation(dict: &Bound<'_, PyDict>) -> PyResult<Annotation> {
     let bbox: Option<[f64; 4]> = opt!(dict, "bbox");
     let area: Option<f64> = opt!(dict, "area");
     let segmentation: Option<Segmentation> = dict
-        .get_item("segmentation")?
+        .get_item(pyo3::intern!(dict.py(), "segmentation"))?
         .map(|v| py_to_segmentation(&v))
         .transpose()?;
     let iscrowd: bool = opt_with!(dict, "iscrowd", extract_flag).unwrap_or(false);
     let keypoints: Option<Vec<f64>> = dict
-        .get_item("keypoints")?
+        .get_item(pyo3::intern!(dict.py(), "keypoints"))?
         .map(|v| {
             // Flat `[x, y, v, …]` is the COCO spelling; an `(N, 3)` array is
             // how the same triplets sit in a tensor. Both mean one thing.
@@ -334,7 +353,7 @@ fn py_to_segmentation(obj: &Bound<'_, PyAny>) -> PyResult<Segmentation> {
     if let Ok(dict) = obj.cast::<PyDict>() {
         let size: [u32; 2] = req!(dict, "size");
         let counts_obj = dict
-            .get_item("counts")?
+            .get_item(pyo3::intern!(dict.py(), "counts"))?
             .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("dict missing 'counts'"))?;
         if let Some(counts) = counts_as_str(&counts_obj)? {
             return Ok(Segmentation::CompressedRle { size, counts });
@@ -460,11 +479,13 @@ pub fn rle_to_coco_py(py: Python<'_>, rle: &Rle) -> PyResult<Py<PyAny>> {
 pub fn py_to_rle(dict: &Bound<'_, PyDict>) -> PyResult<Rle> {
     // Support {"h", "w", "counts": [ints]}, {"size": [h,w], "counts": "string"},
     // and {"size": [h,w], "counts": b"bytes"} (pycocotools format)
-    if let Some(size_obj) = dict.get_item("size")? {
+    if let Some(size_obj) = dict.get_item(pyo3::intern!(dict.py(), "size"))? {
         let size: [u32; 2] = size_obj.extract()?;
-        let counts_obj = dict.get_item("counts")?.ok_or_else(|| {
-            pyo3::exceptions::PyValueError::new_err("RLE dict has 'size' but missing 'counts'")
-        })?;
+        let counts_obj = dict
+            .get_item(pyo3::intern!(dict.py(), "counts"))?
+            .ok_or_else(|| {
+                pyo3::exceptions::PyValueError::new_err("RLE dict has 'size' but missing 'counts'")
+            })?;
         if let Some(s) = counts_as_str(&counts_obj)? {
             return hotcoco_core::mask::rle_from_string(&s, size[0], size[1])
                 .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()));
@@ -478,15 +499,15 @@ pub fn py_to_rle(dict: &Bound<'_, PyDict>) -> PyResult<Rle> {
         });
     }
     let h: u32 = dict
-        .get_item("h")?
+        .get_item(pyo3::intern!(dict.py(), "h"))?
         .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("RLE dict missing 'h'"))?
         .extract()?;
     let w: u32 = dict
-        .get_item("w")?
+        .get_item(pyo3::intern!(dict.py(), "w"))?
         .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("RLE dict missing 'w'"))?
         .extract()?;
     let counts: Vec<u32> = dict
-        .get_item("counts")?
+        .get_item(pyo3::intern!(dict.py(), "counts"))?
         .ok_or_else(|| pyo3::exceptions::PyValueError::new_err("RLE dict missing 'counts'"))?
         .extract()?;
     Ok(Rle { h, w, counts })

crates/hotcoco-pyo3/tests/test_dropin_gaps.py

@@ -110,6 +110,111 @@ def test_slice_by_callable_sees_custom_image_keys(self):
         assert result["rainy"]["num_images"] == 1
 
 
+# ---------------------------------------------------------------------------
+# Every known dict key survives decode (candidate G: interned get_item keys)
+#
+# convert.rs's decode macros (opt!/req!/opt_with!/req_with!) and the standalone
+# get_item calls in py_to_annotation/py_to_segmentation/py_to_rle now fetch
+# each key through pyo3::intern! instead of a bare `&str` literal, to avoid
+# allocating a fresh PyString per key per record. A typo in one of those
+# literals breaks silently: a required key (id, image_id) raises "dict missing
+# '<real name>'" because the interned typo never matches, and an optional key
+# (score, is_group_of, ...) just vanishes instead of raising. This round-trips
+# every field of every record type through COCO to catch either failure mode.
+# ---------------------------------------------------------------------------
+
+
+class TestKnownKeysRoundTrip:
+    def test_every_annotation_field_survives(self):
+        ds = tiny_dataset()
+        ds["annotations"] = [
+            {
+                "id": 7,
+                "image_id": 1,
+                "category_id": 1,
+                "bbox": [1.0, 2.0, 3.0, 4.0],
+                "area": 12.5,
+                "segmentation": {"size": [10, 10], "counts": [100]},
+                "iscrowd": 1,
+                "keypoints": [1.0, 2.0, 2.0],
+                "num_keypoints": 1,
+                "obb": [5.0, 5.0, 2.0, 2.0, 0.3],
+                "score": 0.75,
+                "is_group_of": 1,
+            }
+        ]
+        coco = COCO(ds)
+
+        ann = coco.dataset["annotations"][0]
+
+        assert ann["id"] == 7
+        assert ann["image_id"] == 1
+        assert ann["category_id"] == 1
+        assert ann["bbox"] == [1.0, 2.0, 3.0, 4.0]
+        assert ann["area"] == 12.5
+        assert ann["segmentation"] == {"size": [10, 10], "counts": [100]}
+        assert ann["iscrowd"] == 1
+        assert ann["keypoints"] == [1.0, 2.0, 2.0]
+        assert ann["num_keypoints"] == 1
+        assert ann["obb"] == [5.0, 5.0, 2.0, 2.0, 0.3]
+        assert ann["score"] == 0.75
+        assert ann["is_group_of"] is True
+
+    def test_every_image_field_survives(self):
+        ds = tiny_dataset()
+        ds["images"] = [
+            {
+                "id": 9,
+                "file_name": "x.jpg",
+                "height": 100,
+                "width": 100,
+                "license": 3,
+                "coco_url": "http://a",
+                "flickr_url": "http://b",
+                "date_captured": "2020-01-01",
+                "neg_category_ids": [5, 6],
+                "not_exhaustive_category_ids": [7],
+            }
+        ]
+        ds["annotations"] = [ds["annotations"][0] | {"image_id": 9}]
+        coco = COCO(ds)
+
+        img = next(i for i in coco.dataset["images"] if i["id"] == 9)
+
+        assert img["file_name"] == "x.jpg"
+        assert img["height"] == 100
+        assert img["width"] == 100
+        assert img["license"] == 3
+        assert img["coco_url"] == "http://a"
+        assert img["flickr_url"] == "http://b"
+        assert img["date_captured"] == "2020-01-01"
+        assert img["neg_category_ids"] == [5, 6]
+        assert img["not_exhaustive_category_ids"] == [7]
+
+    def test_every_category_field_survives(self):
+        ds = tiny_dataset()
+        ds["categories"] = [
+            {
+                "id": 1,
+                "name": "person",
+                "supercategory": "animal",
+                "skeleton": [[0, 1], [1, 2]],
+                "keypoints": ["nose", "eye"],
+                "frequency": "f",
+            },
+            ds["categories"][1],
+        ]
+        coco = COCO(ds)
+
+        cat = next(c for c in coco.dataset["categories"] if c["id"] == 1)
+
+        assert cat["name"] == "person"
+        assert cat["supercategory"] == "animal"
+        assert cat["skeleton"] == [[0, 1], [1, 2]]
+        assert cat["keypoints"] == ["nose", "eye"]
+        assert cat["frequency"] == "f"
+
+
 # ---------------------------------------------------------------------------
 # annToRLE returns pycocotools format
 # ---------------------------------------------------------------------------