mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-14 22:56:54 +00:00
Yang Cen
2d86532de9
The vector-search guides do not clearly explain when to use Auto, search effort or a fixed probe count. They also describe `minimum_nprobes` as the exact starting count, although Auto can search more partitions from the start. Explain how to choose public query controls for recall and latency, with short executable examples. Clarify that probe limits apply per segment, calibration targets are not recall guarantees, and more probes cannot remove quantization error. Keep necessary client/server version guidance in a short note. Validation: - The existing example test passes against published `lancedb==0.41.0`, covering all four snippets, result counts, filtering, finite distances and recall against an exact cosine query. The test body takes 0.20 s. Generated snippets match their source. - Vale reports no findings on added prose. Changed-function Python formatting and `git diff --check` pass. Repository-wide Ruff has 137 findings identical to the previous scan, including seven pre-existing findings in the example file. - Mint 4.2.888 strict validation reports the same four `RsExampleBody`/`RsExampleTail` import warnings as pristine base `0b20965d`, with no additional warnings.
797 lines
22 KiB
Python
797 lines
22 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
import random
|
|
import string
|
|
import uuid
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
|
|
def _make_vector_rows(count: int, dim: int, column: str = "vector"):
|
|
rows = [
|
|
{column: np.random.random(dim).astype(np.float32).tolist(), "id": i}
|
|
for i in range(count)
|
|
]
|
|
assert len(rows) == count
|
|
assert all(len(row[column]) == dim for row in rows)
|
|
return rows
|
|
|
|
|
|
def test_vector_index_configure_ivf(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"vector_index_configure_ivf",
|
|
_make_vector_rows(512, 4),
|
|
mode="overwrite",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_configure_ivf]
|
|
table.create_index(metric="l2", num_partitions=16, num_sub_vectors=4)
|
|
# --8<-- [end:vector_index_configure_ivf]
|
|
|
|
assert table.list_indices()
|
|
|
|
|
|
def test_vector_index_setup(tmp_db):
|
|
tmp_db.create_table(
|
|
"vector-index-tbl",
|
|
_make_vector_rows(8, 4),
|
|
mode="overwrite",
|
|
)
|
|
|
|
db = tmp_db
|
|
# --8<-- [start:vector_index_setup]
|
|
table_name = "vector-index-tbl"
|
|
table = db.open_table(table_name)
|
|
# --8<-- [end:vector_index_setup]
|
|
|
|
assert table.name == table_name
|
|
|
|
|
|
def test_vector_index_build_ivf(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"vector-index-build-ivf",
|
|
_make_vector_rows(512, 4, column="keywords_embeddings"),
|
|
mode="overwrite",
|
|
)
|
|
db = tmp_db
|
|
# --8<-- [start:vector_index_build_ivf]
|
|
table_name = "vector-index-build-ivf"
|
|
table = db.open_table(table_name)
|
|
table.create_index(
|
|
metric="cosine",
|
|
vector_column_name="keywords_embeddings",
|
|
)
|
|
# --8<-- [end:vector_index_build_ivf]
|
|
|
|
assert table.list_indices()
|
|
|
|
|
|
def test_vector_index_nested_field(tmp_db):
|
|
dim = 2
|
|
schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.int32()),
|
|
pa.field(
|
|
"image",
|
|
pa.struct([pa.field("embedding", pa.list_(pa.float32(), dim))]),
|
|
),
|
|
]
|
|
)
|
|
data = [
|
|
{
|
|
"id": i,
|
|
"image": {"embedding": np.random.random(dim).astype(np.float32).tolist()},
|
|
}
|
|
for i in range(512)
|
|
]
|
|
table = tmp_db.create_table(
|
|
"vector_index_nested_field", data=data, schema=schema, mode="overwrite"
|
|
)
|
|
|
|
# --8<-- [start:vector_index_nested_field]
|
|
# The vector column `embedding` is nested inside the `image` struct.
|
|
# Pass its full dotted path as `vector_column_name`; the same path is used
|
|
# at query time and is what `list_indices()` reports under `columns`.
|
|
table.create_index(
|
|
vector_column_name="image.embedding",
|
|
num_partitions=1,
|
|
num_sub_vectors=1,
|
|
name="image_embedding_idx",
|
|
)
|
|
|
|
results = (
|
|
table.search([0.0, 1.0], vector_column_name="image.embedding")
|
|
.limit(1)
|
|
.to_list()
|
|
)
|
|
# --8<-- [end:vector_index_nested_field]
|
|
|
|
assert table.index_stats("image_embedding_idx")
|
|
assert len(results) == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_vector_index_async_config(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
# --8<-- [start:vector_index_async_config]
|
|
import lancedb
|
|
import numpy as np
|
|
from lancedb.index import IvfPq
|
|
|
|
async def main():
|
|
data = [
|
|
{"id": i, "vector": np.random.random(8).astype(np.float32).tolist()}
|
|
for i in range(512)
|
|
]
|
|
|
|
db = await lancedb.connect_async("ex_lancedb")
|
|
table = await db.create_table(
|
|
"vector_index_async", data=data, mode="overwrite"
|
|
)
|
|
|
|
await table.create_index(
|
|
"vector",
|
|
config=IvfPq(
|
|
distance_type="cosine",
|
|
num_partitions=16,
|
|
num_sub_vectors=4,
|
|
),
|
|
)
|
|
return await table.list_indices()
|
|
# --8<-- [end:vector_index_async_config]
|
|
|
|
assert await main()
|
|
|
|
|
|
def test_vector_index_query_ivf(tmp_db):
|
|
dim = 1536
|
|
data = [
|
|
{"id": i, "keywords_embeddings": np.random.random(dim).tolist()}
|
|
for i in range(512)
|
|
]
|
|
table = tmp_db.create_table("vector_index_query_ivf", data, mode="overwrite")
|
|
table.create_index(
|
|
metric="cosine",
|
|
vector_column_name="keywords_embeddings",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_query_ivf]
|
|
tbl = table
|
|
tbl.search(np.random.random((1536))).limit(2).nprobes(20).refine_factor(
|
|
10
|
|
).to_pandas()
|
|
# --8<-- [end:vector_index_query_ivf]
|
|
|
|
df = (
|
|
table.search(np.random.random((1536)))
|
|
.limit(2)
|
|
.nprobes(20)
|
|
.refine_factor(10)
|
|
.to_pandas()
|
|
)
|
|
assert len(df) == 2
|
|
|
|
|
|
def test_vector_index_nprobes(tmp_db):
|
|
from lancedb.index import IvfFlat
|
|
|
|
dim = 16
|
|
rng = np.random.default_rng(180)
|
|
data = [
|
|
{"id": i, "vector": rng.random(dim).astype(np.float32).tolist()}
|
|
for i in range(512)
|
|
]
|
|
table = tmp_db.create_table("vector_index_nprobes", data, mode="overwrite")
|
|
table.create_index(
|
|
"vector",
|
|
config=IvfFlat(distance_type="cosine", num_partitions=64, max_iterations=2),
|
|
)
|
|
query_vector = data[123]["vector"]
|
|
|
|
# --8<-- [start:vector_index_auto]
|
|
# Start with Auto.
|
|
auto = table.search(query_vector).limit(10).to_arrow()
|
|
# --8<-- [end:vector_index_auto]
|
|
|
|
# --8<-- [start:vector_index_search_effort]
|
|
# Search more partitions to try to improve recall.
|
|
higher_effort = table.search(query_vector).search_effort(0.75).limit(10).to_arrow()
|
|
# --8<-- [end:vector_index_search_effort]
|
|
|
|
# --8<-- [start:vector_index_fixed_probes]
|
|
# Use a fixed probe count.
|
|
fixed = table.search(query_vector).nprobes(20).limit(10).to_arrow()
|
|
# --8<-- [end:vector_index_fixed_probes]
|
|
|
|
# --8<-- [start:vector_index_nprobes]
|
|
# Let Auto choose between 10 and 50 partitions per segment, if available.
|
|
filtered = (
|
|
table.search(query_vector)
|
|
.minimum_nprobes(10)
|
|
.maximum_nprobes(50)
|
|
.where("id > 100")
|
|
.limit(5)
|
|
.to_arrow()
|
|
)
|
|
# --8<-- [end:vector_index_nprobes]
|
|
|
|
for result, expected_count in [
|
|
(auto, 10),
|
|
(higher_effort, 10),
|
|
(fixed, 10),
|
|
(filtered, 5),
|
|
]:
|
|
assert result.num_rows == expected_count
|
|
assert result["id"][0].as_py() == 123
|
|
assert np.isfinite(result["_distance"].to_numpy()).all()
|
|
exact_query = (
|
|
table.search(query_vector).distance_type("cosine").bypass_vector_index()
|
|
)
|
|
if result is filtered:
|
|
exact_query = exact_query.where("id > 100")
|
|
expected_ids = set(
|
|
exact_query.limit(expected_count).to_arrow()["id"].to_pylist()
|
|
)
|
|
recall = len(set(result["id"].to_pylist()) & expected_ids) / expected_count
|
|
assert recall >= 0.5
|
|
assert all(row_id > 100 for row_id in filtered["id"].to_pylist())
|
|
|
|
|
|
def test_vector_index_distance_range(tmp_db):
|
|
dim = 128
|
|
data = [
|
|
{"id": i, "keywords_embeddings": np.random.random(dim).tolist()}
|
|
for i in range(256)
|
|
]
|
|
table = tmp_db.create_table("vector_index_distance_range", data, mode="overwrite")
|
|
table.create_index(
|
|
metric="cosine",
|
|
vector_column_name="keywords_embeddings",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_distance_range]
|
|
# Only return results whose distance falls within [0.0, 0.5).
|
|
# Useful for near-duplicate detection or thresholded similarity search.
|
|
(
|
|
table.search(np.random.random(128))
|
|
.distance_range(lower_bound=0.0, upper_bound=0.5)
|
|
.limit(10)
|
|
.to_pandas()
|
|
)
|
|
# --8<-- [end:vector_index_distance_range]
|
|
|
|
|
|
def test_vector_index_bypass_recall(tmp_db):
|
|
dim = 128
|
|
data = [
|
|
{"id": i, "keywords_embeddings": np.random.random(dim).tolist()}
|
|
for i in range(256)
|
|
]
|
|
table = tmp_db.create_table("vector_index_bypass_recall", data, mode="overwrite")
|
|
table.create_index(
|
|
metric="cosine",
|
|
vector_column_name="keywords_embeddings",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_bypass_recall]
|
|
query = np.random.random(128)
|
|
k = 10
|
|
|
|
# Ground truth: flat (exhaustive) scan, ignoring the ANN index.
|
|
truth = set(table.search(query).bypass_vector_index().limit(k).to_pandas()["id"])
|
|
|
|
# ANN results with the current nprobes setting.
|
|
ann = set(table.search(query).nprobes(20).limit(k).to_pandas()["id"])
|
|
|
|
recall_at_k = len(truth & ann) / k
|
|
# --8<-- [end:vector_index_bypass_recall]
|
|
assert 0.0 <= recall_at_k <= 1.0
|
|
|
|
|
|
def test_vector_index_custom_name(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"vector_index_custom_name",
|
|
_make_vector_rows(512, 8, column="keywords_embeddings"),
|
|
mode="overwrite",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_custom_name]
|
|
# Override the default `{column}_idx` convention by passing `name=...`.
|
|
table.create_index(
|
|
metric="cosine",
|
|
vector_column_name="keywords_embeddings",
|
|
name="my_custom_index",
|
|
)
|
|
table.wait_for_index(["my_custom_index"])
|
|
print(table.index_stats("my_custom_index"))
|
|
# --8<-- [end:vector_index_custom_name]
|
|
|
|
assert table.index_stats("my_custom_index")
|
|
|
|
|
|
def test_vector_index_hnsw(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"vector_index_hnsw",
|
|
_make_vector_rows(64, 16),
|
|
mode="overwrite",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_build_hnsw]
|
|
table.create_index(index_type="IVF_HNSW_SQ")
|
|
# --8<-- [end:vector_index_build_hnsw]
|
|
|
|
# --8<-- [start:vector_index_query_hnsw]
|
|
tbl = table
|
|
tbl.search(np.random.random((16))).limit(2).to_pandas()
|
|
# --8<-- [end:vector_index_query_hnsw]
|
|
|
|
df = table.search(np.random.random((16))).limit(2).to_pandas()
|
|
assert len(df) == 2
|
|
|
|
|
|
def test_vector_index_binary(tmp_db):
|
|
table_name = "hamming-index-tbl"
|
|
ndim = 256
|
|
schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.int64()),
|
|
pa.field("vector", pa.list_(pa.uint8(), ndim // 8)),
|
|
]
|
|
)
|
|
|
|
# --8<-- [start:vector_index_binary_schema]
|
|
table = tmp_db.create_table(table_name, schema=schema, mode="overwrite")
|
|
# --8<-- [end:vector_index_binary_schema]
|
|
|
|
data = []
|
|
for i in range(64):
|
|
vector = np.random.randint(0, 2, size=ndim)
|
|
vector = np.packbits(vector)
|
|
data.append({"id": i, "vector": vector})
|
|
|
|
# --8<-- [start:vector_index_binary_add_data]
|
|
table.add(data)
|
|
# --8<-- [end:vector_index_binary_add_data]
|
|
|
|
# --8<-- [start:vector_index_binary_build_index]
|
|
table.create_index(
|
|
metric="hamming",
|
|
vector_column_name="vector",
|
|
index_type="IVF_FLAT",
|
|
)
|
|
# --8<-- [end:vector_index_binary_build_index]
|
|
|
|
# --8<-- [start:vector_index_binary_search]
|
|
query = np.random.randint(0, 2, size=ndim)
|
|
query = np.packbits(query)
|
|
df = table.search(query).metric("hamming").limit(10).to_pandas()
|
|
df.vector = df.vector.apply(np.unpackbits)
|
|
# --8<-- [end:vector_index_binary_search]
|
|
|
|
assert not df.empty
|
|
|
|
|
|
def test_quantization_custom_params(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"quantization-custom-params",
|
|
_make_vector_rows(256, 64),
|
|
mode="overwrite",
|
|
)
|
|
|
|
# --8<-- [start:quantization_custom_params]
|
|
table.create_index(
|
|
index_type="IVF_RQ",
|
|
num_bits=2,
|
|
max_iterations=100,
|
|
sample_rate=512,
|
|
)
|
|
# --8<-- [end:quantization_custom_params]
|
|
|
|
assert table.list_indices()
|
|
|
|
|
|
def test_vector_index_check_status(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"vector_index_check_status",
|
|
_make_vector_rows(512, 8, column="keywords_embeddings"),
|
|
mode="overwrite",
|
|
)
|
|
table.create_index(
|
|
metric="cosine",
|
|
vector_column_name="keywords_embeddings",
|
|
)
|
|
|
|
# --8<-- [start:vector_index_check_status]
|
|
index_name = "keywords_embeddings_idx"
|
|
table.wait_for_index([index_name])
|
|
print(table.index_stats(index_name))
|
|
# --8<-- [end:vector_index_check_status]
|
|
|
|
assert table.index_stats(index_name)
|
|
|
|
|
|
def test_scalar_index_build(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"scalar_index_build",
|
|
[
|
|
{"book_id": 1, "publisher": "A", "vector": [0.1, 0.2]},
|
|
{"book_id": 2, "publisher": "B", "vector": [0.2, 0.3]},
|
|
],
|
|
mode="overwrite",
|
|
)
|
|
db = tmp_db
|
|
# --8<-- [start:scalar_index_build]
|
|
tbl = db.open_table("scalar_index_build")
|
|
tbl.create_scalar_index("book_id")
|
|
tbl.create_scalar_index("publisher", index_type="BITMAP")
|
|
# --8<-- [end:scalar_index_build]
|
|
|
|
assert tbl.list_indices()
|
|
|
|
|
|
def test_scalar_index_wait(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"scalar_index_wait",
|
|
[{"label": "fiction"}],
|
|
mode="overwrite",
|
|
)
|
|
table.create_scalar_index("label")
|
|
|
|
# --8<-- [start:scalar_index_wait]
|
|
index_name = "label_idx"
|
|
table.wait_for_index([index_name])
|
|
# --8<-- [end:scalar_index_wait]
|
|
|
|
assert table.list_indices()
|
|
|
|
|
|
def test_scalar_index_optimize(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"scalar_index_optimize",
|
|
[{"vector": [7.0, 8.0], "book_id": 3}],
|
|
mode="overwrite",
|
|
)
|
|
|
|
# --8<-- [start:scalar_index_optimize]
|
|
table.add([{"vector": [7, 8], "book_id": 4}])
|
|
table.optimize()
|
|
# --8<-- [end:scalar_index_optimize]
|
|
|
|
result = table.search().where("book_id = 4").limit(10).to_pandas()
|
|
assert len(result) == 1
|
|
|
|
|
|
def test_scalar_index_filter(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"books",
|
|
[
|
|
{"vector": [1.1, 1.2], "book_id": 1},
|
|
{"vector": [2.1, 2.2], "book_id": 2},
|
|
],
|
|
mode="overwrite",
|
|
)
|
|
db = tmp_db
|
|
# --8<-- [start:scalar_index_filter]
|
|
table = db.open_table("books")
|
|
result = table.search().where("book_id = 2").limit(10).to_pandas()
|
|
# --8<-- [end:scalar_index_filter]
|
|
|
|
assert len(result) == 1
|
|
|
|
|
|
def test_scalar_index_prefilter(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"book_with_embeddings",
|
|
[
|
|
{"vector": [1.2, 1.3], "book_id": 1},
|
|
{"vector": [4.2, 4.3], "book_id": 2},
|
|
],
|
|
mode="overwrite",
|
|
)
|
|
db = tmp_db
|
|
# --8<-- [start:scalar_index_prefilter]
|
|
table = db.open_table("book_with_embeddings")
|
|
table.search([1.2] * 2).where("book_id != 3").limit(10).to_pandas()
|
|
# --8<-- [end:scalar_index_prefilter]
|
|
|
|
result = table.search([1.2] * 2).where("book_id != 3").limit(10).to_pandas()
|
|
assert len(result) == 2
|
|
|
|
|
|
def test_scalar_index_uuid(tmp_db):
|
|
# --8<-- [start:scalar_index_uuid_type]
|
|
import pyarrow as pa
|
|
# --8<-- [end:scalar_index_uuid_type]
|
|
|
|
# --8<-- [start:scalar_index_uuid_data]
|
|
def generate_random_names():
|
|
base_names = ["Alice", "Bob", "Carla", "David", "Eve", "Frank", "Grace"]
|
|
letter = random.choice(string.ascii_uppercase)
|
|
return f"{random.choice(base_names)} {letter}."
|
|
|
|
def generate_uuids(num_items):
|
|
return [uuid.uuid4().bytes for _ in range(num_items)]
|
|
|
|
# Generate some UUIDs and random names
|
|
n = 7
|
|
uuids = generate_uuids(n)
|
|
names = [generate_random_names() for _ in range(n)]
|
|
# --8<-- [end:scalar_index_uuid_data]
|
|
|
|
db = tmp_db
|
|
# --8<-- [start:scalar_index_uuid_table]
|
|
table_name = "index-on-uuid"
|
|
|
|
uuid_array = pa.array(uuids, pa.uuid())
|
|
name_array = pa.array(names, pa.string())
|
|
schema = pa.schema(
|
|
[
|
|
pa.field("id", pa.uuid()),
|
|
pa.field("name", pa.string()),
|
|
]
|
|
)
|
|
data_table = pa.Table.from_arrays([uuid_array, name_array], schema=schema)
|
|
table = db.create_table(table_name, data=data_table, mode="overwrite")
|
|
# --8<-- [end:scalar_index_uuid_table]
|
|
|
|
# --8<-- [start:scalar_index_uuid_wait]
|
|
index_name = "id_idx"
|
|
table.create_scalar_index("id")
|
|
table.wait_for_index([index_name])
|
|
# --8<-- [end:scalar_index_uuid_wait]
|
|
|
|
# --8<-- [start:scalar_index_uuid_upsert]
|
|
new_users = [
|
|
{"id": uuid.uuid4().bytes, "name": "Hannah D."},
|
|
{"id": uuid.uuid4().bytes, "name": "Ian B."},
|
|
]
|
|
# Insert or update using the UUID index
|
|
table.merge_insert(
|
|
"id"
|
|
).when_matched_update_all().when_not_matched_insert_all().execute(new_users)
|
|
# --8<-- [end:scalar_index_uuid_upsert]
|
|
|
|
assert table.list_indices()
|
|
result = table.search().limit(100).to_pandas()
|
|
assert len(result) == n + len(new_users)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_scalar_index_nested_fields(mem_db_async):
|
|
db = mem_db_async
|
|
|
|
# --8<-- [start:scalar_index_nested_fields]
|
|
import pyarrow as pa
|
|
from lancedb.index import BTree
|
|
|
|
metadata_type = pa.struct(
|
|
[
|
|
pa.field("user_id", pa.int32()),
|
|
pa.field("user.id", pa.int32()),
|
|
]
|
|
)
|
|
data = pa.Table.from_arrays(
|
|
[
|
|
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,
|
|
),
|
|
],
|
|
names=["user_id", "metadata"],
|
|
)
|
|
table = await db.create_table("nested_scalar_index", data)
|
|
|
|
# Index a nested struct field.
|
|
await table.create_index(
|
|
"metadata.user_id", config=BTree(), name="nested_user_id_idx"
|
|
)
|
|
|
|
# Escape literal dots inside a segment with backticks.
|
|
await table.create_index(
|
|
"metadata.`user.id`", config=BTree(), name="escaped_user_id_idx"
|
|
)
|
|
|
|
# `columns` is returned as the canonical path you passed in.
|
|
for index in await table.list_indices():
|
|
print(index.name, index.columns)
|
|
# nested_user_id_idx ['metadata.user_id']
|
|
# escaped_user_id_idx ['metadata.`user.id`']
|
|
# --8<-- [end:scalar_index_nested_fields]
|
|
|
|
index_columns = {index.name: index.columns for index in await table.list_indices()}
|
|
assert index_columns["nested_user_id_idx"] == ["metadata.user_id"]
|
|
assert index_columns["escaped_user_id_idx"] == ["metadata.`user.id`"]
|
|
|
|
|
|
def test_fts_index_create(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"fts-index-create",
|
|
[{"text": "hello world", "vector": [0.1, 0.2]}],
|
|
mode="overwrite",
|
|
)
|
|
|
|
db = tmp_db
|
|
# --8<-- [start:fts_index_create]
|
|
table_name = "fts-index-create"
|
|
table = db.open_table(table_name)
|
|
table.create_fts_index("text")
|
|
# --8<-- [end:fts_index_create]
|
|
|
|
assert table.list_indices()
|
|
|
|
|
|
def test_fts_index_wait(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"fts-index-wait",
|
|
[{"text": "full text search"}],
|
|
mode="overwrite",
|
|
)
|
|
|
|
db = tmp_db
|
|
# --8<-- [start:fts_index_wait]
|
|
table_name = "fts-index-wait"
|
|
|
|
table = db.open_table(table_name)
|
|
table.create_fts_index("text")
|
|
|
|
index_name = "text_idx"
|
|
table.wait_for_index([index_name])
|
|
# --8<-- [end:fts_index_wait]
|
|
|
|
assert table.list_indices()
|
|
|
|
|
|
def test_fts_index_nested_field(tmp_db):
|
|
nested_schema = pa.struct([
|
|
pa.field("text", pa.string()),
|
|
pa.field("count", pa.int32()),
|
|
])
|
|
schema = pa.schema([
|
|
pa.field("id", pa.int64()),
|
|
pa.field("payload", nested_schema),
|
|
])
|
|
tmp_db.create_table(
|
|
"fts-index-nested",
|
|
pa.table(
|
|
{
|
|
"id": pa.array([1, 2], pa.int64()),
|
|
"payload": pa.array(
|
|
[
|
|
{"text": "Frodo was a happy puppy", "count": 1},
|
|
{"text": "puppy runs through the meadow", "count": 2},
|
|
],
|
|
type=nested_schema,
|
|
),
|
|
},
|
|
schema=schema,
|
|
),
|
|
mode="overwrite",
|
|
)
|
|
|
|
db = tmp_db
|
|
# --8<-- [start:fts_index_nested]
|
|
from lancedb.query import MatchQuery, PhraseQuery
|
|
|
|
table = db.open_table("fts-index-nested")
|
|
|
|
# Index a text leaf inside a struct column using a dotted path.
|
|
table.create_fts_index("payload.text", with_position=True)
|
|
|
|
# The same dotted path works in MatchQuery and PhraseQuery.
|
|
matches = (
|
|
table.search(MatchQuery("puppy", "payload.text")).limit(5).to_list()
|
|
)
|
|
phrases = (
|
|
table.search(PhraseQuery("puppy runs", "payload.text"))
|
|
.limit(5)
|
|
.to_list()
|
|
)
|
|
# --8<-- [end:fts_index_nested]
|
|
|
|
assert len(matches) > 0
|
|
assert all("puppy" in row["payload"]["text"] for row in matches)
|
|
assert len(phrases) > 0
|
|
assert all("puppy runs" in row["payload"]["text"] for row in phrases)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fts_index_async(tmp_path, monkeypatch):
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
# --8<-- [start:fts_index_async]
|
|
import asyncio
|
|
|
|
import lancedb
|
|
import polars as pl
|
|
from lancedb.index import FTS
|
|
|
|
data = pl.DataFrame(
|
|
{
|
|
"id": [1, 2],
|
|
"text": [
|
|
"His first language is spanish",
|
|
"Her first language is english",
|
|
],
|
|
}
|
|
)
|
|
|
|
async def main(data: pl.DataFrame):
|
|
uri = "ex_lancedb"
|
|
db = await lancedb.connect_async(uri)
|
|
tbl = await db.create_table("my_text", data=data, mode="overwrite")
|
|
|
|
await tbl.create_index("text", config=FTS(language="English"))
|
|
|
|
response = await tbl.search("spanish", query_type="fts")
|
|
result = await response.limit(1).to_polars()
|
|
print(result)
|
|
return result
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main(data))
|
|
# --8<-- [end:fts_index_async]
|
|
|
|
result = await main(data)
|
|
assert result.height == 1
|
|
|
|
|
|
def test_gpu_index_snippets(tmp_db, monkeypatch):
|
|
table = tmp_db.create_table(
|
|
"gpu_index",
|
|
_make_vector_rows(32, 8),
|
|
mode="overwrite",
|
|
)
|
|
|
|
calls = []
|
|
|
|
def fake_create_index(*args, **kwargs):
|
|
calls.append(kwargs)
|
|
return None
|
|
|
|
monkeypatch.setattr(table, "create_index", fake_create_index)
|
|
|
|
# --8<-- [start:gpu_index_cuda]
|
|
table.create_index(
|
|
num_partitions=256,
|
|
num_sub_vectors=96,
|
|
accelerator="cuda",
|
|
)
|
|
# --8<-- [end:gpu_index_cuda]
|
|
|
|
# --8<-- [start:gpu_index_mps]
|
|
table.create_index(
|
|
num_partitions=256,
|
|
num_sub_vectors=96,
|
|
accelerator="mps",
|
|
)
|
|
# --8<-- [end:gpu_index_mps]
|
|
|
|
assert calls[0]["accelerator"] == "cuda"
|
|
assert calls[1]["accelerator"] == "mps"
|
|
|
|
|
|
def test_reindexing_incremental(tmp_db):
|
|
table = tmp_db.create_table(
|
|
"reindexing_incremental",
|
|
[{"vector": [3.1, 4.1], "text": "Frodo was a happy puppy"}],
|
|
mode="overwrite",
|
|
)
|
|
db = tmp_db
|
|
# --8<-- [start:reindexing_incremental]
|
|
table = db.open_table("reindexing_incremental")
|
|
table.add([{"vector": [3.1, 4.1], "text": "Frodo was a happy puppy"}])
|
|
table.optimize()
|
|
# --8<-- [end:reindexing_incremental]
|
|
|
|
result = table.search().limit(10).to_pandas()
|
|
assert len(result) == 2
|