fix: build IVF_HNSW_SQ graphs with CAGRA

python/python/lance/dataset.py

@@ -103,6 +103,42 @@
         np.ndarray,
         Iterable[Union[float, Iterable[float]]],
     ]
+
+
+def _locate_cuvs_library() -> Optional[str]:
+    """Find a cuVS C library that can derive CAGRA params from HNSW params."""
+    import ctypes
+    from importlib.util import find_spec
+
+    try:
+        spec = find_spec("libcuvs")
+    except ModuleNotFoundError:
+        return None
+    if spec is None:
+        return None
+
+    package_dirs = [Path(path) for path in spec.submodule_search_locations or []]
+    if spec.origin is not None:
+        package_dirs.append(Path(spec.origin).parent)
+
+    candidates = {
+        candidate
+        for package_dir in package_dirs
+        for candidate in (package_dir / "lib64").glob("libcuvs_c.so*")
+    }
+    for candidate in sorted(
+        candidates,
+        key=lambda path: (path.name != "libcuvs_c.so", path.name),
+    ):
+        try:
+            library = ctypes.CDLL(str(candidate))
+        except OSError:
+            continue
+        if hasattr(library, "cuvsCagraIndexParamsFromHnswParams"):
+            return str(candidate)
+    return None
+
+
 LANCE_COMMIT_MESSAGE_KEY = "__lance_commit_message"
 # Mirrors Rust's `lance::dataset::DEFAULT_COMMIT_TIMEOUT`; keep the two in sync.
 _DEFAULT_COMMIT_TIMEOUT = timedelta(minutes=30)
@@ -3779,9 +3815,48 @@ def _create_index_impl(
                 f"Only {valid_index_types} index types supported. Got {index_type}"
             )
 
+        require_cagra = kwargs.pop("_require_cagra", False)
+        if not isinstance(require_cagra, bool):
+            raise TypeError(f"_require_cagra must be bool, got {type(require_cagra)}")
+        if require_cagra and (
+            index_type != "IVF_HNSW_SQ"
+            or accelerator is None
+            or str(accelerator).lower() != "cuda"
+        ):
+            raise ValueError(
+                "_require_cagra requires index_type='IVF_HNSW_SQ' "
+                "and accelerator='cuda'"
+            )
+
         # Handle timing for various parts of accelerated builds
         timers = {}
