Compare commits
1 commit
main
...
chore/opti
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
581ddf4f41 |
9 changed files with 601 additions and 25 deletions
3
.gitignore
vendored
3
.gitignore
vendored
|
|
@ -41,3 +41,6 @@ site/
|
|||
.claude/*
|
||||
!.claude/skills/
|
||||
!.claude/skills/**
|
||||
|
||||
# Vacuum investigation scratch: kept on disk, never committed
|
||||
vacuum-findings/
|
||||
|
|
|
|||
|
|
@ -2,6 +2,14 @@
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- `storage.compaction_target_bytes` (default 2 GiB): target size for the fragments compaction writes on the tables holding docling blobs.
|
||||
|
||||
### Changed
|
||||
|
||||
- `Store.vacuum` compacts `documents` and `document_items` through `lance` with a fragment target derived from `compaction_target_bytes`, instead of `AsyncTable.optimize`. The other tables are unchanged. New dependency: `pylance`.
|
||||
|
||||
## [0.82.1] - 2026-09-03
|
||||
|
||||
### Fixed
|
||||
|
|
|
|||
|
|
@ -37,36 +37,44 @@ storage:
|
|||
data_dir: /path/to/data # Empty = use default platform location
|
||||
auto_vacuum: true # Enable automatic vacuuming after operations
|
||||
vacuum_retention_seconds: 86400 # Cleanup threshold in seconds
|
||||
compaction_target_bytes: 2147483648 # Target size for a compacted fragment
|
||||
```
|
||||
|
||||
- **data_dir**: Directory for local database storage. When empty, uses platform-specific default locations
|
||||
- **auto_vacuum**: When enabled (default), automatically runs vacuum after document create/update/delete operations and database rebuilds. Background vacuums are throttled to at most one every 5 minutes, so sustained ingestion does not trigger continuous compaction, and a final vacuum runs when the client closes. Set to `false` to disable automatic vacuuming and rely on manual `haiku-rag vacuum` commands only. Disabling can help avoid potential crashes in high-concurrency scenarios
|
||||
- **vacuum_retention_seconds**: When vacuum runs, old table versions older than this threshold are removed. Default: 86400 seconds (1 day). Set to 0 for aggressive cleanup (removes all old versions immediately)
|
||||
- **compaction_target_bytes**: Target size for the fragments compaction writes on the tables that store docling blobs. Default: 2 GiB. Advisory rather than a cap, see [Vacuum Memory](#vacuum-memory) below
|
||||
|
||||
!!! warning "Vacuum Retention Threshold"
|
||||
The `vacuum_retention_seconds` value should be larger than the typical time it takes to process and write a document. If a concurrent operation is in progress while vacuum runs, setting this value too low can cause race conditions where vacuum removes table versions that an in-flight operation still needs. The default of 86400 seconds (1 day) is conservative and safe for most use cases.
|
||||
|
||||
### Vacuum Memory Requirements
|
||||
### Vacuum Memory
|
||||
|
||||
Vacuum compacts small data files into larger ones. LanceDB targets roughly one million rows per fragment, which a `documents` table holding multi-megabyte docling blobs never reaches, so each vacuum that follows new documents re-merges the whole existing fragment rather than only the new ones. Peak memory therefore scales with the total size of the `documents` table, not with how much was added.
|
||||
Vacuum compacts small data files into larger ones. LanceDB targets roughly one million rows per fragment, which a `documents` table holding multi-megabyte docling blobs never reaches, so an unsized pass re-merges the whole existing fragment rather than only the new rows, and peak memory scales with the table rather than with what was added.
|
||||
|
||||
Measured peak resident memory is about 5x the size of the `documents` table's data files. An 8.8 GB table peaked at 48.7 GB. Plan for **6x the size of `documents/` on disk** as available RAM, or the vacuum will be killed by the OOM killer partway through.
|
||||
The `documents` and `document_items` tables are therefore compacted with an explicit fragment target derived from `compaction_target_bytes`. The target is sized from the widest fragment's bytes per row, taken from table metadata without reading any payload, so a handful of very large documents shrinks it for the whole table. The remaining tables use LanceDB's own optimize, which is already bounded for them because their rows are small.
|
||||
|
||||
Check the current size with:
|
||||
!!! warning "The target is not a memory cap"
|
||||
`compaction_target_bytes` sizes the fragments compaction **writes**. It cannot shrink a fragment that is already larger, and LanceDB rewrites such a fragment in one piece, costing roughly its own size no matter how low the target is set.
|
||||
|
||||
```bash
|
||||
du -sh /path/to/database.lancedb/documents.lance
|
||||
```
|
||||
Oversized fragments come from two places: a single large ingest batch, since each write becomes one fragment, and any database vacuumed before this release. In both cases the cost is paid once per fragment, the first time deletions within it pass LanceDB's 10% threshold, after which it is split to the target and stays there.
|
||||
|
||||
If that number times six exceeds available RAM, use one of:
|
||||
So the practical peak is the larger of `compaction_target_bytes` and your biggest existing fragment. To lower the first, reduce it; to lower the second, ingest in smaller batches.
|
||||
|
||||
A row larger than the target cannot be sized at all: the target floors at one row and the pass costs roughly that row's size. A 300-page PDF at `images_scale: 2.0` produces a row of around 445 MB. To reduce row size rather than raise the target:
|
||||
|
||||
- Reduce `images_scale` (see [Image Settings](processing.md#image-settings)). Rendered page rasters dominate the size of `documents`, and their byte cost falls with the square of the scale factor.
|
||||
- Set `generate_page_images: false` if visual grounding through `visualize_chunk()` is not needed. This removes page rasters entirely.
|
||||
- Set `auto_vacuum: false` and run `haiku-rag vacuum` manually when the machine is otherwise idle, so the peak does not land alongside ingestion.
|
||||
|
||||
Vacuum also folds new rows into the full-text index. Search stays correct without it but scans the uncovered rows on every query. `haiku-rag doctor` reports the coverage.
|
||||
|
||||
This is an upstream limitation rather than a `haiku.rag` setting. Compaction bounds itself by row count instead of bytes, and LanceDB's async API exposes no batch size or fragment target to override it. Tracked at [lancedb/lancedb#2325](https://github.com/lancedb/lancedb/issues/2325). The requirement above will drop once compaction batches by bytes.
|
||||
If fragment sizes are missing from the table metadata, which can happen for databases written by much older versions, compaction is skipped for that table and a warning is logged. Old versions are still pruned.
|
||||
|
||||
#### The first vacuum after upgrading
|
||||
|
||||
A database written before this release may contain one large fragment built by unsized compaction. It is left alone until deletions within it pass the 10% threshold, so early vacuums are cheap but its superseded payload is not yet reclaimed and disk usage can sit above the live data size. The pass that crosses the threshold rewrites it once, splitting it to the target, after which the table stays at the target and disk returns to normal.
|
||||
|
||||
Compaction options are not exposed by LanceDB's async API ([lancedb/lancedb#2325](https://github.com/lancedb/lancedb/issues/2325)), so these tables are compacted through `lance` directly.
|
||||
|
||||
### Placing the Database
|
||||
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ class StorageConfig(ConfigModel):
|
|||
data_dir: Path = Field(default_factory=get_default_data_dir)
|
||||
auto_vacuum: bool = True
|
||||
vacuum_retention_seconds: int = Field(default=86400, ge=0)
|
||||
compaction_target_bytes: int = Field(default=2 * 1024**3, gt=0)
|
||||
|
||||
@field_validator("data_dir", mode="before")
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -1,15 +1,17 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Coroutine
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from enum import Enum
|
||||
from importlib import metadata
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from time import monotonic
|
||||
from typing import Any, cast
|
||||
|
||||
import lance
|
||||
import lancedb
|
||||
from lancedb.index import IvfPq
|
||||
from packaging.version import parse
|
||||
|
|
@ -26,6 +28,7 @@ from haiku.rag.store.schema import (
|
|||
ensure_indexes,
|
||||
get_document_items_arrow_schema,
|
||||
get_documents_arrow_schema,
|
||||
has_payload_columns,
|
||||
query_to_pydantic,
|
||||
)
|
||||
|
||||
|
|
@ -118,6 +121,37 @@ def _stored_embedding(
|
|||
# version; guards against timestamp precision at the boundary.
|
||||
TAG_RETENTION_MARGIN = timedelta(seconds=1)
|
||||
|
||||
|
||||
def compaction_target_rows(
|
||||
dataset: "lance.LanceDataset", target_bytes: int
|
||||
) -> int | None:
|
||||
"""Rows per fragment sized so one compaction task reads about `target_bytes`.
|
||||
|
||||
Taken from the widest fragment's bytes per row rather than the table
|
||||
average: document payloads span kilobytes to hundreds of megabytes, so an
|
||||
average lets a task pick up many outsized rows. Sizes come from the
|
||||
manifest, so no payload is read.
|
||||
|
||||
This sizes fragments compaction *writes*; it does not shrink fragments that
|
||||
are already larger, which lance rewrites whole. Returns None when any
|
||||
fragment cannot be measured, which means leave the table alone rather than
|
||||
fall back to lance's unbounded default: sizing from the fragments that do
|
||||
carry metadata would still hand an unmeasured one to compaction whole.
|
||||
"""
|
||||
widest = 0.0
|
||||
for fragment in dataset.get_fragments():
|
||||
rows = fragment.physical_rows
|
||||
if not rows:
|
||||
continue
|
||||
sizes = [f.file_size_bytes for f in fragment.data_files()]
|
||||
if not all(sizes):
|
||||
return None
|
||||
widest = max(widest, sum(cast(list[int], sizes)) / rows)
|
||||
if widest <= 0:
|
||||
return None
|
||||
return max(1, int(target_bytes // widest))
|
||||
|
||||
|
||||
# Restore order for multi-table restore and its rollback. documents restores
|
||||
# last: writes land in it last on the ingest path, making it the closest
|
||||
# available database commit point.
|
||||
|
|
@ -127,10 +161,12 @@ RESTORE_TABLE_ORDER: tuple[str, ...] = tuple(
|
|||
|
||||
|
||||
async def _wait_protected[T](coro: Coroutine[Any, Any, T]) -> tuple[T, bool]:
|
||||
"""Await a recovery coroutine that a cancellation cannot interrupt.
|
||||
"""Await a coroutine that a cancellation cannot interrupt.
|
||||
|
||||
Runs the coroutine as a task and keeps waiting for it even if this
|
||||
coroutine is cancelled, so a Ctrl-C cannot leave recovery half applied.
|
||||
coroutine is cancelled, so a Ctrl-C cannot leave half-applied state behind.
|
||||
Used for rollback, which is bounded, and for lance maintenance, which is
|
||||
not: a shielded compaction can hold its caller for minutes.
|
||||
Returns the result and whether a cancellation was absorbed; the caller
|
||||
must re-deliver an absorbed cancellation.
|
||||
"""
|
||||
|
|
@ -362,17 +398,94 @@ class Store:
|
|||
# Perform maintenance per table using optimize() with configurable retention
|
||||
retention = timedelta(seconds=retention_seconds)
|
||||
for table in self._tables().values():
|
||||
await table.optimize(
|
||||
cleanup_older_than=await self._tag_safe_retention(
|
||||
table, retention
|
||||
)
|
||||
)
|
||||
cutoff = await self._tag_safe_retention(table, retention)
|
||||
if has_payload_columns(await table.schema()):
|
||||
await self._compact_to_target(table, cutoff)
|
||||
else:
|
||||
await table.optimize(cleanup_older_than=cutoff)
|
||||
except OSError as e:
|
||||
# Resource errors (e.g. disk pressure) skip the pass; lance
|
||||
# errors surface as RuntimeError and must not be swallowed —
|
||||
# a silently skipped cleanup hides tag-interaction bugs.
|
||||
logger.debug(f"Vacuum skipped due to resource constraints: {e}")
|
||||
|
||||
async def _run_lance_maintenance(
|
||||
self, table: lancedb.AsyncTable, step: "Callable[[lance.LanceDataset], Any]"
|
||||
) -> bool:
|
||||
"""Run one lance maintenance step against `table`, then refresh it.
|
||||
|
||||
The step mutates the dataset behind the open handle, so the refresh has
|
||||
to happen whether or not it succeeded: a step that commits and then
|
||||
fails would otherwise leave the handle on a superseded version. The
|
||||
dataset comes from `to_lance()` so object-storage credentials travel
|
||||
with it, and it is re-obtained per step because each commit supersedes
|
||||
the version the previous one pinned.
|
||||
|
||||
A worker thread cannot be cancelled, so the step and the refresh that
|
||||
follows it are shielded as one unit and awaited to completion rather
|
||||
than abandoned while they still hold `_write_lock`. Protecting only the
|
||||
thread would leave a cancellation arriving after the commit but before
|
||||
the refresh with a stale handle. Returns whether a cancellation was
|
||||
absorbed.
|
||||
"""
|
||||
|
||||
async def step_and_refresh() -> None:
|
||||
dataset = await table.to_lance()
|
||||
try:
|
||||
await asyncio.to_thread(step, dataset)
|
||||
finally:
|
||||
await table.checkout_latest()
|
||||
|
||||
_, cancelled = await _wait_protected(step_and_refresh())
|
||||
return cancelled
|
||||
|
||||
async def _compact_to_target(
|
||||
self, table: lancedb.AsyncTable, cutoff: timedelta
|
||||
) -> None:
|
||||
"""Compact a table of inline payloads to an explicit fragment target.
|
||||
|
||||
`AsyncTable.optimize` exposes no compaction options, and its default row
|
||||
target is never reached by a table of multi-megabyte rows, so every pass
|
||||
re-merges the whole table and its peak memory tracks the table rather
|
||||
than the delta. Going through lance directly is the only way to size it.
|
||||
"""
|
||||
dataset = await table.to_lance()
|
||||
target = compaction_target_rows(
|
||||
dataset, self._config.storage.compaction_target_bytes
|
||||
)
|
||||
if target is None:
|
||||
if dataset.get_fragments():
|
||||
logger.warning(
|
||||
f"{table.name}: fragment sizes are absent from the manifest, "
|
||||
"so compaction cannot be sized and is skipped; old versions "
|
||||
"are still pruned"
|
||||
)
|
||||
else:
|
||||
# A compaction cannot be cancelled once it reaches the worker
|
||||
# thread, so a close waits it out. Say what is running, or an
|
||||
# operator cannot tell a long pass from a wedge.
|
||||
logger.info(
|
||||
f"{table.name}: compacting to {target} rows per fragment "
|
||||
f"({len(dataset.get_fragments())} fragments)"
|
||||
)
|
||||
started = monotonic()
|
||||
if await self._run_lance_maintenance(
|
||||
table,
|
||||
lambda ds: ds.optimize.compact_files(target_rows_per_fragment=target),
|
||||
):
|
||||
raise asyncio.CancelledError
|
||||
logger.info(f"{table.name}: compacted in {monotonic() - started:.1f}s")
|
||||
# compact_files leaves indices covering the pre-compaction
|
||||
# fragments and never prunes; optimize() did both.
|
||||
if await self._run_lance_maintenance(
|
||||
table, lambda ds: ds.optimize.optimize_indices()
|
||||
):
|
||||
raise asyncio.CancelledError
|
||||
if await self._run_lance_maintenance(
|
||||
table, lambda ds: ds.cleanup_old_versions(older_than=cutoff)
|
||||
):
|
||||
raise asyncio.CancelledError
|
||||
|
||||
async def _tag_safe_retention(
|
||||
self, table: lancedb.AsyncTable, retention: timedelta
|
||||
) -> timedelta:
|
||||
|
|
|
|||
|
|
@ -52,6 +52,15 @@ class DocumentMetaRecord(LanceModel):
|
|||
updated_at: str = Field(default_factory=lambda: "")
|
||||
|
||||
|
||||
def has_payload_columns(schema: pa.Schema) -> bool:
|
||||
"""Whether a table stores multi-megabyte values inline.
|
||||
|
||||
Vacuum sizes compaction only for these; the rest are small enough that
|
||||
lance's own row target already bounds it.
|
||||
"""
|
||||
return any(pa.types.is_large_binary(field.type) for field in schema)
|
||||
|
||||
|
||||
def get_documents_arrow_schema() -> pa.Schema:
|
||||
"""Generate Arrow schema for documents table with large_binary for docling_document.
|
||||
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ dependencies = [
|
|||
"jinja2>=3.1.0",
|
||||
"fastmcp>=3.3.0",
|
||||
"lancedb==0.37.1",
|
||||
"pylance==10.0.0",
|
||||
"pathspec>=1.0.4",
|
||||
"pydantic>=2.12.5",
|
||||
"pydantic-ai-slim[openai,logfire,ag-ui]>=2.18.0,<3.0.0",
|
||||
|
|
|
|||
413
tests/store/test_vacuum_bounded.py
Normal file
413
tests/store/test_vacuum_bounded.py
Normal file
|
|
@ -0,0 +1,413 @@
|
|||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import threading
|
||||
from typing import cast
|
||||
|
||||
import lance
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from haiku.rag.config.models import AppConfig, StorageConfig
|
||||
from haiku.rag.store.engine import Store, compaction_target_rows
|
||||
|
||||
MB = 1024**2
|
||||
|
||||
|
||||
def _write(path, rows, payload_bytes, per_fragment=10):
|
||||
schema = pa.schema(
|
||||
[pa.field("id", pa.string()), pa.field("blob", pa.large_binary())]
|
||||
)
|
||||
dataset = lance.dataset(path) if os.path.exists(path) else None
|
||||
for start in range(0, rows, per_fragment):
|
||||
batch = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array(
|
||||
[f"d{i}" for i in range(start, start + per_fragment)], pa.string()
|
||||
),
|
||||
pa.array(
|
||||
[os.urandom(payload_bytes) for _ in range(per_fragment)],
|
||||
pa.large_binary(),
|
||||
),
|
||||
],
|
||||
schema=schema,
|
||||
)
|
||||
dataset = lance.write_dataset(
|
||||
batch, path, mode="append" if dataset else "create"
|
||||
)
|
||||
return dataset
|
||||
|
||||
|
||||
def test_target_scales_inversely_with_row_size(tmp_path):
|
||||
dataset = _write(tmp_path / "wide.lance", 20, 1 * MB)
|
||||
|
||||
assert compaction_target_rows(dataset, 8 * MB) == pytest.approx(8, abs=1)
|
||||
# halving the budget halves the rows a task may take
|
||||
assert compaction_target_rows(dataset, 4 * MB) == pytest.approx(4, abs=1)
|
||||
|
||||
|
||||
def test_target_uses_widest_fragment_not_the_average(tmp_path):
|
||||
"""A few outsized rows must shrink the target for the whole table."""
|
||||
path = tmp_path / "skewed.lance"
|
||||
_write(path, 10, 4 * MB)
|
||||
_write(path, 10, 64 * 1024) # a second, much narrower fragment
|
||||
dataset = lance.dataset(path)
|
||||
|
||||
target = compaction_target_rows(dataset, 16 * MB)
|
||||
|
||||
# the average row is ~2 MB, which would allow ~8 rows; the widest is 4 MB
|
||||
assert target is not None and target <= 5
|
||||
|
||||
|
||||
def test_target_floors_at_one_row(tmp_path):
|
||||
dataset = _write(tmp_path / "huge.lance", 10, 2 * MB)
|
||||
|
||||
assert compaction_target_rows(dataset, 1024) == 1
|
||||
|
||||
|
||||
def test_unmeasurable_table_is_not_compacted(tmp_path):
|
||||
"""No size information must mean no compaction, not lance's default.
|
||||
|
||||
Falling back to the default row target is the unbounded behaviour this
|
||||
exists to avoid.
|
||||
"""
|
||||
schema = pa.schema(
|
||||
[pa.field("id", pa.string()), pa.field("blob", pa.large_binary())]
|
||||
)
|
||||
dataset = lance.write_dataset(schema.empty_table(), tmp_path / "empty.lance")
|
||||
|
||||
assert compaction_target_rows(dataset, 1024) is None
|
||||
|
||||
|
||||
class _FakeDataFile:
|
||||
def __init__(self, size):
|
||||
self.file_size_bytes = size
|
||||
|
||||
|
||||
class _FakeFragment:
|
||||
def __init__(self, rows, sizes):
|
||||
self.physical_rows = rows
|
||||
self._files = [_FakeDataFile(s) for s in sizes]
|
||||
|
||||
def data_files(self):
|
||||
return self._files
|
||||
|
||||
|
||||
class _FakeDataset:
|
||||
"""Stands in for a manifest lance 10 will not produce.
|
||||
|
||||
Every file lance writes today records its size, so a mixed manifest can only
|
||||
come from a database written by a much older version.
|
||||
"""
|
||||
|
||||
def __init__(self, fragments):
|
||||
self._fragments = fragments
|
||||
|
||||
def get_fragments(self):
|
||||
return self._fragments
|
||||
|
||||
|
||||
def test_one_unmeasurable_fragment_disables_the_whole_table():
|
||||
"""Sizing from the measurable fragments would still hand the other one over
|
||||
to compaction whole."""
|
||||
dataset = _FakeDataset(
|
||||
[
|
||||
_FakeFragment(rows=10, sizes=[1 * MB]), # measurable and narrow
|
||||
_FakeFragment(rows=10, sizes=[None]), # size missing
|
||||
]
|
||||
)
|
||||
|
||||
assert compaction_target_rows(cast("lance.LanceDataset", dataset), 8 * MB) is None
|
||||
|
||||
|
||||
def test_empty_fragments_do_not_block_sizing():
|
||||
dataset = _FakeDataset(
|
||||
[
|
||||
_FakeFragment(rows=0, sizes=[]),
|
||||
_FakeFragment(rows=10, sizes=[10 * MB]),
|
||||
]
|
||||
)
|
||||
|
||||
assert compaction_target_rows(cast("lance.LanceDataset", dataset), 8 * MB) == 8
|
||||
|
||||
|
||||
def test_budget_must_be_positive():
|
||||
with pytest.raises(ValidationError):
|
||||
StorageConfig(compaction_target_bytes=0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vacuum_leaves_the_handle_usable(temp_db_path):
|
||||
"""Reads and writes must work through the same handle after a vacuum.
|
||||
|
||||
The bounded path mutates the dataset behind the open AsyncTable, so a
|
||||
missing checkout_latest leaves the handle on a stale version.
|
||||
"""
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = False
|
||||
|
||||
async with Store(temp_db_path, config=config, create=True) as store:
|
||||
await store.documents_table.add(
|
||||
[{"id": "a", "content": "x", "docling_document": b"0" * MB}]
|
||||
)
|
||||
await store.vacuum(retention_seconds=0)
|
||||
|
||||
assert await store.documents_table.count_rows() == 1
|
||||
await store.documents_table.add(
|
||||
[{"id": "b", "content": "y", "docling_document": b"1" * MB}]
|
||||
)
|
||||
assert await store.documents_table.count_rows() == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_is_refreshed_when_a_later_step_fails(temp_db_path, monkeypatch):
|
||||
"""A committed compaction followed by a failure must not strand the handle."""
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = False
|
||||
|
||||
async with Store(temp_db_path, config=config, create=True) as store:
|
||||
await store.documents_table.add(
|
||||
[{"id": "a", "content": "x", "docling_document": b"0" * MB}]
|
||||
)
|
||||
await store.documents_table.add(
|
||||
[{"id": "b", "content": "y", "docling_document": b"1" * MB}]
|
||||
)
|
||||
|
||||
def explode(dataset):
|
||||
raise RuntimeError("index optimize failed")
|
||||
|
||||
real = Store._run_lance_maintenance
|
||||
calls = {"n": 0}
|
||||
|
||||
async def fail_on_second_step(self, table, step):
|
||||
calls["n"] += 1
|
||||
# 1 compact, 2 optimize_indices, 3 prune -- fail after the commit
|
||||
return await real(self, table, explode if calls["n"] == 2 else step)
|
||||
|
||||
monkeypatch.setattr(Store, "_run_lance_maintenance", fail_on_second_step)
|
||||
with pytest.raises(RuntimeError, match="index optimize failed"):
|
||||
await store.vacuum(retention_seconds=0)
|
||||
|
||||
# the handle must still serve reads and writes despite the failure
|
||||
assert await store.documents_table.count_rows() == 2
|
||||
await store.documents_table.add(
|
||||
[{"id": "c", "content": "z", "docling_document": b"2" * MB}]
|
||||
)
|
||||
assert await store.documents_table.count_rows() == 3
|
||||
|
||||
|
||||
# 1 compact, 2 optimize_indices, 3 prune -- all three on the first payload table
|
||||
@pytest.mark.parametrize("cancel_at", [1, 2, 3])
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_vacuum_waits_for_the_running_step(
|
||||
temp_db_path, cancel_at, monkeypatch
|
||||
):
|
||||
"""A worker thread cannot be cancelled, so the step must be awaited.
|
||||
|
||||
Releasing the write lock while lance is still mutating the dataset would let
|
||||
the next write race a commit that is still in flight, so a cancellation has
|
||||
to wait for the thread rather than abandon it.
|
||||
"""
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = False
|
||||
|
||||
async with Store(temp_db_path, config=config, create=True) as store:
|
||||
await store.documents_table.add(
|
||||
[{"id": "a", "content": "x", "docling_document": b"0" * MB}]
|
||||
)
|
||||
in_thread = threading.Event()
|
||||
release = threading.Event()
|
||||
completed: list[int] = []
|
||||
real = Store._run_lance_maintenance
|
||||
calls = {"n": 0}
|
||||
|
||||
def blocking(dataset):
|
||||
in_thread.set()
|
||||
release.wait(30)
|
||||
completed.append(1)
|
||||
|
||||
async def block_on_nth(self, table, step):
|
||||
calls["n"] += 1
|
||||
return await real(
|
||||
self, table, blocking if calls["n"] == cancel_at else step
|
||||
)
|
||||
|
||||
monkeypatch.setattr(Store, "_run_lance_maintenance", block_on_nth)
|
||||
try:
|
||||
task = asyncio.create_task(store.vacuum(retention_seconds=0))
|
||||
while not in_thread.is_set():
|
||||
await asyncio.sleep(0.01)
|
||||
task.cancel()
|
||||
release.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
# the step ran to completion despite the cancellation
|
||||
assert completed == [1]
|
||||
assert not store._write_lock.locked()
|
||||
assert await store.documents_table.count_rows() == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_thin_tables_keep_using_optimize(temp_db_path, monkeypatch):
|
||||
"""Only payload-bearing tables take the sized path."""
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = False
|
||||
optimized: list[str] = []
|
||||
compacted: list[str] = []
|
||||
|
||||
async with Store(temp_db_path, config=config, create=True) as store:
|
||||
real_optimize = type(store.chunks_table).optimize
|
||||
|
||||
async def record_optimize(self, **kwargs):
|
||||
optimized.append(self.name)
|
||||
return await real_optimize(self, **kwargs)
|
||||
|
||||
async def record_compact(self, table, cutoff):
|
||||
compacted.append(table.name)
|
||||
|
||||
monkeypatch.setattr(type(store.chunks_table), "optimize", record_optimize)
|
||||
monkeypatch.setattr(Store, "_compact_to_target", record_compact)
|
||||
await store.vacuum(retention_seconds=0)
|
||||
|
||||
assert sorted(compacted) == ["document_items", "documents"]
|
||||
assert sorted(optimized) == ["chunks", "document_meta", "settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unmeasurable_table_still_prunes(temp_db_path, monkeypatch, caplog):
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = False
|
||||
pruned: list[str] = []
|
||||
|
||||
async with Store(temp_db_path, config=config, create=True) as store:
|
||||
await store.documents_table.add(
|
||||
[{"id": "a", "content": "x", "docling_document": b"0" * MB}]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.store.engine.compaction_target_rows", lambda *_: None
|
||||
)
|
||||
real = Store._run_lance_maintenance
|
||||
|
||||
async def record(self, table, step):
|
||||
pruned.append(getattr(step, "__qualname__", "?"))
|
||||
return await real(self, table, step)
|
||||
|
||||
monkeypatch.setattr(Store, "_run_lance_maintenance", record)
|
||||
with caplog.at_level("WARNING"):
|
||||
await store.vacuum(retention_seconds=0)
|
||||
|
||||
# compaction skipped, pruning still ran, and the skip is visible
|
||||
assert len(pruned) == 2 # documents and document_items, prune only
|
||||
assert "sizes are absent" in caplog.text
|
||||
|
||||
|
||||
# The failure is peak RSS while rewriting, and a bounded and an unbounded run
|
||||
# produce identical files, so only a memory measurement can tell them apart.
|
||||
_RSS_PROBE = textwrap.dedent(
|
||||
"""
|
||||
import asyncio, os, sys, threading, time
|
||||
from pathlib import Path
|
||||
import psutil
|
||||
|
||||
db, budget, repo = sys.argv[1], int(sys.argv[2]), sys.argv[3]
|
||||
sys.path.insert(0, repo)
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
proc = psutil.Process()
|
||||
peak = {"v": 0}
|
||||
|
||||
def watch():
|
||||
while True:
|
||||
peak["v"] = max(peak["v"], proc.memory_info().rss)
|
||||
time.sleep(0.01)
|
||||
|
||||
threading.Thread(target=watch, daemon=True).start()
|
||||
|
||||
async def main():
|
||||
config = AppConfig()
|
||||
config.storage.auto_vacuum = False
|
||||
if budget:
|
||||
config.storage.compaction_target_bytes = budget
|
||||
peaks = []
|
||||
async with Store(Path(db), config=config, create=True) as store:
|
||||
row = 0
|
||||
for round_no in range(4):
|
||||
await store.documents_table.add([
|
||||
{"id": f"d{row + i}", "content": "x",
|
||||
"docling_document": os.urandom(4 * 1024 * 1024)}
|
||||
for i in range(8)])
|
||||
row += 8
|
||||
if round_no:
|
||||
# churn: rewrite an old row so deletions accumulate
|
||||
await store.documents_table.delete(f"id = 'd{round_no}'")
|
||||
base = peak["v"] = proc.memory_info().rss
|
||||
await store.vacuum(retention_seconds=0)
|
||||
peaks.append((peak["v"] - base) / (1024 * 1024))
|
||||
print(",".join(f"{p:.1f}" for p in peaks))
|
||||
|
||||
asyncio.run(main())
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def _probe(tmp_path, name, budget):
|
||||
"""Peak RSS in MB per append+churn+vacuum round, in a fresh process."""
|
||||
out = subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
"-c",
|
||||
_RSS_PROBE,
|
||||
str(tmp_path / name),
|
||||
str(budget),
|
||||
os.getcwd(),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return [float(p) for p in out.stdout.strip().splitlines()[-1].split(",")]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_sized_vacuum_does_not_grow_with_the_table(tmp_path):
|
||||
"""Drives Store.vacuum end to end, with deletions, not lance directly."""
|
||||
unbounded = _probe(tmp_path, "unbounded.lancedb", 8 * 1024**3)
|
||||
bounded = _probe(tmp_path, "bounded.lancedb", 24 * MB)
|
||||
|
||||
# the unsized pass re-reads everything written so far, so its cost rises
|
||||
assert unbounded[-1] > unbounded[1]
|
||||
assert max(bounded) < max(unbounded)
|
||||
# and the sized one does not trend upward as the table grows
|
||||
assert max(bounded[2:]) <= max(bounded[1], 1.0) * 1.5
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_a_fragment_larger_than_the_target_is_rewritten_whole(tmp_path):
|
||||
"""The target sizes what compaction writes, not what it may read.
|
||||
|
||||
A fragment written larger than the target -- by a big ingest batch, or by an
|
||||
older unsized vacuum -- is rewritten in one piece the first time deletions
|
||||
make it a candidate, costing roughly its own size regardless of the target.
|
||||
Recorded because it is the ceiling the config knob cannot lower.
|
||||
"""
|
||||
path = tmp_path / "batched.lance"
|
||||
_write(path, 40, 2 * MB, per_fragment=40) # one 80 MB fragment
|
||||
dataset = lance.dataset(path)
|
||||
target = compaction_target_rows(dataset, 8 * MB)
|
||||
assert target is not None and target <= 4
|
||||
|
||||
# deletions must pass lance's 10% threshold for the fragment to be a candidate
|
||||
dataset.delete("id in ('d0','d1','d2','d3','d4','d5')")
|
||||
dataset = lance.dataset(path)
|
||||
metrics = dataset.optimize.compact_files(target_rows_per_fragment=target)
|
||||
|
||||
# the oversized fragment was read whole and split into target-sized pieces
|
||||
assert metrics.fragments_removed == 1
|
||||
assert metrics.fragments_added > 1
|
||||
32
uv.lock
32
uv.lock
|
|
@ -1681,6 +1681,7 @@ dependencies = [
|
|||
{ name = "pydantic" },
|
||||
{ name = "pydantic-ai-slim", extra = ["ag-ui", "logfire", "openai"] },
|
||||
{ name = "pydantic-monty" },
|
||||
{ name = "pylance" },
|
||||
{ name = "pypdfium2" },
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "pyyaml" },
|
||||
|
|
@ -1773,6 +1774,7 @@ requires-dist = [
|
|||
{ name = "pydantic-ai-slim", extras = ["openai", "logfire", "ag-ui"], specifier = ">=2.18.0,<3.0.0" },
|
||||
{ name = "pydantic-ai-slim", extras = ["voyageai"], marker = "extra == 'voyageai'" },
|
||||
{ name = "pydantic-monty", specifier = ">=0.0.19" },
|
||||
{ name = "pylance", specifier = "==10.0.0" },
|
||||
{ name = "pypdfium2", specifier = ">=5.0" },
|
||||
{ name = "python-dotenv", specifier = ">=1.2.2" },
|
||||
{ name = "pyyaml", specifier = ">=6.0.3" },
|
||||
|
|
@ -2234,19 +2236,19 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "lance-namespace"
|
||||
version = "0.6.1"
|
||||
version = "0.8.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace-urllib3-client" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/28/9f/7906ba4117df8d965510285eaf07264a77de2fd283b9d44ec7fc63a4a57a/lance_namespace-0.6.1.tar.gz", hash = "sha256:f0deea442bd3f1056a8e2fed056ae2778e3356517ec2e680db049058b824d131", size = 10666, upload-time = "2026-03-17T17:55:44.977Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/af/12/f7ab93b29be3edbf5fc3610714bf2d06088e7f4524bfb38dfd6852458b08/lance_namespace-0.8.6.tar.gz", hash = "sha256:18232e721c8188145f4ec9389cc2dfbeeabf54a619d94885ea1b3375bee9f4af", size = 11529, upload-time = "2026-06-12T17:36:41.651Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/d1/91/aee1c0a04d17f2810173bd304bd444eb78332045df1b0c1b07cebd01f530/lance_namespace-0.6.1-py3-none-any.whl", hash = "sha256:9699c9e3f12236e5e08ea979cc4e036a8e3c67ed2f37ae6f25c5353ab908e1be", size = 12498, upload-time = "2026-03-17T17:55:44.062Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a0/1b/5b1668ee2dc8910965f390640359112a31157092fcf8e000b89c79b58708/lance_namespace-0.8.6-py3-none-any.whl", hash = "sha256:571eae34f9aad70e5b05020416c2860889b9ec82993ccd0eb015e7b39c3ea309", size = 13383, upload-time = "2026-06-12T17:36:43.456Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lance-namespace-urllib3-client"
|
||||
version = "0.6.1"
|
||||
version = "0.8.6"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pydantic" },
|
||||
|
|
@ -2254,9 +2256,9 @@ dependencies = [
|
|||
{ name = "typing-extensions" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/63/a1/8706a2be25bd184acccc411e48f1a42a4cbf3b6556cba15b9fcf4c15cfcc/lance_namespace_urllib3_client-0.6.1.tar.gz", hash = "sha256:31fbd058ce1ea0bf49045cdeaa756360ece0bc61e9e10276f41af6d217debe87", size = 182567, upload-time = "2026-03-17T17:55:46.87Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/c7/80/fb224b4a89c1c1638cde949cb6cce6c3aca7759effbfea46a3d9c3960b21/lance_namespace_urllib3_client-0.8.6.tar.gz", hash = "sha256:b6fb1d306e74a7576e5309919020be744527de484a63dbf5eed10f8b368548df", size = 228772, upload-time = "2026-06-12T17:36:42.609Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/cd/c7/cb9580602dec25f0fdd6005c1c9ba1d4c8c0c3dc8d543107e5a9f248bba8/lance_namespace_urllib3_client-0.6.1-py3-none-any.whl", hash = "sha256:b9c103e1377ad46d2bd70eec894bfec0b1e2133dae0964d7e4de543c6e16293b", size = 317111, upload-time = "2026-03-17T17:55:45.546Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c5/90/1e27de15cd1b16785a1c7312beb0a59e75c8344a815f600f58173a565bd1/lance_namespace_urllib3_client-0.8.6-py3-none-any.whl", hash = "sha256:9d78249c3fb15aa3d15d668f78f04a275af3d08d800a7027492f37996ac4968b", size = 369950, upload-time = "2026-06-12T17:36:40.438Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4110,6 +4112,24 @@ crypto = [
|
|||
{ name = "cryptography" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pylance"
|
||||
version = "10.0.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "lance-namespace" },
|
||||
{ name = "numpy" },
|
||||
{ name = "pyarrow" },
|
||||
]
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/82/b2/c81de196076c4c8d768f485324a0043e113ce0950a339977d83fee9783ef/pylance-10.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:d4bba56ae829202b7e9cdc82c92c00a4b52f03e679dc3edfad1db09d1285be2e", size = 69279797, upload-time = "2026-08-07T18:25:24.812Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d2/e9/af671a6225740bd70c1e70bd83fa08091d628b960397cbecc477f16aaaf1/pylance-10.0.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:489b944827c0271e16a62b4006f8c75acc3bd7fbc381f35388a2afaaae38d438", size = 72775384, upload-time = "2026-08-07T18:31:29.686Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/22/e07194195bb3bbdf062b0c31690fc92fccb2686d5e768efa1dc379a93350/pylance-10.0.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:018efe7d437d326b9049c1223458bd955e85e48509b4b0bbf8b2d3d94075dbff", size = 76632194, upload-time = "2026-08-07T18:44:29.001Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f2/71/ed9956cf657e86a5fee5d5ddcfa7bfed0d7c2968bee39ed85d12629d8c93/pylance-10.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:254bad5d765c14db6c4eddd5323133dee3c76ff6d40c4777afe0b0e91d4c04d2", size = 72804222, upload-time = "2026-08-07T18:31:24.264Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/94/96/de449c246b2892df9d5a0af775b79a061e9897e6942bcee5c4ab6d6071f8/pylance-10.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:9f0c089e30389d9b7765a4c16f3d9325b05ed4fbf6cc2bbc9418292983472e54", size = 76602015, upload-time = "2026-08-07T18:46:27.724Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/53/81/4a5a9072b6d68c4dbb8c7b5530a381c7a6bea4ba92f9eca58ac885142722/pylance-10.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:9fecf46e6835dc64b2d71767d5eab4f73787543a852432e9ba8ecf8527b3cdab", size = 82799247, upload-time = "2026-08-07T18:47:46.855Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pylatexenc"
|
||||
version = "2.10"
|
||||
|
|
|
|||
Loading…
Reference in a new issue