mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-14 12:15:25 +02:00
perf(checkpoint-postgres): start the delta walk at the target
Reading an older checkpoint walked down from the newest one, so every newer row was read before the target's chain. Ancestors have smaller ids, so the first page now starts at the target itself and later pages continue below the oldest row read.
This commit is contained in:
@@ -451,7 +451,8 @@ class PostgresSaver(BasePostgresSaver):
|
||||
* Stage 1 (paged): dynamic SELECT over `checkpoints` with three
|
||||
columns per channel: its version, an `EXISTS` probe for a stored
|
||||
blob at that version, and its inline value. Pages newest-first by
|
||||
`checkpoint_id` with a cursor; page size is `_DELTA_PAGE_SIZE`.
|
||||
`checkpoint_id`, starting at the target; page size is
|
||||
`_DELTA_PAGE_SIZE`.
|
||||
Stops paging when every channel has found its seed or a page comes
|
||||
back short.
|
||||
|
||||
@@ -476,7 +477,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
|
||||
# Stage 1: paged K-JSONB-lookup scan, walking the parent chain in
|
||||
# Python after each page. Stops as soon as every channel has its seed.
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
||||
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
||||
@@ -486,7 +487,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
seed_inline_by_ch: dict[str, Any] = {}
|
||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
||||
seeded: set[str] = set()
|
||||
cursor: str | None = None
|
||||
cursor: str | None = checkpoint_id
|
||||
|
||||
with self._cursor() as cur:
|
||||
while True:
|
||||
@@ -495,7 +496,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
# ver_i, blob channel, blob version, inline_i
|
||||
stage1_params.extend([ch, ch, ch, ch])
|
||||
stage1_params.extend(
|
||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
||||
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
|
||||
)
|
||||
cur.execute(stage1_sql, stage1_params)
|
||||
page = cur.fetchall()
|
||||
@@ -527,6 +528,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
||||
break
|
||||
cursor = oldest
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
|
||||
# Stage 2: per-channel UNION ALL — one writes branch per channel
|
||||
# with non-empty chain, plus one blob branch per seeded channel.
|
||||
|
||||
@@ -423,7 +423,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
checkpoint_id = target.config["configurable"]["checkpoint_id"]
|
||||
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True, include_cursor=True)
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_by_i_by_cid: list[dict[str, str | None]] = [{} for _ in channels]
|
||||
hb_by_i_by_cid: list[dict[str, bool]] = [{} for _ in channels]
|
||||
@@ -433,7 +433,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
seed_inline_by_ch: dict[str, Any] = {}
|
||||
walk_cursor_by_ch: dict[str, str | None] = {}
|
||||
seeded: set[str] = set()
|
||||
cursor: str | None = None
|
||||
cursor: str | None = checkpoint_id
|
||||
|
||||
async with self._cursor() as cur:
|
||||
while True:
|
||||
@@ -442,7 +442,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
# ver_i, blob channel, blob version, inline_i
|
||||
stage1_params.extend([ch, ch, ch, ch])
|
||||
stage1_params.extend(
|
||||
[thread_id, checkpoint_ns, cursor, cursor, _DELTA_PAGE_SIZE]
|
||||
[thread_id, checkpoint_ns, cursor, _DELTA_PAGE_SIZE]
|
||||
)
|
||||
await cur.execute(stage1_sql, stage1_params)
|
||||
page = await cur.fetchall()
|
||||
@@ -472,6 +472,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
if len(seeded) == len(channels) or len(page) < _DELTA_PAGE_SIZE:
|
||||
break
|
||||
cursor = oldest
|
||||
stage1_sql = _build_delta_stage1_sql(channels, paged=True)
|
||||
|
||||
channels_with_chain = [ch for ch in channels if chain_by_ch[ch]]
|
||||
channels_with_seed = [ch for ch in channels if seed_ver_by_ch[ch] is not None]
|
||||
|
||||
@@ -178,10 +178,13 @@ class _DeltaStage2Row(TypedDict, total=False):
|
||||
# `_build_delta_stage2_sql` document their shapes.
|
||||
|
||||
|
||||
def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
def _build_delta_stage1_sql(
|
||||
channels: Sequence[str], *, paged: bool, include_cursor: bool = False
|
||||
) -> str:
|
||||
"""Build stage 1 SQL with K parallel version lookups + seed probes.
|
||||
|
||||
For channels=["messages", "files"] (with `paged=True`) the result is::
|
||||
For channels=["messages", "files"] (with `paged=True, include_cursor=True`)
|
||||
the result is::
|
||||
|
||||
SELECT checkpoint_id, parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver_0,
|
||||
@@ -197,7 +200,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
checkpoint -> 'channel_values' -> %s AS inline_1
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
AND (%s::text IS NULL OR checkpoint_id < %s)
|
||||
AND checkpoint_id <= %s
|
||||
ORDER BY checkpoint_id DESC
|
||||
LIMIT %s
|
||||
|
||||
@@ -238,9 +241,16 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
and uses safe identifiers).
|
||||
|
||||
Caller must extend params with `[ch_0 x4, ch_1 x4, ..., thread_id, ns,
|
||||
cursor, cursor, page_size]` when `paged=True` — four per channel: the
|
||||
version lookup, the blob's channel, the version the blob must match, and the
|
||||
inline lookup.
|
||||
cursor, page_size]` when `paged=True` — four per channel: the version
|
||||
lookup, the blob's channel, the version the blob must match, and the inline
|
||||
lookup.
|
||||
|
||||
Pages run newest-first from the target down. The first page passes the
|
||||
target as the cursor with `include_cursor=True`, so it opens with the
|
||||
target's own row, whose parent starts the walk; each later page continues
|
||||
below the oldest row read. A checkpoint's ancestors have smaller ids (uuid6
|
||||
is time-ordered, which `get_tuple` also relies on to find the latest
|
||||
checkpoint), so no row newer than the target is part of its chain.
|
||||
|
||||
When `paged=False`, the WHERE has no cursor predicate and there's no
|
||||
LIMIT/ORDER BY — kept as a non-public helper for tests/diagnostics.
|
||||
@@ -264,7 +274,7 @@ def _build_delta_stage1_sql(channels: Sequence[str], *, paged: bool) -> str:
|
||||
)
|
||||
if paged:
|
||||
sql += (
|
||||
" AND (%s::text IS NULL OR checkpoint_id < %s)"
|
||||
f" AND checkpoint_id {'<=' if include_cursor else '<'} %s"
|
||||
" ORDER BY checkpoint_id DESC LIMIT %s"
|
||||
)
|
||||
return sql
|
||||
@@ -413,8 +423,8 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
materialized at this point),
|
||||
(c) the next ancestor cid isn't in `parent_of` yet (waiting for
|
||||
a later page; the cursor stays put), or
|
||||
(d) the target's own row isn't in `parent_of` yet (the walk has
|
||||
not started; no cursor is set, so a later page retries).
|
||||
(d) the target's own row isn't in `parent_of` (the target doesn't
|
||||
exist, so the walk never starts).
|
||||
|
||||
Mutates `chain_by_ch`, `seed_ver_by_ch`, `seed_inline_by_ch`,
|
||||
`walk_cursor_by_ch`, and `seeded` in place.
|
||||
@@ -422,8 +432,9 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
for i, ch in enumerate(channels):
|
||||
if ch in seeded:
|
||||
continue
|
||||
# Pages start at the thread head, so the target may not have
|
||||
# loaded yet; a `None` cursor would read as "target is a root".
|
||||
# The first page opens with the target's row, so it's missing only
|
||||
# when the target doesn't exist; a `None` cursor would read as
|
||||
# "target is a root".
|
||||
if ch not in walk_cursor_by_ch:
|
||||
if target_id not in parent_of:
|
||||
continue
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
@@ -14,7 +15,7 @@ from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE
|
||||
from langgraph.checkpoint.postgres.base import _DELTA_PAGE_SIZE, BasePostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
CHANNEL = "items"
|
||||
@@ -23,8 +24,8 @@ SEED_STEP = 1
|
||||
SEED_VALUE = [10, 20]
|
||||
TARGET_STEP = 4
|
||||
|
||||
# The real page size is the control; the rest leave the target off the first
|
||||
# page (three checkpoints are newer than it).
|
||||
# The real page size is the control; the rest split the walk from the target
|
||||
# to its seed across pages.
|
||||
PAGE_SIZES = [_DELTA_PAGE_SIZE, 3, 2, 1]
|
||||
|
||||
|
||||
@@ -92,7 +93,7 @@ def _assert_history(entry: DeltaChannelHistory, page_size: int) -> None:
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", PAGE_SIZES)
|
||||
async def test_async_target_older_than_the_first_page(
|
||||
async def test_async_walk_continues_across_pages(
|
||||
page_size: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres.aio._DELTA_PAGE_SIZE", page_size)
|
||||
@@ -106,7 +107,7 @@ async def test_async_target_older_than_the_first_page(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_size", PAGE_SIZES)
|
||||
def test_sync_target_older_than_the_first_page(
|
||||
def test_sync_walk_continues_across_pages(
|
||||
page_size: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("langgraph.checkpoint.postgres._DELTA_PAGE_SIZE", page_size)
|
||||
@@ -119,6 +120,50 @@ def test_sync_target_older_than_the_first_page(
|
||||
_assert_history(result[CHANNEL], page_size)
|
||||
|
||||
|
||||
def _record_rows_read(monkeypatch: pytest.MonkeyPatch) -> list[str]:
|
||||
read: list[str] = []
|
||||
ingest = BasePostgresSaver._ingest_stage1_page
|
||||
|
||||
def record(rows: Sequence[Mapping[str, Any]], *args: Any) -> str | None:
|
||||
read.extend(row["checkpoint_id"] for row in rows)
|
||||
return ingest(rows, *args)
|
||||
|
||||
monkeypatch.setattr(BasePostgresSaver, "_ingest_stage1_page", staticmethod(record))
|
||||
return read
|
||||
|
||||
|
||||
def _ids_from_target_down(configs: list[dict]) -> list[str]:
|
||||
return [c["configurable"]["checkpoint_id"] for c in configs[TARGET_STEP::-1]]
|
||||
|
||||
|
||||
async def test_async_walk_reads_nothing_newer_than_the_target(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
read = _record_rows_read(monkeypatch)
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
configs = await _abuild_chain(saver)
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=configs[TARGET_STEP], channels=[CHANNEL]
|
||||
)
|
||||
_assert_history(result[CHANNEL], _DELTA_PAGE_SIZE)
|
||||
assert read == _ids_from_target_down(configs)
|
||||
|
||||
|
||||
def test_sync_walk_reads_nothing_newer_than_the_target(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
read = _record_rows_read(monkeypatch)
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
saver.setup()
|
||||
configs = _build_chain(saver)
|
||||
result = saver.get_delta_channel_history(
|
||||
config=configs[TARGET_STEP], channels=[CHANNEL]
|
||||
)
|
||||
_assert_history(result[CHANNEL], _DELTA_PAGE_SIZE)
|
||||
assert read == _ids_from_target_down(configs)
|
||||
|
||||
|
||||
async def test_root_target_has_no_history_and_still_terminates(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
Reference in New Issue
Block a user