mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 20:18:37 +00:00
lancedb-gatefixer[bot]
c3176a47ce
## Summary - add Python regression coverage for an IVF build that cannot form all requested non-empty partitions - verify hierarchical k-means returns an actionable RuntimeError instead of panicking or silently creating a degenerate index - exercise the current Lance v10.1.0-beta.1 dependency, which contains the upstream error-return fix ## Root cause Hierarchical k-means previously guarded a shortfall in generated clusters with only a debug assertion. Debug builds panicked, while release builds could silently publish an index with many empty partitions. The upstream Lance fix now returns a descriptive error and is already included in the dependency pinned on main; this test locks in propagation through the LanceDB Python API. ## Validation - uv run --extra tests pytest python/tests/test_index.py -q (24 passed) - uv run --extra tests pytest python/tests/test_index.py::test_create_ivf_index_reports_unsplittable_partitions -q (1 passed) - python/.venv/bin/ruff format python/python/tests/test_index.py - python/.venv/bin/ruff check . - git diff --check Fixes #3649 <!-- lance-gatekeeper-fix:v1 agent=a4d34448a9d350a3e2e659f33f5db6f2 generation=1 --> Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>
534 lines
18 KiB
Python
534 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
from datetime import timedelta
|
|
import random
|
|
from typing import get_args, get_type_hints
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
import pytest_asyncio
|
|
from lancedb import AsyncConnection, AsyncTable, connect_async
|
|
from lancedb.index import (
|
|
BTree,
|
|
IvfFlat,
|
|
IvfPq,
|
|
IvfSq,
|
|
IvfHnswPq,
|
|
IvfHnswSq,
|
|
IvfHnswFlat,
|
|
IvfRq,
|
|
Bitmap,
|
|
LabelList,
|
|
Fm,
|
|
HnswPq,
|
|
HnswSq,
|
|
HnswFlat,
|
|
FTS,
|
|
)
|
|
from lancedb.table import IndexStatistics
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def db_async(tmp_path) -> AsyncConnection:
|
|
return await connect_async(tmp_path, read_consistency_interval=timedelta(seconds=0))
|
|
|
|
|
|
def sample_fixed_size_list_array(nrows, dim):
|
|
vector_data = pa.array([float(i) for i in range(dim * nrows)], pa.float32())
|
|
return pa.FixedSizeListArray.from_arrays(vector_data, dim)
|
|
|
|
|
|
DIM = 8
|
|
NROWS = 256
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def some_table(db_async):
|
|
data = pa.Table.from_pydict(
|
|
{
|
|
"id": list(range(NROWS)),
|
|
"vector": sample_fixed_size_list_array(NROWS, DIM),
|
|
"fsb": pa.array([bytes([i]) for i in range(NROWS)], pa.binary(1)),
|
|
"tags": [
|
|
[f"tag{random.randint(0, 8)}" for _ in range(2)] for _ in range(NROWS)
|
|
],
|
|
"is_active": [random.choice([True, False]) for _ in range(NROWS)],
|
|
"data": [random.randbytes(random.randint(0, 128)) for _ in range(NROWS)],
|
|
}
|
|
)
|
|
return await db_async.create_table(
|
|
"some_table",
|
|
data,
|
|
)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def binary_table(db_async):
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"vector": [i] * 128,
|
|
}
|
|
for i in range(NROWS)
|
|
]
|
|
return await db_async.create_table(
|
|
"binary_table",
|
|
data,
|
|
schema=pa.schema(
|
|
[
|
|
pa.field("id", pa.int64()),
|
|
pa.field("vector", pa.list_(pa.uint8(), 128)),
|
|
]
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_index_async_returns_done_job(some_table: AsyncTable):
|
|
job = await some_table.create_index_async("id", config=BTree())
|
|
assert job.id is None
|
|
await job.wait()
|
|
assert len(await some_table.list_indices()) == 1
|
|
await job.cancel()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_scalar_index(some_table: AsyncTable):
|
|
# Can create
|
|
await some_table.create_index("id")
|
|
# Can recreate if replace=True
|
|
await some_table.create_index("id", replace=True)
|
|
indices = await some_table.list_indices()
|
|
assert str(indices).startswith(
|
|
'[IndexConfig(name="id_idx", index_type="BTree", columns=["id"]'
|
|
)
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "BTree"
|
|
assert indices[0].columns == ["id"]
|
|
# Can't recreate if replace=False
|
|
with pytest.raises(RuntimeError, match="already exists"):
|
|
await some_table.create_index("id", replace=False)
|
|
# can also specify index type
|
|
await some_table.create_index("id", config=BTree())
|
|
|
|
await some_table.drop_index("id_idx")
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_config_repr(db_async):
|
|
# Use >= 1000 rows so the thousands separator in the repr is exercised.
|
|
nrows = 1500
|
|
table = await db_async.create_table(
|
|
"repr_table", pa.Table.from_pydict({"id": list(range(nrows))})
|
|
)
|
|
await table.create_index("id", config=BTree())
|
|
indices = await table.list_indices()
|
|
assert len(indices) == 1
|
|
|
|
r = repr(indices[0])
|
|
assert r.startswith('IndexConfig(name="id_idx", index_type="BTree", columns=["id"]')
|
|
# Integer counts use `_` thousands separators (valid Python int syntax).
|
|
assert "num_indexed_rows=1_500" in r
|
|
assert "num_unindexed_rows=0" in r
|
|
# created_at renders as a datetime so the value round-trips.
|
|
assert "created_at=datetime.datetime(" in r
|
|
assert r.endswith(")")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_nested_scalar_index_lists_canonical_paths(db_async):
|
|
metadata_type = pa.struct(
|
|
[
|
|
pa.field("user_id", pa.int32()),
|
|
pa.field("user.id", pa.int32()),
|
|
]
|
|
)
|
|
mixed_case_metadata_type = pa.struct([pa.field("userId", pa.int32())])
|
|
escaped_metadata_type = pa.struct([pa.field("user-id", pa.int32())])
|
|
literal_type = pa.struct([pa.field("a.b", pa.int32())])
|
|
data = pa.Table.from_arrays(
|
|
[
|
|
pa.array([1, 2, 3], type=pa.int32()),
|
|
pa.array([1, 2, 3], type=pa.int32()),
|
|
pa.array([1, 2, 3], type=pa.int32()),
|
|
pa.array([1, 2, 3], type=pa.int32()),
|
|
pa.array(
|
|
[
|
|
{"user_id": 10, "user.id": 100},
|
|
{"user_id": 20, "user.id": 200},
|
|
{"user_id": 30, "user.id": 300},
|
|
],
|
|
type=metadata_type,
|
|
),
|
|
pa.array(
|
|
[{"userId": 10}, {"userId": 20}, {"userId": 30}],
|
|
type=mixed_case_metadata_type,
|
|
),
|
|
pa.array(
|
|
[{"user-id": 10}, {"user-id": 20}, {"user-id": 30}],
|
|
type=escaped_metadata_type,
|
|
),
|
|
pa.array(
|
|
[{"a.b": 10}, {"a.b": 20}, {"a.b": 30}],
|
|
type=literal_type,
|
|
),
|
|
],
|
|
names=[
|
|
"rowId",
|
|
"row-id",
|
|
"userId",
|
|
"user_id",
|
|
"metadata",
|
|
"MetaData",
|
|
"meta-data",
|
|
"literal",
|
|
],
|
|
)
|
|
table = await db_async.create_table("nested_scalar_index", data)
|
|
|
|
await table.create_index("rowId", config=BTree(), name="row_id_idx")
|
|
await table.create_index("`row-id`", config=BTree(), name="row_dash_id_idx")
|
|
await table.create_index("userId", config=BTree(), name="top_user_id_idx")
|
|
await table.create_index("user_id", config=BTree(), name="top_snake_user_id_idx")
|
|
await table.create_index(
|
|
"metadata.user_id", config=BTree(), name="nested_user_id_idx"
|
|
)
|
|
await table.create_index(
|
|
"metadata.`user.id`", config=BTree(), name="escaped_user_id_idx"
|
|
)
|
|
await table.create_index(
|
|
"MetaData.userId", config=BTree(), name="mixed_case_metadata_user_id_idx"
|
|
)
|
|
await table.create_index(
|
|
"`meta-data`.`user-id`", config=BTree(), name="escaped_names_idx"
|
|
)
|
|
await table.create_index("literal.`a.b`", config=BTree(), name="literal_dot_idx")
|
|
|
|
columns_by_name = {
|
|
index.name: index.columns for index in await table.list_indices()
|
|
}
|
|
assert columns_by_name["row_id_idx"] == ["rowId"]
|
|
assert columns_by_name["row_dash_id_idx"] == ["`row-id`"]
|
|
assert columns_by_name["top_user_id_idx"] == ["userId"]
|
|
assert columns_by_name["top_snake_user_id_idx"] == ["user_id"]
|
|
assert columns_by_name["nested_user_id_idx"] == ["metadata.user_id"]
|
|
assert columns_by_name["escaped_user_id_idx"] == ["metadata.`user.id`"]
|
|
assert columns_by_name["mixed_case_metadata_user_id_idx"] == ["MetaData.userId"]
|
|
assert columns_by_name["escaped_names_idx"] == ["`meta-data`.`user-id`"]
|
|
assert columns_by_name["literal_dot_idx"] == ["literal.`a.b`"]
|
|
|
|
for index_name in columns_by_name:
|
|
stats = await table.index_stats(index_name)
|
|
assert stats is not None
|
|
assert stats.num_indexed_rows == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_fixed_size_binary_index(some_table: AsyncTable):
|
|
await some_table.create_index("fsb", config=BTree())
|
|
indices = await some_table.list_indices()
|
|
assert str(indices).startswith(
|
|
'[IndexConfig(name="fsb_idx", index_type="BTree", columns=["fsb"]'
|
|
)
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "BTree"
|
|
assert indices[0].columns == ["fsb"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_fm_index(some_table: AsyncTable):
|
|
# FM-Index accelerates substring search on string/binary columns.
|
|
await some_table.create_index("data", config=Fm())
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "Fm"
|
|
assert indices[0].columns == ["data"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_bitmap_index(some_table: AsyncTable):
|
|
await some_table.create_index("id", config=Bitmap())
|
|
await some_table.create_index("is_active", config=Bitmap())
|
|
await some_table.create_index("data", config=Bitmap())
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 3
|
|
# list_indices returns indices in alphabetical order by name
|
|
assert indices[0].index_type == "Bitmap"
|
|
assert indices[0].columns == ["data"]
|
|
assert indices[1].index_type == "Bitmap"
|
|
assert indices[1].columns == ["id"]
|
|
assert indices[2].index_type == "Bitmap"
|
|
assert indices[2].columns == ["is_active"]
|
|
|
|
index_name = indices[0].name
|
|
stats = await some_table.index_stats(index_name)
|
|
assert stats.index_type == "BITMAP"
|
|
assert stats.distance_type is None
|
|
assert stats.num_indexed_rows == await some_table.count_rows()
|
|
assert stats.num_unindexed_rows == 0
|
|
assert stats.num_indices == 1
|
|
|
|
assert (
|
|
"ScalarIndexQuery"
|
|
in await some_table.query().where("is_active = TRUE").explain_plan()
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_label_list_index(some_table: AsyncTable):
|
|
await some_table.create_index("tags", config=LabelList())
|
|
indices = await some_table.list_indices()
|
|
assert str(indices).startswith(
|
|
'[IndexConfig(name="tags_idx", index_type="LabelList", columns=["tags"]'
|
|
)
|
|
plan = await some_table.query().where("array_has(tags, 'tag0')").explain_plan()
|
|
assert "ScalarIndexQuery" in plan
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_large_list_label_list_index(db_async):
|
|
data = pa.Table.from_pydict(
|
|
{"tags": [[f"tag{i % 2}", "shared"] for i in range(16)]},
|
|
schema=pa.schema([pa.field("tags", pa.large_list(pa.string()))]),
|
|
)
|
|
table = await db_async.create_table("large_list_label_list_index", data)
|
|
|
|
await table.create_index("tags", config=LabelList())
|
|
indices = await table.list_indices()
|
|
assert str(indices).startswith(
|
|
'[IndexConfig(name="tags_idx", index_type="LabelList", columns=["tags"]'
|
|
)
|
|
plan = await table.query().where("array_has(tags, 'shared')").explain_plan()
|
|
assert "ScalarIndexQuery" in plan
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_label_list_index_rejects_list_struct(db_async):
|
|
item_type = pa.struct(
|
|
[
|
|
pa.field("tag", pa.string()),
|
|
pa.field(
|
|
"metadata",
|
|
pa.struct([pa.field("userId", pa.string())]),
|
|
),
|
|
]
|
|
)
|
|
data = pa.Table.from_pylist(
|
|
[
|
|
{
|
|
"items": [
|
|
{"tag": "tag0", "metadata": {"userId": "user0"}},
|
|
{"tag": "shared", "metadata": {"userId": "user1"}},
|
|
]
|
|
}
|
|
],
|
|
schema=pa.schema([pa.field("items", pa.list_(item_type))]),
|
|
)
|
|
table = await db_async.create_table("list_struct_label_list_index", data)
|
|
|
|
with pytest.raises(Exception, match="LabelList index cannot be created"):
|
|
await table.create_index("items", config=LabelList())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_full_text_search_index(some_table: AsyncTable):
|
|
await some_table.create_index("tags", config=FTS(with_position=False))
|
|
indices = await some_table.list_indices()
|
|
assert str(indices).startswith(
|
|
'[IndexConfig(name="tags_idx", index_type="FTS", columns=["tags"]'
|
|
)
|
|
|
|
await some_table.prewarm_index("tags_idx")
|
|
|
|
res = await (await some_table.search("tag0")).to_arrow()
|
|
assert res.num_rows > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_vector_index(some_table: AsyncTable):
|
|
# Can create
|
|
await some_table.create_index("vector")
|
|
# Can recreate if replace=True
|
|
await some_table.create_index("vector", replace=True)
|
|
# Can't recreate if replace=False
|
|
with pytest.raises(RuntimeError, match="already exists"):
|
|
await some_table.create_index("vector", replace=False)
|
|
# Can also specify index type
|
|
await some_table.create_index("vector", config=IvfPq(num_partitions=100))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "IvfPq"
|
|
assert indices[0].columns == ["vector"]
|
|
assert indices[0].name == "vector_idx"
|
|
|
|
stats = await some_table.index_stats("vector_idx")
|
|
assert stats.index_type == "IVF_PQ"
|
|
assert stats.distance_type == "l2"
|
|
assert stats.num_indexed_rows == await some_table.count_rows()
|
|
assert stats.num_unindexed_rows == 0
|
|
assert stats.num_indices == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_ivf_index_reports_unsplittable_partitions(db_async):
|
|
dim = 8
|
|
num_partitions = 300 # More than 256 selects hierarchical k-means.
|
|
base_vectors = [[float(row == column) for column in range(dim)] for row in range(5)]
|
|
vectors = pa.array(base_vectors * 200, pa.list_(pa.float32(), dim))
|
|
table = await db_async.create_table(
|
|
"unsplittable_partitions",
|
|
pa.table({"vector": vectors}),
|
|
)
|
|
|
|
error_pattern = (
|
|
rf"Cannot create {num_partitions} IVF partitions: k-means could only form"
|
|
)
|
|
with pytest.raises(RuntimeError, match=error_pattern):
|
|
await table.create_index(
|
|
"vector",
|
|
config=IvfFlat(
|
|
distance_type="dot",
|
|
num_partitions=num_partitions,
|
|
max_iterations=10,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
|
|
# Can create
|
|
await some_table.create_index("vector", config=IvfPq(num_bits=4))
|
|
# Can recreate if replace=True
|
|
await some_table.create_index("vector", config=IvfPq(num_bits=4), replace=True)
|
|
# Can't recreate if replace=False
|
|
with pytest.raises(RuntimeError, match="already exists"):
|
|
await some_table.create_index("vector", replace=False)
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "IvfPq"
|
|
assert indices[0].columns == ["vector"]
|
|
assert indices[0].name == "vector_idx"
|
|
|
|
stats = await some_table.index_stats("vector_idx")
|
|
assert stats.index_type == "IVF_PQ"
|
|
assert stats.distance_type == "l2"
|
|
assert stats.num_indexed_rows == await some_table.count_rows()
|
|
assert stats.num_unindexed_rows == 0
|
|
assert stats.num_indices == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_ivfrq_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=IvfRq(num_bits=1))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "IvfRq"
|
|
assert indices[0].columns == ["vector"]
|
|
assert indices[0].name == "vector_idx"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_hnswpq_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=HnswPq(num_partitions=10))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_hnswsq_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=HnswSq(num_partitions=10))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_hnswsq_alias_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=IvfHnswSq(num_partitions=5))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type in {"HnswSq", "IvfHnswSq"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_hnswpq_alias_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=IvfHnswPq(num_partitions=5))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type in {"HnswPq", "IvfHnswPq"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_hnswflat_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=HnswFlat(num_partitions=10))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_hnswflat_alias_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=IvfHnswFlat(num_partitions=5))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type in {"HnswFlat", "IvfHnswFlat"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_ivfsq_index(some_table: AsyncTable):
|
|
await some_table.create_index("vector", config=IvfSq(num_partitions=10))
|
|
indices = await some_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "IvfSq"
|
|
stats = await some_table.index_stats(indices[0].name)
|
|
assert stats.index_type == "IVF_SQ"
|
|
assert stats.distance_type == "l2"
|
|
assert stats.num_indexed_rows == await some_table.count_rows()
|
|
assert stats.num_unindexed_rows == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_index_with_binary_vectors(binary_table: AsyncTable):
|
|
await binary_table.create_index(
|
|
"vector", config=IvfFlat(distance_type="hamming", num_partitions=10)
|
|
)
|
|
indices = await binary_table.list_indices()
|
|
assert len(indices) == 1
|
|
assert indices[0].index_type == "IvfFlat"
|
|
assert indices[0].columns == ["vector"]
|
|
assert indices[0].name == "vector_idx"
|
|
|
|
stats = await binary_table.index_stats("vector_idx")
|
|
assert stats.index_type == "IVF_FLAT"
|
|
assert stats.distance_type == "hamming"
|
|
assert stats.num_indexed_rows == await binary_table.count_rows()
|
|
assert stats.num_unindexed_rows == 0
|
|
assert stats.num_indices == 1
|
|
|
|
# the dataset contains vectors with all values from 0 to 255
|
|
for v in range(256):
|
|
res = await binary_table.query().nearest_to([v] * 128).to_arrow()
|
|
assert res["id"][0].as_py() == v
|
|
|
|
|
|
def test_index_statistics_index_type_lists_all_supported_values():
|
|
expected_index_types = {
|
|
"IVF_FLAT",
|
|
"IVF_SQ",
|
|
"IVF_PQ",
|
|
"IVF_RQ",
|
|
"IVF_HNSW_SQ",
|
|
"IVF_HNSW_PQ",
|
|
"IVF_HNSW_FLAT",
|
|
"FTS",
|
|
"BTREE",
|
|
"BITMAP",
|
|
"LABEL_LIST",
|
|
}
|
|
|
|
assert (
|
|
set(get_args(get_type_hints(IndexStatistics)["index_type"]))
|
|
== expected_index_types
|
|
)
|