Add a faster query for get_last_state_changes when the number of states is 1 (#90211)

* Add a faster query for get_last_state_changes when the number of states is 1

related issue #90113

* Add a faster query for get_last_state_changes when the number of states is 1

related issue #90113

* coverage

* Apply suggestions from code review

homeassistant/components/recorder/history/modern.py

@@ -406,16 +406,38 @@ def _get_last_state_changes_stmt(
     stmt, join_attributes = _lambda_stmt_and_join_attributes(
         False, include_last_changed=False
     )
-    stmt += lambda q: q.where(
-        States.state_id
-        == (
-            select(States.state_id)
-            .filter(States.metadata_id == metadata_id)
-            .order_by(States.last_updated_ts.desc())
-            .limit(number_of_states)
-            .subquery()
-        ).c.state_id
-    )
+    if number_of_states == 1:
+        stmt += lambda q: q.join(
+            (
+                lastest_state_for_metadata_id := (
+                    select(
+                        States.metadata_id.label("max_metadata_id"),
+                        # https://github.com/sqlalchemy/sqlalchemy/issues/9189
+                        # pylint: disable-next=not-callable
+                        func.max(States.last_updated_ts).label("max_last_updated"),
+                    )
+                    .filter(States.metadata_id == metadata_id)
+                    .group_by(States.metadata_id)
+                    .subquery()
+                )
+            ),
+            and_(
+                States.metadata_id == lastest_state_for_metadata_id.c.max_metadata_id,
+                States.last_updated_ts
+                == lastest_state_for_metadata_id.c.max_last_updated,
+            ),
+        )
+    else:
+        stmt += lambda q: q.where(
+            States.state_id
+            == (
+                select(States.state_id)
+                .filter(States.metadata_id == metadata_id)
+                .order_by(States.last_updated_ts.desc())
+                .limit(number_of_states)
+                .subquery()
+            ).c.state_id
+        )
     if join_attributes:
         stmt += lambda q: q.outerjoin(
             StateAttributes, States.attributes_id == StateAttributes.attributes_id
@@ -432,6 +454,10 @@ def get_last_state_changes(
     entity_id_lower = entity_id.lower()
     entity_ids = [entity_id_lower]
 
+    # Calling this function with number_of_states > 1 can cause instability
+    # because it has to scan the table to find the last number_of_states states
+    # because the metadata_id_last_updated_ts index is in ascending order.
+
     with session_scope(hass=hass, read_only=True) as session:
         instance = recorder.get_instance(hass)
         if not (

tests/components/recorder/test_history.py

@@ -382,6 +382,42 @@ def test_get_last_state_changes(hass_recorder: Callable[..., HomeAssistant]) ->
     assert_multiple_states_equal_without_context(states, hist[entity_id])
 
 
+def test_get_last_state_change(hass_recorder: Callable[..., HomeAssistant]) -> None:
+    """Test getting the last state change for an entity."""
+    hass = hass_recorder()
+    entity_id = "sensor.test"
+
+    def set_state(state):
+        """Set the state."""
+        hass.states.set(entity_id, state)
+        wait_recording_done(hass)
+        return hass.states.get(entity_id)
+
+    start = dt_util.utcnow() - timedelta(minutes=2)
+    point = start + timedelta(minutes=1)
+    point2 = point + timedelta(minutes=1, seconds=1)
+
+    with patch(
+        "homeassistant.components.recorder.core.dt_util.utcnow", return_value=start
+    ):
+        set_state("1")
+
+    states = []
+    with patch(
+        "homeassistant.components.recorder.core.dt_util.utcnow", return_value=point
+    ):
+        set_state("2")
+
+    with patch(
+        "homeassistant.components.recorder.core.dt_util.utcnow", return_value=point2
+    ):
+        states.append(set_state("3"))
+
+    hist = history.get_last_state_changes(hass, 1, entity_id)
+
+    assert_multiple_states_equal_without_context(states, hist[entity_id])
+
+
 def test_ensure_state_can_be_copied(
     hass_recorder: Callable[..., HomeAssistant]
 ) -> None: