diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 8265ad3f793..7d3b2da561f 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -2797,6 +2797,10 @@ def create_index( accelerator="cuda" ) + Note: GPU acceleration is currently supported only for the ``IVF_PQ`` index + type. Providing an accelerator for other index types will fall back to CPU + index building. + References ---------- * `Faiss Index `_ @@ -2881,6 +2885,13 @@ def create_index( # Handle timing for various parts of accelerated builds timers = {} + if accelerator is not None and index_type != "IVF_PQ": + LOGGER.warning( + "Index type %s does not support GPU acceleration; falling back to CPU", + index_type, + ) + accelerator = None + if accelerator is not None: from .vector import ( one_pass_assign_ivf_pq_on_accelerator, diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index be0ac879fba..8be2a2b1c0e 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors +import logging import platform import random import string @@ -461,6 +462,27 @@ def test_create_index_unsupported_accelerator(tmp_path): ) +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(