mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-22 03:32:28 +00:00
Gatefixer
7ba76d15ff
Derive a row target from estimated row size, link the options reference, and retain first-pass and source-budget caveats. Keep the section hidden until release availability is confirmed. Validate the example across initial compaction and a later append; all seven storage documentation tests pass. Guide, draft, and snippet MDX compilation and targeted Ruff checks pass.
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}
|