mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-29 23:06:51 +00:00
lancedb-gatefixer[bot]andXuanwo
3352bf49df
## Summary - expose typed compaction options on synchronous and asynchronous table optimization - translate validated Python options into native Lance compaction settings - document the row-based steady-state limit and first-pass/source-budget caveats - carry the Lance beta.15 fragment-exclusion option through the typed Python API ## Root cause The native Python optimize binding always constructed default compaction settings, and the public AsyncTable API offered no way to lower the fragment row target. Wide-row fragments therefore remained compaction candidates on every later optimize call, repeatedly rewriting the accumulated table. ## Limitations This is an explicit row-based mitigation, not byte-native task sizing or a hard memory cap. Callers must derive target_rows_per_fragment from expected row size. Existing oversized fragments are rewritten whole on the first pass; max_bytes_per_file limits output file size rather than compaction memory; and source budgets admit whole planned tasks, so a budget below the first task performs no compaction. Splitting tasks by a byte budget or returning a planner error requires a Lance-side change. ## Validation - cargo check --release -p lancedb-python - cargo clippy --quiet -p lancedb-python - cargo fmt --all --check - ruff format --check python tests - ruff check python tests - uv run --no-sync pytest python/tests/test_table.py -k optimize_compaction -q (16 passed) Fixes #2325 <!-- lance-gatekeeper-fix:v1 agent=9b58eee0c79040b19540e259af0e676c generation=1 --> --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo <github@xuanwo.io>
115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import lancedb
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
|
|
def test_compaction_tuning(mem_db):
|
|
payload = b"x" * 1016
|
|
table = mem_db.create_table(
|
|
"compaction_tuning", [{"id": i, "payload": payload} for i in range(32)]
|
|
)
|
|
table.add([{"id": i, "payload": payload} for i in range(32, 64)])
|
|
assert table.stats()["fragment_stats"]["num_fragments"] == 2
|
|
|
|
# --8<-- [start:storage_compaction_tuning]
|
|
from lancedb import CompactionOptions
|
|
|
|
desired_fragment_bytes = 64 * 1024
|
|
average_row_bytes = 1024 # Estimate from representative rows.
|
|
options: CompactionOptions = {
|
|
"target_rows_per_fragment": max(1, desired_fragment_bytes // average_row_bytes),
|
|
"batch_size": 16,
|
|
"num_threads": 1,
|
|
}
|
|
table.optimize(compaction_options=options)
|
|
# --8<-- [end:storage_compaction_tuning]
|
|
|
|
assert table.stats()["fragment_stats"]["num_fragments"] == 1
|
|
assert table.count_rows() == 64
|
|
|
|
table.add([{"id": 64, "payload": payload}])
|
|
table.optimize(compaction_options=options)
|
|
assert table.stats()["fragment_stats"]["num_fragments"] == 2
|
|
assert sorted(table.to_arrow()["id"].to_pylist()) == list(range(65))
|
|
|
|
|
|
def test_rebuild_vector_index(tmp_db):
|
|
rng = np.random.default_rng(7)
|
|
table = tmp_db.create_table(
|
|
"rebuild_index",
|
|
[{"id": i, "vector": v.tolist()} for i, v in enumerate(rng.random((512, 4)))],
|
|
)
|
|
table.create_index(
|
|
"vector", config=lancedb.index.IvfFlat(distance_type="l2", num_partitions=4)
|
|
)
|
|
original_index = table.list_indices()[0]
|
|
table.add(
|
|
[
|
|
{"id": 512 + i, "vector": v.tolist()}
|
|
for i, v in enumerate(rng.random((64, 4)) + 5)
|
|
]
|
|
)
|
|
assert table.index_stats("vector_idx").num_unindexed_rows == 64
|
|
|
|
# --8<-- [start:storage_rebuild_vector_index]
|
|
from lancedb.index import IvfFlat
|
|
|
|
table.create_index(
|
|
"vector",
|
|
config=IvfFlat(distance_type="l2", num_partitions=4),
|
|
replace=True,
|
|
)
|
|
# --8<-- [end:storage_rebuild_vector_index]
|
|
|
|
indices = table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].name == original_index.name
|
|
assert indices[0].index_uuid != original_index.index_uuid
|
|
stats = table.index_stats("vector_idx")
|
|
assert stats.num_indexed_rows == 576
|
|
assert stats.num_unindexed_rows == 0
|
|
|
|
|
|
def test_local_shallow_clone_pins_source_version(tmp_path):
|
|
db = lancedb.connect(str(tmp_path))
|
|
source = db.create_table("source", [{"id": 1}])
|
|
clone = db.clone_table("snapshot", source_uri=str(tmp_path / "source.lance"))
|
|
source.add([{"id": 2}])
|
|
assert source.count_rows() == 2
|
|
assert clone.to_arrow().to_pylist() == [{"id": 1}]
|
|
|
|
|
|
def _reader():
|
|
return pa.RecordBatchReader.from_batches(
|
|
pa.schema([("id", pa.int64())]),
|
|
[pa.record_batch([[1, 2]], names=["id"]), pa.record_batch([[3]], names=["id"])],
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("parallelism", [1, 2])
|
|
def test_streaming_write_parallelism(tmp_db, parallelism):
|
|
table = tmp_db.create_table("streaming", [{"id": 0}])
|
|
progress = []
|
|
table.add(_reader(), write_parallelism=parallelism, progress=progress.append)
|
|
assert sorted(table.to_arrow()["id"].to_pylist()) == [0, 1, 2, 3]
|
|
assert progress[-1]["done"]
|
|
assert progress[-1]["output_rows"] == 3
|
|
assert {p["total_tasks"] for p in progress} == {parallelism}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("parallelism", [1, 2])
|
|
async def test_async_streaming_write_parallelism(mem_db_async, parallelism):
|
|
table = await mem_db_async.create_table("streaming", [{"id": 0}])
|
|
progress = []
|
|
await table.add(_reader(), write_parallelism=parallelism, progress=progress.append)
|
|
data = await table.to_arrow()
|
|
assert sorted(data["id"].to_pylist()) == [0, 1, 2, 3]
|
|
assert progress[-1]["done"]
|
|
assert progress[-1]["output_rows"] == 3
|
|
assert {p["total_tasks"] for p in progress} == {parallelism}
|