perf(api): derive entity edges from unit_entities instead of materializing them

hindsight-api-slim/hindsight_api/alembic/versions/e9b2c7d1f3a4_drop_entity_memory_links.py

@@ -0,0 +1,89 @@
+"""Drop materialized entity rows from memory_links.
+
+Entity edges are no longer stored in ``memory_links``. The /graph endpoint
+derives them on demand from ``unit_entities``, and recall already used the
+``unit_entities`` self-join. Storing entity rows duplicated state we never
+read from the link table — on a 10k-unit benchmark bank, entity rows were
+53% of all link rows (~190 MB after indexes) and recall never touched them.
+
+This migration deletes ``memory_links`` rows with ``link_type = 'entity'``.
+``idx_memory_links_entity_covering`` was already dropped by migration
+``e1b2c3d4f5a6``; we still issue ``DROP INDEX IF EXISTS`` defensively in case
+this migration runs against an older snapshot that predates that one.
+
+Revision ID: e9b2c7d1f3a4
+Revises: e1b2c3d4f5a6
+Create Date: 2026-05-26
+"""
+
+from collections.abc import Sequence
+
+from alembic import context, op
+
+from hindsight_api.alembic._dialect import run_for_dialect
+
+revision: str = "e9b2c7d1f3a4"
+down_revision: str | Sequence[str] | None = "e1b2c3d4f5a6"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def _pg_schema_prefix() -> str:
+    schema = context.config.get_main_option("target_schema")
+    return f'"{schema}".' if schema else ""
+
+
+def _pg_upgrade() -> None:
+    schema = _pg_schema_prefix()
+
+    # Drop the partial covering index first so the bulk DELETE doesn't churn it.
+    # CREATE/DROP INDEX CONCURRENTLY must run outside a transaction block.
+    op.execute("COMMIT")
+    op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {schema}idx_memory_links_entity_covering")
+
+    # Delete entity rows. Chunked to keep individual transactions small on
+    # large banks (the perf-medium bench had ~345k entity rows; production
+    # banks can be much larger).
+    op.execute(
+        f"""
+        DO $$
+        DECLARE
+            deleted INTEGER;
+        BEGIN
+            LOOP
+                DELETE FROM {schema}memory_links
+                WHERE ctid IN (
+                    SELECT ctid FROM {schema}memory_links
+                    WHERE link_type = 'entity'
+                    LIMIT 50000
+                );
+                GET DIAGNOSTICS deleted = ROW_COUNT;
+                EXIT WHEN deleted = 0;
+                COMMIT;
+            END LOOP;
+        END$$;
+        """
+    )
+
+
+def _pg_downgrade() -> None:
+    # Cannot reconstruct deleted entity links — the writer was path-dependent
+    # on retain order. New retains will not produce entity rows either, so the
+    # partial index would stay empty. Leave both no-op.
+    pass
+
+
+def _oracle_upgrade() -> None:
+    op.execute("DELETE FROM memory_links WHERE link_type = 'entity'")
+
+
+def _oracle_downgrade() -> None:
+    pass
+
+
+def upgrade() -> None:
+    run_for_dialect(pg=_pg_upgrade, oracle=_oracle_upgrade)
+
+
+def downgrade() -> None:
+    run_for_dialect(pg=_pg_downgrade, oracle=_oracle_downgrade)

hindsight-api-slim/hindsight_api/engine/db/ops.py

