Rework def ref system to fix various bugs #571

adriangb · 2023-05-02T06:03:11Z

This resolves the bugs in the tests and some others.

adriangb · 2023-05-02T06:04:57Z

src/build_context.rs

+    // Mapping of reference strings to indexes in the definitions Vec
+    // As we recurse through the input when we find:
+    // - A definitions-reference (`{'type': 'definitions-reference', 'schema_ref': '...'}`)
+    //   then we store that `schema_ref` -> `definitions.len()`
+    //   so that the value is the index into definitions.
+    //   We also set a None value in definitions.
+    //   Then we create a DefinitionsReferenceValidator that stores the index (value) from this mapping.
+    //   At runtime if it needs to call into the validator it simply looks it up by index.
+    // - Any schema with a `'ref'`: before recursing we store the ref here and set a None
+    //   value in definitions.
+    //   As we recurse if a definitions-reference to this
+    //   Then after recursing and completing the CombinedValidator we replace that None value with the
+    //   newly built validator. Now anything that
+    refs: AHashMap<String, usize>,
+
+    // Here we hold an *optional* reference to a definition validator
+    // Any time we encounter a schema with a `'ref'` we build a DefinitionsReferenceValidator that points
+    // to an index in this Vector
+    // At the end of building we check if any values are None, which indicates an unsatisfied reference
+    definitions: Vec<Option<T>>,
 }


It isn't too hard to add a basic ref-count system here so we know which references have more than one referrer and hence which can be inlined or not. But I got bogged down in implementation multiple times.

adriangb · 2023-05-02T06:05:39Z

src/build_context.rs

-fn extract_used_refs(schema: &PyAny, refs: &mut AHashSet<String>) -> PyResult<()> {
-    if let Ok(dict) = schema.downcast::<PyDict>() {
-        if is_definition_ref(dict)? {
-            refs.insert(dict.get_as_req(intern!(schema.py(), "schema_ref"))?);
-        } else {
-            for (key, value) in dict.iter() {
-                if !key.eq(intern!(schema.py(), "metadata"))? {
-                    extract_used_refs(value, refs)?;
-                }
-            }
-        }
-    } else if let Ok(list) = schema.downcast::<PyList>() {
-        for item in list.iter() {
-            extract_used_refs(item, refs)?;
-        }
-    }
-    Ok(())
-}


It seems like a lot of the dragons were in here, and it was slow (hence the building schema speedup).

makes sense.

i guess this is how we did the reference counting.

src/recursion_guard.rs

adriangb · 2023-05-02T06:07:42Z

src/serializers/infer.rs

+                return serializer.to_python(
+                    py,
+                    value,
+                    extra.mode.to_object(py).extract(py)?,
+                    include,
+                    exclude,
+                    extra.by_alias,
+                    extra.exclude_unset,
+                    extra.exclude_defaults,
+                    extra.exclude_none,
+                    extra.round_trip,
+                    extra.warnings.is_active(),
+                    extra.fallback,
+                );


This solves a bug where we were calling into the next serializer without any slots, so if it had any def refs it would crash.

src/serializers/type_serializers/definitions.rs

adriangb · 2023-05-02T06:10:43Z

src/validators/definitions.rs

+    fn validate_assignment<'s, 'data: 's>(
+        &'s self,
+        py: Python<'data>,
+        obj: &'data PyAny,
+        field_name: &'data str,
+        field_value: &'data PyAny,
+        extra: &Extra,
+        slots: &'data [CombinedValidator],
+        recursion_guard: &'s mut RecursionGuard,


This solves a bug where validate_assignment straight up didn't work with any def refs because they didn't implement it.

adriangb · 2023-05-02T06:11:00Z

tests/serializers/test_definitions.py

@@ -13,7 +11,7 @@ def test_custom_ser():
        )
    )
    assert s.to_python([1, 2, 3]) == ['1', '2', '3']
-    assert plain_repr(s).endswith('slots=[])')
+    # assert plain_repr(s).endswith('slots=[])')


What to do with these tests? Not sure.

this is testing that slots weren't used, I guess because we previously identified that the validator was only used once, so inlined it.

adriangb · 2023-05-02T06:14:17Z

src/validators/mod.rs

-        if build_context.ref_used(&schema_ref) {
-            // the ref is used somewhere
-            // check the ref is unique
-            if build_context.ref_already_used(&schema_ref) {
-                return py_err!("Duplicate ref: `{}`", schema_ref);
-            }
-
-            return if build_context.ref_used_within(schema_dict, &schema_ref)? {
-                // the ref is used within itself, so we have to store the validator in slots
-                // and return a DefinitionRefValidator
-                let slot_id = build_context.prepare_slot(schema_ref)?;
-                let inner_val = T::build(schema_dict, config, build_context)?;
-                let name = inner_val.get_name().to_string();
-                build_context.complete_slot(slot_id, inner_val)?;
-                Ok(definitions::DefinitionRefValidator::from_id(slot_id, name))
-            } else {
-                // ref is used, but only out side itself, we want to clone it everywhere it's used
-                let validator = T::build(schema_dict, config, build_context)?;
-                build_context.store_reusable(schema_ref, validator.clone());
-                Ok(validator)
-            };
-        }
+        let validator_id = build_context.get_ref_id(schema_ref);
+        let inner_val = T::build(schema_dict, config, build_context)?;
+        build_context.add_definition(validator_id, inner_val)?;
+        return Ok(DefinitionRefValidator::from_id(validator_id).into());


