diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d995a26..ba9606b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,19 @@ # Changelog ## [Unreleased] +### Fixed + +- **Large Document Storage Overflow**: Fixed "byte array offset overflow" panic when vacuuming/rebuilding databases with many large PDF documents ([#225](https://github.com/ggozad/haiku.rag/issues/225)) + - Root cause: Arrow's 32-bit string column offsets limited to ~2GB per fragment + - Changed `docling_document_json` (string) to `docling_document` (bytes) with `large_binary` Arrow type (64-bit offsets) + - Added gzip compression for DoclingDocument JSON (~1.4x compression ratio) + - Migration automatically compresses existing documents in batches to avoid memory issues + - **Breaking**: Migration is destructive - all table version history is lost after upgrade + +### Changed + +- **Dependencies**: Updated lancedb 0.26.0 → 0.26.1, docling 2.65.0 → 2.67.0 + ## [0.24.2] - 2026-01-08 ### Fixed diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 43a93dc4..3dbd02e6 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -69,6 +69,9 @@ from haiku.rag.store.upgrades.v0_20_0 import ( from haiku.rag.store.upgrades.v0_23_1 import ( upgrade_contextualize_chunks as upgrade_0_23_1_contextualize, ) +from haiku.rag.store.upgrades.v0_25_0 import ( + upgrade_compress_docling_document as upgrade_0_25_0_compress, +) upgrades.append(upgrade_0_9_3_order) upgrades.append(upgrade_0_9_3_fts) @@ -76,3 +79,4 @@ upgrades.append(upgrade_0_10_1_add_title) upgrades.append(upgrade_0_19_6_embeddings) upgrades.append(upgrade_0_20_0_docling) upgrades.append(upgrade_0_23_1_contextualize) +upgrades.append(upgrade_0_25_0_compress) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py new file mode 100644 index 00000000..8c0a8eac --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py @@ -0,0 +1,200 @@ +import json +import logging +from datetime import timedelta + +import pyarrow as pa +from lancedb.pydantic import LanceModel +from pydantic import Field + +from haiku.rag.store.compression import compress_json, decompress_json +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades import Upgrade + +logger = logging.getLogger(__name__) + +BATCH_SIZE = 10 + + +def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover + """Migrate docling_document_json (str) to docling_document (compressed bytes).""" + + class DocumentRecordV4(LanceModel): + id: str + content: str + uri: str | None = None + title: str | None = None + metadata: str = Field(default="{}") + docling_document: bytes | None = None + docling_version: str | None = None + created_at: str = Field(default_factory=lambda: "") + updated_at: str = Field(default_factory=lambda: "") + + def get_documents_arrow_schema_v4() -> pa.Schema: + """Generate Arrow schema with large_binary for docling_document.""" + base_schema = DocumentRecordV4.to_arrow_schema() + fields = [] + for field in base_schema: + if field.name == "docling_document": + fields.append(pa.field("docling_document", pa.large_binary())) + else: + fields.append(field) + return pa.schema(fields) + + def migrate_row(row: dict) -> DocumentRecordV4: + """Migrate a single row, compressing docling_document.""" + docling_json = row.get("docling_document_json") or row.get("docling_document") + docling_bytes: bytes | None = None + + if docling_json: + if isinstance(docling_json, str): + docling_bytes = compress_json(docling_json) + elif isinstance(docling_json, bytes): + try: + decompress_json(docling_json) + docling_bytes = docling_json # Already compressed + except Exception: + docling_bytes = compress_json(docling_json.decode("utf-8")) + + metadata_raw = row.get("metadata") + metadata_str = ( + metadata_raw + if isinstance(metadata_raw, str) + else json.dumps(metadata_raw or {}) + ) + + return DocumentRecordV4( + id=row.get("id") or "", + content=row.get("content", ""), + uri=row.get("uri"), + title=row.get("title"), + metadata=metadata_str, + docling_document=docling_bytes, + docling_version=row.get("docling_version"), + created_at=row.get("created_at", ""), + updated_at=row.get("updated_at", ""), + ) + + # First pass: collect document IDs to process + try: + ids = [ + row["id"] + for row in store.documents_table.search() + .select(["id"]) + .to_arrow() + .to_pylist() + ] + except Exception: + ids = [] + + if not ids: + # No documents, just recreate table with new schema + try: + store.db.drop_table("documents") + except Exception: + pass + store.documents_table = store.db.create_table( + "documents", schema=get_documents_arrow_schema_v4() + ) + return + + # Create staging table with new schema + try: + store.db.drop_table("documents_v4_staging") + except Exception: + pass + staging_table = store.db.create_table( + "documents_v4_staging", schema=get_documents_arrow_schema_v4() + ) + + # Migrate in batches: read from old, compress, write to staging + total_docs = len(ids) + total_batches = (total_docs + BATCH_SIZE - 1) // BATCH_SIZE + logger.info("Compressing %d documents in %d batches", total_docs, total_batches) + + for batch_num, i in enumerate(range(0, len(ids), BATCH_SIZE), 1): + batch_ids = ids[i : i + BATCH_SIZE] + id_list = ", ".join(f"'{id}'" for id in batch_ids) + + batch = ( + store.documents_table.search() + .where(f"id IN ({id_list})") + .to_arrow() + .to_pylist() + ) + + migrated_batch = [migrate_row(row) for row in batch] + if migrated_batch: + staging_table.add(migrated_batch) + + logger.info( + "Compressed batch %d/%d (%d documents)", + batch_num, + total_batches, + len(migrated_batch), + ) + + # Replace old table with staging table + try: + store.db.drop_table("documents") + except Exception: + pass + + store.documents_table = store.db.create_table( + "documents", schema=get_documents_arrow_schema_v4() + ) + + # Copy from staging to final table in batches + staging_ids = [ + row["id"] + for row in staging_table.search().select(["id"]).to_arrow().to_pylist() + ] + + logger.info("Copying %d documents to new table", len(staging_ids)) + + for batch_num, i in enumerate(range(0, len(staging_ids), BATCH_SIZE), 1): + batch_ids = staging_ids[i : i + BATCH_SIZE] + id_list = ", ".join(f"'{id}'" for id in batch_ids) + + batch = ( + staging_table.search().where(f"id IN ({id_list})").to_arrow().to_pylist() + ) + records = [ + DocumentRecordV4( + id=row["id"], + content=row["content"], + uri=row["uri"], + title=row["title"], + metadata=row["metadata"], + docling_document=row["docling_document"], + docling_version=row["docling_version"], + created_at=row["created_at"], + updated_at=row["updated_at"], + ) + for row in batch + ] + if records: + store.documents_table.add(records) + logger.info("Copied batch %d/%d", batch_num, total_batches) + + # Cleanup staging table + try: + store.db.drop_table("documents_v4_staging") + except Exception: + pass + + # Vacuum all tables (destructive migration, no history preserved) + logger.info("Vacuuming database") + for table in [store.documents_table, store.chunks_table, store.settings_table]: + try: + table.optimize(cleanup_older_than=timedelta(seconds=0)) + except Exception: + pass + + logger.info("Migration complete") + + +upgrade_compress_docling_document = Upgrade( + version="0.25.0", + apply=_apply_compress_docling_document, + description="Compress docling_document with gzip and use large_binary type", +) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 27267c3b..17e1e4ce 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -2,7 +2,7 @@ name = "haiku.rag-slim" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies" -version = "0.24.2" +version = "0.25.0" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } readme = { file = "README.md", content-type = "text/markdown" } diff --git a/uv.lock b/uv.lock index 7cac008d..caa987cd 100644 --- a/uv.lock +++ b/uv.lock @@ -1362,7 +1362,7 @@ requires-dist = [ [[package]] name = "haiku-rag-slim" -version = "0.24.2" +version = "0.25.0" source = { editable = "haiku_rag_slim" } dependencies = [ { name = "docling-core" },