store: a workspace window keeps its (workspace_id, id) index whatever the planner statistics say

lib/optimal_engine/embed/chunk_backfill.ex

@@ -273,18 +273,23 @@ defmodule OptimalEngine.Embed.ChunkBackfill do
     )
   end
 
+  # A window always reads one workspace through its (workspace_id, id) index.
+  # INDEXED BY keeps that plan whatever the planner statistics say: statistics
+  # that show few chunks per workspace otherwise lead SQLite to range over the
+  # whole table's primary key. PostgreSQL drops the clause (Store.Dialect).
   defp window_end_sql do
     """
     SELECT MAX(w.id) FROM (
-      SELECT id FROM chunks WHERE workspace_id = ?1 AND id > ?2 ORDER BY id LIMIT ?3
+      SELECT id FROM chunks INDEXED BY idx_chunks_workspace_id
+      WHERE workspace_id = ?1 AND id > ?2 ORDER BY id LIMIT ?3
     ) AS w
     """
   end
 
   defp missing_rows_sql(tenant_sql) do
     """
     SELECT c.id, c.tenant_id, c.workspace_id, c.signal_id, c.text, c.modality, c.asset_ref
-    FROM chunks c
+    FROM chunks c INDEXED BY idx_chunks_workspace_id
     WHERE c.workspace_id = ?1 AND c.id > ?2 AND c.id <= ?3 AND #{missing_sql()}#{tenant_sql}
     ORDER BY c.id
     LIMIT ?4

lib/optimal_engine/graph/backfill.ex

@@ -235,13 +235,24 @@ defmodule OptimalEngine.Graph.Backfill do
   defp read_page(opts, phase, after_id, take) do
     {where, params} = scope(opts, phase)
     n = length(params)
-    Store.raw_query(page_sql(phase, where, n), params ++ [after_id, take])
+    Store.raw_query(page_sql(phase, where, n, index_hint(opts, phase)), params ++ [after_id, take])
   end
 
-  defp page_sql(:contexts, where, n) do
+  # A workspace's pages always read through its (workspace_id, id) index.
+  # INDEXED BY keeps that plan whatever the planner statistics say: statistics
+  # that show few rows per workspace otherwise lead SQLite to range over the
+  # whole table's primary key. A tenant or whole-store scan has no such index
+  # and plans freely. PostgreSQL drops the clause (Store.Dialect).
+  defp index_hint(opts, :contexts),
+    do: if(opts[:workspace_id], do: " INDEXED BY idx_contexts_workspace_id", else: "")
+
+  defp index_hint(opts, :claims),
+    do: if(opts[:workspace_id], do: " INDEXED BY idx_claims_workspace_id", else: "")
+
+  defp page_sql(:contexts, where, n, hint) do
     """
     SELECT c.id, c.node, c.entities, c.routed_to, c.supersedes, c.workspace_id, c.tenant_id
-    FROM contexts c
+    FROM contexts c#{hint}
     WHERE c.id > ?#{n + 1} AND #{where}
     ORDER BY c.id
     LIMIT ?#{n + 2}
@@ -250,10 +261,10 @@ defmodule OptimalEngine.Graph.Backfill do
 
   # Claims link to their source context through signal_id; the entities come
   # from that context.
-  defp page_sql(:claims, where, n) do
+  defp page_sql(:claims, where, n, hint) do
     """
     SELECT cl.id, c.entities, cl.workspace_id, cl.tenant_id
-    FROM claims cl
+    FROM claims cl#{hint}
     JOIN contexts c ON c.id = cl.signal_id
     WHERE cl.signal_id IS NOT NULL AND cl.id > ?#{n + 1} AND #{where}
     ORDER BY cl.id
@@ -267,8 +278,11 @@ defmodule OptimalEngine.Graph.Backfill do
   @spec page_statements(String.t()) :: [{atom(), String.t(), [term()]}]
   def page_statements(workspace_id) do
     for phase <- [:contexts, :claims] do
-      {where, params} = scope([workspace_id: workspace_id], phase)
-      {phase, page_sql(phase, where, length(params)), params ++ ["", @page]}
+      opts = [workspace_id: workspace_id]
+      {where, params} = scope(opts, phase)
+
+      {phase, page_sql(phase, where, length(params), index_hint(opts, phase)),
+       params ++ ["", @page]}
     end
   end
 

lib/optimal_engine/store/dialect.ex

@@ -399,9 +399,12 @@ defmodule OptimalEngine.Store.Dialect do
 
   # IFNULL is COALESCE with two arguments, json_valid of a column is
   # PostgreSQL's IS JSON predicate (PostgreSQL 16), and SQLite's INTEGER and
-  # REAL are eight bytes, PostgreSQL's BIGINT and DOUBLE PRECISION.
+  # REAL are eight bytes, PostgreSQL's BIGINT and DOUBLE PRECISION. SQLite's
+  # INDEXED BY and NOT INDEXED pin a plan; PostgreSQL has no index hints and
+  # plans on its own, so the clause is dropped.
   defp rewrite_simple(masked) do
     masked
+    |> then(&Regex.replace(~r/\s+(?:INDEXED\s+BY\s+#{@identifier}|NOT\s+INDEXED)\b/i, &1, ""))
     |> then(&Regex.replace(~r/\bIFNULL\s*\(/i, &1, "COALESCE("))
     |> then(&Regex.replace(~r/\bjson_valid\s*\(\s*(#{@identifier})\s*\)/i, &1, "(\\1 IS JSON)"))
     |> then(&Regex.replace(~r/\bAS\s+INTEGER\s*\)/i, &1, "AS BIGINT)"))

test/storage/dialect_test.exs

@@ -118,6 +118,24 @@ defmodule OptimalEngine.Store.DialectTest do
                )
     end
 
+    test "SQLite index hints are dropped, and a column named like them is not" do
+      assert {:ok, "SELECT id FROM chunks WHERE workspace_id = $1 AND id > $2"} =
+               Dialect.to_postgres(
+                 "SELECT id FROM chunks INDEXED BY idx_chunks_workspace_id WHERE workspace_id = ?1 AND id > ?2"
+               )
+
+      assert {:ok, "SELECT c.id FROM contexts c WHERE c.workspace_id = $1"} =
+               Dialect.to_postgres(
+                 "SELECT c.id FROM contexts c INDEXED BY idx_contexts_workspace_id WHERE c.workspace_id = ?1"
+               )
+
+      assert {:ok, "SELECT a FROM t WHERE b = $1"} =
+               Dialect.to_postgres("SELECT a FROM t NOT INDEXED WHERE b = ?1")
+
+      assert {:ok, "SELECT id FROM contexts WHERE NOT indexed_at IS NULL"} =
+               Dialect.to_postgres("SELECT id FROM contexts WHERE NOT indexed_at IS NULL")
+    end
+
     test "a list bound as one JSON array parameter, and json_each anywhere else refused" do
       assert {:ok,
               "SELECT id FROM t WHERE a IN (SELECT value FROM jsonb_array_elements_text(CAST($2 AS text)::jsonb)) AND tenant_id = $1"} =

test/support/sql_census_allowlist.txt

@@ -306,7 +306,7 @@ rejected	BEGIN/**/IMMEDIATE
 rejected	DELETE FROM "…" WHERE "…" = ?
 rejected	DELETE FROM "…" WHERE "…" IN (?)
 rejected	DELETE FROM data_embeddings WHERE collection_id = ? AND (model != ? OR record_id NOT IN (SELECT "…" FROM "…"))
-rejected	EXPLAIN QUERY PLAN SELECT MAX(w.id) FROM ( SELECT id FROM chunks WHERE workspace_id = ? AND id > ? ORDER BY id LIMIT ? ) AS w
+rejected	EXPLAIN QUERY PLAN SELECT MAX(w.id) FROM ( SELECT id FROM chunks INDEXED BY idx_chunks_workspace_id WHERE workspace_id = ? AND id > ? ORDER BY id LIMIT ? ) AS w
 rejected	EXPLAIN QUERY PLAN SELECT N FROM asset_adapter_runs WHERE asset_id = ?
 rejected	EXPLAIN QUERY PLAN SELECT N FROM asset_embedding_refs WHERE adapter_run_id = ?
 rejected	EXPLAIN QUERY PLAN SELECT N FROM asset_embedding_refs WHERE asset_id = ?
@@ -336,9 +336,9 @@ rejected	EXPLAIN QUERY PLAN SELECT N FROM decomposition_runs WHERE source_packag
 rejected	EXPLAIN QUERY PLAN SELECT N FROM entity_mentions WHERE context_id = ?
 rejected	EXPLAIN QUERY PLAN SELECT N FROM entity_mentions WHERE source_package_id = ?
 rejected	EXPLAIN QUERY PLAN SELECT N FROM resolution_decisions WHERE mention_id = ?
-rejected	EXPLAIN QUERY PLAN SELECT c.id, c.node, c.entities, c.routed_to, c.supersedes, c.workspace_id, c.tenant_id FROM contexts c WHERE c.id > ? AND c.workspace_id = ? ORDER BY c.id LIMIT ?
-rejected	EXPLAIN QUERY PLAN SELECT c.id, c.tenant_id, c.workspace_id, c.signal_id, c.text, c.modality, c.asset_ref FROM chunks c WHERE c.workspace_id = ? AND c.id > ? AND c.id <= ? AND c.text <> '…' AND NOT EXISTS (SELECT N FROM chunk_embeddings e WHERE e.chunk_id = c.id) ORDER BY c.id LIMIT ?
-rejected	EXPLAIN QUERY PLAN SELECT cl.id, c.entities, cl.workspace_id, cl.tenant_id FROM claims cl JOIN contexts c ON c.id = cl.signal_id WHERE cl.signal_id IS NOT NULL AND cl.id > ? AND cl.workspace_id = ? ORDER BY cl.id LIMIT ?
+rejected	EXPLAIN QUERY PLAN SELECT c.id, c.node, c.entities, c.routed_to, c.supersedes, c.workspace_id, c.tenant_id FROM contexts c INDEXED BY idx_contexts_workspace_id WHERE c.id > ? AND c.workspace_id = ? ORDER BY c.id LIMIT ?
+rejected	EXPLAIN QUERY PLAN SELECT c.id, c.tenant_id, c.workspace_id, c.signal_id, c.text, c.modality, c.asset_ref FROM chunks c INDEXED BY idx_chunks_workspace_id WHERE c.workspace_id = ? AND c.id > ? AND c.id <= ? AND c.text <> '…' AND NOT EXISTS (SELECT N FROM chunk_embeddings e WHERE e.chunk_id = c.id) ORDER BY c.id LIMIT ?
+rejected	EXPLAIN QUERY PLAN SELECT cl.id, c.entities, cl.workspace_id, cl.tenant_id FROM claims cl INDEXED BY idx_claims_workspace_id JOIN contexts c ON c.id = cl.signal_id WHERE cl.signal_id IS NOT NULL AND cl.id > ? AND cl.workspace_id = ? ORDER BY cl.id LIMIT ?
 rejected	INSERT INTO "…" ("…", "…", "…", "…", "…", "…", "…") VALUES (?, N, ?)
 rejected	INSERT INTO "…" ("…", "…", "…", "…", "…", "…", "…", "…") VALUES (?, N, ?)
 rejected	INSERT INTO "…" ("…", "…", "…", "…", "…", "…", "…", "…", "…") VALUES (?, N, ?)