Optimize to only return possible to inline. Add bench

Author: MarkusSintonen

python/pydantic_core/__init__.py

@@ -146,6 +146,6 @@ class MultiHostHost(_TypedDict):
 class GatherResult(_TypedDict):
     """Internal result of gathering schemas for cleaning."""
 
-    definition_refs: dict[str, list[DefinitionReferenceSchema]]
+    inlinable_def_refs: dict[str, DefinitionReferenceSchema]
     recursive_refs: set[str]
     schemas_with_meta_keys: dict[str, list[CoreSchema]] | None

src/schema_traverse.rs

@@ -1,12 +1,17 @@
 use crate::tools::py_err;
 use pyo3::exceptions::{PyException, PyKeyError};
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyList, PySet, PyString, PyTuple};
+use pyo3::types::{PyDict, PyList, PyNone, PySet, PyString, PyTuple};
 use pyo3::{create_exception, intern, Bound, PyResult};
-use std::collections::HashSet;
 
 create_exception!(pydantic_core._pydantic_core, GatherInvalidDefinitionError, PyException);
 
+macro_rules! none {
+    ($py: expr) => {
+        PyNone::get_bound($py)
+    };
+}
+
 macro_rules! get {
     ($dict: expr, $key: expr) => {
         $dict.get_item(intern!($dict.py(), $key))?
@@ -47,28 +52,32 @@ fn gather_definition_ref(schema_ref_dict: &Bound<'_, PyDict>, ctx: &mut GatherCt
         return py_err!(PyKeyError; "Invalid definition-ref, missing schema_ref");
     };
     let schema_ref = schema_ref.downcast_exact::<PyString>()?;
+    let py = schema_ref_dict.py();
 
     if !ctx.recursively_seen_refs.contains(schema_ref)? {
-        // Check if we already gathered this definition ref instance.
-        // No need to again process the same def ref instance if we already did it.
-        if ctx.seen_ref_instances.insert(schema_ref_dict.as_ptr() as isize) {
+        // Def ref in no longer consider as inlinable if its re-encountered. Then its used multiple times.
+        // No need to retraverse it either if we already encountered this.
+        if !ctx.inline_def_ref_candidates.contains(schema_ref)? {
             let Some(definition) = ctx.definitions.get_item(schema_ref)? else {
                 return py_err!(GatherInvalidDefinitionError; "{}", schema_ref.to_str()?);
             };
 
-            defaultdict_list_append!(&ctx.def_refs, schema_ref, schema_ref_dict);
-
+            ctx.inline_def_ref_candidates.set_item(schema_ref, schema_ref_dict)?;
             ctx.recursively_seen_refs.add(schema_ref)?;
 
             gather_schema(definition.downcast_exact()?, ctx)?;
             traverse_key_fn!("serialization", gather_schema, schema_ref_dict, ctx);
             gather_meta(schema_ref_dict, ctx)?;
 
             ctx.recursively_seen_refs.discard(schema_ref)?;
+        } else {
+            ctx.inline_def_ref_candidates.set_item(schema_ref, none!(py))?; // Mark not inlinable (used multiple times)
         }
     } else {
+        ctx.inline_def_ref_candidates.set_item(schema_ref, none!(py))?; // Mark not inlinable (used in recursion)
         ctx.recursive_def_refs.add(schema_ref)?;
         for seen_ref in ctx.recursively_seen_refs.iter() {
+            ctx.inline_def_ref_candidates.set_item(&seen_ref, none!(py))?; // Mark not inlinable (used in recursion)
             ctx.recursive_def_refs.add(seen_ref)?;
         }
     }
@@ -160,10 +169,9 @@ fn gather_schema(schema: &Bound<'_, PyDict>, ctx: &mut GatherCtx) -> PyResult<()
 struct GatherCtx<'a, 'py> {
     definitions: &'a Bound<'py, PyDict>,
     meta_with_keys: Option<(Bound<'py, PyDict>, &'a Bound<'py, PySet>)>,
-    def_refs: Bound<'py, PyDict>,
+    inline_def_ref_candidates: Bound<'py, PyDict>,
     recursive_def_refs: Bound<'py, PySet>,
     recursively_seen_refs: Bound<'py, PySet>,
-    seen_ref_instances: HashSet<isize>,
 }
 
 #[pyfunction(signature = (schema, definitions, find_meta_with_keys))]
@@ -179,15 +187,21 @@ pub fn gather_schemas_for_cleaning<'py>(
             true => None,
             false => Some((PyDict::new_bound(py), find_meta_with_keys.downcast_exact::<PySet>()?)),
         },
-        def_refs: PyDict::new_bound(py),
+        inline_def_ref_candidates: PyDict::new_bound(py),
         recursive_def_refs: PySet::empty_bound(py)?,
         recursively_seen_refs: PySet::empty_bound(py)?,
-        seen_ref_instances: HashSet::new(),
     };
     gather_schema(schema.downcast_exact()?, &mut ctx)?;
 
+    let inlinable_def_refs = PyDict::new_bound(py);
+    for (ref_str, def_ref_candidate) in ctx.inline_def_ref_candidates.iter() {
+        if !def_ref_candidate.is_none() {
+            inlinable_def_refs.set_item(ref_str, def_ref_candidate)?;
+        }
+    }
+
     let res = PyDict::new_bound(py);
-    res.set_item(intern!(py, "definition_refs"), ctx.def_refs)?;
+    res.set_item(intern!(py, "inlinable_def_refs"), inlinable_def_refs)?;
     res.set_item(intern!(py, "recursive_refs"), ctx.recursive_def_refs)?;
     res.set_item(intern!(py, "schemas_with_meta_keys"), ctx.meta_with_keys.map(|v| v.0))?;
     Ok(res)

tests/benchmarks/test_gather_schemas_for_cleaning_benchmark.py

@@ -0,0 +1,17 @@
+from typing import Callable
+
+from python.pydantic_core import gather_schemas_for_cleaning
+
+from .nested_schema import inlined_schema, schema_using_defs
+
+
+def test_nested_schema_using_defs(benchmark: Callable[..., None]) -> None:
+    schema = schema_using_defs()
+    definitions = {def_schema['ref']: def_schema for def_schema in schema['definitions']}
+    schema = schema['schema']
+    benchmark(gather_schemas_for_cleaning, schema, definitions, None)
+
+
+def test_nested_schema_inlined(benchmark: Callable[..., None]) -> None:
+    schema = inlined_schema()
+    benchmark(gather_schemas_for_cleaning, schema, {}, {'some_meta_key'})

tests/test_gather_schemas_for_cleaning.py

@@ -8,7 +8,7 @@ def test_no_refs():
     p2 = core_schema.arguments_parameter('b', core_schema.int_schema())
     schema = core_schema.arguments_schema([p1, p2])
     res = gather_schemas_for_cleaning(schema, definitions={}, find_meta_with_keys=None)
-    assert res['definition_refs'] == {}
+    assert res['inlinable_def_refs'] == {}
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None
 
@@ -18,12 +18,12 @@ def test_simple_ref_schema():
     definitions = {'ref1': core_schema.int_schema(ref='ref1')}
 
     res = gather_schemas_for_cleaning(schema, definitions, find_meta_with_keys=None)
-    assert res['definition_refs'] == {'ref1': [schema]} and res['definition_refs']['ref1'][0] is schema
+    assert res['inlinable_def_refs'] == {'ref1': schema} and res['inlinable_def_refs']['ref1'] is schema
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None
 
 
-def test_deep_ref_schema():
+def test_deep_ref_schema_used_multiple_times():
     class Model:
         pass
 
@@ -42,9 +42,7 @@ class Model:
     definitions = {'ref1': core_schema.str_schema(ref='ref1'), 'ref2': core_schema.bytes_schema(ref='ref2')}
 
     res = gather_schemas_for_cleaning(schema, definitions, find_meta_with_keys=None)
-    assert res['definition_refs'] == {'ref1': [ref11, ref12], 'ref2': [ref2]}
-    assert res['definition_refs']['ref1'][0] is ref11 and res['definition_refs']['ref1'][1] is ref12
-    assert res['definition_refs']['ref2'][0] is ref2
+    assert res['inlinable_def_refs'] == {'ref2': ref2} and res['inlinable_def_refs']['ref2'] is ref2
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None
 
@@ -55,15 +53,15 @@ def test_ref_in_serialization_schema():
         serialization=core_schema.plain_serializer_function_ser_schema(lambda v: v, return_schema=ref),
     )
     res = gather_schemas_for_cleaning(schema, definitions={'ref1': core_schema.str_schema()}, find_meta_with_keys=None)
-    assert res['definition_refs'] == {'ref1': [ref]} and res['definition_refs']['ref1'][0] is ref
+    assert res['inlinable_def_refs'] == {'ref1': ref} and res['inlinable_def_refs']['ref1'] is ref
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None
 
 
 def test_recursive_ref_schema():
     ref1 = core_schema.definition_reference_schema('ref1')
     res = gather_schemas_for_cleaning(ref1, definitions={'ref1': ref1}, find_meta_with_keys=None)
-    assert res['definition_refs'] == {'ref1': [ref1]} and res['definition_refs']['ref1'][0] is ref1
+    assert res['inlinable_def_refs'] == {}
     assert res['recursive_refs'] == {'ref1'}
     assert res['schemas_with_meta_keys'] is None
 
@@ -82,11 +80,8 @@ def test_deep_recursive_ref_schema():
         },
         find_meta_with_keys=None,
     )
