# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright The Lance Authors

import logging

import os

import platform

import random

import shutil

import string

import tempfile

import time

import uuid

from pathlib import Path

from typing import Optional

import lance

import numpy as np

import pyarrow as pa

import pyarrow.compute as pc

import pytest

from lance import LanceDataset, LanceFragment

from lance.dataset import Index, VectorIndexReader

from lance.indices import IndexFileVersion, IndicesBuilder

from lance.util import validate_vector_index # noqa: E402

from lance.vector import vec_to_table # noqa: E402

def create_table(nvec=1000, ndim=128, nans=0, nullify=False, dtype=np.float32):

mat = np.random.randn(nvec, ndim)

if nans > 0:

nans_mat = np.empty((nans, ndim))

nans_mat[:] = np.nan

mat = np.concatenate((mat, nans_mat), axis=0)

mat = mat.astype(dtype)

price = np.random.rand(nvec + nans) * 100

def gen_str(n):

return "".join(random.choices(string.ascii_letters + string.digits, k=n))

meta = np.array([gen_str(100) for _ in range(nvec + nans)])

tbl = (

vec_to_table(data=mat)

.append_column("price", pa.array(price))

.append_column("meta", pa.array(meta))

.append_column("id", pa.array(range(nvec + nans)))

)

if nullify:

idx = tbl.schema.get_field_index("vector")

vecs = tbl[idx].to_pylist()

nullified = [vec if i % 2 == 0 else None for i, vec in enumerate(vecs)]

field = tbl.schema.field(idx)

vecs = pa.array(nullified, field.type)

tbl = tbl.set_column(idx, field, vecs)

return tbl

def create_multivec_table(

nvec=1000, nvec_per_row=5, ndim=128, nans=0, nullify=False, dtype=np.float32

):

mat = np.random.randn(nvec, nvec_per_row, ndim)

if nans > 0:

nans_mat = np.empty((nans, ndim))

nans_mat[:] = np.nan

mat = np.concatenate((mat, nans_mat), axis=0)

mat = mat.astype(dtype)

price = np.random.rand(nvec + nans) * 100

def gen_str(n):

return "".join(random.choices(string.ascii_letters + string.digits, k=n))

meta = np.array([gen_str(100) for _ in range(nvec + nans)])

multi_vec_type = pa.list_(pa.list_(pa.float32(), ndim))

tbl = pa.Table.from_arrays(

[

pa.array((mat[i].tolist() for i in range(nvec)), type=multi_vec_type),

],

schema=pa.schema(

[

pa.field("vector", pa.list_(pa.list_(pa.float32(), ndim))),

]

),

)

tbl = (

tbl.append_column("price", pa.array(price))

.append_column("meta", pa.array(meta))

.append_column("id", pa.array(range(nvec + nans)))

)

if nullify:

idx = tbl.schema.get_field_index("vector")

vecs = tbl[idx].to_pylist()

nullified = [vec if i % 2 == 0 else None for i, vec in enumerate(vecs)]

field = tbl.schema.field(idx)

vecs = pa.array(nullified, field.type)

tbl = tbl.set_column(idx, field, vecs)

return tbl

@pytest.fixture()

def dataset(tmp_path):

tbl = create_table()

yield lance.write_dataset(tbl, tmp_path)

@pytest.fixture()

def indexed_dataset(tmp_path):

tbl = create_table()

dataset = lance.write_dataset(tbl, tmp_path)

yield dataset.create_index(

"vector", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=16

)

@pytest.fixture()

def multivec_dataset():

tbl = create_multivec_table()

yield lance.write_dataset(tbl, "memory://")

@pytest.fixture()

def indexed_multivec_dataset(multivec_dataset):

yield multivec_dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

metric="cosine",

)

def run(ds, q=None, assert_func=None):

if q is None:

q = np.random.randn(128)

project = [None, ["price"], ["vector", "price"], ["vector", "meta", "price"]]

refine = [None, 1, 2]

filters = [None, pc.field("price") > 50.0]

times = []

for columns in project:

expected_columns = []

if columns is None:

expected_columns.extend(ds.schema.names)

else:

expected_columns.extend(columns)

# TODO: _distance shouldn't be returned by default either

if "_distance" not in expected_columns:

expected_columns.append("_distance")

for filter_ in filters:

for rf in refine:

start = time.time()

rs = ds.to_table(

columns=columns,

nearest={

"column": "vector",

"q": q,

"k": 15,

"nprobes": 1,

"refine_factor": rf,

},

filter=filter_,

)

end = time.time()

times.append(end - start)

assert rs.column_names == expected_columns

if filter_ is not None:

inmem = pa.dataset.dataset(rs)

assert len(inmem.to_table(filter=filter_)) == len(rs)

else:

assert len(rs) == 15

distances = rs["_distance"].to_numpy()

assert (distances.max() - distances.min()) > 1e-6

if assert_func is not None:

assert_func(rs)

return times

def test_flat(dataset):

run(dataset)

def test_ann(indexed_dataset):

run(indexed_dataset)

@pytest.mark.parametrize(

"fixture_name,index_type,index_params,similarity_threshold",

[

("dataset", "IVF_FLAT", {"num_partitions": 4}, 0.80),

(

"indexed_dataset",

"IVF_PQ",

{"num_partitions": 4, "num_sub_vectors": 16},

0.80,

),

("dataset", "IVF_SQ", {"num_partitions": 4}, 0.80),

],

)

def test_distributed_vector(

request, fixture_name, index_type, index_params, similarity_threshold

):

ds = request.getfixturevalue(fixture_name)

q = np.random.randn(128).astype(np.float32)

assert_distributed_vector_consistency(

ds.to_table(),

"vector",

index_type=index_type,

index_params=index_params,

queries=[q],

topk=10,

world=2,

similarity_metric="recall",

similarity_threshold=similarity_threshold,

)

def test_rowid_order(indexed_dataset):

rs = indexed_dataset.to_table(

columns=["meta"],

with_row_id=True,

nearest={

"column": "vector",

"q": np.random.randn(128),

"k": 10,

"use_index": False,

},

limit=10,

)

assert rs.schema[0].name == "meta"

assert rs.schema[1].name == "_distance"

assert rs.schema[2].name == "_rowid"

def test_ann_append(tmp_path):

tbl = create_table()

dataset = lance.write_dataset(tbl, tmp_path)

dataset = dataset.create_index(

"vector", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=16

)

new_data = create_table(nvec=10)

dataset = lance.write_dataset(new_data, dataset.uri, mode="append")

q = new_data["vector"][0].as_py()

def func(rs: pa.Table):

if "vector" not in rs:

return

assert rs["vector"][0].as_py() == q

run(dataset, q=np.array(q), assert_func=func)

def test_invalid_subvectors(tmp_path):

tbl = create_table()

dataset = lance.write_dataset(tbl, tmp_path)

with pytest.raises(

ValueError,

match="dimension .* must be divisible by num_sub_vectors",

):

dataset.create_index(

"vector", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=15

)

@pytest.mark.cuda

def test_invalid_subvectors_cuda(tmp_path):

tbl = create_table()

dataset = lance.write_dataset(tbl, tmp_path)

with pytest.raises(

ValueError,

match="dimension .* must be divisible by num_sub_vectors",

):

dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=15,

accelerator="cuda",

)

@pytest.mark.cuda

def test_f16_cuda(tmp_path):

tbl = create_table(dtype=np.float16)

dataset = lance.write_dataset(tbl, tmp_path)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="cuda",

one_pass_ivfpq=True,

)

validate_vector_index(dataset, "vector")

@pytest.mark.parametrize(

"index_file_version", [IndexFileVersion.V3, IndexFileVersion.LEGACY]

)

def test_index_with_nans(tmp_path, index_file_version):

# 1024 rows, the entire table should be sampled

tbl = create_table(nvec=1000, nans=24)

dataset = lance.write_dataset(tbl, tmp_path)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

index_file_version=index_file_version,

)

idx_stats = dataset.stats.index_stats("vector_idx")

assert idx_stats["indices"][0]["index_file_version"] == index_file_version

validate_vector_index(dataset, "vector")

@pytest.mark.parametrize(

"index_file_version", [IndexFileVersion.V3, IndexFileVersion.LEGACY]

)

def test_torch_index_with_nans(tmp_path, index_file_version):

torch = pytest.importorskip("torch")

# 1024 rows, the entire table should be sampled

tbl = create_table(nvec=1000, nans=24)

dataset = lance.write_dataset(tbl, tmp_path)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator=torch.device("cpu"),

one_pass_ivfpq=True,

index_file_version=index_file_version,

)

idx_stats = dataset.stats.index_stats("vector_idx")

assert idx_stats["indices"][0]["index_file_version"] == index_file_version

validate_vector_index(dataset, "vector")

def test_index_with_no_centroid_movement(tmp_path):

torch = pytest.importorskip("torch")

# this test makes the centroids essentially [1..]

# this makes sure the early stop condition in the index building code

# doesn't do divide by zero

mat = np.concatenate([np.ones((256, 32))])

tbl = vec_to_table(data=mat)

dataset = lance.write_dataset(tbl, tmp_path)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=1,

num_sub_vectors=4,

accelerator=torch.device("cpu"),

)

validate_vector_index(dataset, "vector")

def test_index_with_pq_codebook(tmp_path):

tbl = create_table(nvec=1024, ndim=128)

dataset = lance.write_dataset(tbl, tmp_path)

pq_codebook = np.random.randn(4, 256, 128 // 4).astype(np.float32)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=1,

num_sub_vectors=4,

ivf_centroids=np.random.randn(1, 128).astype(np.float32),

pq_codebook=pq_codebook,

)

validate_vector_index(dataset, "vector", refine_factor=10, pass_threshold=0.99)

pq_codebook = pa.FixedShapeTensorArray.from_numpy_ndarray(pq_codebook)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=1,

num_sub_vectors=4,

ivf_centroids=np.random.randn(1, 128).astype(np.float32),

pq_codebook=pq_codebook,

replace=True,

)

validate_vector_index(dataset, "vector", refine_factor=10, pass_threshold=0.99)

@pytest.mark.cuda

@pytest.mark.parametrize("nullify", [False, True])

def test_create_index_using_cuda(tmp_path, nullify):

tbl = create_table(nullify=nullify)

dataset = lance.write_dataset(tbl, tmp_path)

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="cuda",

)

q = np.random.randn(128)

expected = dataset.to_table(

columns=["id"],

nearest={

"column": "vector",

"q": q,

"k": 10, # Use non-default k

},

)["id"].to_numpy()

assert len(expected) == 10

dataset = dataset.create_index(

"vector",

index_type="IVF_PQ",

metric="cosine",

num_partitions=4,

num_sub_vectors=16,

accelerator="cuda",

)

q = np.random.randn(128)

expected = dataset.to_table(

columns=["id"],

nearest={

"column": "vector",

"q": q,

"k": 10, # Use non-default k

},

)["id"].to_numpy()

assert len(expected) == 10

def test_create_index_unsupported_accelerator(tmp_path):

# Even attempting to use an accelerator will trigger torch import

# so make sure it's available

pytest.importorskip("torch")

tbl = create_table()

dataset = lance.write_dataset(tbl, tmp_path)

with pytest.raises(ValueError):

dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="no-supported",

)

with pytest.raises(ValueError):

dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="0cuda",

)

with pytest.raises(ValueError):

dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="cuda-0",

)

with pytest.raises(ValueError):

dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="cuda:",

)

with pytest.raises(ValueError):

dataset.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

accelerator="cuda:abc",

)

def test_create_index_accelerator_fallback(tmp_path, caplog):

tbl = create_table()

dataset = lance.write_dataset(tbl, tmp_path)

with caplog.at_level(logging.WARNING):

dataset = dataset.create_index(

"vector",

index_type="IVF_HNSW_SQ",

num_partitions=4,

accelerator="cuda",

)

indices = dataset.list_indices()

assert len(indices) == 1

assert indices[0]["type"] == "IVF_HNSW_SQ"

assert any(

"does not support GPU acceleration; falling back to CPU" in record.message

for record in caplog.records

)

def test_use_index(dataset, tmp_path):

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=16

)

q = np.random.randn(128)

expected = dataset.to_table(

columns=["id"],

nearest={

"column": "vector",

"q": q,

"k": 12, # Use non-default k

},

)["id"].to_numpy()

actual = ann_ds.to_table(

columns=["id"],

nearest={"column": "vector", "q": q, "k": 12, "use_index": False},

)["id"].to_numpy()

assert np.all(expected == actual)

# Can omit k but provide limit