@@ -166,21 +166,6 @@ async def bulk_insert_unit_entities(
 
     # -- LATERAL / fan-out queries ---------------------------------------
 
-    @abstractmethod
-    async def fetch_entity_unit_fanout(
-        self,
-        conn: DatabaseConnection,
-        ue_table: str,
-        entity_id_list: list[UUID],
-        limit_per_entity: int,
-    ) -> list[ResultRow]:
-        """Fetch unit_ids for a list of entities with per-entity row cap.
-
-        PG uses unnest + CROSS JOIN LATERAL with LIMIT.
-        Non-PG queries each entity individually.
-        """
-        ...
-
     @abstractmethod
     async def fetch_unit_dates(
         self,

hindsight-api-slim/hindsight_api/engine/db/ops_oracle.py

@@ -215,30 +215,6 @@ async def bulk_insert_unit_entities(
             list(zip(unit_ids, entity_ids)),
         )
 
-    async def fetch_entity_unit_fanout(
-        self,
-        conn: DatabaseConnection,
-        ue_table: str,
-        entity_id_list: list[UUID],
-        limit_per_entity: int,
-    ) -> list[ResultRow]:
-        # Query each entity individually
-        rows: list[ResultRow] = []
-        for eid in entity_id_list:
-            entity_rows = await conn.fetch(
-                f"""
-                SELECT $1 AS entity_id, ue.unit_id
-                FROM {ue_table} ue
-                WHERE ue.entity_id = $1
-                ORDER BY ue.unit_id DESC
-                LIMIT $2
-                """,
-                eid,
-                limit_per_entity,
-            )
-            rows.extend(entity_rows)
-        return rows
-
     async def fetch_unit_dates(
         self,
         conn: DatabaseConnection,

hindsight-api-slim/hindsight_api/engine/db/ops_postgresql.py

@@ -290,29 +290,6 @@ async def bulk_insert_unit_entities(
             entity_ids,
         )
 
-    async def fetch_entity_unit_fanout(
-        self,
-        conn: DatabaseConnection,
-        ue_table: str,
-        entity_id_list: list[UUID],
-        limit_per_entity: int,
-    ) -> list[ResultRow]:
-        return await conn.fetch(
-            f"""
-            SELECT e.entity_id, n.unit_id
-            FROM unnest($1::uuid[]) AS e(entity_id)
-            CROSS JOIN LATERAL (
-                SELECT ue.unit_id
-                FROM {ue_table} ue
-                WHERE ue.entity_id = e.entity_id
-                ORDER BY ue.unit_id DESC
-                LIMIT $2
-            ) n
-            """,
-            entity_id_list,
-            limit_per_entity,
-        )
-
     async def fetch_unit_dates(
         self,
         conn: DatabaseConnection,

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

@@ -4881,7 +4881,9 @@ async def get_graph_data(
                     source_memory_ids.extend(unit["source_memory_ids"])
             source_memory_ids = list(set(source_memory_ids))  # Deduplicate
 
-            # Fetch links where BOTH endpoints are in the visible set (or source memories)
+            # Fetch non-entity links where BOTH endpoints are in the visible set (or
+            # source memories). Entity edges are derived below from unit_entities so
+            # we don't materialize them in memory_links anymore.
             # Cap at 10k edges — the UI can't usefully render more, and uncapped queries
             # on highly-connected graphs (e.g. 1000 nodes with 500k+ edges) are too slow.
             max_edges = 10000
@@ -4893,10 +4895,11 @@ async def get_graph_data(
                            ml.to_unit_id,
                            ml.link_type,
                            ml.weight,
-                           e.canonical_name as entity_name
+                           NULL::text AS entity_name
                     FROM {fq_table("memory_links")} ml
-                    LEFT JOIN {fq_table("entities")} e ON ml.entity_id = e.id
-                    WHERE ml.from_unit_id = ANY($1::uuid[]) AND ml.to_unit_id = ANY($1::uuid[])
+                    WHERE ml.link_type <> 'entity'
+                      AND ml.from_unit_id = ANY($1::uuid[])
+                      AND ml.to_unit_id = ANY($1::uuid[])
                     ORDER BY ml.weight DESC NULLS LAST
                     LIMIT $2
                 """,
@@ -5034,16 +5037,21 @@ async def get_graph_data(
                 }
             )
 
-        # Build observation-inferred links from inherited entities and shared source memories.
-        # Observations never have direct memory_links rows, so all their links must be derived.
+        # Build derived links: entity edges for all visible units (from unit_entities)
+        # and observation semantic edges via shared source memories.
+        # Observations never have direct memory_links rows, so all their links are derived.
         observation_units = [unit for unit in units if unit["fact_type"] == "observation"]
-        observation_ids = {unit["id"] for unit in observation_units}
 
-        # Entity links: pair observations that share at least one inherited entity
-        entity_to_observations: dict[str, list] = {}
-        for obs_id in observation_ids:
-            for entity_name in entity_map.get(obs_id, []):
-                entity_to_observations.setdefault(entity_name, []).append(obs_id)
+        # Entity links: pair any visible units that share at least one entity.
+        # Each unit links to up to max_neighbors_per_unit subsequent units in the
+        # per-entity list, so every unit that shares an entity with another visible
+        # unit gets edges (matches the historical writer cap, which was per-unit).
+        # Bounds total edges to ~N * cap per entity instead of N² for hot entities.
+        max_neighbors_per_unit = 10
+        entity_to_units_visible: dict[str, list] = {}
+        for unit_id in unit_ids:
+            for entity_name in entity_map.get(unit_id, []):
+                entity_to_units_visible.setdefault(entity_name, []).append(unit_id)
 
         # Semantic links: pair observations that share at least one source memory
         source_to_obs_for_semantic: dict = {}
@@ -5055,16 +5063,22 @@ async def get_graph_data(
         observation_inferred_links = []
         seen_inferred: set[tuple] = set()
 
-        for entity_name, obs_ids in entity_to_observations.items():
-            for i, obs_a in enumerate(obs_ids):
-                for obs_b in obs_ids[i + 1 :]:
-                    pair = (min(str(obs_a), str(obs_b)), max(str(obs_a), str(obs_b)), "entity", entity_name)
+        for entity_name, ent_unit_ids in entity_to_units_visible.items():
+            n = len(ent_unit_ids)
+            for i, unit_a in enumerate(ent_unit_ids):
+                # Sliding window: link unit_a to its next max_neighbors_per_unit
+                # in the list. Each pair is also "incoming" for the later unit,
+                # so every unit ends up with up to ~2*max_neighbors_per_unit edges
+                # for this entity (its successors + its predecessors via their pairs).
+                for j in range(i + 1, min(i + 1 + max_neighbors_per_unit, n)):
+                    unit_b = ent_unit_ids[j]
+                    pair = (min(str(unit_a), str(unit_b)), max(str(unit_a), str(unit_b)), "entity", entity_name)
                     if pair not in seen_inferred:
                         seen_inferred.add(pair)
                         observation_inferred_links.append(
                             {
-                                "from_unit_id": obs_a,
-                                "to_unit_id": obs_b,
+                                "from_unit_id": unit_a,
+                                "to_unit_id": unit_b,
                                 "link_type": "entity",
                                 "weight": 1.0,
                                 "entity_name": entity_name,
@@ -6996,19 +7010,56 @@ async def get_bank_stats(
                 bank_id,
             )
 
-            # Link stats — filter on ml.bank_id directly instead of joining through mu.bank_id.
-            # link_counts and link_counts_by_fact_type are derived in Python from the breakdown.
-            link_breakdown_stats = await conn.fetch(
+            # Link stats — filter on ml.bank_id (indexed) instead of joining through mu.bank_id.
+            # With the idx_memory_links_bank_link_type index this turns a full-table hash join
+            # into an indexed scan + PK lookups.  link_counts and link_counts_by_fact_type are
+            # derived in Python from the breakdown.
+            non_entity_breakdown = await conn.fetch(
                 f"""
                 SELECT mu.fact_type, ml.link_type, COUNT(*) as count
                 FROM {fq_table("memory_links")} ml
                 JOIN {fq_table("memory_units")} mu ON ml.from_unit_id = mu.id
-                WHERE ml.bank_id = $1
+                WHERE ml.bank_id = $1 AND ml.link_type <> 'entity'
                 GROUP BY mu.fact_type, ml.link_type
                 """,
                 bank_id,
             )
 
+            # Entity links are derived from unit_entities (no longer stored in memory_links).
+            # Replicate the historical writer cap: each unit linked bidirectionally to up to
+            # MAX_LINKS_PER_ENTITY others sharing each entity. Count per (unit, entity) the
+            # outgoing rows the writer would have produced — by fact_type of the source unit.
+            max_links_per_entity = 10
+            entity_breakdown = await conn.fetch(
+                f"""
+                WITH per_entity AS (
+                    SELECT ue.entity_id, COUNT(*) AS n
+                    FROM {fq_table("unit_entities")} ue
+                    JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
+                    WHERE mu.bank_id = $1
+                    GROUP BY ue.entity_id
+                )
+                SELECT mu.fact_type, SUM(LEAST(pe.n - 1, $2))::bigint AS count
+                FROM {fq_table("unit_entities")} ue
+                JOIN {fq_table("memory_units")} mu ON mu.id = ue.unit_id
+                JOIN per_entity pe ON pe.entity_id = ue.entity_id
+                WHERE mu.bank_id = $1
+                GROUP BY mu.fact_type
+                """,
+                bank_id,
+                max_links_per_entity,
+            )
+
+            link_breakdown_stats = [
+                {"fact_type": row["fact_type"], "link_type": row["link_type"], "count": row["count"]}
+                for row in non_entity_breakdown
+            ]
+            link_breakdown_stats.extend(
+                {"fact_type": row["fact_type"], "link_type": "entity", "count": int(row["count"] or 0)}
+                for row in entity_breakdown
+                if (row["count"] or 0) > 0
+            )
+
             link_counts: dict[str, int] = {}
             link_counts_by_fact_type: dict[str, int] = {}
             for row in link_breakdown_stats: