highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-05-12 12:53:01 +02:00
Code Issues Packages Projects Releases Wiki Activity

Speedup orphan doc cleanup script (#3596)

Browse Source 
* Speedup orphan doc cleanup script

* Fix mypy
This commit is contained in:
Chris Weaver 2025-01-05 06:28:25 -08:00 committed by GitHub
parent 2fc58252f4
commit f895e5f7d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 88 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

127
backend/scripts/orphan_doc_cleanup_script.py
View File

@ -1,3 +1,4 @@
import concurrent.futures
import os
import sys
@ -8,16 +9,17 @@ from sqlalchemy.orm import Session
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(parent_dir)
from onyx.context.search.models import IndexFilters # noqa: E402
from onyx.document_index.interfaces import VespaChunkRequest # noqa: E402
from onyx.db.engine import get_session_context_manager # noqa: E402
from onyx.db.document import delete_documents_complete__no_commit # noqa: E402
from onyx.db.search_settings import get_current_search_settings # noqa: E402
from onyx.document_index.vespa.index import VespaIndex # noqa: E402
from onyx.background.celery.tasks.shared.RetryDocumentIndex import ( # noqa: E402
RetryDocumentIndex,
)
BATCH_SIZE = 100
def _get_orphaned_document_ids(db_session: Session) -> list[str]:
def _get_orphaned_document_ids(db_session: Session, limit: int) -> list[str]:
"""Get document IDs that don't have any entries in document_by_connector_credential_pair"""
query = text(
"""
@ -25,54 +27,101 @@ def _get_orphaned_document_ids(db_session: Session) -> list[str]:
FROM document d
LEFT JOIN document_by_connector_credential_pair dbcc ON d.id = dbcc.id
WHERE dbcc.id IS NULL
LIMIT :limit
"""
)
orphaned_ids = [doc_id[0] for doc_id in db_session.execute(query)]
print(f"Found {len(orphaned_ids)} orphaned documents")
orphaned_ids = [doc_id[0] for doc_id in db_session.execute(query, {"limit": limit})]
print(f"Found {len(orphaned_ids)} orphaned documents in this batch")
return orphaned_ids
def main() -> None:
with get_session_context_manager() as db_session:
# Get orphaned document IDs
orphaned_ids = _get_orphaned_document_ids(db_session)
if not orphaned_ids:
print("No orphaned documents found")
return
total_processed = 0
while True:
# Get orphaned document IDs in batches
orphaned_ids = _get_orphaned_document_ids(db_session, BATCH_SIZE)
if not orphaned_ids:
if total_processed == 0:
print("No orphaned documents found")
else:
print(
f"Finished processing all batches. Total documents "
f"processed: {total_processed}"
)
return
# Setup Vespa index
search_settings = get_current_search_settings(db_session)
index_name = search_settings.index_name
vespa_index = VespaIndex(index_name=index_name, secondary_index_name=None)
retry_index = RetryDocumentIndex(vespa_index)
# Setup Vespa index
search_settings = get_current_search_settings(db_session)
index_name = search_settings.index_name
vespa_index = VespaIndex(index_name=index_name, secondary_index_name=None)
# Delete chunks from Vespa first
print("Deleting orphaned document chunks from Vespa")
successfully_vespa_deleted_doc_ids = []
for doc_id in orphaned_ids:
# Delete chunks from Vespa first
print("Deleting orphaned document chunks from Vespa")
successfully_vespa_deleted_doc_ids: list[str] = []
# Process documents in parallel using ThreadPoolExecutor
with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
def process_doc(doc_id: str) -> str | None:
# Check if document exists in Vespa first
try:
chunks = vespa_index.id_based_retrieval(
chunk_requests=[
VespaChunkRequest(document_id=doc_id, max_chunk_ind=2)
],
filters=IndexFilters(access_control_list=None),
batch_retrieval=True,
)
if not chunks:
print(f"Document {doc_id} not found in Vespa")
return doc_id
except Exception as e:
print(
f"Error checking if document {doc_id} exists in Vespa: {e}"
)
return None
try:
print(f"Deleting document {doc_id} in Vespa")
chunks_deleted = vespa_index.delete_single(doc_id)
if chunks_deleted > 0:
print(
f"Deleted {chunks_deleted} chunks for document {doc_id}"
)
return doc_id
except Exception as e:
print(
f"Error deleting document {doc_id} in Vespa and "
f"will not delete from Postgres: {e}"
)
return None
# Submit all tasks and gather results
futures = [
executor.submit(process_doc, doc_id) for doc_id in orphaned_ids
]
for future in concurrent.futures.as_completed(futures):
doc_id = future.result()
if doc_id:
successfully_vespa_deleted_doc_ids.append(doc_id)
# Delete documents from Postgres
print("Deleting orphaned documents from Postgres")
try:
chunks_deleted = retry_index.delete_single(doc_id)
successfully_vespa_deleted_doc_ids.append(doc_id)
if chunks_deleted > 0:
print(f"Deleted {chunks_deleted} chunks for document {doc_id}")
except Exception as e:
print(
f"Error deleting document {doc_id} in Vespa and will not delete from Postgres: {e}"
delete_documents_complete__no_commit(
db_session, successfully_vespa_deleted_doc_ids
)
db_session.commit()
except Exception as e:
print(f"Error deleting documents from Postgres: {e}")
break
# Delete documents from Postgres
print("Deleting orphaned documents from Postgres")
try:
delete_documents_complete__no_commit(
db_session, successfully_vespa_deleted_doc_ids
total_processed += len(successfully_vespa_deleted_doc_ids)
print(
f"Successfully cleaned up {len(successfully_vespa_deleted_doc_ids)}"
f" orphaned documents in this batch"
)
db_session.commit()
except Exception as e:
print(f"Error deleting documents from Postgres: {e}")
print(
f"Successfully cleaned up {len(successfully_vespa_deleted_doc_ids)} orphaned documents"
)
print(f"Total documents processed so far: {total_processed}")
if __name__ == "__main__":