actual = ann_ds.to_table(

columns=["id"],

nearest={"column": "vector", "q": q, "use_index": False},

limit=12,

)["id"].to_numpy()

assert np.all(expected == actual)

def test_nearest_errors(dataset, tmp_path):

import pandas as pd

with pytest.raises(ValueError, match="does not match index column size"):

dataset.to_table(

columns=["id"],

nearest={"column": "vector", "q": np.random.randn(127), "k": 10},

)

df = pd.DataFrame({"a": [5], "b": [10]})

ds = lance.write_dataset(pa.Table.from_pandas(df), tmp_path / "dataset.lance")

with pytest.raises(TypeError, match="must be a vector"):

ds.to_table(nearest={"column": "a", "q": np.random.randn(128), "k": 10})

def test_has_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector", index_type="IVF_PQ", num_partitions=4, num_sub_vectors=16

)

assert ann_ds.has_index

assert ann_ds.list_indices()[0]["fields"] == ["vector"]

def test_index_type(dataset, tmp_path):

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

replace=True,

)

assert ann_ds.list_indices()[0]["type"] == "IVF_PQ"

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_HNSW_SQ",

num_partitions=4,

num_sub_vectors=16,

replace=True,

)

assert ann_ds.list_indices()[0]["type"] == "IVF_HNSW_SQ"

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_HNSW_PQ",

num_partitions=4,

num_sub_vectors=16,

replace=True,

)

assert ann_ds.list_indices()[0]["type"] == "IVF_HNSW_PQ"

def test_create_dot_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=16,

metric="dot",

)

assert ann_ds.has_index

def test_create_4bit_ivf_pq_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=1,

num_sub_vectors=16,

num_bits=4,

metric="l2",

)

index = ann_ds.stats.index_stats("vector_idx")

assert index["indices"][0]["sub_index"]["nbits"] == 4

def test_create_ivf_pq_with_target_partition_size(dataset, tmp_path):

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_PQ",

num_sub_vectors=16,

target_partition_size=1000,

)

assert ann_ds.stats.index_stats("vector_idx")["indices"][0]["num_partitions"] == 1

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_PQ",

num_sub_vectors=16,

target_partition_size=500,

replace=True,

)

assert ann_ds.stats.index_stats("vector_idx")["indices"][0]["num_partitions"] == 2

# setting both num_partitions and target_partition_size will use num_partitions

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_PQ",

num_sub_vectors=16,

num_partitions=2,

target_partition_size=1000,

replace=True,

)

assert ann_ds.stats.index_stats("vector_idx")["indices"][0]["num_partitions"] == 2

def test_index_size_stats(tmp_path: Path):

num_rows = 512

dims = 32

schema = pa.schema([pa.field("a", pa.list_(pa.float32(), dims), False)])

values = pc.random(num_rows * dims).cast("float32")

table = pa.Table.from_pydict(

{"a": pa.FixedSizeListArray.from_arrays(values, dims)}, schema=schema

)

base_dir = tmp_path / "test"

dataset = lance.write_dataset(table, base_dir)

index_name = "vec_idx"

dataset.create_index(

"a", "IVF_PQ", name=index_name, num_partitions=2, num_sub_vectors=1

)

# Expect to see non-zero sizes here but all sizes are zero

stats = dataset.stats.index_stats(index_name)

stats = stats["indices"][0]

assert stats["partitions"][0]["size"] + stats["partitions"][1]["size"] == num_rows

def test_ivf_flat_over_binary_vector(tmp_path):

dim = 128

nvec = 1000

