Repository: pydantic/pydantic
Pull request: Fix discriminated union schema gen bug
Fix https://github.com/pydantic/pydantic/issues/8709

With this code snippet:

```py
from __future__ import annotations
from typing import Literal, Annotated
from pydantic import Field, TypeAdapter, BaseModel

class NestedState(BaseModel):
    state_type: Literal["nested"]
    substate: AnyState

# If this type is left out, the model behaves normally again
class LoopState(BaseModel):
    state_type: Literal["loop"]
    substate: AnyState

class LeafState(BaseModel):
    state_type: Literal["leaf"]

AnyState = Annotated[NestedState | LoopState | LeafState, Field(..., discriminator="state_type")]

def build_nested_state(n):
    if n <= 0:
        return {"state_type": "leaf"}
    else:
        return {"state_type": "loop", "substate": {"state_type": "nested", "substate": build_nested_state(n-1)}}
        
adapter = TypeAdapter(AnyState)
```

Previously:

```py
# the next statement takes around 0.8s
adapter.validate_python(build_nested_state(9))

# the next statement takes around 3.5s
adapter.validate_python(build_nested_state(10))

# the next statement takes around 12.5s
adapter.validate_python(build_nested_state(11))
```

Now, note the speedups:

```py
start = time.time()
adapter.validate_python(build_nested_state(9))
print(time.time() - start)
#> 0.0063190460205078125

start = time.time()
adapter.validate_python(build_nested_state(10))
print(time.time() - start)
#> 0.02297687530517578

start = time.time()
adapter.validate_python(build_nested_state(11))
print(time.time() - start)
#> 0.057311058044433594
```

Problem statement:
Recursive model with discriminated union has exponential time and space complexity for validation
### Initial Checks

- [X] I confirm that I'm using Pydantic V2

### Description

When validating the recursive model given in the example, Pydantic has exponential time and space complexity, meaning that validation time and memory usage may grow indefinitely and rapidly depending on model depth, even when using a discriminator field.

Given the presence of a discriminator field, this should not be the case.

A model like this can potentially be abused in a denial-of-service attack, if the data passed to the validator can be controlled by the user.

### Example Code

```Python
from __future__ import annotations
from typing import Literal, Annotated
from pydantic import Field, TypeAdapter, BaseModel

class NestedState(BaseModel):
    state_type: Literal["nested"]
    substate: AnyState

# If this type is left out, the model behaves normally again
class LoopState(BaseModel):
    state_type: Literal["loop"]
    substate: AnyState

class LeafState(BaseModel):
    state_type: Literal["leaf"]

AnyState = Annotated[NestedState | LoopState | LeafState, Field(..., discriminator="state_type")]

def build_nested_state(n):
    if n <= 0:
        return {"state_type": "leaf"}
    else:
        return {"state_type": "loop", "substate": {"state_type": "nested", "substate": build_nested_state(n-1)}}
        
adapter = TypeAdapter(AnyState)

# the next statement takes around 0.8s
adapter.validate_python(build_nested_state(9))

# the next statement takes around 3.5s
adapter.validate_python(build_nested_state(10))

# the next statement takes around 12.5s
adapter.validate_python(build_nested_state(11))
```


### Python, Pydantic & OS Version

```Text
pydantic version: 2.6.0
        pydantic-core version: 2.16.1
          pydantic-core build: profile=release pgo=true
                 install path: /<redacted>/venv-3.11/lib/python3.11/site-packages/pydantic
               python version: 3.11.7 (main, Dec  4 2023, 18:10:11) [Clang 15.0.0 (clang-1500.1.0.2.5)]
                     platform: macOS-14.3-arm64-arm-64bit
             related packages: fastapi-0.105.0 typing_extensions-4.8.0 pydantic-settings-2.1.0
                       commit: unknown
```

Hints from issue discussion:
@palle-k,

Thanks for reporting this. @davidhewitt, could you please take a look at this performance issue when you have a chance? Thank you!
It looks like there's a potential schema-building bug here.

The core schema looks like this:

```
{'definitions': [{'cls': <class '__main__.NestedState'>,
                  'config': {'title': 'NestedState'},
                  'custom_init': False,
                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False,
                               'pydantic_js_annotation_functions': [],
                               'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x7fba5440fce0>, cls=<class '__main__.NestedState'>),
                                                         <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.NestedState'>>]},
                  'ref': '__main__.NestedState:94515929304032',
                  'root_model': False,
                  'schema': {'computed_fields': [],
                             'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x7fba54486fc0>],
                                                                    'pydantic_js_functions': []},
                                                       'schema': {'expected': ['nested'],
                                                                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                  'type': 'literal'},
                                                       'type': 'model-field'},
                                        'substate': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x7fba54487100>],
                                                                  'pydantic_js_functions': []},
                                                     'schema': {'default': Ellipsis,
                                                                'schema': {'choices': [{'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.NestedState:94515929304032',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LoopState:94515929369184',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LeafState:94515929326848',
                                                                                        'type': 'definition-ref'}],
                                                                           'metadata': {'pydantic.internal.needs_apply_discriminated_union': True,
                                                                                        'pydantic.internal.union_discriminator': 'state_type'},
                                                                           'type': 'union'},
                                                                'type': 'default'},
                                                     'type': 'model-field'}},
                             'model_name': 'NestedState',
                             'type': 'model-fields'},
                  'type': 'model'},
                 {'cls': <class '__main__.LoopState'>,
                  'config': {'title': 'LoopState'},
                  'custom_init': False,
                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False,
                               'pydantic_js_annotation_functions': [],
                               'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x7fba5440fce0>, cls=<class '__main__.LoopState'>),
                                                         <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.LoopState'>>]},
                  'ref': '__main__.LoopState:94515929369184',
                  'root_model': False,
                  'schema': {'computed_fields': [],
                             'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x7fba544871a0>],
                                                                    'pydantic_js_functions': []},
                                                       'schema': {'expected': ['loop'],
                                                                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                  'type': 'literal'},
                                                       'type': 'model-field'},
                                        'substate': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x7fba54487380>],
                                                                  'pydantic_js_functions': []},
                                                     'schema': {'default': Ellipsis,
                                                                'schema': {'choices': [{'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.NestedState:94515929304032',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LoopState:94515929369184',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LeafState:94515929326848',
                                                                                        'type': 'definition-ref'}],
                                                                           'metadata': {'pydantic.internal.needs_apply_discriminated_union': True,
                                                                                        'pydantic.internal.union_discriminator': 'state_type'},
                                                                           'type': 'union'},
                                                                'type': 'default'},
                                                     'type': 'model-field'}},
                             'model_name': 'LoopState',
                             'type': 'model-fields'},
                  'type': 'model'},
                 {'cls': <class '__main__.LeafState'>,
                  'config': {'title': 'LeafState'},
                  'custom_init': False,
                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False,
                               'pydantic_js_annotation_functions': [],
                               'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x7fba5440fce0>, cls=<class '__main__.LeafState'>),
                                                         <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.LeafState'>>]},
                  'ref': '__main__.LeafState:94515929326848',
                  'root_model': False,
                  'schema': {'computed_fields': [],
                             'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x7fba54484fe0>],
                                                                    'pydantic_js_functions': []},
                                                       'schema': {'expected': ['leaf'],
                                                                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                  'type': 'literal'},
                                                       'type': 'model-field'}},
                             'model_name': 'LeafState',
                             'type': 'model-fields'},
                  'type': 'model'}],
 'schema': {'choices': {'leaf': {'schema_ref': '__main__.LeafState:94515929326848',
                                 'type': 'definition-ref'},
                        'loop': {'schema_ref': '__main__.LoopState:94515929369184',
                                 'type': 'definition-ref'},
                        'nested': {'schema_ref': '__main__.NestedState:94515929304032',
                                   'type': 'definition-ref'}},
            'discriminator': 'state_type',
            'from_attributes': True,
            'metadata': {'pydantic.internal.needs_apply_discriminated_union': True,
                         'pydantic.internal.union_discriminator': 'state_type'},
            'strict': False,
            'type': 'tagged-union'},
 'type': 'definitions'}
 ```

In particular the interesting bit to me is that we have 3 definitions instead of 4; the `AnyState` union is inlined in each model. But in the inlined model copies it is `'type': 'union'` instead of `'type': 'tagged-union'`, which would explain where there's linear search rather than tagged dispatch at each level.

Hence the exponential blowup.

cc @adriangb as the schema building expert ;)
The schema simplifies to:

```json
{
    "type": "definitions",
    "schema": {
        "type": "tagged-union",
        "choices": {
            "nested": {
                "type": "definition-ref",
                "schema_ref": "__main__.NestedState:4587657680"
            },
            "loop": {
                "type": "definition-ref",
                "schema_ref": "__main__.LoopState:4587670848"
            },
            "leaf": {
                "type": "definition-ref",
                "schema_ref": "__main__.LeafState:5124597648"
            }
        },
        "discriminator": "state_type"
    },
    "definitions": [
        {
            "type": "model",
            "cls": "__main__.LeafState",
            "schema": {
                "type": "model-fields",
                "fields": {
                    "state_type": {
                        "type": "model-field",
                        "schema": {
                            "type": "literal",
                            "expected": [
                                "leaf"
                            ]
                        }
                    }
                },
                "model_name": "LeafState"
            },
            "ref": "__main__.LeafState:5124597648"
        },
        {
            "type": "model",
            "cls": "__main__.LoopState",
            "schema": {
                "type": "model-fields",
                "fields": {
                    "state_type": {
                        "type": "model-field",
                        "schema": {
                            "type": "literal",
                            "expected": [
                                "loop"
                            ]
                        }
                    },
                    "substate": {
                        "type": "model-field",
                        "schema": {
                            "type": "default",
                            "schema": {
                                "type": "union",
                                "choices": [
                                    {
                                        "type": "definition-ref",
                                        "schema_ref": "__main__.NestedState:4587657680"
                                    },
                                    {
                                        "type": "definition-ref",
                                        "schema_ref": "__main__.LoopState:4587670848"
                                    },
                                    {
                                        "type": "definition-ref",
                                        "schema_ref": "__main__.LeafState:5124597648"
                                    }
                                ]
                            }
                        }
                    }
                }
            },
            "ref": "__main__.LoopState:4587670848"
        },
        {
            "type": "model",
            "cls": "__main__.NestedState",
            "schema": {
                "type": "model-fields",
                "fields": {
                    "state_type": {
                        "type": "model-field",
                        "schema": {
                            "type": "literal",
                            "expected": [
                                "nested"
                            ]
                        }
                    },
                    "substate": {
                        "type": "model-field",
                        "schema": {
                            "type": "default",
                            "schema": {
                                "type": "union",
                                "choices": [
                                    {
                                        "type": "definition-ref",
                                        "schema_ref": "__main__.NestedState:4587657680"
                                    },
                                    {
                                        "type": "definition-ref",
                                        "schema_ref": "__main__.LoopState:4587670848"
                                    },
                                    {
                                        "type": "definition-ref",
                                        "schema_ref": "__main__.LeafState:5124597648"
                                    }
                                ]
                            }
                        }
                    }
                }
            },
            "ref": "__main__.NestedState:4587657680"
        }
    ]
}
```

The issue is that the tagged union is not carried along so it's loss after the first union and thus performance degrades.
So yes it's a schema generation bug.

@sydney-runkle is that enough to go off of? I'm guessing we're not doing the `apply_discrimnators` stuff correctly.
Hmm, this looks similar to this issue: https://github.com/pydantic/pydantic/issues/8688, which we fixed here https://github.com/pydantic/pydantic/pull/8702 for the 2.6.1 release. @palle-k, could you please confirm whether or not this issue has been resolved with 2.6.1?
Afraid I'm reproducing this on main 😢 
@sydney-runkle unfortunately still reproducible with 2.6.1, as David already noticed.
@davidhewitt, ah rats, ok 🐀.
+1 this seems pretty critical 😬, is there any workaround?
This is a regression because I don't see this issue with
```
pydantic==2.4.2
pydantic_core==2.10.1
```

Unfortunately I can't downgrade because my dependencies have updated to >2.5.0
@rossmacarthur,

