[Improvement] Celery worker OOMKilled

[iziang]

🧠 Summary

celery-worker pods are being OOMKilled repeatedly due to excessive memory usage when loading data from very large PostgreSQL tables (lightrag_vdb_relation and lightrag_vdb_entity).

---

🐛 Problem Description

The celeryworker pod in the aperag deployment frequently restarts because of Out of Memory (OOMKilled) events.
After investigation, it appears that the worker is loading entire tables from PostgreSQL into memory during the execution of the query_lightrag_vdb_relation_all function.

These two tables have grown significantly and now exceed the pod’s memory limit (4 GiB):

postgres=# SELECT pg_size_pretty(pg_total_relation_size('lightrag_vdb_relation'));
 pg_size_pretty 
----------------
 4228 MB
(1 row)

postgres=# SELECT pg_size_pretty(pg_total_relation_size('lightrag_vdb_entity'));
 pg_size_pretty 
----------------
 3533 MB
(1 row)

postgres=# 

---

📋 Environment Details

Pod:

Name:        celeryworker-799fc9f787-65jmf
Namespace:   default
Node:        cn-hongkong.10.231.119.38
Status:      Running (frequently OOMKilled)
RestartCount: 412
Memory Limit: 4Gi

Image:

apecloud-registry.cn-zhangjiakou.cr.aliyuncs.com/apecloud/aperag:v0.7.0-alpha.5

Command:

/opt/venv/bin/celery -A config.celery worker -l INFO --concurrency=16 -Q 10.231.119.38,celery --pool=threads

---

🔍 Investigation

Using py-spy to inspect the running worker process:

root@:~# py-spy dump -p 1956455
Process 1956455: /opt/venv/bin/python3 /opt/venv/bin/celery -A config.celery worker -l INFO --concurrency=16 -Q 10.231.119.38,celery --pool=threads
Python v3.11.13 (/usr/local/bin/python3.11)

Thread 11 (active): "MainThread"
    poll (kombu/utils/eventio.py:83)
    create_loop (kombu/asynchronous/hub.py:317)
    asynloop (celery/worker/loops.py:97)
    start (celery/worker/consumer/consumer.py:772)
    start (celery/bootsteps.py:116)
    start (celery/worker/consumer/consumer.py:341)
    start (celery/bootsteps.py:365)
    start (celery/bootsteps.py:116)
    start (celery/worker/worker.py:203)
    worker (celery/bin/worker.py:356)
    caller (celery/bin/base.py:135)
    new_func (click/decorators.py:33)
    invoke (click/core.py:788)
    invoke (click/core.py:1443)
    invoke (click/core.py:1697)
    main (click/core.py:1082)
    __call__ (click/core.py:1161)
    main (celery/bin/celery.py:231)
    main (celery/__main__.py:15)
    <module> (celery:10)
Thread 30 (active): "ThreadPoolExecutor-3_0"
    _read_ready__get_buffer (asyncio/selector_events.py:974)
    _read_ready (asyncio/selector_events.py:956)
    _run (asyncio/events.py:84)
    _run_once (asyncio/base_events.py:1936)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    process_document_for_celery (lightrag_manager.py:118)
    create_index (aperag/tasks/document.py:114)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 31 (active): "ThreadPoolExecutor-3_1"
    read (ssl.py:1168)
    recv (ssl.py:1295)
    read (httpcore/_backends/sync.py:128)
    _receive_event (httpcore/_sync/http11.py:217)
    _receive_response_headers (httpcore/_sync/http11.py:177)
    handle_request (httpcore/_sync/http11.py:106)
    handle_request (httpcore/_sync/connection.py:103)
    handle_request (httpcore/_sync/connection_pool.py:236)
    handle_request (httpx/_transports/default.py:250)
    _send_single_request (httpx/_client.py:1014)
    _send_handling_redirects (httpx/_client.py:979)
    _send_handling_auth (httpx/_client.py:942)
    send (httpx/_client.py:914)
    post (http_handler.py:761)
    _make_common_sync_call (llm_http_handler.py:175)
    completion (llm_http_handler.py:471)
    completion (litellm/main.py:2626)
    wrapper (litellm/utils.py:1219)
    _completion_core (aperag/llm/completion/completion_service.py:173)
    generate (aperag/llm/completion/completion_service.py:211)
    _summarize_text (aperag/index/summary_index.py:368)
    _generate_document_summary (aperag/index/summary_index.py:305)
    create_index (aperag/index/summary_index.py:75)
    update_index (aperag/index/summary_index.py:196)
    update_index (aperag/tasks/document.py:325)
    update_index_task (config/celery_tasks.py:393)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 32 (idle): "ThreadPoolExecutor-3_2"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    delete_document_for_celery (lightrag_manager.py:126)
    delete_index (aperag/tasks/document.py:216)
    delete_index_task (config/celery_tasks.py:334)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 33 (active): "ThreadPoolExecutor-3_3"
    acquire (logging/__init__.py:927)
    handle (logging/__init__.py:976)
    callHandlers (logging/__init__.py:1706)
    handle (logging/__init__.py:1644)
    _log (logging/__init__.py:1634)
    info (logging/__init__.py:1489)
    cleanup_expired_documents (aperag/tasks/collection.py:211)
    reconcile_all (aperag/tasks/reconciler.py:649)
    cleanup_expired_documents_task (config/celery_tasks.py:841)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 34 (idle): "ThreadPoolExecutor-3_4"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    delete_document_for_celery (lightrag_manager.py:126)
    delete_index (aperag/tasks/document.py:216)
    delete_index_task (config/celery_tasks.py:334)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 35 (active): "ThreadPoolExecutor-3_5"
    _read_ready__get_buffer (asyncio/selector_events.py:974)
    _read_ready (asyncio/selector_events.py:956)
    _run (asyncio/events.py:84)
    _run_once (asyncio/base_events.py:1936)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    process_document_for_celery (lightrag_manager.py:118)
    create_index (aperag/tasks/document.py:114)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 36 (active): "ThreadPoolExecutor-3_6"
    read (ssl.py:1168)
    recv (ssl.py:1295)
    read (httpcore/_backends/sync.py:128)
    _receive_event (httpcore/_sync/http11.py:217)
    _receive_response_headers (httpcore/_sync/http11.py:177)
    handle_request (httpcore/_sync/http11.py:106)
    handle_request (httpcore/_sync/connection.py:103)
    handle_request (httpcore/_sync/connection_pool.py:236)
    handle_request (httpx/_transports/default.py:250)
    _send_single_request (httpx/_client.py:1014)
    _send_handling_redirects (httpx/_client.py:979)
    _send_handling_auth (httpx/_client.py:942)
    send (httpx/_client.py:914)
    post (http_handler.py:761)
    _make_common_sync_call (llm_http_handler.py:175)
    completion (llm_http_handler.py:471)
    completion (litellm/main.py:2626)
    wrapper (litellm/utils.py:1219)
    _completion_core (aperag/llm/completion/completion_service.py:173)
    generate (aperag/llm/completion/completion_service.py:211)
    _summarize_text (aperag/index/summary_index.py:368)
    _generate_document_summary (aperag/index/summary_index.py:305)
    create_index (aperag/index/summary_index.py:75)
    create_index (aperag/tasks/document.py:135)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 37 (idle): "ThreadPoolExecutor-3_7"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    process_document_for_celery (lightrag_manager.py:118)
    create_index (aperag/tasks/document.py:114)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 38 (idle): "ThreadPoolExecutor-3_8"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    delete_document_for_celery (lightrag_manager.py:126)
    delete_index (aperag/tasks/document.py:216)
    delete_index_task (config/celery_tasks.py:334)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 39 (active): "ThreadPoolExecutor-3_9"
    read (ssl.py:1168)
    recv (ssl.py:1295)
    read (httpcore/_backends/sync.py:128)
    _receive_event (httpcore/_sync/http11.py:217)
    _receive_response_headers (httpcore/_sync/http11.py:177)
    handle_request (httpcore/_sync/http11.py:106)
    handle_request (httpcore/_sync/connection.py:103)
    handle_request (httpcore/_sync/connection_pool.py:236)
    handle_request (httpx/_transports/default.py:250)
    _send_single_request (httpx/_client.py:1014)
    _send_handling_redirects (httpx/_client.py:979)
    _send_handling_auth (httpx/_client.py:942)
    send (httpx/_client.py:914)
    post (http_handler.py:761)
    _make_common_sync_call (llm_http_handler.py:175)
    completion (llm_http_handler.py:471)
    completion (litellm/main.py:2626)
    wrapper (litellm/utils.py:1219)
    _completion_core (aperag/llm/completion/completion_service.py:173)
    generate (aperag/llm/completion/completion_service.py:211)
    _summarize_text (aperag/index/summary_index.py:368)
    _generate_document_summary (aperag/index/summary_index.py:305)
    create_index (aperag/index/summary_index.py:75)
    create_index (aperag/tasks/document.py:135)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 40 (idle): "ThreadPoolExecutor-3_10"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    delete_document_for_celery (lightrag_manager.py:126)
    delete_index (aperag/tasks/document.py:216)
    delete_index_task (config/celery_tasks.py:334)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 41 (active): "ThreadPoolExecutor-3_11"
    read (ssl.py:1168)
    recv (ssl.py:1295)
    read (httpcore/_backends/sync.py:128)
    _receive_event (httpcore/_sync/http11.py:217)
    _receive_response_headers (httpcore/_sync/http11.py:177)
    handle_request (httpcore/_sync/http11.py:106)
    handle_request (httpcore/_sync/connection.py:103)
    handle_request (httpcore/_sync/connection_pool.py:236)
    handle_request (httpx/_transports/default.py:250)
    _send_single_request (httpx/_client.py:1014)
    _send_handling_redirects (httpx/_client.py:979)
    _send_handling_auth (httpx/_client.py:942)
    send (httpx/_client.py:914)
    post (http_handler.py:761)
    _make_common_sync_call (llm_http_handler.py:175)
    completion (llm_http_handler.py:471)
    completion (litellm/main.py:2626)
    wrapper (litellm/utils.py:1219)
    _completion_core (aperag/llm/completion/completion_service.py:173)
    generate (aperag/llm/completion/completion_service.py:211)
    _summarize_text (aperag/index/summary_index.py:368)
    _generate_document_summary (aperag/index/summary_index.py:305)
    create_index (aperag/index/summary_index.py:75)
    create_index (aperag/tasks/document.py:135)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 42 (idle): "ThreadPoolExecutor-3_12"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    delete_document_for_celery (lightrag_manager.py:126)
    delete_index (aperag/tasks/document.py:216)
    delete_index_task (config/celery_tasks.py:334)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 43 (active): "ThreadPoolExecutor-3_13"
    emit (logging/__init__.py:1113)
    handle (logging/__init__.py:978)
    callHandlers (logging/__init__.py:1706)
    handle (logging/__init__.py:1644)
    _log (logging/__init__.py:1634)
    info (logging/__init__.py:1489)
    cleanup_expired_documents (aperag/tasks/collection.py:211)
    reconcile_all (aperag/tasks/reconciler.py:649)
    cleanup_expired_documents_task (config/celery_tasks.py:841)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 44 (idle): "ThreadPoolExecutor-3_14"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    process_document_for_celery (lightrag_manager.py:118)
    create_index (aperag/tasks/document.py:114)
    create_index_task (config/celery_tasks.py:267)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 45 (idle): "ThreadPoolExecutor-3_15"
    select (selectors.py:468)
    _run_once (asyncio/base_events.py:1898)
    run_forever (asyncio/base_events.py:608)
    run_until_complete (asyncio/base_events.py:641)
    _run_in_new_loop (lightrag_manager.py:216)
    delete_document_for_celery (lightrag_manager.py:126)
    delete_index (aperag/tasks/document.py:216)
    delete_index_task (config/celery_tasks.py:334)
    run (celery/app/autoretry.py:38)
    __protected_call__ (celery/app/trace.py:736)
    trace_task (celery/app/trace.py:453)
    fast_trace_task (celery/app/trace.py:651)
    apply_target (celery/concurrency/base.py:30)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 46 (active): "ThreadPoolExecutor-1_0"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3476 (active): "asyncio_0"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3478 (active): "asyncio_0"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3480 (active): "asyncio_0"
    do_execute (sqlalchemy/engine/default.py:945)
    _exec_single_context (sqlalchemy/engine/base.py:1964)
    _execute_context (sqlalchemy/engine/base.py:1843)
    _execute_clauseelement (sqlalchemy/engine/base.py:1638)
    _execute_on_connection (sqlalchemy/sql/elements.py:523)
    execute (sqlalchemy/engine/base.py:1416)
    orm_execute_statement (sqlalchemy/orm/bulk_persistence.py:1294)
    _execute_internal (sqlalchemy/orm/session.py:2251)
    execute (sqlalchemy/orm/session.py:2365)
    _upsert_node (aperag/db/repositories/graph.py:65)
    _execute_transaction (aperag/db/repositories/base.py:65)
    upsert_graph_node (aperag/db/repositories/graph.py:69)
    _sync_upsert_node (lightrag/kg/pg_ops_sync_graph_storage.py:68)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3482 (active): "asyncio_0"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3484 (active): "asyncio_0"
    do_execute (sqlalchemy/engine/default.py:945)
    _exec_single_context (sqlalchemy/engine/base.py:1964)
    _execute_context (sqlalchemy/engine/base.py:1843)
    _execute_clauseelement (sqlalchemy/engine/base.py:1638)
    _execute_on_connection (sqlalchemy/sql/elements.py:523)
    execute (sqlalchemy/engine/base.py:1416)
    orm_execute_statement (sqlalchemy/orm/context.py:306)
    _execute_internal (sqlalchemy/orm/session.py:2251)
    execute (sqlalchemy/orm/session.py:2365)
    _query (aperag/db/repositories/lightrag.py:576)
    _execute_query (aperag/db/repositories/base.py:55)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3485 (active): "asyncio_0"
    do_execute (sqlalchemy/engine/default.py:945)
    _exec_single_context (sqlalchemy/engine/base.py:1964)
    _execute_context (sqlalchemy/engine/base.py:1843)
    _execute_clauseelement (sqlalchemy/engine/base.py:1638)
    _execute_on_connection (sqlalchemy/sql/elements.py:523)
    execute (sqlalchemy/engine/base.py:1416)
    orm_execute_statement (sqlalchemy/orm/context.py:306)
    _execute_internal (sqlalchemy/orm/session.py:2251)
    execute (sqlalchemy/orm/session.py:2365)
    _query (aperag/db/repositories/lightrag.py:576)
    _execute_query (aperag/db/repositories/base.py:55)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3486 (active): "asyncio_0"
    <listcomp> (sqlalchemy/orm/loading.py:223)
    chunks (sqlalchemy/orm/loading.py:223)
    _fetchall_impl (sqlalchemy/engine/result.py:2268)
    _fetchall_impl (sqlalchemy/engine/result.py:1674)
    _allrows (sqlalchemy/engine/result.py:548)
    all (sqlalchemy/engine/result.py:1767)
    _query (aperag/db/repositories/lightrag.py:577)
    _execute_query (aperag/db/repositories/base.py:55)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3487 (active+gil): "asyncio_0"
    fetchall (sqlalchemy/engine/cursor.py:1137)
    _fetchall_impl (sqlalchemy/engine/cursor.py:2135)
    _raw_all_rows (sqlalchemy/engine/result.py:540)
    chunks (sqlalchemy/orm/loading.py:219)
    _fetchall_impl (sqlalchemy/engine/result.py:2268)
    _fetchall_impl (sqlalchemy/engine/result.py:1674)
    _allrows (sqlalchemy/engine/result.py:548)
    all (sqlalchemy/engine/result.py:1767)
    _query (aperag/db/repositories/lightrag.py:577)
    _execute_query (aperag/db/repositories/base.py:55)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3488 (active): "asyncio_0"
    do_execute (sqlalchemy/engine/default.py:945)
    _exec_single_context (sqlalchemy/engine/base.py:1964)
    _execute_context (sqlalchemy/engine/base.py:1843)
    _execute_clauseelement (sqlalchemy/engine/base.py:1638)
    _execute_on_connection (sqlalchemy/sql/elements.py:523)
    execute (sqlalchemy/engine/base.py:1416)
    orm_execute_statement (sqlalchemy/orm/context.py:306)
    _execute_internal (sqlalchemy/orm/session.py:2251)
    execute (sqlalchemy/orm/session.py:2365)
    _query (aperag/db/repositories/lightrag.py:576)
    _execute_query (aperag/db/repositories/base.py:55)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3489 (active): "asyncio_0"
    do_execute (sqlalchemy/engine/default.py:945)
    _exec_single_context (sqlalchemy/engine/base.py:1964)
    _execute_context (sqlalchemy/engine/base.py:1843)
    _execute_clauseelement (sqlalchemy/engine/base.py:1638)
    _execute_on_connection (sqlalchemy/sql/elements.py:523)
    execute (sqlalchemy/engine/base.py:1416)
    orm_execute_statement (sqlalchemy/orm/context.py:306)
    _execute_internal (sqlalchemy/orm/session.py:2251)
    execute (sqlalchemy/orm/session.py:2365)
    _query (aperag/db/repositories/lightrag.py:576)
    _execute_query (aperag/db/repositories/base.py:55)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)
    run (concurrent/futures/thread.py:58)
    _worker (concurrent/futures/thread.py:83)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3490 (active): "asyncio_1"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3491 (active): "asyncio_1"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3492 (active): "asyncio_1"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3493 (active): "asyncio_1"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3497 (active): "ThreadPoolExecutor-1_1"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3498 (active): "ThreadPoolExecutor-1_2"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)
Thread 3499 (active): "ThreadPoolExecutor-1_3"
    _worker (concurrent/futures/thread.py:81)
    run (threading.py:982)
    _bootstrap_inner (threading.py:1045)
    _bootstrap (threading.py:1002)

Stack traces show multiple active calls to:

query_lightrag_vdb_relation_all

This function attempts to load all records from lightrag_vdb_relation, causing high memory pressure.

---

💥 Root Cause

• The tables lightrag_vdb_relation and lightrag_vdb_entity are very large (4 GB and 3.5 GB respectively).
• The query_lightrag_vdb_relation_all function does not paginate or stream results.
• As a result, the entire dataset is loaded into memory, exceeding the Celery worker’s 4 GiB memory limit and triggering OOMKilled events.

---

✅ Suggested Fixes / Improvements

1. Implement pagination or streaming queries when fetching relation/entity data.
   Example: use LIMIT/OFFSET or server-side cursors (psycopg2.extras.DictCursor(name='cursor_name')).
2. Avoid loading entire tables into memory for processing — process items in batches.
3. Consider offloading large queries to background jobs with higher memory limits or move data aggregation logic into SQL.
4. Optionally, increase the worker memory limit temporarily as a mitigation (from 4Gi → 8Gi), but the core issue is unbounded memory usage.

---

📈 Impact

• Frequent Celery worker restarts.
• Tasks interrupted or retried repeatedly.
• Reduced system reliability and performance.

[github-actions]

This issue has been marked as stale because it has been open for 30 days with no activity

[earayu]

📊 Root Cause Analysis and Solution Proposal for OOM Issue #1398

1️⃣ Current Code Purpose and Logic

The adelete_by_doc_id() method in lightrag.py (lines 874-1110) is responsible for deleting a document and cleaning up all its related references in the knowledge graph. The deletion process follows these steps:

Step 1: Find Document Chunks (Lines 884-896)

all_chunks = await self.text_chunks.get_all()  # Load ALL chunks
related_chunks = {
    chunk_id: chunk_data
    for chunk_id, chunk_data in all_chunks.items()
    if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
}

Problem: Loads ALL chunks from the database into memory, then filters in Python to find chunks belonging to the target document.

Step 2: Clean Vector Storage References (Lines 898-960)

2.1 Process Entities:

all_entities = await self.entities_vdb.get_all()  # Line 904 - Load ALL entities
for entity_id, entity_data in all_entities.items():
    old_chunk_ids = set(entity_data.get("chunk_ids", []))
    new_chunk_ids = old_chunk_ids - chunk_ids  # Filter in Python
    
    if not new_chunk_ids:
        # Mark entity for deletion if no chunks remain
        entities_to_delete_from_vdb.append(entity_name)
    elif len(new_chunk_ids) != len(old_chunk_ids):
        # Update entity if some chunks remain
        entity_data["chunk_ids"] = list(new_chunk_ids)
        entities_to_update_in_vdb[entity_id] = entity_data

Problem: Loads ALL entities from lightrag_vdb_entity table (3533 MB on disk) into memory, then iterates through each one to check if it references any deleted chunks.

2.2 Process Relationships:

all_relationships = await self.relationships_vdb.get_all()  # Line 935 - Load ALL relationships
for rel_id, rel_data in all_relationships.items():
    old_chunk_ids = set(rel_data.get("chunk_ids", []))
    new_chunk_ids = old_chunk_ids - chunk_ids  # Filter in Python
    
    if not new_chunk_ids:
        # Mark relationship for deletion if no chunks remain
        relationships_to_delete_from_vdb.append((src_id, tgt_id))
    elif len(new_chunk_ids) != len(old_chunk_ids):
        # Update relationship if some chunks remain
        rel_data["chunk_ids"] = list(new_chunk_ids)
        relationships_to_update_in_vdb[rel_id] = rel_data

Problem: Loads ALL relationships from lightrag_vdb_relation table (4228 MB on disk) into memory, then iterates through each one.

Step 3: Clean Graph Storage References (Lines 962-1016)

Iterates through all graph nodes and edges to update source_id references.

Step 4: Execute Deletions and Updates (Lines 1018-1076)

Performs the actual database operations.

Step 5: Verification (Lines 1083-1099)

remaining_chunks = await self.text_chunks.get_all()  # Load ALL chunks again

Loads all chunks again to verify deletion was successful.

---

2️⃣ The Root Problem

Problem Pattern: "Load Everything, Filter in Memory"

The code follows an anti-pattern:

1. Execute SELECT * FROM table WHERE workspace = ?
2. Load entire result set into Python memory
3. Iterate through every record in Python
4. Filter and process in application layer

Underlying Implementation

File: aperag/db/repositories/lightrag.py (Lines 571-579)

def query_lightrag_vdb_relation_all(self, workspace: str):
    """Query all LightRAG VDB Relation records for workspace"""
    def _query(session):
        stmt = select(LightRAGVDBRelationModel).where(
            LightRAGVDBRelationModel.workspace == workspace
        )
        result = session.execute(stmt)
        return {relation.id: relation for relation in result.scalars().all()}  # ❌ .all()
    return self._execute_query(_query)

The .all() method forces SQLAlchemy to load the entire result set into memory.

