fix(vectors): harden large-repo merge_insert (mem pool + PK scalar index)

java_codebase_rag/config.py

@@ -47,15 +47,39 @@
 COCOINDEX_MAX_INFLIGHT_COMPONENTS_ENV = "COCOINDEX_MAX_INFLIGHT_COMPONENTS"
 COCOINDEX_DEFAULT_MAX_INFLIGHT_COMPONENTS = "256"
 
+# Lance native DataFusion hash-join memory pool ceiling (FairSpillPool). The
+# lance default is ~100 MiB, tuned for query workloads — too small for the
+# single big ``merge_insert`` cocoindex emits at the end of a flow component.
+# On ``--full-reprocess`` (all rows match the existing table → bulk-update
+# path) the hash join builds on a large side and exhausts the pool somewhere
+# around 75k-100k chunks: "Resources exhausted: Failed to allocate ... for
+# HashJoinInput ... N MiB remain available for the total pool". cocoindex is a
+# bare pass-through to lancedb (it never sets a Session/memory_limit), so it
+# inherits this default — we raise it here. FairSpillPool is a *reservation
+# ceiling*, not a pre-allocation: setting 1 GiB does not reserve 1 GiB upfront,
+# it just allows the join to grow before spilling/erroring, so it is safe on
+# memory-constrained hosts. An operator can still override via their own
+# ``LANCE_MEM_POOL_SIZE`` (subprocess_env copies os.environ, and apply is via
+# ``setdefault`` so the operator value wins). Increment is unaffected (tiny
+# batch → tiny hash table); only the full-reprocess write path is at risk.
+LANCE_MEM_POOL_SIZE_ENV = "LANCE_MEM_POOL_SIZE"
+LANCE_DEFAULT_MEM_POOL_SIZE = "1073741824"  # 1 GiB
+
 
 def cocoindex_subprocess_env_defaults() -> dict[str, str]:
-    """Env defaults applied to every CocoIndex subprocess to bound concurrency.
+    """Env defaults applied to every CocoIndex subprocess.
+
+    Bounds CocoIndex concurrency (``COCOINDEX_MAX_INFLIGHT_COMPONENTS``; see
+    :issue:`306`) and raises the Lance hash-join memory ceiling
+    (``LANCE_MEM_POOL_SIZE``) so a large full-reprocess does not exhaust the
+    default ~100 MiB pool mid-``merge_insert``.
 
     Apply with ``env.setdefault(...)`` so a caller-provided (operator) value
