mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-22 01:25:06 +02:00
refactor(postgres): replace recursive CTE with two-query ancestor walk for DeltaChannel
Instead of a recursive SQL CTE, collect the ancestor checkpoint ID chain in Python by fetching all (checkpoint_id, parent_checkpoint_id) for the thread in one query, then fetch writes with a plain WHERE checkpoint_id = ANY(...). Simpler, avoids recursive query planner overhead, and uses well-indexed lookups.
This commit is contained in:
@@ -400,36 +400,42 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
|||||||
channel: str,
|
channel: str,
|
||||||
cur: Any,
|
cur: Any,
|
||||||
) -> list[Any]:
|
) -> list[Any]:
|
||||||
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async)."""
|
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest (async).
|
||||||
|
|
||||||
|
Two queries instead of a recursive CTE — see sync version for rationale.
|
||||||
|
"""
|
||||||
await cur.execute(
|
await cur.execute(
|
||||||
"""
|
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
|
||||||
WITH RECURSIVE chain(cid, depth) AS (
|
"WHERE thread_id = %s AND checkpoint_ns = %s",
|
||||||
SELECT parent_checkpoint_id, 0
|
(thread_id, checkpoint_ns),
|
||||||
FROM checkpoints
|
)
|
||||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
|
parent_map: dict[str, str | None] = {
|
||||||
UNION ALL
|
row["checkpoint_id"]: row["parent_checkpoint_id"]
|
||||||
SELECT c.parent_checkpoint_id, ch.depth + 1
|
for row in await cur.fetchall()
|
||||||
FROM checkpoints c
|
}
|
||||||
JOIN chain ch ON c.checkpoint_id = ch.cid
|
ancestor_ids: list[str] = []
|
||||||
WHERE ch.cid IS NOT NULL
|
cid: str | None = parent_map.get(checkpoint_id)
|
||||||
)
|
while cid is not None:
|
||||||
SELECT cw.type, cw.blob
|
ancestor_ids.append(cid)
|
||||||
FROM checkpoint_writes cw
|
cid = parent_map.get(cid)
|
||||||
JOIN chain ON cw.checkpoint_id = chain.cid
|
if not ancestor_ids:
|
||||||
WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s
|
return []
|
||||||
ORDER BY chain.depth DESC, cw.task_id, cw.idx
|
await cur.execute(
|
||||||
""",
|
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
|
||||||
(
|
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
||||||
thread_id,
|
" AND checkpoint_id = ANY(%s) "
|
||||||
checkpoint_ns,
|
"ORDER BY task_id, idx",
|
||||||
checkpoint_id,
|
(thread_id, checkpoint_ns, channel, ancestor_ids),
|
||||||
thread_id,
|
|
||||||
checkpoint_ns,
|
|
||||||
channel,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
rows = await cur.fetchall()
|
rows = await cur.fetchall()
|
||||||
return [self.serde.loads_typed((row["type"], row["blob"])) for row in rows]
|
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
|
||||||
|
for row in rows:
|
||||||
|
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
|
||||||
|
result = []
|
||||||
|
for cid in reversed(ancestor_ids):
|
||||||
|
for type_tag, blob in writes_by_cp.get(cid, []):
|
||||||
|
result.append(self.serde.loads_typed((type_tag, blob)))
|
||||||
|
return result
|
||||||
|
|
||||||
async def aget_channel_writes(
|
async def aget_channel_writes(
|
||||||
self, config: RunnableConfig, channel: str
|
self, config: RunnableConfig, channel: str
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import random
|
import random
|
||||||
import warnings
|
import warnings
|
||||||
|
from collections import defaultdict
|
||||||
from collections.abc import Sequence
|
from collections.abc import Sequence
|
||||||
from importlib.metadata import version as get_version
|
from importlib.metadata import version as get_version
|
||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
@@ -222,37 +223,46 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
|||||||
channel: str,
|
channel: str,
|
||||||
cur: Any,
|
cur: Any,
|
||||||
) -> list[Any]:
|
) -> list[Any]:
|
||||||
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest."""
|
"""Fetch writes for `channel` across the checkpoint ancestor chain, oldest→newest.
|
||||||
|
|
||||||
|
Two queries instead of a recursive CTE:
|
||||||
|
1. Fetch all (checkpoint_id, parent_checkpoint_id) for the thread — cheap, just IDs.
|
||||||
|
2. Walk the ancestor chain in Python, then fetch writes with a plain ANY() filter.
|
||||||
|
"""
|
||||||
cur.execute(
|
cur.execute(
|
||||||
"""
|
"SELECT checkpoint_id, parent_checkpoint_id FROM checkpoints "
|
||||||
WITH RECURSIVE chain(cid, depth) AS (
|
"WHERE thread_id = %s AND checkpoint_ns = %s",
|
||||||
SELECT parent_checkpoint_id, 0
|
(thread_id, checkpoint_ns),
|
||||||
FROM checkpoints
|
|
||||||
WHERE thread_id = %s AND checkpoint_ns = %s AND checkpoint_id = %s
|
|
||||||
UNION ALL
|
|
||||||
SELECT c.parent_checkpoint_id, ch.depth + 1
|
|
||||||
FROM checkpoints c
|
|
||||||
JOIN chain ch ON c.checkpoint_id = ch.cid
|
|
||||||
WHERE ch.cid IS NOT NULL
|
|
||||||
)
|
|
||||||
SELECT cw.type, cw.blob
|
|
||||||
FROM checkpoint_writes cw
|
|
||||||
JOIN chain ON cw.checkpoint_id = chain.cid
|
|
||||||
WHERE cw.thread_id = %s AND cw.checkpoint_ns = %s AND cw.channel = %s
|
|
||||||
ORDER BY chain.depth DESC, cw.task_id, cw.idx
|
|
||||||
""",
|
|
||||||
(
|
|
||||||
thread_id,
|
|
||||||
checkpoint_ns,
|
|
||||||
checkpoint_id,
|
|
||||||
thread_id,
|
|
||||||
checkpoint_ns,
|
|
||||||
channel,
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
return [
|
parent_map: dict[str, str | None] = {
|
||||||
self.serde.loads_typed((row["type"], row["blob"])) for row in cur.fetchall()
|
row["checkpoint_id"]: row["parent_checkpoint_id"] for row in cur.fetchall()
|
||||||
]
|
}
|
||||||
|
# Walk newest→oldest starting from the current checkpoint's parent.
|
||||||
|
# Writes stored under checkpoint C produced the state *after* C, so we
|
||||||
|
# want ancestors of the current checkpoint (not the checkpoint itself).
|
||||||
|
ancestor_ids: list[str] = []
|
||||||
|
cid: str | None = parent_map.get(checkpoint_id)
|
||||||
|
while cid is not None:
|
||||||
|
ancestor_ids.append(cid)
|
||||||
|
cid = parent_map.get(cid)
|
||||||
|
if not ancestor_ids:
|
||||||
|
return []
|
||||||
|
cur.execute(
|
||||||
|
"SELECT checkpoint_id, type, blob FROM checkpoint_writes "
|
||||||
|
"WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s "
|
||||||
|
" AND checkpoint_id = ANY(%s) "
|
||||||
|
"ORDER BY task_id, idx",
|
||||||
|
(thread_id, checkpoint_ns, channel, ancestor_ids),
|
||||||
|
)
|
||||||
|
writes_by_cp: dict[str, list[tuple[str, bytes]]] = defaultdict(list)
|
||||||
|
for row in cur.fetchall():
|
||||||
|
writes_by_cp[row["checkpoint_id"]].append((row["type"], row["blob"]))
|
||||||
|
# ancestor_ids is newest→oldest; replay oldest→newest
|
||||||
|
result = []
|
||||||
|
for cid in reversed(ancestor_ids):
|
||||||
|
for type_tag, blob in writes_by_cp.get(cid, []):
|
||||||
|
result.append(self.serde.loads_typed((type_tag, blob)))
|
||||||
|
return result
|
||||||
|
|
||||||
def _dump_blobs(
|
def _dump_blobs(
|
||||||
self,
|
self,
|
||||||
|
|||||||
Reference in New Issue
Block a user