data = np.random.randint(0, 256, (nvec, dim // 8)).tolist()

array = pa.array(data, type=pa.list_(pa.uint8(), dim // 8))

tbl = pa.Table.from_pydict({"vector": array})

ds = lance.write_dataset(tbl, tmp_path)

ds.create_index("vector", index_type="IVF_FLAT", num_partitions=4, metric="hamming")

stats = ds.stats.index_stats("vector_idx")

assert stats["indices"][0]["metric_type"] == "hamming"

assert stats["index_type"] == "IVF_FLAT"

query = np.random.randint(0, 256, dim // 8).astype(np.uint8)

ds.to_table(

nearest={

"column": "vector",

"q": query,

"k": 10,

"metric": "hamming",

}

)

def test_ivf_flat_respects_index_metric_binary(tmp_path):

# Searching with binary vectors should default to hamming distance

table = pa.Table.from_pydict(

{

"vector": pa.array([[0], [128], [255]], type=pa.list_(pa.uint8(), 1)),

"id": pa.array([0, 1, 2], type=pa.int32()),

}

)

ds = lance.write_dataset(table, tmp_path)

ds = ds.create_index(

"vector",

index_type="IVF_FLAT",

num_partitions=1,

metric="hamming",

)

query = np.array([128], dtype=np.uint8)

# Search should succeed and use the index's Hamming metric.

indexed = ds.scanner(

columns=["id"],

nearest={

"column": "vector",

"q": query,

"k": 3,

},

)

plan = indexed.explain_plan()

indexed = indexed.to_table()

# Should succeed even though user asked for L2 (index metric is used).

assert indexed["id"].to_pylist() == [1, 0, 2]

assert "metric=Hamming" in plan

assert "metric=L2" not in plan

def test_bruteforce_uses_user_metric(tmp_path):

# Even if an index exists, a brute-force scan (use_index=False) should

# respect the user-specified metric instead of the index metric.

vectors = np.array(

[

[10.0, 10.0], # Large magnitude, best under dot product

[-1.0, -1.0],

[1.0, 1.0], # Closest under L2

],

dtype=np.float32,

)

table = pa.Table.from_pydict(

{

"vector": pa.array(vectors.tolist(), type=pa.list_(pa.float32(), 2)),

"id": pa.array([0, 1, 2], type=pa.int32()),

}

)

ds = lance.write_dataset(table, tmp_path)

# Build an index with L2 metric.

ds = ds.create_index(

"vector",

index_type="IVF_FLAT",

num_partitions=1,

metric="l2",

)

query = np.array([1.0, 1.0], dtype=np.float32)

# Brute-force search should honor the requested dot metric (not the index's L2).

brute_force = ds.to_table(

columns=["id"],

nearest={

"column": "vector",

"q": query,

"k": 3,

"metric": "dot",

"use_index": False,

},

)

# Under dot product the largest magnitude vector ranks first; under L2 it is last.

assert brute_force["id"].to_pylist() == [0, 2, 1]

def test_create_ivf_sq_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_SQ",

num_partitions=4,

)

assert ann_ds.list_indices()[0]["fields"] == ["vector"]

def test_create_ivf_rq_index():

ds = lance.write_dataset(create_table(), "memory://")

ds = ds.create_index(

"vector",

index_type="IVF_RQ",

num_partitions=4,

num_bits=1,

)

assert ds.list_indices()[0]["fields"] == ["vector"]

with pytest.raises(

NotImplementedError,

match="Creating empty vector indices with train=False is not yet implemented",

):

ds.delete("id>=0")

ds = ds.create_index(

"vector",

index_type="IVF_RQ",

num_partitions=4,

num_bits=1,

replace=True,

)

zero_vectors = np.zeros((1000, 128)).astype(np.float32).tolist()

tbl = pa.Table.from_pydict(

{"vector": pa.array(zero_vectors, type=pa.list_(pa.float32(), 128))}

)

ds = lance.write_dataset(tbl, "memory://", mode="overwrite")

ds = ds.create_index(

"vector",

index_type="IVF_RQ",

num_partitions=4,

num_bits=1,

)

res = ds.to_table(

nearest={

"column": "vector",

"q": np.zeros(128),

"k": 10,

}

)

assert res.num_rows == 10

assert res["_distance"].to_numpy().min() == 0.0

assert res["_distance"].to_numpy().max() == 0.0

def test_create_ivf_hnsw_pq_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_HNSW_PQ",

num_partitions=4,

num_sub_vectors=16,

)

assert ann_ds.list_indices()[0]["fields"] == ["vector"]

def test_create_ivf_hnsw_sq_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_HNSW_SQ",

num_partitions=4,

num_sub_vectors=16,

)

assert ann_ds.list_indices()[0]["fields"] == ["vector"]

def test_create_ivf_hnsw_flat_index(dataset, tmp_path):

assert not dataset.has_index

ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")

ann_ds = ann_ds.create_index(

"vector",

index_type="IVF_HNSW_FLAT",

num_partitions=4,

num_sub_vectors=16,

)

assert ann_ds.list_indices()[0]["fields"] == ["vector"]

def test_multivec_ann(indexed_multivec_dataset: lance.LanceDataset):

query = np.random.rand(5, 128)

results = indexed_multivec_dataset.scanner(

nearest={"column": "vector", "q": query, "k": 100}

).to_table()

assert results.num_rows == 100

assert results["vector"].type == pa.list_(pa.list_(pa.float32(), 128))

assert len(results["vector"][0]) == 5

# query with single vector also works

query = np.random.rand(128)

results = indexed_multivec_dataset.to_table(

nearest={"column": "vector", "q": query, "k": 100}

)

# we don't verify the number of results here,

# because for multivector, it's not guaranteed to return k results

assert results["vector"].type == pa.list_(pa.list_(pa.float32(), 128))

assert len(results["vector"][0]) == 5

query = [query, query]

doubled_results = indexed_multivec_dataset.to_table(

nearest={"column": "vector", "q": query, "k": 100}

)

assert len(results) == len(doubled_results)

for i in range(len(results)):

assert (

results["_distance"][i].as_py() * 2

== doubled_results["_distance"][i].as_py()

)

# query with a vector that dim not match

query = np.random.rand(256)

with pytest.raises(ValueError, match="does not match index column size"):

indexed_multivec_dataset.to_table(

nearest={"column": "vector", "q": query, "k": 100}

)

# query with a list of vectors that some dim not match

query = [np.random.rand(128)] * 5 + [np.random.rand(256)]

with pytest.raises(ValueError, match="All query vectors must have the same length"):

indexed_multivec_dataset.to_table(

nearest={"column": "vector", "q": query, "k": 100}

)

def test_pre_populated_ivf_centroids(dataset, tmp_path: Path):

centroids = np.random.randn(5, 128).astype(np.float32) # IVF5

dataset_with_index = dataset.create_index(

["vector"],

index_type="IVF_PQ",

metric="cosine",

ivf_centroids=centroids,

num_partitions=5,

num_sub_vectors=8,

)

q = np.random.randn(128)

actual = dataset_with_index.to_table(

columns=["id"],

nearest={"column": "vector", "q": q, "k": 10, "use_index": False},

)["id"].to_numpy()

assert len(actual) == 10

index_meta = dataset_with_index.list_indices()[0]

index_uuid = index_meta["uuid"]

assert len(index_uuid) == 36

assert index_meta["fragment_ids"] == {0}

expected_filepath = str(tmp_path / "_indices" / index_uuid / "index.idx")

if platform.system() == "Windows":

expected_filepath = expected_filepath.replace("\\", "/")

expected_statistics = {

"index_type": "IVF_PQ",

"uuid": index_uuid,

"uri": expected_filepath,

"metric_type": "cosine",

"num_partitions": 5,

"sub_index": {

"dimension": 128,

"index_type": "PQ",

"metric_type": "l2",

"nbits": 8,

"num_sub_vectors": 8,

"transposed": True,

},

"index_file_version": IndexFileVersion.V3,

}

with pytest.raises(KeyError, match='Index "non-existent_idx" not found'):

# increase 1 miss of index_cache.metadata_cache

assert dataset_with_index.stats.index_stats("non-existent_idx")

with pytest.raises(KeyError, match='Index "" not found'):

# increase 1 miss of index_cache.metadata_cache

assert dataset_with_index.stats.index_stats("")

with pytest.raises(TypeError):

dataset_with_index.stats.index_stats()

# increase 1 hit of index_cache.metadata_cache

actual_statistics = dataset_with_index.stats.index_stats("vector_idx")

assert actual_statistics["num_indexed_rows"] == 1000

assert actual_statistics["num_unindexed_rows"] == 0

idx_stats = actual_statistics["indices"][0]

partitions = idx_stats.pop("partitions")

idx_stats.pop("centroids")

idx_stats.pop("loss")

assert idx_stats == expected_statistics

assert len(partitions) == 5

partition_keys = {"size"}

assert all([partition_keys == set(p.keys()) for p in partitions])

def test_optimize_index(dataset, tmp_path):

dataset_uri = tmp_path / "dataset.lance"

assert not dataset.has_index

ds = lance.write_dataset(dataset.to_table(), dataset_uri)

ds = ds.create_index(

"vector",

index_type="IVF_PQ",

num_partitions=4,

num_sub_vectors=2,

)