Repository: streamlit/streamlit
Issue category: Performance Issue

Exponential slowdown in Streamlit as the number of widgets increases (starting in v1.34)
### Checklist

- [X] I have searched the [existing issues](https://github.com/streamlit/streamlit/issues) for similar issues.
- [X] I added a very descriptive title to this issue.
- [X] I have provided sufficient information below to help reproduce this issue.

### Summary

There is an exponential slowdown in Streamlit as the number of input widgets is increases. This started happening with the 1.34 release and is tied to PR #8529. Calling [`get_session_state()`](https://github.com/streamlit/streamlit/blob/develop/lib/streamlit/elements/widgets/number_input.py#L390) in the widgets is triggering this [dict comp method](https://github.com/streamlit/streamlit/blob/develop/lib/streamlit/runtime/state/session_state.py#L370) exponentially more times than previously. For a simple app containing 200 number inputs, there is a 50x increase in calls to `_reverse_key_wid_map` which causes the app to take over 1s to re-load locally (vs ~0.2s with release 1.33). Here is a plot of the number of calls vs the number of inputs:

![image](https://github.com/user-attachments/assets/4b0ef9c9-49c4-482b-85d9-9684e41f523d)

### Reproducible Code Example

```Python
import streamlit as st
import time
import random

NUM_TABS = 10
NUM_INPUTS = 20


def create_inputs(t):
    for i in range(NUM_INPUTS):
        key = f"tab_{t + 1}_input_{i + 1}_value"
        if key in st.session_state:
            value = st.session_state[key]
        else:
            value = random.randint(0, 100000)
        st.number_input(
            label=f"Number Input {i + 1}",
            value=value,
            min_value=0,
            step=1,
            format="%d",
            key=key,
        )


def main():
    st.title("My dummy app")
    with st.container(border=True, height=500):
        tabs = st.tabs([f"Tab {i + 1}" for i in range(NUM_TABS)])
        for t, tab in enumerate(tabs):
            with tab:
                create_inputs(t)


if __name__ == "__main__":
    import cProfile
    from pstats import Stats

    pr = cProfile.Profile()
    pr.enable()

    start_time = time.time()
    main()
    st.info(f"run time: {time.time() - start_time:.3f} seconds")

    pr.disable()
    stats = Stats(pr)
    stats.sort_stats("tottime").print_stats(10)
```

### Steps To Reproduce

1. Run the app using Streamlit version 1.33
2. Adjust `NUM_TABS` value from 1 to 10 and monitor the `cProfile` output
3. Switch to Streamlit version 1.34 or above
4. Repeat step 2

here is a screenshot of the cProfile output with Streamlit 1.33 using 200 widgets:
![image](https://github.com/user-attachments/assets/2c39e939-5cc9-407a-99f1-e1c9de0257b2)

and here's a screenshot using Streamlit 1.34:
![image](https://github.com/user-attachments/assets/3da40e2d-025c-4752-9a96-6b10b233a0bd)

### Expected Behavior

_No response_

### Current Behavior

_No response_

### Is this a regression?

- [x] Yes, this used to work in a previous version.

### Debug info

- Streamlit version: 1.33 vs 1.34+
- Python version: 3.11
- Operating System: Windows 10
- Browser: Chrome


### Additional Information

_No response_


Functions to edit:
lib/streamlit/runtime/state/session_state.py:SessionState
lib/streamlit/runtime/state/session_state.py:SessionState.clear
lib/streamlit/runtime/state/session_state.py:SessionState.filtered_state
lib/streamlit/runtime/state/session_state.py:SessionState._reverse_key_wid_map
lib/streamlit/runtime/state/session_state.py:SessionState.__getitem__
lib/streamlit/runtime/state/session_state.py:SessionState.__setitem__
lib/streamlit/runtime/state/session_state.py:SessionState.__delitem__
lib/streamlit/runtime/state/session_state.py:SessionState._get_widget_id
lib/streamlit/runtime/state/session_state.py:SessionState._set_key_widget_mapping
Functions added: lib/streamlit/runtime/state/session_state.py:KeyIdMapper
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.__contains__
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.__setitem__
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.__delitem__
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.id_key_mapping
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.set_key_id_mapping
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.get_id_from_key
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.get_key_from_id
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.update
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.clear
lib/streamlit/runtime/state/session_state.py:KeyIdMapper.delete
lib/streamlit/runtime/state/session_state.py:SessionState

Patch:
diff --git a/lib/streamlit/runtime/state/session_state.py b/lib/streamlit/runtime/state/session_state.py
index 5481cd2be27a..b9f016e19ad7 100644
--- a/lib/streamlit/runtime/state/session_state.py
+++ b/lib/streamlit/runtime/state/session_state.py
@@ -280,6 +280,56 @@ def _missing_key_error_message(key: str) -> str:
     )
 
 
+@dataclass
+class KeyIdMapper:
+    """A mapping of user-provided keys to element IDs.
+    It also maps element IDs to user-provided keys so that this reverse mapping
+    does not have to be computed ad-hoc.
+    All built-in dict-operations such as setting and deleting expect the key as the
+    argument, not the element ID.
+    """
+
+    _key_id_mapping: dict[str, str] = field(default_factory=dict)
+    _id_key_mapping: dict[str, str] = field(default_factory=dict)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._key_id_mapping
+
+    def __setitem__(self, key: str, widget_id: Any) -> None:
+        self._key_id_mapping[key] = widget_id
+        self._id_key_mapping[widget_id] = key
+
+    def __delitem__(self, key: str) -> None:
+        self.delete(key)
+
+    @property
+    def id_key_mapping(self) -> dict[str, str]:
+        return self._id_key_mapping
+
+    def set_key_id_mapping(self, key_id_mapping: dict[str, str]) -> None:
+        self._key_id_mapping = key_id_mapping
+        self._id_key_mapping = {v: k for k, v in key_id_mapping.items()}
+
+    def get_id_from_key(self, key: str, default: Any = None) -> str:
+        return self._key_id_mapping.get(key, default)
+
+    def get_key_from_id(self, widget_id: str) -> str:
+        return self._id_key_mapping[widget_id]
+
+    def update(self, other: KeyIdMapper) -> None:
+        self._key_id_mapping.update(other._key_id_mapping)
+        self._id_key_mapping.update(other._id_key_mapping)
+
+    def clear(self):
+        self._key_id_mapping.clear()
+        self._id_key_mapping.clear()
+
+    def delete(self, key: str):
+        widget_id = self._key_id_mapping[key]
+        del self._key_id_mapping[key]
+        del self._id_key_mapping[widget_id]
+
+
 @dataclass
 class SessionState:
     """SessionState allows users to store values that persist between app
@@ -308,10 +358,11 @@ class SessionState:
     # Widget values from the frontend, usually one changing prompted the script rerun
     _new_widget_state: WStates = field(default_factory=WStates)
 
-    # Keys used for widgets will be eagerly converted to the matching widget id
-    _key_id_mapping: dict[str, str] = field(default_factory=dict)
+    # Keys used for widgets will be eagerly converted to the matching element id
+    _key_id_mapper: KeyIdMapper = field(default_factory=KeyIdMapper)
 
-    # query params are stored in session state because query params will be tied with widget state at one point.
+    # query params are stored in session state because query params will be tied with
+    # widget state at one point.
     query_params: QueryParams = field(default_factory=QueryParams)
 
     def __repr__(self):
@@ -338,13 +389,13 @@ def clear(self) -> None:
         self._old_state.clear()
         self._new_session_state.clear()
         self._new_widget_state.clear()
-        self._key_id_mapping.clear()
+        self._key_id_mapper.clear()
 
     @property
     def filtered_state(self) -> dict[str, Any]:
         """The combined session and widget state, excluding keyless widgets."""
 
-        wid_key_map = self._reverse_key_wid_map
+        wid_key_map = self._key_id_mapper.id_key_mapping
 
         state: dict[str, Any] = {}
 
@@ -366,12 +417,6 @@ def filtered_state(self) -> dict[str, Any]:
 
         return state
 
-    @property
-    def _reverse_key_wid_map(self) -> dict[str, str]:
-        """Return a mapping of widget_id : widget_key."""
-        wid_key_map = {v: k for k, v in self._key_id_mapping.items()}
-        return wid_key_map
-
     def _keys(self) -> set[str]:
         """All keys active in Session State, with widget keys converted
         to widget ids when one is known. (This includes autogenerated keys
@@ -400,7 +445,7 @@ def __len__(self) -> int:
         return len(self._keys())
 
     def __getitem__(self, key: str) -> Any:
-        wid_key_map = self._reverse_key_wid_map
+        wid_key_map = self._key_id_mapper.id_key_mapping
         widget_id = self._get_widget_id(key)
 
         if widget_id in wid_key_map and widget_id == key:
@@ -463,7 +508,7 @@ def __setitem__(self, user_key: str, value: Any) -> None:
         ctx = get_script_run_ctx()
 
         if ctx is not None:
-            widget_id = self._key_id_mapping.get(user_key, None)
+            widget_id = self._key_id_mapper.get_id_from_key(user_key, None)
             widget_ids = ctx.widget_ids_this_run
             form_ids = ctx.form_ids_this_run
 
@@ -487,8 +532,8 @@ def __delitem__(self, key: str) -> None:
         if key in self._old_state:
             del self._old_state[key]
 
-        if key in self._key_id_mapping:
-            del self._key_id_mapping[key]
+        if key in self._key_id_mapper:
+            self._key_id_mapper.delete(key)
 
         if widget_id in self._new_widget_state:
             del self._new_widget_state[widget_id]
@@ -608,10 +653,10 @@ def _get_widget_id(self, k: str) -> str:
         """Turns a value that might be a widget id or a user provided key into
         an appropriate widget id.
         """
-        return self._key_id_mapping.get(k, k)
+        return self._key_id_mapper.get_id_from_key(k, k)
 
     def _set_key_widget_mapping(self, widget_id: str, user_key: str) -> None:
-        self._key_id_mapping[user_key] = widget_id
+        self._key_id_mapper[user_key] = widget_id
 
     def register_widget(
         self, metadata: WidgetMetadata[T], user_key: str | None

Test patch: diff --git a/lib/tests/streamlit/runtime/state/session_state_test.py b/lib/tests/streamlit/runtime/state/session_state_test.py
index a9408a209a9b..6eac40a58a64 100644
--- a/lib/tests/streamlit/runtime/state/session_state_test.py
+++ b/lib/tests/streamlit/runtime/state/session_state_test.py
@@ -38,6 +38,7 @@
 from streamlit.runtime.state import SessionState, get_session_state
 from streamlit.runtime.state.common import GENERATED_ELEMENT_ID_PREFIX
 from streamlit.runtime.state.session_state import (
+    KeyIdMapper,
     Serialized,
     Value,
     WidgetMetadata,
@@ -656,7 +657,9 @@ def test_setitem_disallows_setting_created_widget(self):
             return_value=mock_ctx,
         ):
             with pytest.raises(StreamlitAPIException) as e:
-                self.session_state._key_id_mapping = {"widget_id": "widget_id"}
+                self.session_state._key_id_mapper.set_key_id_mapping(
+                    {"widget_id": "widget_id"}
+                )
                 self.session_state["widget_id"] = "blah"
             assert "`st.session_state.widget_id` cannot be modified" in str(e.value)
 
@@ -800,8 +803,8 @@ def test_map_set_del(m, key, value1):
 
 
 @given(state=stst.session_state())
-def test_key_wid_lookup_equiv(state):
-    k_wid_map = state._key_id_mapping
+def test_key_wid_lookup_equiv(state: SessionState):
+    k_wid_map = state._key_id_mapper._key_id_mapping
     for k, wid in k_wid_map.items():
         assert state[k] == state[wid]
 
@@ -885,13 +888,17 @@ def test_session_state_stats(self):
         stat = state.get_stats()[0]
         assert stat.category_name == "st_session_state"
 
+        # The expected size of the session state in bytes.
+        # It composes of the session_state's fields.
+        expected_session_state_size_bytes = 3000
+
         init_size = stat.byte_length
-        assert init_size < 2500
+        assert init_size < expected_session_state_size_bytes
 
         state["foo"] = 2
         new_size = state.get_stats()[0].byte_length
         assert new_size > init_size
-        assert new_size < 2500
+        assert new_size < expected_session_state_size_bytes
 
         state["foo"] = 1
         new_size_2 = state.get_stats()[0].byte_length
@@ -900,8 +907,74 @@ def test_session_state_stats(self):
         st.checkbox("checkbox", key="checkbox")
         new_size_3 = state.get_stats()[0].byte_length
         assert new_size_3 > new_size_2
-        assert new_size_3 - new_size_2 < 2500
+        assert new_size_3 - new_size_2 < expected_session_state_size_bytes
 
         state._compact_state()
         new_size_4 = state.get_stats()[0].byte_length
         assert new_size_4 <= new_size_3
+
+
+class KeyIdMapperTest(unittest.TestCase):
+    def test_key_id_mapping(self):
+        key_id_mapper = KeyIdMapper()
+        key_id_mapper.set_key_id_mapping({"key": "wid"})
+        assert key_id_mapper.get_id_from_key("key") == "wid"
+        assert key_id_mapper.get_key_from_id("wid") == "key"
+
+    def test_key_id_mapping_errors(self):
+        key_id_mapper = KeyIdMapper()
+        key_id_mapper.set_key_id_mapping({"key": "wid"})
+        assert key_id_mapper.get_id_from_key("nonexistent") is None
+        with pytest.raises(KeyError):
+            key_id_mapper.get_key_from_id("nonexistent")
+
+    def test_key_id_mapping_clear(self):
+        key_id_mapper = KeyIdMapper()
+        key_id_mapper.set_key_id_mapping({"key": "wid"})
+        assert key_id_mapper.get_id_from_key("key") == "wid"
+        assert key_id_mapper.get_key_from_id("wid") == "key"
+        key_id_mapper.clear()
+        assert key_id_mapper.get_id_from_key("key") is None
+        with pytest.raises(KeyError):
+            key_id_mapper.get_key_from_id("wid")
+
+    def test_key_id_mapping_delete(self):
+        key_id_mapper = KeyIdMapper()
+        key_id_mapper.set_key_id_mapping({"key": "wid"})
+        assert key_id_mapper.get_id_from_key("key") == "wid"
+        assert key_id_mapper.get_key_from_id("wid") == "key"
+        del key_id_mapper["key"]
+        assert key_id_mapper.get_id_from_key("key") is None
+        with pytest.raises(KeyError):
+            key_id_mapper.get_key_from_id("wid")
+
+    def test_key_id_mapping_set_key_id_mapping(self):
+        key_id_mapper = KeyIdMapper()
+        key_id_mapper.set_key_id_mapping({"key": "wid"})
+        key_id_mapper["key2"] = "wid2"
+        assert key_id_mapper.get_id_from_key("key") == "wid"
+        assert key_id_mapper.get_key_from_id("wid") == "key"
+        assert key_id_mapper.get_id_from_key("key2") == "wid2"
+        assert key_id_mapper.get_key_from_id("wid2") == "key2"
+
+    def test_key_id_mapping_update(self):
+        key_id_mapper = KeyIdMapper()
+        key_id_mapper.set_key_id_mapping({"key": "wid"})
+        assert key_id_mapper.get_id_from_key("key") == "wid"
+        assert key_id_mapper.get_key_from_id("wid") == "key"
+
+        key_id_mapper2 = KeyIdMapper()
+        key_id_mapper2.set_key_id_mapping({"key2": "wid2"})
+        key_id_mapper.update(key_id_mapper2)
+        assert key_id_mapper.get_id_from_key("key2") == "wid2"
+        assert key_id_mapper.get_key_from_id("wid2") == "key2"
+
+        key_id_mapper3 = KeyIdMapper()
+        key_id_mapper3.set_key_id_mapping({"key": "wid3"})
+        key_id_mapper.update(key_id_mapper3)
+        assert key_id_mapper.get_id_from_key("key") == "wid3"
+        assert key_id_mapper.get_key_from_id("wid3") == "key"
+        assert key_id_mapper.get_id_from_key("key2") == "wid2"
+        assert key_id_mapper.get_key_from_id("wid2") == "key2"
+        assert key_id_mapper.get_id_from_key("key") == "wid3"
+        assert key_id_mapper.get_key_from_id("wid") == "key"
diff --git a/lib/tests/streamlit/runtime/state/strategies.py b/lib/tests/streamlit/runtime/state/strategies.py
index 42c2fc3c20d0..2a4702aa70f4 100644
--- a/lib/tests/streamlit/runtime/state/strategies.py
+++ b/lib/tests/streamlit/runtime/state/strategies.py
@@ -82,7 +82,7 @@ def _merge_states(a: SessionState, b: SessionState) -> None:
     a._new_session_state.update(b._new_session_state)
     a._new_widget_state.update(b._new_widget_state)
     a._old_state.update(b._old_state)
-    a._key_id_mapping.update(b._key_id_mapping)
+    a._key_id_mapper.update(b._key_id_mapper)
 
 
 # TODO: don't generate states where there is a k-wid mapping where the key exists but the widget doesn't