perf(retain): store-owned backend writes facts once with entities inline

hindsight-api-slim/hindsight_api/engine/entity_resolver.py

@@ -1246,6 +1246,7 @@ async def record_unit_entity_postings(
         self,
         unit_entity_pairs: list[tuple[str, str]] | list[tuple[str, str, datetime | None]],
         bank_id: str | None = None,
+        store_write: bool = True,
     ):
         """Store-owned variant of :meth:`link_units_to_entities_batch` that touches NO
         Postgres connection.
@@ -1257,6 +1258,13 @@ async def record_unit_entity_postings(
         its connection-free store phase and never hold the data-plane connection across the
         object-store write. NOT for the Postgres store, whose posting is a real ``unit_entities``
         INSERT that requires the connection.
+
+        ``store_write=False`` skips the store-side posting and does ONLY the co-occurrence
+        accumulation. The caller uses this when it has already attached entity ids to the memories
+        as part of the same write (a single deferred write with entities inline, instead of
+        write-then-reattach) — so the store row is already correct and a second store write would
+        be redundant. Co-occurrence still runs: it references only ``entities`` and is needed by the
+        entity-graph endpoint and resolution's disambiguation signal regardless of who wrote the row.
         """
         if not unit_entity_pairs:
             return
@@ -1265,10 +1273,16 @@ async def record_unit_entity_postings(
             (t[0], t[1], t[2] if len(t) >= 3 else None)  # type: ignore[misc]
             for t in unit_entity_pairs
         ]
-        return await self._link_units_to_entities_batch_impl(None, normalized, bank_id)
+        return await self._link_units_to_entities_batch_impl(
+            None, normalized, bank_id, store_write=store_write
+        )
 
     async def _link_units_to_entities_batch_impl(
-        self, conn, unit_entity_pairs: list[tuple[str, str, datetime | None]], bank_id: str | None = None
+        self,
+        conn,
+        unit_entity_pairs: list[tuple[str, str, datetime | None]],
+        bank_id: str | None = None,
+        store_write: bool = True,
     ):
         # Sorted bulk insert to prevent deadlocks from inconsistent lock ordering
         # across concurrent transactions on the unit_entities unique index.
