diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 6e884e578..149254702 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -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": { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 7dbefcf51..06f809087 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -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( { diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index c5837ca61..f6a0d8645 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -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, diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index 0247714eb..c9b7a837f 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -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 diff --git a/notes/delta_channel_query_bench.md b/notes/delta_channel_query_bench.md new file mode 100644 index 000000000..8dd69728a --- /dev/null +++ b/notes/delta_channel_query_bench.md @@ -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.