-    always wins. See :issue:`306`.
+    always wins.
     """
     return {
-        COCOINDEX_MAX_INFLIGHT_COMPONENTS_ENV: COCOINDEX_DEFAULT_MAX_INFLIGHT_COMPONENTS
+        COCOINDEX_MAX_INFLIGHT_COMPONENTS_ENV: COCOINDEX_DEFAULT_MAX_INFLIGHT_COMPONENTS,
+        LANCE_MEM_POOL_SIZE_ENV: LANCE_DEFAULT_MEM_POOL_SIZE,
     }
 
 _DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

java_codebase_rag/lance_optimize.py

@@ -184,6 +184,31 @@ async def optimize_lance_tables(
 
                 if last_exc is None:
                     results[name] = "ok"
+                    # Best-effort BTREE scalar index on the primary key ("id").
+                    # cocoindex's merge_insert defaults to use_index=True but
+                    # never creates a scalar PK index itself (declaring
+                    # primary_key in the schema does NOT auto-build a lance
+                    # index), so without this every merge_insert — increment
+                    # included — is a forced full scan of the PK column,
+                    # O(existing rows). On a large repo that scan dominates
+                    # increment wall-clock; with the index present the join does
+                    # lookups (~O(batch*log N)). Failure is non-fatal (the table
+                    # is still correct, just un-indexed) and never alters the
+                    # "ok" status, mirroring the FTS block below. ``replace=True``
+                    # keeps it idempotent across runs; table.optimize() above
+                    # maintains it on subsequent runs.
+                    try:
+                        from lancedb.index import BTree
+                        await table.create_index("id", config=BTree(), replace=True)
+                    except Exception as exc:
+                        low = str(exc).lower()
+                        if not any(
+                            w in low for w in ("exist", "duplicate", "already", "same name")
+                        ) and not quiet:
+                            print(
+                                f"java-codebase-rag: optimize: {name} id-index skipped: {exc}",
+                                file=sys.stderr,
+                            )
                     # Best-effort FTS index at index time (PR-SEARCH-3) so hybrid
                     # search works on all tables (java/sql/yaml) without a
                     # first-query race. Failure is non-fatal — the lazy

tests/test_config.py

@@ -506,6 +506,10 @@ def test_cocoindex_subprocess_env_defaults_uses_real_inflight_env_var() -> None:
     assert defaults["COCOINDEX_MAX_INFLIGHT_COMPONENTS"] == "256"
     # The bogus name from the broken #293 fix must NOT leak back in.
     assert "COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS" not in defaults
+    # Lance native hash-join pool ceiling, raised from the ~100 MiB default so a
+    # large full-reprocess merge_insert does not exhaust it mid-commit. Applied
+    # via setdefault, so an operator's own value still wins.
+    assert defaults["LANCE_MEM_POOL_SIZE"] == "1073741824"
 
 
 class TestConfigSourcePointer:

tests/test_lance_optimize.py

@@ -123,31 +123,92 @@ async def test_optimize_builds_fts_index_after_success(monkeypatch, tmp_path) ->
     table = _FakeTable(name, [None])  # optimize succeeds on first try
     conn = _FakeConnection(table_names={name}, tables={name: table})
     _install_fake_lancedb(monkeypatch, conn)
-    # Make ``from lancedb.index import FTS`` resolve against a stand-in module.
-    index_mod = types.ModuleType("lancedb.index")
-
-    class FTS:  # stands in for lancedb.index.FTS config object
-        pass
-
-    index_mod.FTS = FTS
-    monkeypatch.setitem(sys.modules, "lancedb.index", index_mod)
+    fts_cls, _btree_cls = _install_fake_lancedb_index(monkeypatch)
 
     results = await lance_optimize.optimize_lance_tables(tmp_path, quiet=True)
 
     assert results[name] == "ok"
-    assert len(table.create_index_calls) == 1, (
-        f"expected FTS create_index once after optimize, got {table.create_index_calls}"
-    )
-    call = table.create_index_calls[0]
-    assert call["args"] and call["args"][0] == "text", (
-        f"FTS index must target the 'text' column, got args={call['args']}"
+    # optimize now builds TWO indices (BTREE on "id" + FTS on "text"); find the
+    # FTS one by column rather than asserting on call count / order.
+    fts_calls = [
+        c for c in table.create_index_calls if c["args"] and c["args"][0] == "text"
+    ]
+    assert len(fts_calls) == 1, (
+        f"expected one FTS create_index on 'text', got {table.create_index_calls}"
     )
+    call = fts_calls[0]
     assert call["kwargs"].get("replace") is True, (
         f"FTS index must be built with replace=True, got kwargs={call['kwargs']}"
     )
+    assert isinstance(call["kwargs"].get("config"), fts_cls), (
+        f"FTS index config must be a lancedb.index.FTS, got kwargs={call['kwargs']}"
+    )
     assert conn.closed is True
 
 
+async def test_optimize_builds_btree_pk_index_after_success(monkeypatch, tmp_path) -> None:
+    """A successful optimize builds a BTREE scalar index on the ``id`` PK.
+
+    cocoindex's ``merge_insert`` defaults to ``use_index=True`` but never creates
+    a scalar PK index itself (declaring ``primary_key`` does NOT auto-build a
+    lance index), so without this every ``merge_insert`` — increment included —
+    is a forced full scan of the PK column (O(existing rows)). The BTREE index
+    flips that to ~O(batch*log N) lookups. Guards the path against silent
+    regression under the broad ``except`` that swallows index-build failures.
+    """
+    import types
+
+    from java_codebase_rag import lance_optimize
+
+    name = lance_optimize.LANCE_TABLE_NAMES[0]
+    table = _FakeTable(name, [None])  # optimize succeeds on first try
+    conn = _FakeConnection(table_names={name}, tables={name: table})
+    _install_fake_lancedb(monkeypatch, conn)
+    _fts_cls, btree_cls = _install_fake_lancedb_index(monkeypatch)
+
+    results = await lance_optimize.optimize_lance_tables(tmp_path, quiet=True)
+
+    assert results[name] == "ok"
+    id_calls = [
+        c for c in table.create_index_calls if c["args"] and c["args"][0] == "id"
+    ]
+    assert len(id_calls) == 1, (
+        f"expected one BTREE create_index on 'id', got {table.create_index_calls}"
+    )
+    call = id_calls[0]
+    assert call["kwargs"].get("replace") is True, (
+        f"BTREE PK index must be built with replace=True, got kwargs={call['kwargs']}"
+    )
+    assert isinstance(call["kwargs"].get("config"), btree_cls), (
+        f"BTREE PK index config must be a lancedb.index.BTree, got kwargs={call['kwargs']}"
+    )
+
+
+def _install_fake_lancedb_index(monkeypatch):
+    """Make ``from lancedb.index import FTS, BTree`` resolve to stand-in classes.
+
+    Returns the ``(FTS, BTree)`` config stand-ins so a test can assert the right
+    config object was passed. Both the FTS path (PR-SEARCH-3) and the BTREE PK
+    path build their index config via a local ``from lancedb.index import ...``;
+    without a stand-in those imports hit the broad ``except`` and silently no-op,
+    leaving the path unexercised.
+    """
+    import types
+
+    index_mod = types.ModuleType("lancedb.index")
+
+    class FTS:  # stands in for lancedb.index.FTS config object
+        pass
+
+    class BTree:  # stands in for lancedb.index.BTree config object
+        pass
+
+    index_mod.FTS = FTS
+    index_mod.BTree = BTree
+    monkeypatch.setitem(sys.modules, "lancedb.index", index_mod)
+    return FTS, BTree
+
+
 async def test_optimize_does_not_retry_non_conflict_error(monkeypatch, tmp_path) -> None:
     """A non-conflict exception is re-raised (captured per-table), never retried."""
     from java_codebase_rag import lance_optimize