-        if accelerator is not None and index_type != "IVF_PQ":
+        cagra_library = None
+        if accelerator is not None and index_type == "IVF_HNSW_SQ":
+            if str(accelerator).lower() != "cuda":
+                LOGGER.warning(
+                    "IVF_HNSW_SQ CAGRA acceleration supports only accelerator='cuda'; "
+                    "falling back to CPU"
+                )
+            else:
+                cagra_library = _locate_cuvs_library()
+                if cagra_library is None:
+                    if require_cagra:
+                        raise RuntimeError(
+                            "CAGRA was required, but a compatible cuVS library "
+                            "was not found"
+                        )
+                    LOGGER.warning(
+                        "A compatible cuVS CAGRA library was not found; falling back "
+                        "to CPU for IVF_HNSW_SQ"
+                    )
+                else:
+                    kwargs["cuvs_library"] = cagra_library
+                    if require_cagra:
+                        kwargs["_require_cagra"] = True
+            # CAGRA accelerates only the per-partition HNSW graph build. The
+            # one-pass Torch path below applies exclusively to IVF_PQ.
+            accelerator = None
+        elif 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,
@@ -3790,10 +3865,10 @@ def _create_index_impl(
 
         # IMPORTANT: Distributed indexing is CPU-only. Enforce single-node when
         # accelerator or torch-related paths are detected.
-        torch_detected = False
+        accelerator_or_torch_detected = cagra_library is not None
         try:
             if accelerator is not None:
-                torch_detected = True
+                accelerator_or_torch_detected = True
             else:
                 impl = kwargs.get("implementation")
                 use_torch_flag = kwargs.get("use_torch") is True
@@ -3807,24 +3882,24 @@ def _create_index_impl(
                     or torch_centroids
                     or torch_codebook
                 ):
-                    torch_detected = True
+                    accelerator_or_torch_detected = True
         except Exception:
             # Be conservative: if detection fails, do not modify behavior
             pass
 
-        if torch_detected:
+        if accelerator_or_torch_detected:
             if require_commit:
                 if fragment_ids is not None or index_uuid is not None:
                     LOGGER.info(
-                        "Torch detected; "
+                        "Accelerator or Torch input detected; "
                         "enforce single-node indexing (distributed is CPU-only)."
                     )
                 fragment_ids = None
                 index_uuid = None
             else:
                 if index_uuid is not None:
                     LOGGER.info(
-                        "Torch detected; "
+                        "Accelerator or Torch input detected; "
                         "enforce single-node indexing (distributed is CPU-only)."
                     )
                 index_uuid = None
@@ -4126,9 +4201,11 @@ def create_index(
         num_sub_vectors : int, optional
             The number of sub-vectors for PQ (Product Quantization).
         accelerator : str or ``torch.Device``, optional
-            If set, use an accelerator to speed up the training process.
-            Accepted accelerator: "cuda" (Nvidia GPU) and "mps" (Apple Silicon GPU).
-            If not set, use the CPU.
+            If set, use an accelerator for supported index build stages.
+            ``IVF_PQ`` accepts "cuda" (Nvidia GPU) and "mps" (Apple Silicon GPU)
+            and requires PyTorch. ``IVF_HNSW_SQ`` accepts "cuda" and uses cuVS
+            CAGRA to construct each HNSW graph when the ``libcuvs`` Python package
+            is installed. Unsupported or unavailable accelerators fall back to CPU.
         index_cache_size : int, optional
             The size of the index cache in number of entries. Default value is 256.
         shuffle_partition_batches : int, optional
@@ -4262,9 +4339,9 @@ def create_index(
 
         Experimental Accelerator (GPU) support:
 
-        - *accelerate*: use GPU to train IVF partitions.
-            Only supports CUDA (Nvidia) or MPS (Apple) currently.
-            Requires PyTorch being installed.
+        - ``IVF_PQ`` uses CUDA or MPS through PyTorch for IVF/PQ training.
+        - ``IVF_HNSW_SQ`` uses CUDA through cuVS CAGRA for HNSW graph
+          construction and requires the ``libcuvs`` Python package.
 
         .. code-block:: python
 
@@ -4279,9 +4356,9 @@ 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.
+        Other index types fall back to CPU index building when an accelerator is
+        provided. ``IVF_HNSW_SQ`` also falls back to CPU when compatible cuVS CAGRA
+        support is unavailable.
 
         References
         ----------

python/python/tests/test_vector_index.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright The Lance Authors
 
+import importlib
 import logging
 import os
 import platform
@@ -738,7 +739,9 @@ def test_create_index_unsupported_accelerator(tmp_path):
         )
 
 
-def test_create_index_accelerator_fallback(tmp_path, caplog):
+def test_create_index_accelerator_fallback(tmp_path, caplog, monkeypatch):
+    dataset_module = importlib.import_module("lance.dataset")
+    monkeypatch.setattr(dataset_module, "_locate_cuvs_library", lambda: None)
     tbl = create_table()
     dataset = lance.write_dataset(tbl, tmp_path)
 
@@ -753,11 +756,94 @@ def test_create_index_accelerator_fallback(tmp_path, caplog):
     stats = dataset.stats.index_stats("vector_idx")
     assert stats["index_type"] == "IVF_HNSW_SQ"
     assert any(
-        "does not support GPU acceleration; falling back to CPU" in record.message
+        "compatible cuVS CAGRA library was not found" in record.message
         for record in caplog.records
     )
 
 
+def test_create_index_cagra_accelerator_dispatch(tmp_path, caplog, monkeypatch):
+    dataset_module = importlib.import_module("lance.dataset")
+    monkeypatch.setattr(
+        dataset_module, "_locate_cuvs_library", lambda: "/missing/libcuvs_c.so"
+    )
+    dataset = lance.write_dataset(create_table(), tmp_path)
+
+    with caplog.at_level(logging.WARNING):
+        dataset = dataset.create_index(
+            "vector",
+            index_type="IVF_HNSW_SQ",
+            num_partitions=4,
+            accelerator="cuda",
+        )
+
+    stats = dataset.stats.index_stats("vector_idx")
+    assert stats["index_type"] == "IVF_HNSW_SQ"
+    assert not any(
+        "does not support GPU acceleration" in record.message
+        for record in caplog.records
+    )
+
+
+def test_create_index_cagra_required_disables_fallback(tmp_path, monkeypatch):
+    dataset_module = importlib.import_module("lance.dataset")
+    monkeypatch.setattr(
+        dataset_module, "_locate_cuvs_library", lambda: "/missing/libcuvs_c.so"
+    )
+    dataset = lance.write_dataset(create_table(), tmp_path)
+
+    with pytest.raises(
+        OSError,
+        match=(
+            "CAGRA HNSW acceleration is currently supported only on Unix"
+            "|failed to load cuVS library"
+        ),
+    ):
+        dataset.create_index(
+            "vector",
+            index_type="IVF_HNSW_SQ",
+            num_partitions=1,
+            accelerator="cuda",
+            _require_cagra=True,
+        )
+
+
+@pytest.mark.cuda
+@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"])
+def test_create_index_cagra_accelerated_recall(tmp_path, metric):
+    dataset_module = importlib.import_module("lance.dataset")
+    if dataset_module._locate_cuvs_library() is None:
+        pytest.skip(
+            "compatible libcuvs is unavailable; "
+            "https://github.com/lance-format/lance/issues/5061"
+        )
+
+    rng = np.random.default_rng(42)
+    vectors = rng.standard_normal((2048, 32)).astype(np.float32)
+    table = vec_to_table(data=vectors).append_column(
+        "id", pa.array(np.arange(len(vectors)))
+    )
+    dataset = lance.write_dataset(table, tmp_path)
+    query = vectors[17]
+    nearest = {"column": "vector", "q": query, "k": 10, "metric": metric}
+    ground_truth = dataset.to_table(columns=["id"], nearest=nearest)["id"].to_numpy()
+
+    indexed = dataset.create_index(
+        "vector",
+        index_type="IVF_HNSW_SQ",
+        metric=metric,
+        num_partitions=1,
+        accelerator="cuda",
+        _require_cagra=True,
+    )
+
+    result = indexed.to_table(columns=["id"], nearest=nearest)["id"].to_numpy()
+    recall = len(set(ground_truth) & set(result)) / len(ground_truth)
+    assert recall >= 0.5, (
+        f"metric={metric}, recall={recall}, "
+        f"ground_truth={ground_truth}, result={result}"
+    )
+
+
 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(

python/src/dataset.rs

@@ -5151,12 +5151,24 @@ fn prepare_vector_index_params(
     }?;
     params.version(index_file_version);
     params.skip_transpose(skip_transpose);
-    if let Some(kwargs) = kwargs
-        && let Some(acc) = kwargs.get_item("accelerator")?
-    {
-        params
-            .runtime_hints
-            .insert("lancedb.accelerator".to_string(), acc.to_string());
+    if let Some(kwargs) = kwargs {
+        if let Some(acc) = kwargs.get_item("accelerator")? {
+            params
+                .runtime_hints
+                .insert("lancedb.accelerator".to_string(), acc.to_string());
+        }
+        if let Some(library_path) = kwargs.get_item("cuvs_library")? {
+            params
+                .runtime_hints
+                .insert("lancedb.cuvs_library".to_string(), library_path.extract()?);
+        }
+        if let Some(required) = kwargs.get_item("_require_cagra")?
+            && required.extract::<bool>()?
+        {
+            params
+                .runtime_hints
+                .insert("lancedb.cuvs_required".to_string(), "true".to_string());
+        }
     }
     Ok(params)
 }

rust/lance-index/Cargo.toml

@@ -52,6 +52,7 @@ lance-linalg.workspace = true
 lance-select.workspace = true
 lance-tokenizer.workspace = true
 lance-table.workspace = true
+libc.workspace = true
 log.workspace = true
 ndarray.workspace = true
 num-traits.workspace = true

rust/lance-index/src/vector/hnsw.rs

@@ -16,6 +16,7 @@ use super::graph::OrderedNode;
 use super::storage::VectorStore;
 
 pub mod builder;
+mod cagra;
 pub mod index;
 pub mod online;