Skip to content

Commit 8aac8b4

Browse files
georgelichenChen Li
georgelichen
and
Chen Li
authored
Repair stale Pandas Arrow registrations before GraphRAG workers start (#1414)
Pandas registers the period and interval Arrow extension types as an import side effect, while reload or dynamic import hooks can remove the module without clearing pyarrow's process-local registry. Prepare the extension module inside the isolated GraphRAG worker and recover only the known stale-registration error before retrying the import. Constraint: GraphRAG is an optional dependency and its imports must remain inside the isolated worker thread Rejected: warning filters and BLAS thread limits | the failure is a raised ArrowKeyError during module registration, not a warning or numeric thread race Confidence: high Scope-risk: narrow Reversibility: clean Directive: Keep the recovery limited to pandas.period/pandas.interval and the exact duplicate-registration error Tested: targeted GraphRAG tests (87 passed, 2 skipped); subprocess regression on pandas 3.0.0/pyarrow 25.0.1 and pandas 2.3.3/pyarrow 22.0.0; Ruff, format, compileall, git diff --check Not-tested: full GraphRAG indexing with external model services; pyright without optional GraphRAG dependencies Co-authored-by: Chen Li <[email protected]>
1 parent 695414d commit 8aac8b4

4 files changed

Lines changed: 156 additions & 0 deletions

File tree

‎deeptutor/services/rag/pipelines/graphrag/engine.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
classify_embedding_error,
3535
classify_model_error,
3636
)
37+
from .pandas_compat import prepare_pandas_arrow_extensions
3738
from .provider import (
3839
COMPLETION_TYPE,
3940
resolve_completion_call_args,
@@ -106,6 +107,7 @@ async def _run_isolated(work: Callable[[], Awaitable[_T]]) -> _T:
106107
"""
107108

108109
def _runner() -> _T:
110+
prepare_pandas_arrow_extensions()
109111
loop = asyncio.new_event_loop()
110112
asyncio.set_event_loop(loop)
111113
try:

‎deeptutor/services/rag/pipelines/graphrag/pandas_compat.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
"""Scoped compatibility for stale Pandas/Arrow extension registrations."""
2+
3+
from __future__ import annotations
4+
5+
import importlib
6+
import threading
7+
from typing import Any
8+
9+
_PANDAS_ARROW_EXTENSION_MODULE = "pandas.core.arrays.arrow.extension_types"
10+
_PANDAS_ARROW_EXTENSION_NAMES = ("pandas.period", "pandas.interval")
11+
_extension_registration_lock = threading.Lock()
12+
13+
14+
def _is_duplicate_registration_error(error: BaseException) -> bool:
15+
message = str(error)
16+
return (
17+
"A type extension with name" in message
18+
and "already defined" in message
19+
and any(name in message for name in _PANDAS_ARROW_EXTENSION_NAMES)
20+
)
21+
22+
23+
def _unregister_stale_extension_types(
24+
pyarrow: Any,
25+
arrow_key_error: type[BaseException],
26+
) -> None:
27+
for name in _PANDAS_ARROW_EXTENSION_NAMES:
28+
try:
29+
pyarrow.unregister_extension_type(name)
30+
except arrow_key_error:
31+
# The extension may not have been registered yet.
32+
continue
33+
34+
35+
def prepare_pandas_arrow_extensions() -> None:
36+
"""Import Pandas Arrow extensions, repairing stale registry entries.
37+
38+
Pandas registers ``pandas.period`` and ``pandas.interval`` as a module import
39+
side effect. If a test runner, plugin, or reload hook removes that module from
40+
``sys.modules`` without unregistering the pyarrow types, importing it again
41+
raises ``pyarrow.lib.ArrowKeyError``. GraphRAG runs in isolated worker threads,
42+
so this repair is performed immediately inside each worker before GraphRAG
43+
imports its pandas users.
44+
45+
The helper is a no-op when pyarrow is not installed. Other import or registry
46+
errors are allowed to propagate so a real dependency failure is not hidden.
47+
"""
48+
try:
49+
import pyarrow
50+
from pyarrow.lib import ArrowKeyError
51+
except ImportError:
52+
return
53+
54+
with _extension_registration_lock:
55+
try:
56+
importlib.import_module(_PANDAS_ARROW_EXTENSION_MODULE)
57+
except ModuleNotFoundError as error:
58+
if error.name == "pandas":
59+
return
60+
raise
61+
except ArrowKeyError as error:
62+
if not _is_duplicate_registration_error(error):
63+
raise
64+
_unregister_stale_extension_types(pyarrow, ArrowKeyError)
65+
importlib.import_module(_PANDAS_ARROW_EXTENSION_MODULE)
66+
67+
68+
__all__ = ["prepare_pandas_arrow_extensions"]

‎tests/services/rag/test_graphrag_pandas_compat.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
"""Regression tests for Pandas/Arrow extension registration in GraphRAG."""
2+
3+
from __future__ import annotations
4+
5+
import importlib
6+
import os
7+
from pathlib import Path
8+
import subprocess
9+
import sys
10+
11+
import pytest
12+
13+
14+
def test_prepare_pandas_arrow_extensions_repairs_stale_registration() -> None:
15+
"""A re-import after module removal must not hit ArrowKeyError."""
16+
pytest.importorskip("pandas", reason="pandas is required for this regression test")
17+
pytest.importorskip("pyarrow", reason="pyarrow is required for this regression test")
18+
root = Path(__file__).parents[3]
19+
script = """
20+
import importlib
21+
import sys
22+
23+
import pandas.core.arrays.arrow.extension_types
24+
25+
module_name = "pandas.core.arrays.arrow.extension_types"
26+
del sys.modules[module_name]
27+
28+
from deeptutor.services.rag.pipelines.graphrag.pandas_compat import (
29+
prepare_pandas_arrow_extensions,
30+
)
31+
32+
prepare_pandas_arrow_extensions()
33+
importlib.import_module(module_name)
34+
print("ok")
35+
"""
36+
env = os.environ.copy()
37+
env["PYTHONPATH"] = os.pathsep.join(
38+
part for part in (str(root), env.get("PYTHONPATH", "")) if part
39+
)
40+
result = subprocess.run(
41+
[sys.executable, "-c", script],
42+
cwd=root,
43+
env=env,
44+
capture_output=True,
45+
text=True,
46+
check=False,
47+
)
48+
49+
assert result.returncode == 0, result.stderr
50+
assert result.stdout.strip().endswith("ok")
51+
52+
53+
def test_prepare_pandas_arrow_extensions_does_not_hide_unrelated_arrow_errors(
54+
monkeypatch: pytest.MonkeyPatch,
55+
) -> None:
56+
"""Only the known duplicate-registration error is recovered."""
57+
pytest.importorskip("pyarrow", reason="pyarrow is required for this regression test")
58+
59+
from pyarrow.lib import ArrowKeyError
60+
61+
from deeptutor.services.rag.pipelines.graphrag import pandas_compat
62+
63+
def raise_unrelated_error(_name: str) -> None:
64+
raise ArrowKeyError("unrelated Arrow registry failure")
65+
66+
monkeypatch.setattr(importlib, "import_module", raise_unrelated_error)
67+
68+
with pytest.raises(ArrowKeyError, match="unrelated Arrow registry failure"):
69+
pandas_compat.prepare_pandas_arrow_extensions()

‎tests/services/rag/test_graphrag_pipeline.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,23 @@ async def caller() -> str:
382382
assert seen["current"] is seen["running"]
383383

384384

385+
def test_engine_isolated_runtime_prepares_pandas_arrow_extensions(monkeypatch) -> None:
386+
calls: list[str] = []
387+
388+
monkeypatch.setattr(
389+
engine,
390+
"prepare_pandas_arrow_extensions",
391+
lambda: calls.append("prepared"),
392+
)
393+
394+
async def work() -> str:
395+
calls.append("work")
396+
return "done"
397+
398+
assert asyncio.run(engine._run_isolated(work)) == "done"
399+
assert calls == ["prepared", "work"]
400+
401+
385402
def test_compatibility_probe_resolves_candidate_without_mutating_active_model(
386403
monkeypatch, tmp_path: Path
387404
) -> None:

0 commit comments

Comments
 (0)