diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index db7b4919d..8d841fa70 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -897,6 +897,9 @@ class IvfRq: target_partition_size: int, default is 8192 Target size of each partition. + accelerator: Optional[str], default None + Reserved for accelerator support. IVF RQ currently builds on CPU, + preserving all configured index options. """ distance_type: Literal["l2", "cosine", "dot"] = "l2" @@ -905,8 +908,7 @@ class IvfRq: max_iterations: int = 50 sample_rate: int = 256 target_partition_size: Optional[int] = None - # Name of the accelerator (e.g. "cuda") to use for IVF training. When set, - # create_index() dispatches to pylance to build the index on the accelerator. + # IVF RQ currently uses the native CPU builder even when this is set. accelerator: Optional[str] = None layered: Optional[bool] = None diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 68f4b95eb..d8325e20e 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -3609,8 +3609,13 @@ class LanceTable(Table): # New API: metric is the column name column = metric - # Check if config has accelerator set and dispatch to pylance - if config is not None and hasattr(config, "accelerator"): + # IVF_RQ currently falls back to CPU in pylance. Keep its build here + # so the fallback preserves the native builder's layout options. + if ( + config is not None + and not isinstance(config, IvfRq) + and hasattr(config, "accelerator") + ): acc = getattr(config, "accelerator", None) if acc is not None: # Dispatch to pylance for GPU acceleration diff --git a/python/python/tests/test_index.py b/python/python/tests/test_index.py index 19bddacd2..174743b96 100644 --- a/python/python/tests/test_index.py +++ b/python/python/tests/test_index.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors from datetime import timedelta +import json import random from typing import get_args, get_type_hints +import numpy as np import pyarrow as pa import pytest import pytest_asyncio -from lancedb import AsyncConnection, AsyncTable, connect_async +from lancedb import AsyncConnection, AsyncTable, connect, connect_async from lancedb.index import ( BTree, IvfFlat, @@ -446,6 +448,41 @@ async def test_create_ivfrq_index(some_table: AsyncTable, num_bits, layered): assert len(set(actual["id"].to_pylist()) & set(expected["id"].to_pylist())) >= 5 +@pytest.mark.parametrize("num_bits", [1, 2, 5]) +@pytest.mark.parametrize("layered", [None, False, True]) +def test_create_ivfrq_accelerator_preserves_layout(tmp_path, num_bits, layered): + from lance.file import LanceFileReader + + vectors = np.random.default_rng(4252).normal(size=(256, 32)).astype("float32") + data = pa.table( + { + "id": range(len(vectors)), + "vector": pa.FixedSizeListArray.from_arrays(pa.array(vectors.ravel()), 32), + } + ) + db = connect(tmp_path) + table = db.create_table("test", data.slice(0, 128)) + table.add(data.slice(128)) + table.create_index( + "vector", + config=IvfRq( + num_partitions=2, + num_bits=num_bits, + layered=layered, + accelerator="cuda", + ), + ) + + path = next((tmp_path / "test.lance" / "_indices").glob("*/auxiliary.idx")) + schema = LanceFileReader(str(path)).metadata().schema + metadata = json.loads(json.loads(schema.metadata[b"storage_metadata"])[0]) + assert metadata.get("layered", False) is (layered is True and num_bits > 2) + + expected = table.search(vectors[0]).bypass_vector_index().limit(10).to_arrow() + actual = table.search(vectors[0]).nprobes(2).refine_factor(26).limit(10).to_arrow() + assert len(set(actual["id"].to_pylist()) & set(expected["id"].to_pylist())) >= 5 + + @pytest.mark.asyncio async def test_create_hnswpq_index(some_table: AsyncTable): await some_table.create_index("vector", config=HnswPq(num_partitions=10))