From 28db1837df4f8b19d33e9e1b0b453900fb07b09f Mon Sep 17 00:00:00 2001 From: Dmitry Teryaev Date: Mon, 6 Jul 2026 22:24:47 +0300 Subject: [PATCH 1/2] fix(config): raise LANCE_MEM_POOL_SIZE default to 1 GiB for cocoindex subprocess MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cocoindex is a bare pass-through to lancedb — it never sets a Session or memory_limit, and has zero references to memory_pool / LANCE_MEM_POOL_SIZE across its Python package and 56 MB Rust engine. So it inherits lance's ~100 MiB FairSpillPool default, which is tuned for query workloads and is too small for the single big merge_insert a flow component emits at commit. On --full-reprocess (all rows match the existing table → bulk-update path) the hash-join build side can exceed the pool around 75k–100k chunks, raising: RuntimeError: lance error: Resources exhausted: Failed to allocate ~200 MB for HashJoinInput ... 100.0 MB remain available for the total pool The crash is allocation-granularity dependent and non-deterministic (observed once under sustained bench load; does not reproduce on isolated reruns up to 200k), but the mechanism and the LANCE_MEM_POOL_SIZE knob are source-confirmed (the error string matches DataFusion 52.4.0 embedded in lancedb._lancedb.abi3.so; the var is read once at session init). Raise the default to 1 GiB via cocoindex_subprocess_env_defaults() — applied with setdefault, so an operator's own LANCE_MEM_POOL_SIZE still wins. The pool is a reservation ceiling, not a pre-allocation, so 1 GiB costs nothing upfront and is safe on constrained hosts. Covers ~500k-chunk bulk-update build sides; larger repos can raise it further. Increment is unaffected (tiny batch → tiny hash table); only the full-reprocess write path is at risk. Co-Authored-By: Claude --- java_codebase_rag/config.py | 30 +++++++++++++++++++++++++++--- tests/test_config.py | 4 ++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/java_codebase_rag/config.py b/java_codebase_rag/config.py index 7f294c1..c6e6e47 100644 --- a/java_codebase_rag/config.py +++ b/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" diff --git a/tests/test_config.py b/tests/test_config.py index 716bb09..8bb3303 100644 --- a/tests/test_config.py +++ b/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: From 9a53f94d885a9832935d97aba236c67d8edd7dec Mon Sep 17 00:00:00 2001 From: Dmitry Teryaev Date: Mon, 6 Jul 2026 22:24:47 +0300 Subject: [PATCH 2/2] perf(optimize): build a BTREE scalar index on the id PK each run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cocoindex's merge_insert defaults to use_index=True but never creates a scalar index on the primary key (declaring primary_key in the schema does NOT auto- build a lance index), so every merge_insert — increment included — is a forced full scan of the id column, O(existing rows). Measured ~flat through 100k (<5 ms) because the constant is tiny (columnar PK-only scan), but it is O(N) and would dominate increment wall-clock on a multi-million-chunk repo. With the index present the join does lookups (~O(batch*log N)). Add a best-effort BTREE scalar index on "id" in optimize_lance_tables(), right after table.optimize() succeeds, mirroring the existing FTS block: local import of the config class (lancedb.index.BTree), replace=True for idempotency, and a non-fatal broad except that never flips the "ok" status. Runs after every init/increment/reprocess, so the index exists before the next merge_insert; table.optimize() maintains it on subsequent runs. Verified end-to-end (real lancedb): optimize on a temp index yields [Index(BTree, id), Index(FTS, text)] on all three tables, and empty tables (sql/yaml at count=0) create the index without error. Tests: the existing test_optimize_builds_fts_index_after_success defined only FTS in its fake lancedb.index module, so the BTree import silently no-op'd under the broad except — fixed by also defining BTree (via a shared _install_fake_lancedb_index helper) and asserting by column rather than call count. Add test_optimize_builds_btree_pk_index_after_success, mutation-checked (fails without this change). Co-Authored-By: Claude --- java_codebase_rag/lance_optimize.py | 25 ++++++++ tests/test_lance_optimize.py | 89 ++++++++++++++++++++++++----- 2 files changed, 100 insertions(+), 14 deletions(-) diff --git a/java_codebase_rag/lance_optimize.py b/java_codebase_rag/lance_optimize.py index 85d1db3..520de9a 100644 --- a/java_codebase_rag/lance_optimize.py +++ b/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 diff --git a/tests/test_lance_optimize.py b/tests/test_lance_optimize.py index 69e1c33..0676fdf 100644 --- a/tests/test_lance_optimize.py +++ b/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