The big question here is: what is the cost of this change? In theory, it's large. As you can see anything with a ref permanently becomes a DefinitionRefValidator, which also always use slots so they'll be "slow".

However this is not what needs to happen in practice: it would be quite easy to, in Python, remove any unused refs. Certainly much much easier than maintaining the complex system of refs and slots and whatnot that we had here. And in the end, it will be just as efficient because we can give pydantic-core the most efficient arrangement of references to begin with instead of trying to jump through hoops in Rust to get the borrow checker to accept circular references.
If we do that in Pydantic we can even enforce it here if we want to.

seems fine to me, but are we sure it's okay for core-schema's validator? If that's not effected, I'm fine to move it to pydantic.

The other option is to have another pass at reference counting in pydantic-core after this is merged?

In particular, I think we can copy this stuff into pydantic-core: https://github.com/pydantic/pydantic/blob/de0cf4bca502aad9f7e747cae60298e654bc5a75/pydantic/_internal/_core_utils.py#L222-L423. At which point we're just moving some complex logic from Rust into Python. I don't think it was much faster in Rust to begin with (in fact, I think it may have been slower because it was recursing into Python objects more) and it's certainly easier to debug in Python. And we wouldn't need to do any fancy stuff in Rust.

That said, yes I see the possibility of coming back after this PR and adding back complexity to do the schema inlining / reference counting.

adriangb · 2023-05-02T06:15:17Z

By the way, this is not based on the current main because I needed to do a lot of testing against pydantic proper and the current main is quite broken when used with pydantic 😞

samuelcolvin

overall I think this change is good.

I'm marginally in favour of implementing reference counting in pydantic-core rather than changing pydantic, but:

I can see both arguments
I think it should be done after this PR.

src/build_context.rs

src/serializers/infer.rs

src/serializers/type_serializers/definitions.rs

samuelcolvin · 2023-05-02T12:40:27Z

src/validators/mod.rs

-        if build_context.ref_used(&schema_ref) {
-            // the ref is used somewhere
-            // check the ref is unique
-            if build_context.ref_already_used(&schema_ref) {
-                return py_err!("Duplicate ref: `{}`", schema_ref);
-            }
-
-            return if build_context.ref_used_within(schema_dict, &schema_ref)? {
-                // the ref is used within itself, so we have to store the validator in slots
-                // and return a DefinitionRefValidator
-                let slot_id = build_context.prepare_slot(schema_ref)?;
-                let inner_val = T::build(schema_dict, config, build_context)?;
-                let name = inner_val.get_name().to_string();
-                build_context.complete_slot(slot_id, inner_val)?;
-                Ok(definitions::DefinitionRefValidator::from_id(slot_id, name))
-            } else {
-                // ref is used, but only out side itself, we want to clone it everywhere it's used
-                let validator = T::build(schema_dict, config, build_context)?;
-                build_context.store_reusable(schema_ref, validator.clone());
-                Ok(validator)
-            };
-        }
+        let validator_id = build_context.get_ref_id(schema_ref);
+        let inner_val = T::build(schema_dict, config, build_context)?;
+        build_context.add_definition(validator_id, inner_val)?;
+        return Ok(DefinitionRefValidator::from_id(validator_id).into());


seems fine to me, but are we sure it's okay for core-schema's validator? If that's not effected, I'm fine to move it to pydantic.

The other option is to have another pass at reference counting in pydantic-core after this is merged?

samuelcolvin · 2023-05-02T12:41:27Z

tests/serializers/test_definitions.py

@@ -13,7 +11,7 @@ def test_custom_ser():
        )
    )
    assert s.to_python([1, 2, 3]) == ['1', '2', '3']
-    assert plain_repr(s).endswith('slots=[])')
+    # assert plain_repr(s).endswith('slots=[])')


this is testing that slots weren't used, I guess because we previously identified that the validator was only used once, so inlined it.

tests/validators/test_union.py

codspeed-hq · 2023-05-02T19:01:02Z

CodSpeed Performance Report

Merging #571 rework-references (04c2936) will not alter performances.

Summary

🔥 0 improvements
❌ 1 regressions
✅ 111 untouched benchmarks

🆕 0 new benchmarks
⁉️ 0 dropped benchmarks

Benchmarks breakdown

	Benchmark	`main`	`rework-references`	Change
❌	`test_definition_out_of_tree`	28.3 ms	32.1 ms	-13.38%

