Batch look-ups to see if rooms are partial stated.

synapse/storage/database.py

@@ -1819,7 +1819,7 @@ async def simple_select_many_batch(
         keyvalues: Optional[Dict[str, Any]] = None,
         desc: str = "simple_select_many_batch",
         batch_size: int = 100,
-    ) -> List[Any]:
+    ) -> List[Dict[str, Any]]:
         """Executes a SELECT query on the named table, which may return zero or
         more rows, returning the result as a list of dicts.
 

synapse/storage/databases/main/room.py

@@ -23,6 +23,7 @@
     Collection,
     Dict,
     List,
+    Literal,
     Mapping,
     Optional,
     Sequence,
@@ -1285,17 +1286,15 @@ async def is_partial_state_room_batched(
         complete.
         """
 
-        entries = set(
-            await self.db_pool.simple_select_many_batch(
-                table="partial_state_rooms",
-                column="room_id",
-                iterable=room_ids,
-                retcols=("room_id",),
-                desc="is_partial_state_room",
-            )
+        rows: List[Dict[str, Literal[1]]] = await self.db_pool.simple_select_many_batch(
+            table="partial_state_rooms",
+            column="1",
+            iterable=room_ids,
+            retcols=("room_id",),
+            desc="is_partial_state_room",
         )
-
-        return {room_id: room_id in entries for room_id in room_ids}
+        partial_state_rooms = {room_id for row_dict in rows for room_id in row_dict}
+        return {room_id: room_id in partial_state_rooms for room_id in room_ids}
 
     async def get_join_event_id_and_device_lists_stream_id_for_partial_state(
         self, room_id: str