File: aperag/graph/lightrag/kg/pg_ops_sync_vector_storage.py (Lines 61-121)

async def get_all(self) -> dict[str, Any]:
    def _sync_get_all():
        from aperag.db.ops import db_ops
        
        if is_namespace(self.namespace, NameSpace.VECTOR_STORE_ENTITIES):
            models = db_ops.query_lightrag_vdb_entity_all(self.workspace)  # Load all
            return {
                entity_id: {
                    "id": entity_id,
                    "entity_name": model.entity_name,
                    "content": model.content or "",
                    "content_vector": model.content_vector,  # Large vector data
                    "chunk_ids": model.chunk_ids or [],
                    # ... more fields ...
                }
                for entity_id, model in models.items()
            }

Memory Impact Analysis

Data Structure (from aperag/db/models.py lines 733-764):

class LightRAGVDBEntityModel(Base):
    __tablename__ = "lightrag_vdb_entity"
    
    id = Column(String(255), primary_key=True)
    workspace = Column(String(255), primary_key=True)
    entity_name = Column(String(255), nullable=True)
    content = Column(Text, nullable=True)  # Few KB
    content_vector = Column(Vector(), nullable=True)  # ~6KB (1536-dim float32)
    chunk_ids = Column(ARRAY(String), nullable=True)  # Variable size
    # ... more fields ...

class LightRAGVDBRelationModel(Base):
    __tablename__ = "lightrag_vdb_relation"
    
    id = Column(String(255), primary_key=True)
    workspace = Column(String(255), primary_key=True)
    source_id = Column(String(256), nullable=True)
    target_id = Column(String(256), nullable=True)
    content = Column(Text, nullable=True)  # Few KB
    content_vector = Column(Vector(), nullable=True)  # ~6KB (1536-dim float32)
    chunk_ids = Column(ARRAY(String), nullable=True)  # Variable size
    # ... more fields ...

Memory Per Record:

• content_vector: ~6 KB (1536-dimensional float32 vector)
• content: 1-5 KB (text description)
• chunk_ids: 0.5-2 KB (array of UUIDs)
• Other fields: ~1 KB
• Total: ~8-14 KB per record

Your Database Size:

• lightrag_vdb_relation table: 4228 MB on disk
• lightrag_vdb_entity table: 3533 MB on disk
• Total: 7,761 MB = ~7.6 GB on disk

In Python Memory:

• Python object overhead: ~1.5-2x of raw data
• Estimated memory usage: 10-15 GB when fully loaded
• With multiple concurrent Celery workers calling adelete_by_doc_id() simultaneously: Memory multiplies

Complexity Analysis

Memory Complexity: O(N)

• N = Total number of records in the database
• Every deletion operation loads all N records regardless of how many are actually affected

Time Complexity: O(N)

• Must iterate through all N records even if only 0.1% are relevant

Scalability: Linear degradation

• As data grows from 7 GB → 14 GB → 28 GB, memory requirements grow proportionally
• Upgrading from 4GB → 8GB only delays the inevitable

Why It Causes OOMKilled

From the py-spy stack trace in issue #1398:

Thread 3487 (active+gil): "asyncio_0"
    fetchall (sqlalchemy/engine/cursor.py:1137)
    _fetchall_impl (sqlalchemy/engine/cursor.py:2135)
    _raw_all_rows (sqlalchemy/engine/result.py:540)
    chunks (sqlalchemy/orm/loading.py:219)
    _fetchall_impl (sqlalchemy/engine/result.py:2268)
    all (sqlalchemy/engine/result.py:1767)  # ← .all() loading everything
    _query (aperag/db/repositories/lightrag.py:577)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)

Multiple threads are simultaneously calling query_lightrag_vdb_relation_all(), each attempting to load 7+ GB into memory, causing the 4 GB pod limit to be exceeded.

---

3️⃣ Proposed Solution: Database-Level Filtering

Core Principle

Push filtering logic from Python application layer to PostgreSQL database layer.

Instead of:

# ❌ Load everything, filter in Python
all_data = load_all_from_database()
filtered_data = [x for x in all_data if condition(x)]

Do this:

# ✅ Filter in database, load only matches
filtered_data = load_from_database_where(condition)

Benefits

• ✅ Memory: O(M) where M = matching records (typically 0.01-1% of N)
• ✅ Performance: Leverage database indexes and query optimizer
• ✅ Correctness: Identical semantic behavior, just optimized execution
• ✅ Scalability: Handles data growth without memory explosion

Implementation Plan

Phase 1: Add Database-Level Filtering Methods

File: aperag/db/repositories/lightrag.py

Add three new methods after line 580:

def query_lightrag_doc_chunks_by_doc_id(self, workspace: str, doc_id: str):
    """Query chunks by document ID - database-level filtering.
    
    This is a memory-optimized alternative to query_lightrag_doc_chunks_all()
    that only loads chunks for a specific document.
    
    Args:
        workspace: Workspace identifier
        doc_id: Document ID to filter by
        
    Returns:
        Dictionary mapping chunk_id to chunk model
    """
    def _query(session):
        stmt = select(LightRAGDocChunksModel).where(
            LightRAGDocChunksModel.workspace == workspace,
            LightRAGDocChunksModel.full_doc_id == doc_id  # ✅ Database-level filter
        )
        result = session.execute(stmt)
        return {chunk.id: chunk for chunk in result.scalars().all()}
    
    return self._execute_query(_query)


def query_lightrag_vdb_entity_by_chunk_ids(self, workspace: str, chunk_ids: list[str]):
    """Query entities that reference any of the specified chunk_ids.
    
    Uses PostgreSQL array overlap operator (&&) for efficient filtering at database level.
    This is a memory-optimized alternative to query_lightrag_vdb_entity_all().
    
    Args:
        workspace: Workspace identifier
        chunk_ids: List of chunk IDs to search for
        
    Returns:
        Dictionary mapping entity_id to entity model
        
    Example:
        If chunk_ids = ['chunk1', 'chunk2'], this returns all entities where:
        entity.chunk_ids && ['chunk1', 'chunk2'] = true
        (i.e., entity.chunk_ids contains at least one of the specified chunks)
    """
    def _query(session):
        from sqlalchemy import cast
        from sqlalchemy.dialects.postgresql import ARRAY
        from sqlalchemy.types import String
        
        # PostgreSQL array overlap operator: ARRAY[...] && ARRAY[...]
        # Returns true if arrays have any elements in common
        stmt = select(LightRAGVDBEntityModel).where(
            LightRAGVDBEntityModel.workspace == workspace,
            LightRAGVDBEntityModel.chunk_ids.op("&&")(  # ✅ Array overlap
                cast(chunk_ids, ARRAY(String))
            )
        )
        result = session.execute(stmt)
        return {entity.id: entity for entity in result.scalars().all()}
    
    return self._execute_query(_query)


def query_lightrag_vdb_relation_by_chunk_ids(self, workspace: str, chunk_ids: list[str]):
    """Query relationships that reference any of the specified chunk_ids.
    
    Uses PostgreSQL array overlap operator (&&) for efficient filtering at database level.
    This is a memory-optimized alternative to query_lightrag_vdb_relation_all().
    
    Args:
        workspace: Workspace identifier
        chunk_ids: List of chunk IDs to search for
        
    Returns:
        Dictionary mapping relation_id to relation model
        
    Example:
        If chunk_ids = ['chunk1', 'chunk2'], this returns all relationships where:
        relation.chunk_ids && ['chunk1', 'chunk2'] = true
        (i.e., relation.chunk_ids contains at least one of the specified chunks)
    """
    def _query(session):
        from sqlalchemy import cast
        from sqlalchemy.dialects.postgresql import ARRAY
        from sqlalchemy.types import String
        
        stmt = select(LightRAGVDBRelationModel).where(
            LightRAGVDBRelationModel.workspace == workspace,
            LightRAGVDBRelationModel.chunk_ids.op("&&")(  # ✅ Array overlap
                cast(chunk_ids, ARRAY(String))
            )
        )
        result = session.execute(stmt)
        return {relation.id: relation for relation in result.scalars().all()}
    
    return self._execute_query(_query)

PostgreSQL Array Operator && Explanation:

• Syntax: array1 && array2
• Returns: true if arrays have any common elements
• Example: ARRAY['a','b','c'] && ARRAY['c','d'] = true
• Performance: Can use GIN index for fast lookups

Phase 2: Add Vector Storage Layer Methods

File: aperag/graph/lightrag/kg/pg_ops_sync_vector_storage.py

Add new method after line 121:

async def get_by_doc_id(self, doc_id: str) -> dict[str, Any]:
    """Get chunks for a specific document - memory-optimized.
    
    This is an alternative to get_all() that only loads chunks for one document.
    
    Args:
        doc_id: Document ID to filter by
        
    Returns:
        Dictionary mapping chunk_id to chunk data
    """
    def _sync_get_by_doc_id():
        from aperag.db.ops import db_ops
        from aperag.graph.lightrag.namespace import NameSpace, is_namespace
        
        if is_namespace(self.namespace, NameSpace.VECTOR_STORE_CHUNKS):
            models = db_ops.query_lightrag_doc_chunks_by_doc_id(self.workspace, doc_id)
            return {
                chunk_id: {
                    "id": chunk_id,
                    "tokens": model.tokens,
                    "content": model.content or "",
                    "chunk_order_index": model.chunk_order_index,
                    "full_doc_id": model.full_doc_id,
                    "content_vector": model.content_vector,
                    "file_path": model.file_path,
                    "created_at": int(model.create_time.timestamp()) if model.create_time else None,
                }
                for chunk_id, model in models.items()
            }
        else:
            logger.error(f"get_by_doc_id only supports VECTOR_STORE_CHUNKS namespace")
            return {}
    
    return await asyncio.to_thread(_sync_get_by_doc_id)


async def get_by_chunk_ids(self, chunk_ids: list[str]) -> dict[str, Any]:
    """Get entities/relationships that reference the specified chunk_ids - memory-optimized.
    
    This is an alternative to get_all() that only loads records referencing specific chunks.
    
    Args:
        chunk_ids: List of chunk IDs to search for
        
    Returns:
        Dictionary mapping id to entity/relationship data
    """
    def _sync_get_by_chunk_ids():
        from aperag.db.ops import db_ops
        from aperag.graph.lightrag.namespace import NameSpace, is_namespace
        
        if is_namespace(self.namespace, NameSpace.VECTOR_STORE_ENTITIES):
            models = db_ops.query_lightrag_vdb_entity_by_chunk_ids(self.workspace, chunk_ids)
            return {
                entity_id: {
                    "id": entity_id,
                    "entity_name": model.entity_name,
                    "content": model.content or "",
                    "content_vector": model.content_vector,
                    "chunk_ids": model.chunk_ids or [],
                    "file_path": model.file_path,
                    "created_at": int(model.create_time.timestamp()) if model.create_time else None,
                }
                for entity_id, model in models.items()
            }
        elif is_namespace(self.namespace, NameSpace.VECTOR_STORE_RELATIONSHIPS):
            models = db_ops.query_lightrag_vdb_relation_by_chunk_ids(self.workspace, chunk_ids)
            return {
                relation_id: {
                    "id": relation_id,
                    "source_id": model.source_id,
                    "target_id": model.target_id,
                    "content": model.content or "",
                    "content_vector": model.content_vector,
                    "chunk_ids": model.chunk_ids or [],
                    "file_path": model.file_path,
                    "created_at": int(model.create_time.timestamp()) if model.create_time else None,
                    "src_id": model.source_id,
                    "tgt_id": model.target_id,
                }
                for relation_id, model in models.items()
            }
        else:
            logger.error(f"get_by_chunk_ids not supported for namespace: {self.namespace}")
            return {}
    
    return await asyncio.to_thread(_sync_get_by_chunk_ids)

Phase 3: Update Document Deletion Logic

File: aperag/graph/lightrag/lightrag.py

Replace lines 884-960 with optimized version:

async def adelete_by_doc_id(self, doc_id: str) -> None:
    """Delete a document and all its related data - MEMORY OPTIMIZED VERSION
    
    Args:
        doc_id: Document ID to delete
    """
    try:
        self.lightrag_logger.info(f"Starting deletion for document {doc_id}")
        
        # ========== STEP 1: Get chunks for this document (OPTIMIZED) ==========
        # OLD CODE:
        # all_chunks = await self.text_chunks.get_all()  # ❌ Loads ALL chunks
        # related_chunks = {
        #     chunk_id: chunk_data
        #     for chunk_id, chunk_data in all_chunks.items()
        #     if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
        # }
        
        # NEW CODE: Only load chunks for this specific document
        if hasattr(self.text_chunks, "get_by_doc_id"):
            related_chunks = await self.text_chunks.get_by_doc_id(doc_id)  # ✅ Filtered
        else:
            # Fallback to old behavior if method not available
            all_chunks = await self.text_chunks.get_all()
            related_chunks = {
                chunk_id: chunk_data
                for chunk_id, chunk_data in all_chunks.items()
                if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
            }
        
        if not related_chunks:
            logger.warning(f"No chunks found for document {doc_id}")
            return
        
        chunk_ids = set(related_chunks.keys())
        chunk_ids_list = list(chunk_ids)
        self.lightrag_logger.info(f"Found {len(chunk_ids)} chunks to delete for document {doc_id}")
        
        # ========== STEP 2: Handle Vector Storage References (OPTIMIZED) ==========
        # Process entities in vector storage
        entities_to_delete_from_vdb = []
        entities_to_update_in_vdb = {}
        
        # OLD CODE:
        # if hasattr(self.entities_vdb, "get_all"):
        #     all_entities = await self.entities_vdb.get_all()  # ❌ Loads ALL entities
        #     for entity_id, entity_data in all_entities.items():
        #         ...
        
        # NEW CODE: Only load entities that reference these chunks
        if hasattr(self.entities_vdb, "get_by_chunk_ids"):
            affected_entities = await self.entities_vdb.get_by_chunk_ids(
                chunk_ids_list
            )  # ✅ Only affected records
            self.lightrag_logger.debug(
                f"Loaded {len(affected_entities)} affected entities (optimized query)"
            )
        elif hasattr(self.entities_vdb, "get_all"):
            # Fallback to old behavior if optimized method not available
            all_entities = await self.entities_vdb.get_all()
            affected_entities = {
                entity_id: entity_data
                for entity_id, entity_data in all_entities.items()
                if isinstance(entity_data, dict) 
                and "chunk_ids" in entity_data
                and set(entity_data.get("chunk_ids", [])) & chunk_ids  # Has overlap
            }
            self.lightrag_logger.debug(
                f"Loaded {len(all_entities)} entities, {len(affected_entities)} affected (fallback mode)"
            )
        else:
            affected_entities = {}
        
        for entity_id, entity_data in affected_entities.items():
            if not isinstance(entity_data, dict) or "chunk_ids" not in entity_data:
                continue
            
            # Remove deleted chunks from entity's chunk_ids array
            old_chunk_ids = set(entity_data.get("chunk_ids", []))
            new_chunk_ids = old_chunk_ids - chunk_ids
            
            if not new_chunk_ids:
                # Entity has no remaining chunks, mark for deletion
                entity_name = entity_data.get("entity_name")
                if entity_name:
                    entities_to_delete_from_vdb.append(entity_name)
                    self.lightrag_logger.debug(
                        f"Entity {entity_name} marked for deletion from VDB - no remaining chunks"
                    )
            elif len(new_chunk_ids) != len(old_chunk_ids):
                # Entity has some remaining chunks, update chunk_ids array
                entity_data["chunk_ids"] = list(new_chunk_ids)
                entity_data["source_id"] = GRAPH_FIELD_SEP.join(new_chunk_ids)
                entities_to_update_in_vdb[entity_id] = entity_data
                self.lightrag_logger.debug(
                    f"Entity {entity_data.get('entity_name')} chunk_ids updated: {len(old_chunk_ids)} -> {len(new_chunk_ids)}"
                )
        
        # Process relationships in vector storage
        relationships_to_delete_from_vdb = []
        relationships_to_update_in_vdb = {}
        
        # OLD CODE:
        # if hasattr(self.relationships_vdb, "get_all"):
        #     all_relationships = await self.relationships_vdb.get_all()  # ❌ Loads ALL
        #     for rel_id, rel_data in all_relationships.items():
        #         ...
        
        # NEW CODE: Only load relationships that reference these chunks
        if hasattr(self.relationships_vdb, "get_by_chunk_ids"):
            affected_relationships = await self.relationships_vdb.get_by_chunk_ids(
                chunk_ids_list
            )  # ✅ Only affected records
            self.lightrag_logger.debug(
                f"Loaded {len(affected_relationships)} affected relationships (optimized query)"
            )
        elif hasattr(self.relationships_vdb, "get_all"):
            # Fallback to old behavior if optimized method not available
            all_relationships = await self.relationships_vdb.get_all()
            affected_relationships = {
                rel_id: rel_data
                for rel_id, rel_data in all_relationships.items()
                if isinstance(rel_data, dict)
                and "chunk_ids" in rel_data
                and set(rel_data.get("chunk_ids", [])) & chunk_ids  # Has overlap
            }
            self.lightrag_logger.debug(
                f"Loaded {len(all_relationships)} relationships, {len(affected_relationships)} affected (fallback mode)"
            )
        else:
            affected_relationships = {}
        
        for rel_id, rel_data in affected_relationships.items():
            if not isinstance(rel_data, dict) or "chunk_ids" not in rel_data:
                continue
            
            # Remove deleted chunks from relationship's chunk_ids array
            old_chunk_ids = set(rel_data.get("chunk_ids", []))
            new_chunk_ids = old_chunk_ids - chunk_ids
            
            if not new_chunk_ids:
                # Relationship has no remaining chunks, mark for deletion
                src_id = rel_data.get("src_id") or rel_data.get("source_id")
                tgt_id = rel_data.get("tgt_id") or rel_data.get("target_id")
                if src_id and tgt_id:
                    relationships_to_delete_from_vdb.append((src_id, tgt_id))
                    self.lightrag_logger.debug(
                        f"Relationship {src_id}-{tgt_id} marked for deletion from VDB - no remaining chunks"
                    )
            elif len(new_chunk_ids) != len(old_chunk_ids):
                # Relationship has some remaining chunks, update chunk_ids array
                rel_data["chunk_ids"] = list(new_chunk_ids)
                rel_data["source_id"] = GRAPH_FIELD_SEP.join(new_chunk_ids)
                relationships_to_update_in_vdb[rel_id] = rel_data
                self.lightrag_logger.debug(
                    f"Relationship {rel_data.get('src_id')}-{rel_data.get('tgt_id')} chunk_ids updated: {len(old_chunk_ids)} -> {len(new_chunk_ids)}"
                )
        
        # ========== STEP 3-5: Rest remains the same ==========
        # ... (continue with existing graph storage cleanup and deletion logic)

Also update the verification step (around line 1085):

# ========== STEP 5: Simple verification ==========
# Verify chunks were actually deleted
# OLD CODE:
# remaining_chunks = await self.text_chunks.get_all()  # ❌ Loads all again
# remaining_related_chunks = {
#     chunk_id: chunk_data
#     for chunk_id, chunk_data in remaining_chunks.items()
#     if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
# }

# NEW CODE: Only load chunks for this document
if hasattr(self.text_chunks, "get_by_doc_id"):
    remaining_related_chunks = await self.text_chunks.get_by_doc_id(doc_id)
else:
    remaining_chunks = await self.text_chunks.get_all()
    remaining_related_chunks = {
        chunk_id: chunk_data
        for chunk_id, chunk_data in remaining_chunks.items()
        if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
    }

if remaining_related_chunks:
    self.lightrag_logger.warning(
        f"Verification failed: {len(remaining_related_chunks)} chunks still exist for document {doc_id}"
    )
else:
    self.lightrag_logger.info(
        f"Verification successful: All chunks for document {doc_id} have been deleted"
    )

Phase 4: Add Database Indexes (Optional but Recommended)

Create a migration script to add GIN indexes for better query performance:

-- File: migrations/add_chunk_ids_indexes.sql

-- Add GIN index for entity chunk_ids array queries
-- This dramatically speeds up the && (array overlap) operator
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_lightrag_vdb_entity_chunk_ids 
ON lightrag_vdb_entity USING GIN (chunk_ids)
WHERE workspace IS NOT NULL;

-- Add GIN index for relationship chunk_ids array queries
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_lightrag_vdb_relation_chunk_ids 
ON lightrag_vdb_relation USING GIN (chunk_ids)
WHERE workspace IS NOT NULL;

-- Add B-tree index for document ID filtering
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_lightrag_doc_chunks_full_doc_id 
ON lightrag_doc_chunks (workspace, full_doc_id)
WHERE full_doc_id IS NOT NULL;

-- Verify indexes were created
\d+ lightrag_vdb_entity
\d+ lightrag_vdb_relation
\d+ lightrag_doc_chunks

Index Performance Impact:

• GIN indexes on array columns enable O(log N) lookups instead of O(N) table scans
• CONCURRENTLY option allows index creation without locking the table
• Expected index size: ~5-10% of table size

---

4️⃣ Expected Impact

Memory Reduction

Current Behavior (Delete 1 document):

1. Load ALL chunks: ~X GB
2. Load ALL entities: ~3.5 GB
3. Load ALL relationships: ~4.2 GB
4. Load ALL chunks again for verification: ~X GB
5. Total peak memory: 8-15 GB

After Optimization (Delete 1 document):

1. Load chunks for 1 document: ~10-100 MB (depending on document size)
2. Load affected entities: ~5-50 MB (only entities referencing those chunks)
3. Load affected relationships: ~5-50 MB (only relationships referencing those chunks)
4. Load chunks for 1 document again: ~10-100 MB
5. Total peak memory: 30-300 MB

Memory Savings: 97-99%

Performance Improvement

Query Time:

• Current: O(N) table scan of millions of records
• Optimized: O(log N + M) with index lookup, where M = matching records
• Expected speedup: 10-100x depending on data distribution

Database Load:

• Reduced data transfer between PostgreSQL and Python
• Less network bandwidth usage
• Lower disk I/O

Python Performance:

• Reduced GC pressure
• Lower CPU usage for iteration
• Faster response times

Scalability

Current System:

• 7.6 GB data → requires 8+ GB RAM
• 15 GB data → requires 16+ GB RAM
• 30 GB data → requires 32+ GB RAM
• Linear scaling: Double data = Double RAM

Optimized System:

• 7.6 GB data → requires 500 MB RAM
• 15 GB data → requires 500 MB RAM
• 30 GB data → requires 500 MB RAM
• Constant scaling: Memory independent of total data size

Correctness Guarantee

✅ Identical Semantic Behavior:

• All references are still properly cleaned up
• Same deletion logic, just filtered at database level
• No data loss or corruption

✅ Backward Compatible:

• Includes fallback to old behavior if new methods not available
• Uses hasattr() checks before calling new methods
• Graceful degradation

✅ Testable:

• Can verify identical results between old and new implementation
• Easy to A/B test with feature flags

---

5️⃣ Testing Plan

Unit Tests

Create tests/test_lightrag_optimized_queries.py:

import pytest
from aperag.db.ops import db_ops

class TestOptimizedQueries:
    """Test new optimized query methods"""
    
    def test_query_by_doc_id_returns_correct_chunks(self):
        """Verify get_by_doc_id returns same results as filtered get_all"""
        doc_id = "test-doc-123"
        
        # Old method
        all_chunks = db_ops.query_lightrag_doc_chunks_all(workspace)
        expected = {k: v for k, v in all_chunks.items() if v.full_doc_id == doc_id}
        
        # New method
        actual = db_ops.query_lightrag_doc_chunks_by_doc_id(workspace, doc_id)
        
        assert actual.keys() == expected.keys()
    
    def test_query_entities_by_chunk_ids(self):
        """Verify entity filtering by chunk_ids works correctly"""
        chunk_ids = ["chunk1", "chunk2", "chunk3"]
        
        # Old method
        all_entities = db_ops.query_lightrag_vdb_entity_all(workspace)
        expected = {
            k: v for k, v in all_entities.items()
            if v.chunk_ids and set(v.chunk_ids) & set(chunk_ids)
        }
        
        # New method
        actual = db_ops.query_lightrag_vdb_entity_by_chunk_ids(workspace, chunk_ids)
        
        assert actual.keys() == expected.keys()
    
    def test_query_relationships_by_chunk_ids(self):
        """Verify relationship filtering by chunk_ids works correctly"""
        chunk_ids = ["chunk1", "chunk2"]
        
        # Old method
        all_rels = db_ops.query_lightrag_vdb_relation_all(workspace)
        expected = {
            k: v for k, v in all_rels.items()
            if v.chunk_ids and set(v.chunk_ids) & set(chunk_ids)
        }
        
        # New method
        actual = db_ops.query_lightrag_vdb_relation_by_chunk_ids(workspace, chunk_ids)
        
        assert actual.keys() == expected.keys()

Integration Tests

Create tests/test_document_deletion_optimized.py:

import pytest
from aperag.graph.lightrag.lightrag import LightRAG

class TestDocumentDeletionOptimized:
    """Test optimized document deletion"""
    
    @pytest.mark.asyncio
    async def test_delete_document_memory_usage(self, lightrag_instance):
        """Verify optimized deletion uses less memory"""
        import tracemalloc
        
        doc_id = "test-doc-to-delete"
        
        # Measure memory usage
        tracemalloc.start()
        await lightrag_instance.adelete_by_doc_id(doc_id)
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        
        # Should use significantly less than 1 GB
        assert peak < 1_000_000_000, f"Peak memory {peak} exceeds 1GB"
    
    @pytest.mark.asyncio
    async def test_delete_document_correctness(self, lightrag_instance):
        """Verify optimized deletion produces correct results"""
        doc_id = "test-doc-correctness"
        
        # Record state before deletion
        chunks_before = await lightrag_instance.text_chunks.get_by_doc_id(doc_id)
        chunk_ids = set(chunks_before.keys())
        
        # Perform deletion
        await lightrag_instance.adelete_by_doc_id(doc_id)
        
        # Verify chunks deleted
        chunks_after = await lightrag_instance.text_chunks.get_by_doc_id(doc_id)
        assert len(chunks_after) == 0
        
        # Verify no entities reference deleted chunks
        all_entities = await lightrag_instance.entities_vdb.get_all()
        for entity in all_entities.values():
            if entity.get("chunk_ids"):
                assert not (set(entity["chunk_ids"]) & chunk_ids)

Performance Tests

Create tests/perf/test_deletion_performance.py:

import time
import pytest

class TestDeletionPerformance:
    """Benchmark deletion performance"""
    
    @pytest.mark.benchmark
    @pytest.mark.asyncio
    async def test_deletion_speed(self, lightrag_instance, benchmark_doc_id):
        """Measure deletion speed"""
        start = time.time()
        await lightrag_instance.adelete_by_doc_id(benchmark_doc_id)
        elapsed = time.time() - start
        
        # Should complete in reasonable time even with large database
        assert elapsed < 60, f"Deletion took {elapsed}s, expected < 60s"
    
    @pytest.mark.benchmark
    def test_query_performance_with_index(self, db_connection):
        """Verify indexes are being used"""
        import sqlparse
        
        # Run EXPLAIN ANALYZE
        result = db_connection.execute("""
            EXPLAIN ANALYZE
            SELECT * FROM lightrag_vdb_entity
            WHERE workspace = 'test' AND chunk_ids && ARRAY['chunk1']
        """)
        
        plan = result.fetchall()
        plan_text = "\n".join([row[0] for row in plan])
        
        # Should use index scan, not sequential scan
        assert "Index Scan" in plan_text or "Bitmap Index Scan" in plan_text
        assert "Seq Scan" not in plan_text

Load Testing

@pytest.mark.load_test
@pytest.mark.asyncio
async def test_concurrent_deletions(lightrag_instances):
    """Test multiple concurrent deletions don't cause OOM"""
    import asyncio
    
    # Simulate 10 concurrent deletion tasks
    tasks = [
        instance.adelete_by_doc_id(f"doc-{i}")
        for i, instance in enumerate(lightrag_instances[:10])
    ]
    
    # Should all complete without OOM
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # Check no OOM errors
    for result in results:
        if isinstance(result, Exception):
            assert "MemoryError" not in str(result)
            assert "OOMKilled" not in str(result)

---

6️⃣ Rollout Plan

Phase 1: Development (Week 1-2)

1. Implement new database methods
2. Implement vector storage methods
3. Write unit tests
4. Local testing

Phase 2: Staging Testing (Week 3)

1. Deploy to staging environment
2. Run integration tests
3. Performance benchmarking
4. Memory profiling

Phase 3: Production Rollout (Week 4)

1. Deploy database indexes (use CONCURRENTLY to avoid locks)
2. Deploy new code with feature flag
3. Enable for 10% of traffic
4. Monitor metrics (memory usage, deletion times, error rates)
5. Gradually increase to 100%

Phase 4: Cleanup (Week 5)

1. Remove old code paths after verification
2. Update documentation
3. Remove feature flags

---

7️⃣ Monitoring and Metrics

Key Metrics to Track

Memory Usage:

# Add to celery worker monitoring
celery_worker_memory_usage_bytes
celery_worker_memory_limit_bytes
celery_worker_oom_kills_total

Query Performance:

# Add timing metrics
lightrag_deletion_duration_seconds
lightrag_query_chunks_duration_seconds
lightrag_query_entities_duration_seconds
lightrag_query_relationships_duration_seconds

Database Performance:

-- Monitor query plans
SELECT * FROM pg_stat_statements 
WHERE query LIKE '%chunk_ids%' 
ORDER BY total_exec_time DESC;

Alerts

# Alert if memory usage exceeds 75% of limit
alert: CeleryWorkerHighMemory
expr: celery_worker_memory_usage_bytes / celery_worker_memory_limit_bytes > 0.75

# Alert if deletion takes too long
alert: SlowDocumentDeletion
expr: lightrag_deletion_duration_seconds > 300

---

8️⃣ References

• Issue: [Improvement] Celery worker OOMKilled #1398
• PostgreSQL Array Functions: https://www.postgresql.org/docs/current/functions-array.html
• PostgreSQL Array Operators: https://www.postgresql.org/docs/current/functions-array.html#ARRAY-OPERATORS-TABLE
• PostgreSQL GIN Indexes: https://www.postgresql.org/docs/current/gin-intro.html
• SQLAlchemy Array Operations: https://docs.sqlalchemy.org/en/20/core/type_basics.html#sqlalchemy.types.ARRAY
• Python Memory Profiling: https://docs.python.org/3/library/tracemalloc.html

[earayu]

📊 OOM 问题 #1398 的根本原因分析和解决方案

1️⃣ 当前代码的目的和逻辑

lightrag.py 中的 adelete_by_doc_id() 方法（第 874-1110 行）负责删除一个文档并清理知识图谱中所有相关的引用。删除过程包括以下步骤：

步骤 1：查找文档的 chunks（第 884-896 行）

all_chunks = await self.text_chunks.get_all()  # 加载所有 chunks
related_chunks = {
    chunk_id: chunk_data
    for chunk_id, chunk_data in all_chunks.items()
    if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
}

问题： 将数据库中的所有 chunks 加载到内存，然后在 Python 中过滤出属于目标文档的 chunks。

步骤 2：清理向量存储中的引用（第 898-960 行）

2.1 处理实体（Entities）：

all_entities = await self.entities_vdb.get_all()  # 第 904 行 - 加载所有实体
for entity_id, entity_data in all_entities.items():
    old_chunk_ids = set(entity_data.get("chunk_ids", []))
    new_chunk_ids = old_chunk_ids - chunk_ids  # 在 Python 中过滤
    
    if not new_chunk_ids:
        # 如果没有剩余 chunks，标记实体为删除
        entities_to_delete_from_vdb.append(entity_name)
    elif len(new_chunk_ids) != len(old_chunk_ids):
        # 如果还有一些 chunks 剩余，更新实体
        entity_data["chunk_ids"] = list(new_chunk_ids)
        entities_to_update_in_vdb[entity_id] = entity_data

问题： 将 lightrag_vdb_entity 表中的所有实体（磁盘上 3533 MB）加载到内存，然后遍历每一个实体，检查其 chunk_ids 数组是否包含任何待删除的 chunks。

2.2 处理关系（Relationships）：

all_relationships = await self.relationships_vdb.get_all()  # 第 935 行 - 加载所有关系
for rel_id, rel_data in all_relationships.items():
    old_chunk_ids = set(rel_data.get("chunk_ids", []))
    new_chunk_ids = old_chunk_ids - chunk_ids  # 在 Python 中过滤
    
    if not new_chunk_ids:
        # 如果没有剩余 chunks，标记关系为删除
        relationships_to_delete_from_vdb.append((src_id, tgt_id))
    elif len(new_chunk_ids) != len(old_chunk_ids):
        # 如果还有一些 chunks 剩余，更新关系
        rel_data["chunk_ids"] = list(new_chunk_ids)
        relationships_to_update_in_vdb[rel_id] = rel_data

问题： 将 lightrag_vdb_relation 表中的所有关系（磁盘上 4228 MB）加载到内存，然后遍历每一个关系。

步骤 3：清理图存储中的引用（第 962-1016 行）

遍历所有图节点和边，更新 source_id 引用。

步骤 4：执行删除和更新操作（第 1018-1076 行）

执行实际的数据库操作。

步骤 5：验证（第 1083-1099 行）

remaining_chunks = await self.text_chunks.get_all()  # 再次加载所有 chunks

再次加载所有 chunks 来验证删除是否成功。

---

2️⃣ 根本问题

问题模式："先全部加载，再在内存中过滤"

代码采用了一种反模式：

1. 执行 SELECT * FROM table WHERE workspace = ?
2. 将整个结果集加载到 Python 内存
3. 在 Python 中遍历每一条记录
4. 在应用层进行过滤和处理

底层实现

文件： aperag/db/repositories/lightrag.py（第 571-579 行）

def query_lightrag_vdb_relation_all(self, workspace: str):
    """查询工作空间的所有 LightRAG VDB 关系记录"""
    def _query(session):
        stmt = select(LightRAGVDBRelationModel).where(
            LightRAGVDBRelationModel.workspace == workspace
        )
        result = session.execute(stmt)
        return {relation.id: relation for relation in result.scalars().all()}  # ❌ .all()
    return self._execute_query(_query)

.all() 方法强制 SQLAlchemy 将整个结果集加载到内存。

文件： aperag/graph/lightrag/kg/pg_ops_sync_vector_storage.py（第 61-121 行）

async def get_all(self) -> dict[str, Any]:
    def _sync_get_all():
        from aperag.db.ops import db_ops
        
        if is_namespace(self.namespace, NameSpace.VECTOR_STORE_ENTITIES):
            models = db_ops.query_lightrag_vdb_entity_all(self.workspace)  # 加载全部
            return {
                entity_id: {
                    "id": entity_id,
                    "entity_name": model.entity_name,
                    "content": model.content or "",
                    "content_vector": model.content_vector,  # 大型向量数据
                    "chunk_ids": model.chunk_ids or [],
                    # ... 更多字段 ...
                }
                for entity_id, model in models.items()
            }

内存影响分析

数据结构（来自 aperag/db/models.py 第 733-764 行）：

class LightRAGVDBEntityModel(Base):
    __tablename__ = "lightrag_vdb_entity"
    
    id = Column(String(255), primary_key=True)
    workspace = Column(String(255), primary_key=True)
    entity_name = Column(String(255), nullable=True)
    content = Column(Text, nullable=True)  # 几 KB
    content_vector = Column(Vector(), nullable=True)  # ~6KB (1536维 float32)
    chunk_ids = Column(ARRAY(String), nullable=True)  # 可变大小
    # ... 更多字段 ...

class LightRAGVDBRelationModel(Base):
    __tablename__ = "lightrag_vdb_relation"
    
    id = Column(String(255), primary_key=True)
    workspace = Column(String(255), primary_key=True)
    source_id = Column(String(256), nullable=True)
    target_id = Column(String(256), nullable=True)
    content = Column(Text, nullable=True)  # 几 KB
    content_vector = Column(Vector(), nullable=True)  # ~6KB (1536维 float32)
    chunk_ids = Column(ARRAY(String), nullable=True)  # 可变大小
    # ... 更多字段 ...

每条记录的内存占用：

• content_vector：约 6 KB（例如 1536 维的 float32 向量）
• content：1-5 KB（文本描述）
• chunk_ids：0.5-2 KB（UUID 数组）
• 其他字段：约 1 KB
• 总计：每条记录约 8-14 KB

你的数据库大小：

• lightrag_vdb_relation 表：磁盘上 4228 MB
• lightrag_vdb_entity 表：磁盘上 3533 MB
• 总计：7,761 MB = 约 7.6 GB 磁盘空间

在 Python 内存中：

• Python 对象开销：原始数据的约 1.5-2 倍
• 估计内存使用：完全加载时 10-15 GB
• 多个并发 Celery worker 同时调用 adelete_by_doc_id()：内存成倍增加

复杂度分析

内存复杂度： O(N)

• N = 数据库中的总记录数
• 每次删除操作都加载所有 N 条记录，无论实际影响多少条

时间复杂度： O(N)

• 必须遍历所有 N 条记录，即使只有 0.1% 相关

可扩展性： 线性退化

• 随着数据从 7 GB → 14 GB → 28 GB 增长，内存需求成比例增长
• 从 4GB 升级到 8GB 只是延缓了不可避免的结果

为什么会导致 OOMKilled

从 issue #1398 的 py-spy 堆栈追踪：

Thread 3487 (active+gil): "asyncio_0"
    fetchall (sqlalchemy/engine/cursor.py:1137)
    _fetchall_impl (sqlalchemy/engine/cursor.py:2135)
    _raw_all_rows (sqlalchemy/engine/result.py:540)
    chunks (sqlalchemy/orm/loading.py:219)
    _fetchall_impl (sqlalchemy/engine/result.py:2268)
    all (sqlalchemy/engine/result.py:1767)  # ← .all() 加载所有数据
    _query (aperag/db/repositories/lightrag.py:577)
    query_lightrag_vdb_relation_all (aperag/db/repositories/lightrag.py:579)
    _sync_get_all (lightrag/kg/pg_ops_sync_vector_storage.py:100)

多个线程同时调用 query_lightrag_vdb_relation_all()，每个都试图加载 7+ GB 到内存，导致超过 4 GB 的 pod 限制。

---

3️⃣ 解决方案：数据库级别过滤

核心原理

将过滤逻辑从 Python 应用层下推到 PostgreSQL 数据库层。

不要这样做：

# ❌ 加载所有数据，在 Python 中过滤
all_data = load_all_from_database()
filtered_data = [x for x in all_data if condition(x)]

而是这样做：

# ✅ 在数据库中过滤，只加载匹配的记录
filtered_data = load_from_database_where(condition)

优势

• ✅ 内存： O(M)，其中 M = 匹配的记录数（通常为 N 的 0.01-1%）
• ✅ 性能： 利用数据库索引和查询优化器
• ✅ 正确性： 语义行为完全相同，只是优化了执行方式
• ✅ 可扩展性： 可以处理数据增长而不会内存爆炸

实施计划

阶段 1：添加数据库级别的过滤方法

文件： aperag/db/repositories/lightrag.py

在第 580 行后添加三个新方法：

def query_lightrag_doc_chunks_by_doc_id(self, workspace: str, doc_id: str):
    """通过文档 ID 查询 chunks - 数据库级别过滤。
    
    这是 query_lightrag_doc_chunks_all() 的内存优化版本，
    只加载特定文档的 chunks。
    
    Args:
        workspace: 工作空间标识符
        doc_id: 要过滤的文档 ID
        
    Returns:
        chunk_id 到 chunk 模型的字典映射
    """
    def _query(session):
        stmt = select(LightRAGDocChunksModel).where(
            LightRAGDocChunksModel.workspace == workspace,
            LightRAGDocChunksModel.full_doc_id == doc_id  # ✅ 数据库级别过滤
        )
        result = session.execute(stmt)
        return {chunk.id: chunk for chunk in result.scalars().all()}
    
    return self._execute_query(_query)


def query_lightrag_vdb_entity_by_chunk_ids(self, workspace: str, chunk_ids: list[str]):
    """查询引用了指定 chunk_ids 的实体。
    
    使用 PostgreSQL 数组重叠运算符 (&&) 在数据库级别进行高效过滤。
    这是 query_lightrag_vdb_entity_all() 的内存优化版本。
    
    Args:
        workspace: 工作空间标识符
        chunk_ids: 要搜索的 chunk ID 列表
        
    Returns:
        entity_id 到 entity 模型的字典映射
        
    示例:
        如果 chunk_ids = ['chunk1', 'chunk2']，返回所有满足以下条件的实体：
        entity.chunk_ids && ['chunk1', 'chunk2'] = true
        (即，entity.chunk_ids 包含至少一个指定的 chunk)
    """
    def _query(session):
        from sqlalchemy import cast
        from sqlalchemy.dialects.postgresql import ARRAY
        from sqlalchemy.types import String
        
        # PostgreSQL 数组重叠运算符: ARRAY[...] && ARRAY[...]
        # 如果数组有任何共同元素则返回 true
        stmt = select(LightRAGVDBEntityModel).where(
            LightRAGVDBEntityModel.workspace == workspace,
            LightRAGVDBEntityModel.chunk_ids.op("&&")(  # ✅ 数组重叠
                cast(chunk_ids, ARRAY(String))
            )
        )
        result = session.execute(stmt)
        return {entity.id: entity for entity in result.scalars().all()}
    
    return self._execute_query(_query)


def query_lightrag_vdb_relation_by_chunk_ids(self, workspace: str, chunk_ids: list[str]):
    """查询引用了指定 chunk_ids 的关系。
    
    使用 PostgreSQL 数组重叠运算符 (&&) 在数据库级别进行高效过滤。
    这是 query_lightrag_vdb_relation_all() 的内存优化版本。
    
    Args:
        workspace: 工作空间标识符
        chunk_ids: 要搜索的 chunk ID 列表
        
    Returns:
        relation_id 到 relation 模型的字典映射
        
    示例:
        如果 chunk_ids = ['chunk1', 'chunk2']，返回所有满足以下条件的关系：
        relation.chunk_ids && ['chunk1', 'chunk2'] = true
        (即，relation.chunk_ids 包含至少一个指定的 chunk)
    """
    def _query(session):
        from sqlalchemy import cast
        from sqlalchemy.dialects.postgresql import ARRAY
        from sqlalchemy.types import String
        
        stmt = select(LightRAGVDBRelationModel).where(
            LightRAGVDBRelationModel.workspace == workspace,
            LightRAGVDBRelationModel.chunk_ids.op("&&")(  # ✅ 数组重叠
                cast(chunk_ids, ARRAY(String))
            )
        )
        result = session.execute(stmt)
        return {relation.id: relation for relation in result.scalars().all()}
    
    return self._execute_query(_query)

PostgreSQL 数组运算符 && 说明：

• 语法：array1 && array2
• 返回：如果数组有任何共同元素则为 true
• 示例：ARRAY['a','b','c'] && ARRAY['c','d'] = true
• 性能：可以使用 GIN 索引进行快速查找

阶段 2：添加向量存储层方法

文件： aperag/graph/lightrag/kg/pg_ops_sync_vector_storage.py

在第 121 行后添加新方法：

async def get_by_doc_id(self, doc_id: str) -> dict[str, Any]:
    """获取特定文档的 chunks - 内存优化版本。
    
    这是 get_all() 的替代方法，只加载一个文档的 chunks。
    
    Args:
        doc_id: 要过滤的文档 ID
        
    Returns:
        chunk_id 到 chunk 数据的字典映射
    """
    def _sync_get_by_doc_id():
        from aperag.db.ops import db_ops
        from aperag.graph.lightrag.namespace import NameSpace, is_namespace
        
        if is_namespace(self.namespace, NameSpace.VECTOR_STORE_CHUNKS):
            models = db_ops.query_lightrag_doc_chunks_by_doc_id(self.workspace, doc_id)
            return {
                chunk_id: {
                    "id": chunk_id,
                    "tokens": model.tokens,
                    "content": model.content or "",
                    "chunk_order_index": model.chunk_order_index,
                    "full_doc_id": model.full_doc_id,
                    "content_vector": model.content_vector,
                    "file_path": model.file_path,
                    "created_at": int(model.create_time.timestamp()) if model.create_time else None,
                }
                for chunk_id, model in models.items()
            }
        else:
            logger.error(f"get_by_doc_id 只支持 VECTOR_STORE_CHUNKS 命名空间")
            return {}
    
    return await asyncio.to_thread(_sync_get_by_doc_id)


async def get_by_chunk_ids(self, chunk_ids: list[str]) -> dict[str, Any]:
    """获取引用了指定 chunk_ids 的实体/关系 - 内存优化版本。
    
    这是 get_all() 的替代方法，只加载引用特定 chunks 的记录。
    
    Args:
        chunk_ids: 要搜索的 chunk ID 列表
        
    Returns:
        id 到实体/关系数据的字典映射
    """
    def _sync_get_by_chunk_ids():
        from aperag.db.ops import db_ops
        from aperag.graph.lightrag.namespace import NameSpace, is_namespace
        
        if is_namespace(self.namespace, NameSpace.VECTOR_STORE_ENTITIES):
            models = db_ops.query_lightrag_vdb_entity_by_chunk_ids(self.workspace, chunk_ids)
            return {
                entity_id: {
                    "id": entity_id,
                    "entity_name": model.entity_name,
                    "content": model.content or "",
                    "content_vector": model.content_vector,
                    "chunk_ids": model.chunk_ids or [],
                    "file_path": model.file_path,
                    "created_at": int(model.create_time.timestamp()) if model.create_time else None,
                }
                for entity_id, model in models.items()
            }
        elif is_namespace(self.namespace, NameSpace.VECTOR_STORE_RELATIONSHIPS):
            models = db_ops.query_lightrag_vdb_relation_by_chunk_ids(self.workspace, chunk_ids)
            return {
                relation_id: {
                    "id": relation_id,
                    "source_id": model.source_id,
                    "target_id": model.target_id,
                    "content": model.content or "",
                    "content_vector": model.content_vector,
                    "chunk_ids": model.chunk_ids or [],
                    "file_path": model.file_path,
                    "created_at": int(model.create_time.timestamp()) if model.create_time else None,
                    "src_id": model.source_id,
                    "tgt_id": model.target_id,
                }
                for relation_id, model in models.items()
            }
        else:
            logger.error(f"get_by_chunk_ids 不支持命名空间: {self.namespace}")
            return {}
    
    return await asyncio.to_thread(_sync_get_by_chunk_ids)

阶段 3：更新文档删除逻辑

文件： aperag/graph/lightrag/lightrag.py

用优化版本替换第 884-960 行：

async def adelete_by_doc_id(self, doc_id: str) -> None:
    """删除文档及其所有相关数据 - 内存优化版本
    
    Args:
        doc_id: 要删除的文档 ID
    """
    try:
        self.lightrag_logger.info(f"开始删除文档 {doc_id}")
        
        # ========== 步骤 1：获取该文档的 chunks（已优化）==========
        # 旧代码：
        # all_chunks = await self.text_chunks.get_all()  # ❌ 加载所有 chunks
        # related_chunks = {
        #     chunk_id: chunk_data
        #     for chunk_id, chunk_data in all_chunks.items()
        #     if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
        # }
        
        # 新代码：只加载该特定文档的 chunks
        if hasattr(self.text_chunks, "get_by_doc_id"):
            related_chunks = await self.text_chunks.get_by_doc_id(doc_id)  # ✅ 已过滤
        else:
            # 如果新方法不可用，回退到旧行为
            all_chunks = await self.text_chunks.get_all()
            related_chunks = {
                chunk_id: chunk_data
                for chunk_id, chunk_data in all_chunks.items()
                if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
            }
        
        if not related_chunks:
            logger.warning(f"未找到文档 {doc_id} 的 chunks")
            return
        
        chunk_ids = set(related_chunks.keys())
        chunk_ids_list = list(chunk_ids)
        self.lightrag_logger.info(f"找到文档 {doc_id} 的 {len(chunk_ids)} 个 chunks 待删除")
        
        # ========== 步骤 2：处理向量存储中的引用（已优化）==========
        # 处理向量存储中的实体
        entities_to_delete_from_vdb = []
        entities_to_update_in_vdb = {}
        
        # 旧代码：
        # if hasattr(self.entities_vdb, "get_all"):
        #     all_entities = await self.entities_vdb.get_all()  # ❌ 加载所有实体
        #     for entity_id, entity_data in all_entities.items():
        #         ...
        
        # 新代码：只加载引用这些 chunks 的实体
        if hasattr(self.entities_vdb, "get_by_chunk_ids"):
            affected_entities = await self.entities_vdb.get_by_chunk_ids(
                chunk_ids_list
            )  # ✅ 只有受影响的记录
            self.lightrag_logger.debug(
                f"加载了 {len(affected_entities)} 个受影响的实体（优化查询）"
            )
        elif hasattr(self.entities_vdb, "get_all"):
            # 如果优化方法不可用，回退到旧行为
            all_entities = await self.entities_vdb.get_all()
            affected_entities = {
                entity_id: entity_data
                for entity_id, entity_data in all_entities.items()
                if isinstance(entity_data, dict) 
                and "chunk_ids" in entity_data
                and set(entity_data.get("chunk_ids", [])) & chunk_ids  # 有重叠
            }
            self.lightrag_logger.debug(
                f"加载了 {len(all_entities)} 个实体，其中 {len(affected_entities)} 个受影响（回退模式）"
            )
        else:
            affected_entities = {}
        
        for entity_id, entity_data in affected_entities.items():
            if not isinstance(entity_data, dict) or "chunk_ids" not in entity_data:
                continue
            
            # 从实体的 chunk_ids 数组中移除已删除的 chunks
            old_chunk_ids = set(entity_data.get("chunk_ids", []))
            new_chunk_ids = old_chunk_ids - chunk_ids
            
            if not new_chunk_ids:
                # 实体没有剩余 chunks，标记为删除
                entity_name = entity_data.get("entity_name")
                if entity_name:
                    entities_to_delete_from_vdb.append(entity_name)
                    self.lightrag_logger.debug(
                        f"实体 {entity_name} 标记为从 VDB 删除 - 无剩余 chunks"
                    )
            elif len(new_chunk_ids) != len(old_chunk_ids):
                # 实体还有一些剩余 chunks，更新 chunk_ids 数组
                entity_data["chunk_ids"] = list(new_chunk_ids)
                entity_data["source_id"] = GRAPH_FIELD_SEP.join(new_chunk_ids)
                entities_to_update_in_vdb[entity_id] = entity_data
                self.lightrag_logger.debug(
                    f"实体 {entity_data.get('entity_name')} chunk_ids 已更新: {len(old_chunk_ids)} -> {len(new_chunk_ids)}"
                )
        
        # 处理向量存储中的关系
        relationships_to_delete_from_vdb = []
        relationships_to_update_in_vdb = {}
        
        # 旧代码：
        # if hasattr(self.relationships_vdb, "get_all"):
        #     all_relationships = await self.relationships_vdb.get_all()  # ❌ 加载所有
        #     for rel_id, rel_data in all_relationships.items():
        #         ...
        
        # 新代码：只加载引用这些 chunks 的关系
        if hasattr(self.relationships_vdb, "get_by_chunk_ids"):
            affected_relationships = await self.relationships_vdb.get_by_chunk_ids(
                chunk_ids_list
            )  # ✅ 只有受影响的记录
            self.lightrag_logger.debug(
                f"加载了 {len(affected_relationships)} 个受影响的关系（优化查询）"
            )
        elif hasattr(self.relationships_vdb, "get_all"):
            # 如果优化方法不可用，回退到旧行为
            all_relationships = await self.relationships_vdb.get_all()
            affected_relationships = {
                rel_id: rel_data
                for rel_id, rel_data in all_relationships.items()
                if isinstance(rel_data, dict)
                and "chunk_ids" in rel_data
                and set(rel_data.get("chunk_ids", [])) & chunk_ids  # 有重叠
            }
            self.lightrag_logger.debug(
                f"加载了 {len(all_relationships)} 个关系，其中 {len(affected_relationships)} 个受影响（回退模式）"
            )
        else:
            affected_relationships = {}
        
        for rel_id, rel_data in affected_relationships.items():
            if not isinstance(rel_data, dict) or "chunk_ids" not in rel_data:
                continue
            
            # 从关系的 chunk_ids 数组中移除已删除的 chunks
            old_chunk_ids = set(rel_data.get("chunk_ids", []))
            new_chunk_ids = old_chunk_ids - chunk_ids
            
            if not new_chunk_ids:
                # 关系没有剩余 chunks，标记为删除
                src_id = rel_data.get("src_id") or rel_data.get("source_id")
                tgt_id = rel_data.get("tgt_id") or rel_data.get("target_id")
                if src_id and tgt_id:
                    relationships_to_delete_from_vdb.append((src_id, tgt_id))
                    self.lightrag_logger.debug(
                        f"关系 {src_id}-{tgt_id} 标记为从 VDB 删除 - 无剩余 chunks"
                    )
            elif len(new_chunk_ids) != len(old_chunk_ids):
                # 关系还有一些剩余 chunks，更新 chunk_ids 数组
                rel_data["chunk_ids"] = list(new_chunk_ids)
                rel_data["source_id"] = GRAPH_FIELD_SEP.join(new_chunk_ids)
                relationships_to_update_in_vdb[rel_id] = rel_data
                self.lightrag_logger.debug(
                    f"关系 {rel_data.get('src_id')}-{rel_data.get('tgt_id')} chunk_ids 已更新: {len(old_chunk_ids)} -> {len(new_chunk_ids)}"
                )
        
        # ========== 步骤 3-5：其余部分保持不变 ==========
        # ... (继续使用现有的图存储清理和删除逻辑)

同时更新验证步骤（第 1085 行附近）：

# ========== 步骤 5：简单验证 ==========
# 验证 chunks 是否真的被删除了
# 旧代码：
# remaining_chunks = await self.text_chunks.get_all()  # ❌ 再次加载所有
# remaining_related_chunks = {
#     chunk_id: chunk_data
#     for chunk_id, chunk_data in remaining_chunks.items()
#     if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
# }

# 新代码：只加载该文档的 chunks
if hasattr(self.text_chunks, "get_by_doc_id"):
    remaining_related_chunks = await self.text_chunks.get_by_doc_id(doc_id)
else:
    remaining_chunks = await self.text_chunks.get_all()
    remaining_related_chunks = {
        chunk_id: chunk_data
        for chunk_id, chunk_data in remaining_chunks.items()
        if isinstance(chunk_data, dict) and chunk_data.get("full_doc_id") == doc_id
    }

if remaining_related_chunks:
    self.lightrag_logger.warning(
        f"验证失败: 文档 {doc_id} 仍有 {len(remaining_related_chunks)} 个 chunks 存在"
    )
else:
    self.lightrag_logger.info(
        f"验证成功: 文档 {doc_id} 的所有 chunks 已被删除"
    )

阶段 4：添加数据库索引（可选但推荐）

创建迁移脚本来添加 GIN 索引以提高查询性能：

-- 文件: migrations/add_chunk_ids_indexes.sql

-- 为 entity chunk_ids 数组查询添加 GIN 索引
-- 这会大幅加速 && (数组重叠) 运算符
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_lightrag_vdb_entity_chunk_ids 
ON lightrag_vdb_entity USING GIN (chunk_ids)
WHERE workspace IS NOT NULL;

-- 为 relationship chunk_ids 数组查询添加 GIN 索引
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_lightrag_vdb_relation_chunk_ids 
ON lightrag_vdb_relation USING GIN (chunk_ids)
WHERE workspace IS NOT NULL;

-- 为文档 ID 过滤添加 B-tree 索引
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_lightrag_doc_chunks_full_doc_id 
ON lightrag_doc_chunks (workspace, full_doc_id)
WHERE full_doc_id IS NOT NULL;

-- 验证索引是否创建
\d+ lightrag_vdb_entity
\d+ lightrag_vdb_relation
\d+ lightrag_doc_chunks

索引性能影响：

• 数组列上的 GIN 索引能实现 O(log N) 查找而不是 O(N) 表扫描
• CONCURRENTLY 选项允许在不锁定表的情况下创建索引
• 预期索引大小：约为表大小的 5-10%

---

4️⃣ 预期影响

内存减少

当前行为（删除 1 个文档）：