codecov-commenter · 2023-05-02T21:45:32Z

Codecov Report

Merging #571 (7b33040) into main (817af72) will decrease coverage by 0.11%.
The diff coverage is 91.82%.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #571      +/-   ##
==========================================
- Coverage   94.24%   94.14%   -0.11%     
==========================================
  Files          95       95              
  Lines       12904    12877      -27     
  Branches       24       24              
==========================================
- Hits        12162    12123      -39     
- Misses        736      748      +12     
  Partials        6        6

Impacted Files	Coverage Δ
src/lib.rs	`100.00% <ø> (ø)`
src/validators/dataclass.rs	`94.11% <66.66%> (+0.22%)`	⬆️
src/validators/is_subclass.rs	`80.48% <66.66%> (ø)`
src/validators/url.rs	`93.41% <66.66%> (ø)`
src/validators/lax_or_strict.rs	`80.76% <71.42%> (ø)`
src/validators/chain.rs	`85.07% <75.00%> (ø)`
src/validators/custom_error.rs	`87.32% <75.00%> (ø)`
src/validators/function.rs	`90.59% <76.92%> (ø)`
src/validators/bytes.rs	`90.24% <80.00%> (+1.21%)`	⬆️
src/validators/json.rs	`97.87% <80.00%> (ø)`
... and 59 more

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 817af72...7b33040. Read the comment docs.

samuelcolvin

a few small things, after that I think this is ready to merge.

samuelcolvin · 2023-05-02T20:38:34Z

src/definitions.rs

+    // Get a ReferenceId for the given reference string.
+    // This ReferenceId can later be used to retrieve a definition


Suggested change

// Get a ReferenceId for the given reference string.

// This ReferenceId can later be used to retrieve a definition

/// Get a ReferenceId for the given reference string.

/// This ReferenceId can later be used to retrieve a definition

and below.

https://doc.rust-lang.org/reference/comments.html#doc-comments TIL, thank you!

samuelcolvin · 2023-05-02T20:39:00Z

src/definitions.rs

+    // Get a ReferenceId for the given reference string.
+    // This ReferenceId can later be used to retrieve a definition
+    pub fn get_reference_id(&mut self, reference: &str) -> ReferenceId {
+        let next_id = self.definitions.len();


can this be moved down to line 45?

I think not because we borrow definitions as mutable. I can do it here because usize will just get copied freeing up definitions to be borrowed

samuelcolvin · 2023-05-02T20:39:09Z

src/definitions.rs

+
+    // Add a definition, returning the ReferenceId that maps to it
+    pub fn add_definition(&mut self, reference: String, value: T) -> PyResult<ReferenceId> {
+        let next_id = self.definitions.len();


Idem above, I don’t think we can, or at least I wasn’t able to; I tried

src/definitions.rs

src/serializers/extra.rs

src/serializers/infer.rs

src/serializers/mod.rs

adriangb self-assigned this May 2, 2023

adriangb commented May 2, 2023

View reviewed changes

adriangb requested review from samuelcolvin and dmontagu and removed request for samuelcolvin May 2, 2023 06:14

adriangb marked this pull request as ready for review May 2, 2023 06:14

adriangb requested a review from samuelcolvin May 2, 2023 06:15

samuelcolvin reviewed May 2, 2023

View reviewed changes

adriangb force-pushed the rework-references branch 2 times, most recently from 7ce340e to 061b721 Compare May 2, 2023 18:55

samuelcolvin reviewed May 2, 2023

View reviewed changes

adriangb added 9 commits May 2, 2023 21:21

Rework def ref system to fix various bugs

fd4d5ae

add comment

d4658b6

better test

dfe212e

bulk rename slots -> definitions

5b54998

use finish() instead of build()

2b915b0

remove format!

fef4174

use proper comment format

6f28ae5

use a Definitions type for documentation / discoverability

58da950

better doc

9d5ad36

adriangb force-pushed the rework-references branch from 4cd975a to 9d5ad36 Compare May 3, 2023 02:21

Fix test (add URL)

3f7fc28

adriangb enabled auto-merge (squash) May 3, 2023 02:23

fix errors

04c2936

adriangb merged commit c655030 into main May 3, 2023

adriangb deleted the rework-references branch May 3, 2023 04:14

		// Get a ReferenceId for the given reference string.
		// This ReferenceId can later be used to retrieve a definition

Rework def ref system to fix various bugs #571

Rework def ref system to fix various bugs #571

Uh oh!

Conversation

adriangb commented May 2, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

adriangb commented May 2, 2023

Uh oh!

samuelcolvin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codspeed-hq bot commented May 2, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Benchmarks breakdown

Uh oh!

codecov-commenter commented May 2, 2023

Codecov Report

Uh oh!

samuelcolvin left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

adriangb May 2, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

adriangb commented May 2, 2023 •

edited

Loading

codspeed-hq bot commented May 2, 2023 •

edited

Loading

adriangb May 2, 2023 •

edited

Loading