@@ -1280,16 +1294,19 @@ async def _link_units_to_entities_batch_impl(
         # memories store records it. Co-occurrence below is separate and unaffected:
         # it references only `entities`, which stays in Postgres either way, and is
         # read by the entity-graph endpoint and by resolution's disambiguation signal.
+        # `store_write=False` means the caller already wrote the postings inline with the
+        # memories, so we skip the (redundant) second store write and keep only co-occurrence.
         from .memories import get_memories
 
-        await get_memories().record_unit_entities(
-            conn=conn,
-            ops=self._ops,
-            fq_table=fq_table,
-            bank_id=bank_id,
-            unit_ids=unit_ids,
-            entity_ids=entity_ids,
-        )
+        if store_write:
+            await get_memories().record_unit_entities(
+                conn=conn,
+                ops=self._ops,
+                fq_table=fq_table,
+                bank_id=bank_id,
+                unit_ids=unit_ids,
+                entity_ids=entity_ids,
+            )
 
         # Build maps keyed by unit_id:
         #   unit_to_entities: entity set per unit (for the co-occurrence cross-product)

hindsight-api-slim/hindsight_api/engine/retain/fact_storage.py

@@ -109,16 +109,21 @@ async def index_facts(
     facts: list[ProcessedFact],
     document_id: str | None = None,
     unit_entity_ids: dict[str, list[str]] | None = None,
+    txn=None,
 ) -> None:
     """Complete a deferred `insert_facts_batch`, now that the edges are known.
 
     ``unit_entity_ids`` is the unit→entity posting and each fact's causal
     relations are its edges; both travel with the memory for a store that owns
     them. A no-op for the Postgres store, which wrote all of it already.
+
+    ``txn`` rides a cross-store write-group handle so this single, entity-bearing
+    write commits (and becomes visible) atomically with the rest of the group —
+    the store-owned retain path writes facts ONCE here rather than write-then-reattach.
     """
     from ..memories import get_memories
 
-    await get_memories().index_facts(bank_id, unit_ids, facts, document_id, unit_entity_ids)
+    await get_memories().index_facts(bank_id, unit_ids, facts, document_id, unit_entity_ids, txn=txn)
 
 
 async def ensure_bank_exists(conn, bank_id: str, ops=None) -> None:

hindsight-api-slim/hindsight_api/engine/retain/orchestrator.py

@@ -670,14 +670,18 @@ async def _streaming_batch_write_ext(
             if cid:
                 processed_fact.chunk_id = cid
 
-    # Stage the memory records to the store (conn unused by a store-owned backend), tagged with
-    # ext_txn. This is the slow object-store write we are keeping OUT of the connection window.
-    unit_ids = await fact_storage.insert_facts_batch(None, bank_id, batch_processed, ops=pool.ops, txn=ext_txn)
+    # Mint the unit ids WITHOUT writing (defer_index): entities can only be resolved onto real ids
+    # after they exist, so we write the memories to the store ONCE below — with their entity ids
+    # already attached — instead of writing them here and then re-upserting each one with entities.
+    # That reattach cost a SECOND full object-store write per memory (plus a read-back of the just-
+    # written records), doubling the store round-trips on the slow path. Connection-free either way.
+    unit_ids = await fact_storage.insert_facts_batch(
+        None, bank_id, batch_processed, ops=pool.ops, txn=ext_txn, defer_index=True
+    )
     batch_result_ids = _map_results_to_contents(batch_contents, batch_processed, unit_ids if unit_ids else [])
 
     if unit_ids:
-        # Remap Phase-1 placeholder ids onto the real unit ids, then re-write each memory with its
-        # entity ids attached — also connection-free for a store-owned backend.
+        # Remap Phase-1 placeholder ids onto the real unit ids.
         resolved_entity_ids = [entity.entity_id for entity in phase1.entities.resolved_entities]
         remapped_entity_to_unit, _remapped_unit_to_entity_ids, _remapped_semantic = _remap_phase1_results(
             resolved_entity_ids, phase1.entities.entity_to_unit, phase1.entities.unit_to_entity_ids, [], unit_ids
@@ -686,7 +690,25 @@ async def _streaming_batch_write_ext(
             (unit_id, resolved_entity_ids[idx], fact_date)
             for idx, (unit_id, _local_idx, fact_date) in enumerate(remapped_entity_to_unit)
         ]
-        await entity_resolver.record_unit_entity_postings(unit_entity_pairs, bank_id=bank_id)
+        # The single, entity-bearing store write — connection-free, and tagged with ext_txn so it
+        # commits (and becomes visible) atomically with the rest of the write-group. This replaces
+        # the earlier insert-then-reattach pair with one write.
+        unit_entity_ids: dict[str, list[str]] = {}
+        for unit_id, entity_id, _fd in unit_entity_pairs:
+            unit_entity_ids.setdefault(unit_id, []).append(entity_id)
+        await fact_storage.index_facts(
+            bank_id,
+            unit_ids,
+            batch_processed,
+            document_id=effective_doc_id,
+            unit_entity_ids=unit_entity_ids,
+            txn=ext_txn,
+        )
+        # The store row was just written with its entities inline, so skip the (now redundant)
+        # second store write and keep ONLY the co-occurrence accumulation the entity graph needs.
+        await entity_resolver.record_unit_entity_postings(
+            unit_entity_pairs, bank_id=bank_id, store_write=False
+        )
 
     # ---- CONNECTION PHASE (short transaction: local metadata + commit witness) ----
     try:

hindsight-api-slim/tests/test_retain_ext_writegroup.py

@@ -63,12 +63,21 @@ def _make_common(monkeypatch, tracker, *, calls):
     """Patch the module-level collaborators the helper reaches for."""
     monkeypatch.setattr(orch, "acquire_with_retry", lambda pool: tracker.acquire())
 
-    async def _insert_facts_batch(conn, bank_id, processed, ops=None, txn=None):
-        calls.append(("insert_facts", conn, tracker.open))
-        tracker.store_writes_saw_open.append(tracker.open)
+    async def _insert_facts_batch(conn, bank_id, processed, ops=None, txn=None, defer_index=False):
+        # defer_index=True mints ids WITHOUT writing — the entity-bearing write happens in
+        # index_facts below, so this call is a no-op store-side (records no write).
+        calls.append(("insert_facts", conn, tracker.open, defer_index))
         assert conn is None, "ext store write must not receive a connection"
+        if not defer_index:
+            tracker.store_writes_saw_open.append(tracker.open)
         return ["u1"]
 
+    async def _index_facts(bank_id, unit_ids, facts, document_id=None, unit_entity_ids=None, txn=None):
+        # The single, entity-bearing store write — must be connection-free and carry the txn.
+        calls.append(("index_facts", tracker.open))
+        tracker.store_writes_saw_open.append(tracker.open)
+        assert txn is not None, "the deferred store write must ride the write-group txn"
+
     async def _store_chunks_batch(conn, bank_id, doc_id, meta, ops=None, store_document_text=True):
         calls.append(("store_chunks", tracker.open))
         return {}
@@ -77,6 +86,7 @@ async def _handle_doc_tracking(conn, *a, **k):
         calls.append(("handle_doc_tracking", tracker.open))
 
     monkeypatch.setattr(orch.fact_storage, "insert_facts_batch", _insert_facts_batch)
+    monkeypatch.setattr(orch.fact_storage, "index_facts", _index_facts)
     monkeypatch.setattr(orch.chunk_storage, "store_chunks_batch", _store_chunks_batch)
     monkeypatch.setattr(orch.fact_storage, "handle_document_tracking", _handle_doc_tracking)
     monkeypatch.setattr(orch, "_map_results_to_contents", lambda contents, pf, uids: [list(uids)])
@@ -109,10 +119,11 @@ def __init__(self, tracker):
         self.postings = []
         self.reasserts = []
 
-    async def record_unit_entity_postings(self, pairs, bank_id=None):
-        # THE contract: the store re-posting runs with no connection held.
+    async def record_unit_entity_postings(self, pairs, bank_id=None, store_write=True):
+        # THE contract: co-occurrence accumulation runs with no connection held. store_write=False
+        # here because the entity ids were already written inline by index_facts (single write).
         assert self._t.open is False, "entity posting must run connection-free"
-        self.postings.append(pairs)
+        self.postings.append((pairs, store_write))
 
     async def reassert_entities_batch(self, bank_id, resolved, conn):
         assert conn is not None
@@ -174,10 +185,11 @@ async def test_store_writes_are_connection_free_and_witness_is_in_txn(monkeypatc
 
     assert result.aborted is False
     assert result.batch_result_ids == [["u1"]]
-    # The fact write happened with no connection held.
+    # The fact write happened with no connection held (a single deferred write via index_facts).
     assert tracker.store_writes_saw_open == [False]
-    # Entity re-posting happened (also connection-free — asserted inside the fake).
-    assert er.postings == [[("u1", "e1", None)]]
+    assert ("index_facts", False) in calls  # the entity-bearing store write ran connection-free
+    # Co-occurrence ran connection-free, with store_write=False (entities already written inline).
+    assert er.postings == [([("u1", "e1", None)], False)]
     # Witness written with a real connection, exactly once; commit published after release.
     assert len(provider.witnesses) == 1 and provider.witnesses[0][1] is not None
     assert provider.decisions == [True]
@@ -333,10 +345,11 @@ async def test_delta_store_writes_are_connection_free(monkeypatch):
 
     assert result.fell_back is False
     assert result.result_unit_ids == [["u1"]]
-    # Fact write + body store + entity re-posting all happened connection-free.
+    # Fact write + body store + entity re-posting all happened connection-free. The delta path
+    # still writes-then-reposts (store_write=True) — only the streaming path uses the single write.
     assert tracker.store_writes_saw_open == [False]
     assert ("store_document_bodies", False) in calls
-    assert er.postings == [[("u1", "e1", None)]]
+    assert er.postings == [([("u1", "e1", None)], True)]
     # Witness inside the txn; publish after release.
     assert len(provider.witnesses) == 1 and provider.witnesses[0][1] is not None
     assert provider.decisions == [True]