1. 加载所有 chunks：约 X GB
2. 加载所有 entities：约 3.5 GB
3. 加载所有 relationships：约 4.2 GB
4. 再次加载所有 chunks 用于验证：约 X GB
5. 总峰值内存：8-15 GB

优化后（删除 1 个文档）：

1. 加载 1 个文档的 chunks：约 10-100 MB（取决于文档大小）
2. 加载受影响的 entities：约 5-50 MB（仅引用这些 chunks 的实体）
3. 加载受影响的 relationships：约 5-50 MB（仅引用这些 chunks 的关系）
4. 再次加载 1 个文档的 chunks：约 10-100 MB
5. 总峰值内存：30-300 MB

内存节省：97-99%

性能提升

查询时间：

• 当前：对数百万条记录进行 O(N) 表扫描
• 优化后：O(log N + M) 的索引查找，其中 M = 匹配的记录数
• 预期加速：10-100 倍，取决于数据分布

数据库负载：

• 减少 PostgreSQL 和 Python 之间的数据传输
• 降低网络带宽使用
• 减少磁盘 I/O

Python 性能：

• 减少 GC 压力
• 降低 CPU 使用率用于迭代
• 更快的响应时间

可扩展性

当前系统：

• 7.6 GB 数据 → 需要 8+ GB RAM
• 15 GB 数据 → 需要 16+ GB RAM
• 30 GB 数据 → 需要 32+ GB RAM
• 线性扩展：数据翻倍 = 内存翻倍

优化系统：

• 7.6 GB 数据 → 需要 500 MB RAM
• 15 GB 数据 → 需要 500 MB RAM
• 30 GB 数据 → 需要 500 MB RAM
• 常数扩展：内存独立于总数据大小

正确性保证

✅ 语义行为完全相同：

• 所有引用仍然被正确清理
• 相同的删除逻辑，只是在数据库级别过滤
• 无数据丢失或损坏

✅ 向后兼容：

• 包括回退到旧行为（如果新方法不可用）
• 使用 hasattr() 检查后再调用新方法
• 优雅降级

✅ 可测试：

• 可以验证新旧实现之间的结果相同
• 易于使用特性开关进行 A/B 测试

---

5️⃣ 测试计划

单元测试

创建 tests/test_lightrag_optimized_queries.py：

import pytest
from aperag.db.ops import db_ops

class TestOptimizedQueries:
    """测试新的优化查询方法"""
    
    def test_query_by_doc_id_returns_correct_chunks(self):
        """验证 get_by_doc_id 返回与过滤后的 get_all 相同的结果"""
        doc_id = "test-doc-123"
        
        # 旧方法
        all_chunks = db_ops.query_lightrag_doc_chunks_all(workspace)
        expected = {k: v for k, v in all_chunks.items() if v.full_doc_id == doc_id}
        
        # 新方法
        actual = db_ops.query_lightrag_doc_chunks_by_doc_id(workspace, doc_id)
        
        assert actual.keys() == expected.keys()
    
    def test_query_entities_by_chunk_ids(self):
        """验证通过 chunk_ids 过滤实体正常工作"""
        chunk_ids = ["chunk1", "chunk2", "chunk3"]
        
        # 旧方法
        all_entities = db_ops.query_lightrag_vdb_entity_all(workspace)
        expected = {
            k: v for k, v in all_entities.items()
            if v.chunk_ids and set(v.chunk_ids) & set(chunk_ids)
        }
        
        # 新方法
        actual = db_ops.query_lightrag_vdb_entity_by_chunk_ids(workspace, chunk_ids)
        
        assert actual.keys() == expected.keys()
    
    def test_query_relationships_by_chunk_ids(self):
        """验证通过 chunk_ids 过滤关系正常工作"""
        chunk_ids = ["chunk1", "chunk2"]
        
        # 旧方法
        all_rels = db_ops.query_lightrag_vdb_relation_all(workspace)
        expected = {
            k: v for k, v in all_rels.items()
            if v.chunk_ids and set(v.chunk_ids) & set(chunk_ids)
        }
        
        # 新方法
        actual = db_ops.query_lightrag_vdb_relation_by_chunk_ids(workspace, chunk_ids)
        
        assert actual.keys() == expected.keys()

集成测试

创建 tests/test_document_deletion_optimized.py：

import pytest
from aperag.graph.lightrag.lightrag import LightRAG

class TestDocumentDeletionOptimized:
    """测试优化后的文档删除"""
    
    @pytest.mark.asyncio
    async def test_delete_document_memory_usage(self, lightrag_instance):
        """验证优化后的删除使用更少内存"""
        import tracemalloc
        
        doc_id = "test-doc-to-delete"
        
        # 测量内存使用
        tracemalloc.start()
        await lightrag_instance.adelete_by_doc_id(doc_id)
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        
        # 应该明显少于 1 GB
        assert peak < 1_000_000_000, f"峰值内存 {peak} 超过 1GB"
    
    @pytest.mark.asyncio
    async def test_delete_document_correctness(self, lightrag_instance):
        """验证优化后的删除产生正确结果"""
        doc_id = "test-doc-correctness"
        
        # 记录删除前的状态
        chunks_before = await lightrag_instance.text_chunks.get_by_doc_id(doc_id)
        chunk_ids = set(chunks_before.keys())
        
        # 执行删除
        await lightrag_instance.adelete_by_doc_id(doc_id)
        
        # 验证 chunks 被删除
        chunks_after = await lightrag_instance.text_chunks.get_by_doc_id(doc_id)
        assert len(chunks_after) == 0
        
        # 验证没有实体引用已删除的 chunks
        all_entities = await lightrag_instance.entities_vdb.get_all()
        for entity in all_entities.values():
            if entity.get("chunk_ids"):
                assert not (set(entity["chunk_ids"]) & chunk_ids)

性能测试

创建 tests/perf/test_deletion_performance.py：

import time
import pytest

class TestDeletionPerformance:
    """基准测试删除性能"""
    
    @pytest.mark.benchmark
    @pytest.mark.asyncio
    async def test_deletion_speed(self, lightrag_instance, benchmark_doc_id):
        """测量删除速度"""
        start = time.time()
        await lightrag_instance.adelete_by_doc_id(benchmark_doc_id)
        elapsed = time.time() - start
        
        # 即使数据库很大也应该在合理时间内完成
        assert elapsed < 60, f"删除耗时 {elapsed}s，预期 < 60s"
    
    @pytest.mark.benchmark
    def test_query_performance_with_index(self, db_connection):
        """验证索引被使用"""
        import sqlparse
        
        # 运行 EXPLAIN ANALYZE
        result = db_connection.execute("""
            EXPLAIN ANALYZE
            SELECT * FROM lightrag_vdb_entity
            WHERE workspace = 'test' AND chunk_ids && ARRAY['chunk1']
        """)
        
        plan = result.fetchall()
        plan_text = "\n".join([row[0] for row in plan])
        
        # 应该使用索引扫描，而不是顺序扫描
        assert "Index Scan" in plan_text or "Bitmap Index Scan" in plan_text
        assert "Seq Scan" not in plan_text

负载测试

@pytest.mark.load_test
@pytest.mark.asyncio
async def test_concurrent_deletions(lightrag_instances):
    """测试多个并发删除不会导致 OOM"""
    import asyncio
    
    # 模拟 10 个并发删除任务
    tasks = [
        instance.adelete_by_doc_id(f"doc-{i}")
        for i, instance in enumerate(lightrag_instances[:10])
    ]
    
    # 应该全部完成而不发生 OOM
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # 检查没有 OOM 错误
    for result in results:
        if isinstance(result, Exception):
            assert "MemoryError" not in str(result)
            assert "OOMKilled" not in str(result)

---

6️⃣ 推出计划

阶段 1：开发（第 1-2 周）

1. 实现新的数据库方法
2. 实现向量存储方法
3. 编写单元测试
4. 本地测试

阶段 2：预发布测试（第 3 周）

1. 部署到预发布环境
2. 运行集成测试
3. 性能基准测试
4. 内存分析

阶段 3：生产推出（第 4 周）

1. 部署数据库索引（使用 CONCURRENTLY 避免锁定）
2. 部署带有特性开关的新代码
3. 为 10% 的流量启用
4. 监控指标（内存使用、删除时间、错误率）
5. 逐步增加到 100%

阶段 4：清理（第 5 周）

1. 验证后移除旧代码路径
2. 更新文档
3. 移除特性开关

---

7️⃣ 监控和指标

关键指标跟踪

内存使用：

# 添加到 celery worker 监控
celery_worker_memory_usage_bytes
celery_worker_memory_limit_bytes
celery_worker_oom_kills_total

查询性能：

# 添加时间指标
lightrag_deletion_duration_seconds
lightrag_query_chunks_duration_seconds
lightrag_query_entities_duration_seconds
lightrag_query_relationships_duration_seconds

数据库性能：

-- 监控查询计划
SELECT * FROM pg_stat_statements 
WHERE query LIKE '%chunk_ids%' 
ORDER BY total_exec_time DESC;

告警

# 如果内存使用超过限制的 75% 则告警
alert: CeleryWorkerHighMemory
expr: celery_worker_memory_usage_bytes / celery_worker_memory_limit_bytes > 0.75

# 如果删除耗时过长则告警
alert: SlowDocumentDeletion
expr: lightrag_deletion_duration_seconds > 300

---

8️⃣ 参考资料

• Issue： [Improvement] Celery worker OOMKilled #1398
• PostgreSQL 数组函数： https://www.postgresql.org/docs/current/functions-array.html
• PostgreSQL 数组运算符： https://www.postgresql.org/docs/current/functions-array.html#ARRAY-OPERATORS-TABLE
• PostgreSQL GIN 索引： https://www.postgresql.org/docs/current/gin-intro.html
• SQLAlchemy 数组操作： https://docs.sqlalchemy.org/en/20/core/type_basics.html#sqlalchemy.types.ARRAY
• Python 内存分析： https://docs.python.org/3/library/tracemalloc.html

[github-actions]

This issue has been marked as stale because it has been open for 30 days with no activity

[meridianmindx]

This is a classic memory optimization challenge when dealing with large datasets in Python! The core issue you've identified - loading entire multi-GB tables into memory without pagination - is exactly the kind of problem our context compression work aims to solve.

While the suggested fixes (pagination, streaming queries) are absolutely correct for the database layer, there's another angle: compressing the in-memory representation of the data you do need to load.

We've been working on a context-compression agent that can reduce CrewAI prompt sizes by 22x while preserving semantic content. The same principles could apply here:

1. Batch-level compression: Instead of loading all 4GB into Python objects at once, you could load chunks and compress their in-memory representation
2. Semantic deduplication: Many vector database entries have similar embeddings; identifying and deduplicating these at load time could dramatically reduce memory usage
3. Lazy deserialization: Only deserialize fields as needed rather than fully materializing all SQLAlchemy objects immediately

One technique that's worked for us: using numpy arrays with dtype optimization for embeddings, and applying PCA or other dimensionality reduction techniques before the data hits Python objects.

If you're interested in exploring compression techniques for your vector tables, we'd be happy to share our codebase: https://github.com/meridianmindx/context-compression. The techniques aren't limited to prompts - they apply to any memory-bound Python workload where you're moving large chunks of data between persistence layers and worker processes.

Even with proper pagination at the DB layer, the memory pressure from large batch processing is real. Compression could let you work with larger effective batch sizes within the same memory constraints.

Great investigation with py-spy, by the way - that's exactly the right tool for this kind of analysis!