Thanks for the +1. I've added to the list of issues we'll attempt to fix for 2.6.2. Hopefully we can get that release out next week. I'll keep you updated as we progress. I don't currently know of a workaround for this performance issue.

Commit message: fix for earlier python versions

Patch to review:
diff --git a/pydantic/_internal/_discriminated_union.py b/pydantic/_internal/_discriminated_union.py
index e3806b8f2fe..6611404e599 100644
--- a/pydantic/_internal/_discriminated_union.py
+++ b/pydantic/_internal/_discriminated_union.py
@@ -54,6 +54,7 @@ def inner(s: core_schema.CoreSchema, recurse: _core_utils.Recurse) -> core_schem
         if discriminator is not None:
             if definitions is None:
                 definitions = collect_definitions(schema)
+                definitions = {k: recurse(v, inner) for k, v in definitions.items()}
             s = apply_discriminator(s, discriminator, definitions)
         return s
 
diff --git a/tests/test_discriminated_union.py b/tests/test_discriminated_union.py
index 0e3a0eeadba..383a767ca51 100644
--- a/tests/test_discriminated_union.py
+++ b/tests/test_discriminated_union.py
@@ -1852,3 +1852,27 @@ class SubModel(MyModel):
         'title': 'MyModel',
         'type': 'object',
     }
+
+
+def test_nested_schema_gen_uses_tagged_union_in_ref() -> None:
+    class NestedState(BaseModel):
+        state_type: Literal['nested']
+        substate: 'AnyState'
+
+    # If this type is left out, the model behaves normally again
+    class LoopState(BaseModel):
+        state_type: Literal['loop']
+        substate: 'AnyState'
+
+    class LeafState(BaseModel):
+        state_type: Literal['leaf']
+
+    AnyState = Annotated[Union[NestedState, LoopState, LeafState], Field(..., discriminator='state_type')]
+    NestedState.model_rebuild()
+    LoopState.model_rebuild()
+    adapter = TypeAdapter(AnyState)
+
+    assert adapter.core_schema['schema']['type'] == 'tagged-union'
+    for definition in adapter.core_schema['definitions']:
+        if definition['schema']['model_name'] in ['NestedState', 'LoopState']:
+            assert definition['schema']['fields']['substate']['schema']['schema']['type'] == 'tagged-union'


Code review comments:
diff_hunk: @@ -54,6 +54,7 @@ def inner(s: core_schema.CoreSchema, recurse: _core_utils.Recurse) -> core_schem
         if discriminator is not None:
             if definitions is None:
                 definitions = collect_definitions(schema)
+                definitions = {k: recurse(v, inner) for k, v in definitions.items()}
             s = apply_discriminator(s, discriminator, definitions)
path: pydantic/_internal/_discriminated_union.py
text: @user1:
🤯 care to explain how this fixes things?

@author:
Sure thing! I'll add some more comments in the code to explain things 😄!

I'll also add a more in-depth note to this PR, so we can reference that for more detail, if needed 👍.

@author:
So previously, with the above snippet of code, the resultant `core_schema` was the following:

```py
{'definitions': [{'cls': <class '__main__.LeafState'>,
                  'config': {'title': 'LeafState'},
                  'custom_init': False,
                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False,
                               'pydantic_js_annotation_functions': [],
                               'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x101bcd900>, cls=<class '__main__.LeafState'>),
                                                         <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.LeafState'>>]},
                  'ref': '__main__.LeafState:5484303040',
                  'root_model': False,
                  'schema': {'computed_fields': [],
                             'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101c09d80>],
                                                                    'pydantic_js_functions': []},
                                                       'schema': {'expected': ['leaf'],
                                                                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                  'type': 'literal'},
                                                       'type': 'model-field'}},
                             'model_name': 'LeafState',
                             'type': 'model-fields'},
                  'type': 'model'},
                 {'cls': <class '__main__.LoopState'>,
                  'config': {'title': 'LoopState'},
                  'custom_init': False,
                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False,
                               'pydantic_js_annotation_functions': [],
                               'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x101bcd900>, cls=<class '__main__.LoopState'>),
                                                         <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.LoopState'>>]},
                  'ref': '__main__.LoopState:5484288784',
                  'root_model': False,
                  'schema': {'computed_fields': [],
                             'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf8310>],
                                                                    'pydantic_js_functions': []},
                                                       'schema': {'expected': ['loop'],
                                                                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                  'type': 'literal'},
                                                       'type': 'model-field'},
                                        'substate': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf83a0>],
                                                                  'pydantic_js_functions': []},
                                                     'schema': {'default': Ellipsis,
                                                                'schema': {'choices': [{'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.NestedState:5484231952',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LoopState:5484288784',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LeafState:5484303040',
                                                                                        'type': 'definition-ref'}],
                                                                           'metadata': {'pydantic.internal.needs_apply_discriminated_union': True,
                                                                                        'pydantic.internal.union_discriminator': 'state_type'},
                                                                           'type': 'union'},
                                                                'type': 'default'},
                                                     'type': 'model-field'}},
                             'model_name': 'LoopState',
                             'type': 'model-fields'},
                  'type': 'model'},
                 {'cls': <class '__main__.NestedState'>,
                  'config': {'title': 'NestedState'},
                  'custom_init': False,
                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False,
                               'pydantic_js_annotation_functions': [],
                               'pydantic_js_functions': [functools.partial(<function modify_model_json_schema at 0x101bcd900>, cls=<class '__main__.NestedState'>),
                                                         <bound method BaseModel.__get_pydantic_json_schema__ of <class '__main__.NestedState'>>]},
                  'ref': '__main__.NestedState:5484231952',
                  'root_model': False,
                  'schema': {'computed_fields': [],
                             'fields': {'state_type': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf80d0>],
                                                                    'pydantic_js_functions': []},
                                                       'schema': {'expected': ['nested'],
                                                                  'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                  'type': 'literal'},
                                                       'type': 'model-field'},
                                        'substate': {'metadata': {'pydantic_js_annotation_functions': [<function get_json_schema_update_func.<locals>.json_schema_update_func at 0x101cf8160>],
                                                                  'pydantic_js_functions': []},
                                                     'schema': {'default': Ellipsis,
                                                                'schema': {'choices': [{'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.NestedState:5484231952',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LoopState:5484288784',
                                                                                        'type': 'definition-ref'},
                                                                                       {'metadata': {'pydantic.internal.needs_apply_discriminated_union': False},
                                                                                        'schema_ref': '__main__.LeafState:5484303040',
                                                                                        'type': 'definition-ref'}],
                                                                           'metadata': {'pydantic.internal.needs_apply_discriminated_union': True,
                                                                                        'pydantic.internal.union_discriminator': 'state_type'},
                                                                           'type': 'union'},
                                                                'type': 'default'},
                                                     'type': 'model-field'}},
                             'model_name': 'NestedState',
                             'type': 'model-fields'},
                  'type': 'model'}],
 'schema': {'choices': {'leaf': {'schema_ref': '__main__.LeafState:5484303040',
                                 'type': 'definition-ref'},
                        'loop': {'schema_ref': '__main__.LoopState:5484288784',
                                 'type': 'definition-ref'},
                        'nested': {'schema_ref': '__main__.NestedState:5484231952',
                                   'type': 'definition-ref'}},
            'discriminator': 'state_type',
            'from_attributes': True,
            'metadata': {'pydantic.internal.needs_apply_discriminated_union': True,
                         'pydantic.internal.union_discriminator': 'state_type'},
            'strict': False,
            'type': 'tagged-union'},
 'type': 'definitions'}
```

We already identified that the schemas for `NestedState` and `LoopState` within the `definitions` list were marked as having `union` type `substate` schemas, rather than `tagged-union` `substate` schemas. Interestingly, though, the `metadata` for those schemas had the necessary discriminated-union related information:

```
'metadata': {'pydantic.internal.needs_apply_discriminated_union': True, 'pydantic.internal.union_discriminator': 'state_type'}
``` 

So the question, then, is why aren't we "applying" the discriminators to those schemas? This is where my change comes in - after we collect the definitions schemas, we also need to pass over them + apply discriminators where appropriate. Within the `inner` function, we were already calling `recurse`, but we weren't calling it on any definitions schemas, which was the missing piece.

@user1:
Totally makes sense!

@user1:
Totally makes sense!