-    assert res['definition_refs'] == {'ref1': [ref1], 'ref2': [ref2], 'ref3': [ref3]}
+    assert res['inlinable_def_refs'] == {}
     assert res['recursive_refs'] == {'ref1', 'ref2', 'ref3'}
-    assert res['definition_refs']['ref1'][0] is ref1
-    assert res['definition_refs']['ref2'][0] is ref2
-    assert res['definition_refs']['ref3'][0] is ref3
     assert res['schemas_with_meta_keys'] is None
 
 
@@ -108,7 +103,7 @@ class Model:
     res = gather_schemas_for_cleaning(
         schema, definitions={'ref1': field2}, find_meta_with_keys={'find_meta1', 'find_meta2'}
     )
-    assert res['definition_refs'] == {'ref1': [ref1]} and res['definition_refs']['ref1'][0] is ref1
+    assert res['inlinable_def_refs'] == {'ref1': ref1} and res['inlinable_def_refs']['ref1'] is ref1
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] == {'find_meta1': [field1, field2], 'find_meta2': [field2]}
     assert res['schemas_with_meta_keys']['find_meta1'][0] is field1
@@ -134,9 +129,6 @@ def test_no_duplicate_ref_instances_gathered():
     definitions = {'ref1': schema1, 'ref2': schema2}
 
     res = gather_schemas_for_cleaning(schema3, definitions=definitions, find_meta_with_keys=None)
-    assert res['definition_refs'] == {
-        'ref1': [schema2['items_schema'][0], schema2['items_schema'][1]],
-        'ref2': [schema3['items_schema'][0], schema3['items_schema'][1]],
-    }
+    assert res['inlinable_def_refs'] == {}
     assert res['recursive_refs'] == set()
     assert res['schemas_with_meta_keys'] is None