mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-28 18:59:42 +02:00
refactor(delta-channel): plain SELECT WHERE replaces recursive CTE
The recursive CTE was bottlenecked by a JSON-expression join (`bl.version = checkpoint->'channel_versions'->>bl.channel`) that the planner could not index, producing an O(ancestors x blobs) nested-loop. At depth 1000 it ran ~275 ms and removed ~2M filter rows; the recursion itself was 2.4 ms. Switch to three plain indexed SELECTs per delta channel (checkpoints, checkpoint_writes, checkpoint_blobs); a pure helper on BasePostgresSaver walks the parent chain and assembles DeltaChannelWrites. Sync (__init__.py) and async (aio.py) each own their three-roundtrip I/O wrappers. Bench numbers (notes/delta_channel_query_bench.md): 3x at depth 50, 15x at depth 200, ~100x at depth 1000. Plain over-fetches sibling rows when the thread branches but still wins at every realistic depth on both local and remote postgres. Multi-channel coalescing dropped — reconstruction is per-channel now. Same shape as InMemorySaver. Can come back as a SQL-level optimization later if needed.
This commit is contained in:
Sydney Runkle
@@ -4,7 +4,7 @@ import threading
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, cast
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import (
|
||||
@@ -24,7 +24,12 @@ from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import ConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _internal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_BLOBS_SQL,
|
||||
SELECT_DELTA_PARENTS_SQL,
|
||||
SELECT_DELTA_WRITES_SQL,
|
||||
BasePostgresSaver,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import ShallowPostgresSaver
|
||||
|
||||
Conn = _internal.Conn # For backward compatibility
|
||||
@@ -431,6 +436,35 @@ class PostgresSaver(BasePostgresSaver):
|
||||
with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
def _reconstruct_delta_channel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
channel: str,
|
||||
target_id: str,
|
||||
cur: Cursor[DictRow],
|
||||
) -> Any:
|
||||
"""Run the three reconstruction SELECTs and assemble `DeltaChannelWrites`.
|
||||
|
||||
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
|
||||
`checkpoint_blobs`) each filtered by `(thread_id, checkpoint_ns)` and
|
||||
the per-table key. Plain SELECTs let the planner pick straight index
|
||||
scans; rationale + benchmark in `notes/delta_channel_query_bench.md`.
|
||||
"""
|
||||
cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns))
|
||||
parents_rows = cur.fetchall()
|
||||
cur.execute(SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel))
|
||||
writes_rows = cur.fetchall()
|
||||
cur.execute(SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel))
|
||||
blobs_rows = cur.fetchall()
|
||||
return self._build_delta_channel_writes(
|
||||
target_id=target_id,
|
||||
parents_rows=parents_rows,
|
||||
writes_rows=writes_rows,
|
||||
blobs_rows=blobs_rows,
|
||||
)
|
||||
|
||||
def _load_checkpoint_tuple(
|
||||
self, value: DictRow, cur: Cursor[DictRow]
|
||||
) -> CheckpointTuple:
|
||||
@@ -448,18 +482,15 @@ class PostgresSaver(BasePostgresSaver):
|
||||
and pending writes.
|
||||
"""
|
||||
channel_values = self._load_blobs(value["channel_values"])
|
||||
if any(v is DELTA_SENTINEL for v in channel_values.values()):
|
||||
cp_config = cast(
|
||||
RunnableConfig,
|
||||
{
|
||||
"configurable": {
|
||||
"thread_id": value["thread_id"],
|
||||
"checkpoint_ns": value["checkpoint_ns"],
|
||||
"checkpoint_id": value["checkpoint_id"],
|
||||
}
|
||||
},
|
||||
delta_channels = [ch for ch, v in channel_values.items() if v is DELTA_SENTINEL]
|
||||
for channel in delta_channels:
|
||||
channel_values[channel] = self._reconstruct_delta_channel(
|
||||
thread_id=value["thread_id"],
|
||||
checkpoint_ns=value["checkpoint_ns"],
|
||||
channel=channel,
|
||||
target_id=value["checkpoint_id"],
|
||||
cur=cur,
|
||||
)
|
||||
self._resolve_delta_channels(cp_config, channel_values, cur)
|
||||
return CheckpointTuple(
|
||||
{
|
||||
"configurable": {
|
||||
|
||||
@@ -15,7 +15,6 @@ from langgraph.checkpoint.base import (
|
||||
CheckpointMetadata,
|
||||
CheckpointTuple,
|
||||
DeltaChannelWrites,
|
||||
_overwrite_types,
|
||||
get_checkpoint_id,
|
||||
get_serializable_checkpoint_metadata,
|
||||
)
|
||||
@@ -26,7 +25,12 @@ from psycopg.types.json import Jsonb
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
|
||||
from langgraph.checkpoint.postgres import _ainternal
|
||||
from langgraph.checkpoint.postgres.base import BasePostgresSaver
|
||||
from langgraph.checkpoint.postgres.base import (
|
||||
SELECT_DELTA_BLOBS_SQL,
|
||||
SELECT_DELTA_PARENTS_SQL,
|
||||
SELECT_DELTA_WRITES_SQL,
|
||||
BasePostgresSaver,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.shallow import AsyncShallowPostgresSaver
|
||||
|
||||
Conn = _ainternal.Conn # For backward compatibility
|
||||
@@ -394,147 +398,35 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
|
||||
yield cur
|
||||
|
||||
async def _areconstruct_delta_channels_cur(
|
||||
async def _areconstruct_delta_channel(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channels: Sequence[str],
|
||||
cur: Any,
|
||||
) -> dict[str, DeltaChannelWrites]:
|
||||
"""Async mirror of `_reconstruct_delta_channels_cur`.
|
||||
channel: str,
|
||||
target_id: str,
|
||||
cur: AsyncCursor[DictRow],
|
||||
) -> Any:
|
||||
"""Async mirror of `PostgresSaver._reconstruct_delta_channel`.
|
||||
|
||||
Single recursive CTE enumerates on-path ancestors, LEFT-joined once
|
||||
against `checkpoint_writes` and once against `checkpoint_blobs`. Per
|
||||
channel the walk stops at the first terminator — a user `Overwrite`
|
||||
in writes or a pre-delta blob (captured as `seed`). See the sync
|
||||
docstring for the full rationale.
|
||||
Three indexed roundtrips (`checkpoints`, `checkpoint_writes`,
|
||||
`checkpoint_blobs`); rows assembled by the shared pure helper on
|
||||
`BasePostgresSaver`. Rationale + benchmark in
|
||||
`notes/delta_channel_query_bench.md`.
|
||||
"""
|
||||
if not channels:
|
||||
return {}
|
||||
|
||||
overwrite_types = _overwrite_types()
|
||||
channels_list = list(channels)
|
||||
|
||||
await cur.execute(
|
||||
"""
|
||||
WITH RECURSIVE ancestors(cid, parent, depth) AS (
|
||||
SELECT parent_checkpoint_id, NULL::text, 0
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
AND checkpoint_id = %s AND parent_checkpoint_id IS NOT NULL
|
||||
UNION ALL
|
||||
SELECT c.parent_checkpoint_id, a.cid, a.depth + 1
|
||||
FROM checkpoints c
|
||||
JOIN ancestors a ON c.checkpoint_id = a.cid
|
||||
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
|
||||
AND c.parent_checkpoint_id IS NOT NULL
|
||||
),
|
||||
walk AS (
|
||||
SELECT a.cid, a.depth, c.checkpoint
|
||||
FROM ancestors a
|
||||
JOIN checkpoints c
|
||||
ON c.thread_id = %s AND c.checkpoint_ns = %s
|
||||
AND c.checkpoint_id = a.cid
|
||||
)
|
||||
SELECT w.cid, w.depth,
|
||||
cw.channel AS write_channel, cw.type AS write_type,
|
||||
cw.blob AS write_blob, cw.task_id, cw.idx,
|
||||
bl.channel AS blob_channel, bl.type AS blob_type,
|
||||
bl.blob AS blob_blob
|
||||
FROM walk w
|
||||
LEFT JOIN checkpoint_writes cw
|
||||
ON cw.thread_id = %s AND cw.checkpoint_ns = %s
|
||||
AND cw.checkpoint_id = w.cid AND cw.channel = ANY(%s)
|
||||
LEFT JOIN checkpoint_blobs bl
|
||||
ON bl.thread_id = %s AND bl.checkpoint_ns = %s
|
||||
AND bl.channel = ANY(%s)
|
||||
AND bl.version = (w.checkpoint->'channel_versions'->>bl.channel)
|
||||
ORDER BY w.depth ASC, cw.task_id DESC, cw.idx DESC
|
||||
""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channels_list,
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channels_list,
|
||||
),
|
||||
await cur.execute(SELECT_DELTA_PARENTS_SQL, (channel, thread_id, checkpoint_ns))
|
||||
parents_rows = await cur.fetchall()
|
||||
await cur.execute(SELECT_DELTA_WRITES_SQL, (thread_id, checkpoint_ns, channel))
|
||||
writes_rows = await cur.fetchall()
|
||||
await cur.execute(SELECT_DELTA_BLOBS_SQL, (thread_id, checkpoint_ns, channel))
|
||||
blobs_rows = await cur.fetchall()
|
||||
return self._build_delta_channel_writes(
|
||||
target_id=target_id,
|
||||
parents_rows=parents_rows,
|
||||
writes_rows=writes_rows,
|
||||
blobs_rows=blobs_rows,
|
||||
)
|
||||
|
||||
rows_by_cid: dict[str, dict[str, Any]] = {}
|
||||
cid_order: list[str] = []
|
||||
seen_blob: set[tuple[str, str]] = set()
|
||||
seen_write: set[tuple[str, str, str, int]] = set()
|
||||
for row in await cur.fetchall():
|
||||
cid = row["cid"]
|
||||
if cid not in rows_by_cid:
|
||||
rows_by_cid[cid] = {"writes_per_channel": {}, "blob_per_channel": {}}
|
||||
cid_order.append(cid)
|
||||
ch_w = row["write_channel"]
|
||||
if ch_w is not None:
|
||||
key = (cid, ch_w, row["task_id"], row["idx"])
|
||||
if key not in seen_write:
|
||||
seen_write.add(key)
|
||||
rows_by_cid[cid]["writes_per_channel"].setdefault(ch_w, []).append(
|
||||
(row["write_type"], row["write_blob"])
|
||||
)
|
||||
ch_b = row["blob_channel"]
|
||||
if ch_b is not None and (cid, ch_b) not in seen_blob:
|
||||
seen_blob.add((cid, ch_b))
|
||||
rows_by_cid[cid]["blob_per_channel"][ch_b] = (
|
||||
row["blob_type"],
|
||||
row["blob_blob"],
|
||||
)
|
||||
|
||||
collected: dict[str, list[Any]] = {ch: [] for ch in channels_list}
|
||||
done: set[str] = set()
|
||||
seeds: dict[str, Any] = {}
|
||||
|
||||
for cid in cid_order:
|
||||
bucket = rows_by_cid[cid]
|
||||
# Pre-delta blob check first — subsumes any writes at this cp.
|
||||
for ch in list(channels_list):
|
||||
if ch in done:
|
||||
continue
|
||||
blob = bucket["blob_per_channel"].get(ch)
|
||||
if blob is None or blob[0] == "empty":
|
||||
continue
|
||||
blob_value = self.serde.loads_typed(blob)
|
||||
if blob_value is DELTA_SENTINEL:
|
||||
continue
|
||||
seeds[ch] = blob_value
|
||||
done.add(ch)
|
||||
for ch in list(channels_list):
|
||||
if ch in done:
|
||||
continue
|
||||
for type_tag, blob in bucket["writes_per_channel"].get(ch, []):
|
||||
val = self.serde.loads_typed((type_tag, blob))
|
||||
collected[ch].append(val)
|
||||
if isinstance(val, overwrite_types):
|
||||
done.add(ch)
|
||||
break
|
||||
if len(done) == len(channels_list):
|
||||
break
|
||||
|
||||
result: dict[str, DeltaChannelWrites] = {}
|
||||
for ch in channels_list:
|
||||
ch_writes = collected[ch]
|
||||
ch_writes.reverse()
|
||||
if ch in seeds:
|
||||
result[ch] = DeltaChannelWrites(writes=ch_writes, seed=seeds[ch])
|
||||
else:
|
||||
result[ch] = DeltaChannelWrites(writes=ch_writes)
|
||||
return result
|
||||
|
||||
async def aget_channel_writes(
|
||||
self, config: RunnableConfig, channel: str
|
||||
) -> DeltaChannelWrites:
|
||||
@@ -542,14 +434,13 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = config["configurable"]["checkpoint_id"]
|
||||
async with self._cursor() as cur:
|
||||
result = await self._areconstruct_delta_channels_cur(
|
||||
return await self._areconstruct_delta_channel(
|
||||
thread_id=thread_id,
|
||||
checkpoint_ns=checkpoint_ns,
|
||||
checkpoint_id=checkpoint_id,
|
||||
channels=[channel],
|
||||
channel=channel,
|
||||
target_id=checkpoint_id,
|
||||
cur=cur,
|
||||
)
|
||||
return result.get(channel, DeltaChannelWrites(writes=[]))
|
||||
|
||||
async def _load_checkpoint_tuple(
|
||||
self, value: DictRow, cur: AsyncCursor[DictRow]
|
||||
@@ -570,7 +461,6 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
"""
|
||||
thread_id = value["thread_id"]
|
||||
checkpoint_ns = value["checkpoint_ns"]
|
||||
checkpoint_id = value["checkpoint_id"]
|
||||
blob_values = value["channel_values"]
|
||||
|
||||
channel_values: dict[str, Any] = {}
|
||||
@@ -579,16 +469,14 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
delta_channels = [
|
||||
ch for ch, v in channel_values.items() if v is DELTA_SENTINEL
|
||||
]
|
||||
if delta_channels:
|
||||
reconstructed = await self._areconstruct_delta_channels_cur(
|
||||
for channel in delta_channels:
|
||||
channel_values[channel] = await self._areconstruct_delta_channel(
|
||||
thread_id=thread_id,
|
||||
checkpoint_ns=checkpoint_ns,
|
||||
checkpoint_id=checkpoint_id,
|
||||
channels=delta_channels,
|
||||
channel=channel,
|
||||
target_id=value["checkpoint_id"],
|
||||
cur=cur,
|
||||
)
|
||||
for ch, writes in reconstructed.items():
|
||||
channel_values[ch] = writes
|
||||
|
||||
return CheckpointTuple(
|
||||
{
|
||||
|
||||
@@ -155,6 +155,30 @@ INSERT_CHECKPOINT_WRITES_SQL = """
|
||||
ON CONFLICT (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) DO NOTHING
|
||||
"""
|
||||
|
||||
# DeltaChannel reconstruction: three plain indexed SELECTs per channel.
|
||||
# Bench (notes/delta_channel_query_bench.md) showed the prior recursive CTE
|
||||
# carried a hidden O(ancestors x blobs_in_thread) join; plain SELECTs are
|
||||
# 3x-100x faster in the realistic depth range and the Python walk is O(n).
|
||||
SELECT_DELTA_PARENTS_SQL = """
|
||||
SELECT checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> %s AS ver
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
"""
|
||||
|
||||
SELECT_DELTA_WRITES_SQL = """
|
||||
SELECT checkpoint_id, type, blob, task_id, idx
|
||||
FROM checkpoint_writes
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
|
||||
"""
|
||||
|
||||
SELECT_DELTA_BLOBS_SQL = """
|
||||
SELECT version, type, blob
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s AND channel = %s
|
||||
"""
|
||||
|
||||
|
||||
class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
SELECT_SQL = SELECT_SQL
|
||||
@@ -200,190 +224,94 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
|
||||
result[k.decode()] = self.serde.loads_typed((type_tag, v))
|
||||
return result
|
||||
|
||||
def _resolve_delta_channels(
|
||||
self,
|
||||
config: RunnableConfig,
|
||||
channel_values: dict[str, Any],
|
||||
cur: Any,
|
||||
) -> None:
|
||||
delta_channels = [ch for ch, v in channel_values.items() if v is DELTA_SENTINEL]
|
||||
if not delta_channels:
|
||||
return
|
||||
reconstructed = self._reconstruct_delta_channels_cur(
|
||||
thread_id=config["configurable"]["thread_id"],
|
||||
checkpoint_ns=config["configurable"].get("checkpoint_ns", ""),
|
||||
checkpoint_id=config["configurable"]["checkpoint_id"],
|
||||
channels=delta_channels,
|
||||
cur=cur,
|
||||
)
|
||||
for ch, writes in reconstructed.items():
|
||||
channel_values[ch] = writes
|
||||
|
||||
def _reconstruct_delta_channels_cur(
|
||||
def _build_delta_channel_writes(
|
||||
self,
|
||||
*,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channels: Sequence[str],
|
||||
cur: Any,
|
||||
) -> dict[str, DeltaChannelWrites]:
|
||||
"""Reconstruct `DeltaChannelWrites` for every `channel` in `channels`.
|
||||
target_id: str,
|
||||
parents_rows: Sequence[Any],
|
||||
writes_rows: Sequence[Any],
|
||||
blobs_rows: Sequence[Any],
|
||||
) -> DeltaChannelWrites:
|
||||
"""Reconstruct one delta channel from rows of the three SELECTs.
|
||||
|
||||
A single recursive CTE enumerates on-path ancestors (never siblings),
|
||||
left-joined once against `checkpoint_writes` and once against
|
||||
`checkpoint_blobs`. One roundtrip covers every delta channel in the
|
||||
`get_tuple` — avoids the N-channels × 3-queries blowup and fits the
|
||||
intent of the schema (`parent_checkpoint_id` is an indexed ancestor
|
||||
pointer; postgres's planner handles this CTE shape well for the
|
||||
ancestor depths seen in practice).
|
||||
Pure data transform shared by sync (`PostgresSaver`) and async
|
||||
(`AsyncPostgresSaver`); both paths run the queries themselves and
|
||||
feed the rows here.
|
||||
|
||||
The walk newest→oldest stops per channel at the first terminator:
|
||||
Walk is newest → oldest from the target's parent. Stops at the first
|
||||
terminator:
|
||||
|
||||
* a user-emitted `Overwrite` in `checkpoint_writes` — replaces
|
||||
prior history;
|
||||
* a non-sentinel blob in `checkpoint_blobs` — a pre-delta snapshot;
|
||||
bound as `DeltaChannelWrites.seed` so replay starts from it.
|
||||
|
||||
Writes stored AT `checkpoint_id` are pending for the next step and
|
||||
excluded; the recursion starts from the target's parent.
|
||||
Writes stored at `target_id` itself are pending writes for the next
|
||||
step and are excluded — the walk begins at the target's parent.
|
||||
"""
|
||||
if not channels:
|
||||
return {}
|
||||
|
||||
overwrite_types = _overwrite_types()
|
||||
channels_list = list(channels)
|
||||
|
||||
# depth=0 → target's parent (we never read writes or blob for the
|
||||
# target itself). A NULL parent stops the recursion.
|
||||
cur.execute(
|
||||
"""
|
||||
WITH RECURSIVE ancestors(cid, parent, depth) AS (
|
||||
SELECT parent_checkpoint_id, NULL::text, 0
|
||||
FROM checkpoints
|
||||
WHERE thread_id = %s AND checkpoint_ns = %s
|
||||
AND checkpoint_id = %s AND parent_checkpoint_id IS NOT NULL
|
||||
UNION ALL
|
||||
SELECT c.parent_checkpoint_id, a.cid, a.depth + 1
|
||||
FROM checkpoints c
|
||||
JOIN ancestors a ON c.checkpoint_id = a.cid
|
||||
WHERE c.thread_id = %s AND c.checkpoint_ns = %s
|
||||
AND c.parent_checkpoint_id IS NOT NULL
|
||||
),
|
||||
walk AS (
|
||||
-- Each ancestor plus the channel_versions mapping for blob join.
|
||||
SELECT a.cid, a.depth, c.checkpoint
|
||||
FROM ancestors a
|
||||
JOIN checkpoints c
|
||||
ON c.thread_id = %s AND c.checkpoint_ns = %s
|
||||
AND c.checkpoint_id = a.cid
|
||||
parent_of: dict[str, str | None] = {}
|
||||
ver_of: dict[str, str | None] = {}
|
||||
for r in parents_rows:
|
||||
cid = r["checkpoint_id"]
|
||||
parent_of[cid] = r["parent_checkpoint_id"]
|
||||
ver_of[cid] = r["ver"]
|
||||
|
||||
ancestors: list[str] = []
|
||||
cid = parent_of.get(target_id)
|
||||
while cid is not None:
|
||||
ancestors.append(cid)
|
||||
cid = parent_of.get(cid)
|
||||
if not ancestors:
|
||||
return DeltaChannelWrites(writes=[])
|
||||
ancestor_set = set(ancestors)
|
||||
|
||||
# Group writes by ancestor cid; sort within (task_id DESC, idx DESC)
|
||||
# to match the prior CTE ordering — newest write first per ancestor.
|
||||
writes_by_cid: dict[str, list[tuple[str, bytes, str, int]]] = {}
|
||||
for r in writes_rows:
|
||||
cid = r["checkpoint_id"]
|
||||
if cid not in ancestor_set:
|
||||
continue
|
||||
writes_by_cid.setdefault(cid, []).append(
|
||||
(r["type"], r["blob"], r["task_id"], r["idx"])
|
||||
)
|
||||
SELECT w.cid, w.depth,
|
||||
cw.channel AS write_channel, cw.type AS write_type,
|
||||
cw.blob AS write_blob, cw.task_id, cw.idx,
|
||||
bl.channel AS blob_channel, bl.type AS blob_type,
|
||||
bl.blob AS blob_blob
|
||||
FROM walk w
|
||||
LEFT JOIN checkpoint_writes cw
|
||||
ON cw.thread_id = %s AND cw.checkpoint_ns = %s
|
||||
AND cw.checkpoint_id = w.cid AND cw.channel = ANY(%s)
|
||||
LEFT JOIN checkpoint_blobs bl
|
||||
ON bl.thread_id = %s AND bl.checkpoint_ns = %s
|
||||
AND bl.channel = ANY(%s)
|
||||
AND bl.version = (w.checkpoint->'channel_versions'->>bl.channel)
|
||||
ORDER BY w.depth ASC, cw.task_id DESC, cw.idx DESC
|
||||
""",
|
||||
(
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
checkpoint_id, # anchor
|
||||
thread_id,
|
||||
checkpoint_ns, # recursion
|
||||
thread_id,
|
||||
checkpoint_ns, # walk join
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channels_list, # writes join
|
||||
thread_id,
|
||||
checkpoint_ns,
|
||||
channels_list, # blobs join
|
||||
),
|
||||
)
|
||||
for ws in writes_by_cid.values():
|
||||
ws.sort(key=lambda w: (w[2], w[3]), reverse=True)
|
||||
|
||||
# Group incoming rows by `cid` so we can, per ancestor, decide whether
|
||||
# to terminate (pre-delta blob) BEFORE processing writes at that
|
||||
# ancestor. Rows within a cid arrive in (task_id DESC, idx DESC)
|
||||
# order; we preserve that when building per-cid writes lists.
|
||||
rows_by_cid: dict[str, dict[str, Any]] = {}
|
||||
cid_order: list[str] = [] # newest → oldest
|
||||
seen_blob: set[tuple[str, str]] = set()
|
||||
seen_write: set[tuple[str, str, str, int]] = set()
|
||||
for row in cur.fetchall():
|
||||
cid = row["cid"]
|
||||
if cid not in rows_by_cid:
|
||||
rows_by_cid[cid] = {"writes_per_channel": {}, "blob_per_channel": {}}
|
||||
cid_order.append(cid)
|
||||
# Writes: dedupe via (cid, channel, task_id, idx).
|
||||
ch_w = row["write_channel"]
|
||||
if ch_w is not None:
|
||||
key = (cid, ch_w, row["task_id"], row["idx"])
|
||||
if key not in seen_write:
|
||||
seen_write.add(key)
|
||||
rows_by_cid[cid]["writes_per_channel"].setdefault(ch_w, []).append(
|
||||
(row["write_type"], row["write_blob"])
|
||||
)
|
||||
# Blobs: dedupe via (cid, channel).
|
||||
ch_b = row["blob_channel"]
|
||||
if ch_b is not None and (cid, ch_b) not in seen_blob:
|
||||
seen_blob.add((cid, ch_b))
|
||||
rows_by_cid[cid]["blob_per_channel"][ch_b] = (
|
||||
row["blob_type"],
|
||||
row["blob_blob"],
|
||||
)
|
||||
blob_by_ver: dict[str, tuple[str, bytes]] = {
|
||||
r["version"]: (r["type"], r["blob"]) for r in blobs_rows
|
||||
}
|
||||
|
||||
# Per-channel state. `collected[ch]` is newest→oldest during the walk.
|
||||
collected: dict[str, list[Any]] = {ch: [] for ch in channels_list}
|
||||
done: set[str] = set()
|
||||
seeds: dict[str, Any] = {}
|
||||
|
||||
for cid in cid_order: # newest → oldest
|
||||
bucket = rows_by_cid[cid]
|
||||
# At each ancestor, check the blob FIRST — a pre-delta blob
|
||||
# subsumes any writes stored under the same checkpoint, so we
|
||||
# must not fold those writes in before terminating.
|
||||
for ch in list(channels_list):
|
||||
if ch in done:
|
||||
continue
|
||||
blob = bucket["blob_per_channel"].get(ch)
|
||||
if blob is None or blob[0] == "empty":
|
||||
continue
|
||||
blob_value = self.serde.loads_typed(blob)
|
||||
if blob_value is DELTA_SENTINEL:
|
||||
continue
|
||||
seeds[ch] = blob_value
|
||||
done.add(ch)
|
||||
# Then process per-channel writes for any channel still live.
|
||||
for ch in list(channels_list):
|
||||
if ch in done:
|
||||
continue
|
||||
for type_tag, blob in bucket["writes_per_channel"].get(ch, []):
|
||||
val = self.serde.loads_typed((type_tag, blob))
|
||||
collected[ch].append(val)
|
||||
if isinstance(val, overwrite_types):
|
||||
done.add(ch)
|
||||
collected: list[Any] = [] # newest first; reversed at the end
|
||||
seed: Any = None
|
||||
found_seed = False
|
||||
for cid in ancestors:
|
||||
# Pre-delta blob terminator: subsumes any writes at this ancestor.
|
||||
ver = ver_of.get(cid)
|
||||
if ver is not None:
|
||||
seed_blob = blob_by_ver.get(ver)
|
||||
if seed_blob is not None and seed_blob[0] != "empty":
|
||||
blob_value = self.serde.loads_typed(seed_blob)
|
||||
if blob_value is not DELTA_SENTINEL:
|
||||
seed = blob_value
|
||||
found_seed = True
|
||||
break
|
||||
if len(done) == len(channels_list):
|
||||
terminated = False
|
||||
for type_tag, write_blob, _task_id, _idx in writes_by_cid.get(cid, []):
|
||||
val = self.serde.loads_typed((type_tag, write_blob))
|
||||
collected.append(val)
|
||||
if isinstance(val, overwrite_types):
|
||||
terminated = True
|
||||
break
|
||||
if terminated:
|
||||
break
|
||||
|
||||
result: dict[str, DeltaChannelWrites] = {}
|
||||
for ch in channels_list:
|
||||
ch_writes = collected[ch]
|
||||
ch_writes.reverse() # oldest → newest
|
||||
if ch in seeds:
|
||||
result[ch] = DeltaChannelWrites(writes=ch_writes, seed=seeds[ch])
|
||||
else:
|
||||
result[ch] = DeltaChannelWrites(writes=ch_writes)
|
||||
return result
|
||||
collected.reverse() # oldest → newest
|
||||
if found_seed:
|
||||
return DeltaChannelWrites(writes=collected, seed=seed)
|
||||
return DeltaChannelWrites(writes=collected)
|
||||
|
||||
def _dump_blobs(
|
||||
self,
|
||||
|
||||
@@ -257,7 +257,7 @@ def run_benchmark() -> None:
|
||||
import psycopg
|
||||
|
||||
psycopg.connect(_POSTGRES_URI).close()
|
||||
checkpointers.append(("Postgres (recursive CTE)", "postgres"))
|
||||
checkpointers.append(("Postgres (plain SELECT)", "postgres"))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
# Delta-channel reconstruction: query strategy benchmark
|
||||
|
||||
**Branch:** `delta-channel-writes-based`
|
||||
**Question (Nuno):** Is the recursive CTE the right query shape for reconstructing a delta channel inside `get_tuple`, or would a plain `SELECT WHERE` be cheaper even though it returns more rows?
|
||||
**Answer:** Plain `SELECT WHERE` wins at every realistic depth. The recursion isn't the problem — the JSON-expression join inside the CTE is.
|
||||
|
||||
## Setup
|
||||
|
||||
- Postgres 16 on `localhost:5441` (the `compose-postgres.yml` instance, run directly without docker for this round)
|
||||
- Single delta channel `messages`, one write per checkpoint, `DELTA_SENTINEL` blob per checkpoint
|
||||
- Linear chain (`branch=1`) and 5-way branching at every step (`branch=5`) — branching is the case where plain over-fetches sibling rows
|
||||
- Median of 20 timed runs after 3 warmups, fresh psycopg cursor per strategy
|
||||
- Bench script: `bench_get_tuple_strategies.py` at repo root
|
||||
|
||||
Three strategies compared:
|
||||
|
||||
| name | roundtrips | shape |
|
||||
|------|-----------|-------|
|
||||
| `cte` | 1 | Current prod: recursive CTE walks ancestors, LEFT JOINs writes + blobs |
|
||||
| `plain` | 3 | Nuno's suggestion: thread-wide `SELECT WHERE` per table, Python walks parent chain and filters |
|
||||
| `cte+narrow` | 2 | CTE returns ancestor IDs only, then one `UNION ALL` of writes + blobs filtered by `ANY(ids)` |
|
||||
|
||||
## Results (ms per get_tuple, median of 20)
|
||||
|
||||
```
|
||||
depth branch cte plain cte+narrow rows_cte rows_plain plain/cte
|
||||
10 1 0.14ms 0.26ms 0.24ms 9 30 1.89x
|
||||
10 5 0.21ms 0.23ms 0.17ms 9 110 1.11x
|
||||
50 1 0.89ms 0.27ms 0.33ms 49 150 0.31x
|
||||
50 5 2.35ms 0.66ms 0.51ms 49 550 0.28x
|
||||
200 1 11.61ms 0.78ms 1.30ms 199 600 0.07x
|
||||
200 5 34.79ms 2.33ms 3.07ms 199 2200 0.07x
|
||||
1000 1 274.60ms 2.59ms 13.29ms 999 3000 0.01x
|
||||
1000 5 856.01ms 10.14ms 15.31ms 999 11000 0.01x
|
||||
```
|
||||
|
||||
Lower is better. `plain/cte < 1` means plain is faster.
|
||||
|
||||
### Headline numbers
|
||||
|
||||
- depth 50: plain is **3x** faster
|
||||
- depth 200: plain is **15x** faster
|
||||
- depth 1000: plain is **~100x** faster
|
||||
- Branching makes plain over-fetch (3000 rows → 11000 rows at d=1000), but it remains ~85x faster than the CTE
|
||||
|
||||
## Why the CTE collapses
|
||||
|
||||
`EXPLAIN (ANALYZE, BUFFERS)` of the CTE at depth 1000 (linear). Excerpt with the load-bearing nodes:
|
||||
|
||||
```
|
||||
Sort ... actual time=137.798..137.827 rows=999
|
||||
CTE ancestors
|
||||
-> Recursive Union ... actual time=0.005..2.443 rows=999
|
||||
^^^^^^
|
||||
recursion is 2.4 ms — fine
|
||||
-> Nested Loop Left Join ... actual time=2.676..137.529 rows=999
|
||||
Join Filter: (cw.checkpoint_id = a.cid)
|
||||
Rows Removed by Join Filter: 998001
|
||||
^^^^^^^
|
||||
999 ancestors x ~1000 writes
|
||||
-> Nested Loop Left Join ... actual time=2.669..85.061 rows=999
|
||||
Join Filter: (bl.version = ((c.checkpoint -> 'channel_versions'::text) ->> bl.channel))
|
||||
Rows Removed by Join Filter: 998001
|
||||
^^^^^^^
|
||||
same quadratic blow-up on the blob join
|
||||
```
|
||||
|
||||
Two pathological things are happening:
|
||||
|
||||
1. **The blob join filter is on a JSON expression**: `bl.version = (c.checkpoint -> 'channel_versions' ->> bl.channel)`. The planner cannot push this into an index lookup, so it materializes `checkpoint_blobs` for the thread and does a nested-loop comparison against every ancestor — a Cartesian product that grows as `O(ancestors × blobs_in_thread)`.
|
||||
2. **The writes join is similar**: writes for the thread are materialized once, then nested-loop joined against ancestors with a `Join Filter` rather than a hash/merge join over the indexed `checkpoint_id`.
|
||||
|
||||
At depth 1000 that's **~2 million rows evaluated, 99.9% of them discarded**. The recursion itself is a rounding error.
|
||||
|
||||
For comparison, the plain Q1 (`SELECT … FROM checkpoints WHERE thread_id=? AND checkpoint_ns=?`) at depth 1000:
|
||||
|
||||
```
|
||||
Seq Scan on checkpoints ... actual time=0.012..0.121 rows=1000
|
||||
Execution Time: 0.140 ms
|
||||
```
|
||||
|
||||
A simple seq scan over 57 buffers. Q2 and Q3 follow the same shape and complete in well under 1 ms each.
|
||||
|
||||
## Crossover and remote-DB reasoning
|
||||
|
||||
- Pure local Postgres: plain wins from depth ~30 onward; CTE wins by fractions of a ms below that
|
||||
- Remote Postgres at ~5 ms RTT adds ~10 ms to plain (3 roundtrips vs 1). Crossover shifts to ~depth 30. Above that, the CTE's quadratic SQL cost still dominates the RTT savings.
|
||||
|
||||
There is no realistic conversation depth where the CTE wins on a remote DB. At depth 200+ (anything resembling a real multi-turn agent run) plain is faster regardless of network.
|
||||
|
||||
## Recommendation
|
||||
|
||||
**Switch to plain SELECT WHERE, one delta channel at a time.**
|
||||
|
||||
Three indexed queries per delta channel:
|
||||
|
||||
```sql
|
||||
-- Q1: parent chain + per-checkpoint version of this channel
|
||||
SELECT checkpoint_id,
|
||||
parent_checkpoint_id,
|
||||
checkpoint -> 'channel_versions' ->> 'channel_name' AS ver
|
||||
FROM checkpoints
|
||||
WHERE thread_id = ? AND checkpoint_ns = ?;
|
||||
|
||||
-- Q2: writes for this channel, anywhere in the thread
|
||||
SELECT checkpoint_id, type, blob, task_id, idx
|
||||
FROM checkpoint_writes
|
||||
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
|
||||
|
||||
-- Q3: blobs for this channel, anywhere in the thread
|
||||
SELECT version, type, blob
|
||||
FROM checkpoint_blobs
|
||||
WHERE thread_id = ? AND checkpoint_ns = ? AND channel = ?;
|
||||
```
|
||||
|
||||
Python then:
|
||||
- Builds `parent_of: dict[cid, parent_cid]` from Q1
|
||||
- Walks from target's parent newest → oldest
|
||||
- Filters Q2 rows by `ancestor_set`, processes oldest → newest, applies overwrite-terminator
|
||||
- Picks seed blob via the per-ancestor `ver` map, terminates at first non-sentinel blob
|
||||
|
||||
All O(n) on n = thread checkpoints, with tight constants (dict lookups). No recursion, no JSON-expression joins, no quadratic plans.
|
||||
|
||||
If the 3-roundtrip cost ever shows up on remote-DB benchmarks, fold Q2 + Q3 into one `UNION ALL` to get back to 2 roundtrips. Bench says it isn't worth the SQL complexity right now.
|
||||
|
||||
## Bonus: code simplification from single-channel scope
|
||||
|
||||
Multi-channel reconstruction in the current `_reconstruct_delta_channels_cur` carries:
|
||||
|
||||
- `rows_by_cid` nested dicts, keyed by cid then channel
|
||||
- `seen_blob: set[(cid, channel)]` and `seen_write: set[(cid, channel, task_id, idx)]` dedup
|
||||
- `collected: dict[channel, list]`, `done: set[channel]`, `seeds: dict[channel, value]`
|
||||
- Inner `for ch in channels_list` loops and an early-exit `if len(done) == len(channels_list)`
|
||||
|
||||
Single-channel collapses these to a single list, a single bool, and one `Optional[Any]`. Roughly half the Python in that function, plus an obvious shape for splitting pure post-processing into `base.py` so sync and async stop duplicating it.
|
||||
|
||||
If multi-channel coalescing turns out to matter later, it can come back as a SQL-level optimization without re-introducing the bookkeeping in Python.
|
||||
Reference in New Issue
Block a user