From fd3798042553dce6426470a96e890b997b79899d Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 15:11:05 +0000 Subject: [PATCH 1/4] fix: build IVF_HNSW_SQ graphs with CAGRA --- Cargo.lock | 1 + java/lance-jni/Cargo.lock | 1 + python/Cargo.lock | 1 + python/python/lance/dataset.py | 89 +++- python/python/tests/test_vector_index.py | 30 +- python/src/dataset.rs | 17 +- rust/lance-index/Cargo.toml | 1 + rust/lance-index/src/vector/hnsw.rs | 1 + rust/lance-index/src/vector/hnsw/builder.rs | 179 ++++++- rust/lance-index/src/vector/hnsw/cagra.rs | 555 ++++++++++++++++++++ rust/lance-index/src/vector/sq/storage.rs | 32 ++ rust/lance-index/src/vector/v3/subindex.rs | 66 ++- rust/lance/src/index/vector.rs | 49 +- rust/lance/src/index/vector/builder.rs | 50 +- rust/lance/src/index/vector/details.rs | 14 +- 15 files changed, 1035 insertions(+), 51 deletions(-) create mode 100644 rust/lance-index/src/vector/hnsw/cagra.rs diff --git a/Cargo.lock b/Cargo.lock index 48b60c7b514..f8cd83ad23b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4852,6 +4852,7 @@ dependencies = [ "lance-table", "lance-testing", "lance-tokenizer", + "libc", "libsais-rs", "log", "ndarray", diff --git a/java/lance-jni/Cargo.lock b/java/lance-jni/Cargo.lock index 1d8c42337a7..2c5c5b1a12b 100644 --- a/java/lance-jni/Cargo.lock +++ b/java/lance-jni/Cargo.lock @@ -4029,6 +4029,7 @@ dependencies = [ "lance-select", "lance-table", "lance-tokenizer", + "libc", "libsais-rs", "log", "ndarray", diff --git a/python/Cargo.lock b/python/Cargo.lock index 10dc6efc545..b7cf93a9567 100644 --- a/python/Cargo.lock +++ b/python/Cargo.lock @@ -4357,6 +4357,7 @@ dependencies = [ "lance-select", "lance-table", "lance-tokenizer", + "libc", "libsais-rs", "log", "ndarray", diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index d387880be48..ecce451b6ea 100644 --- a/python/python/lance/dataset.py +++ b/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) @@ -3781,7 +3817,26 @@ def _create_index_impl( # 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: + LOGGER.warning( + "A compatible cuVS CAGRA library was not found; falling back " + "to CPU for IVF_HNSW_SQ" + ) + else: + kwargs["cuvs_library"] = cagra_library + # 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 +3845,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,16 +3862,16 @@ 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 @@ -3824,7 +3879,7 @@ def _create_index_impl( 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 +4181,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 +4319,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 +4336,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 ---------- diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 21ae2aaac8e..93e8715c6a4 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 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,7 +756,30 @@ 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 ) diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 69cdc3c8130..8c6dee34ec7 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -5151,12 +5151,17 @@ 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()?); + } } Ok(params) } diff --git a/rust/lance-index/Cargo.toml b/rust/lance-index/Cargo.toml index ec128ee1b85..bb1212b39cd 100644 --- a/rust/lance-index/Cargo.toml +++ b/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 diff --git a/rust/lance-index/src/vector/hnsw.rs b/rust/lance-index/src/vector/hnsw.rs index 0b86f4d2bf7..838c3f6cbdc 100644 --- a/rust/lance-index/src/vector/hnsw.rs +++ b/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; diff --git a/rust/lance-index/src/vector/hnsw/builder.rs b/rust/lance-index/src/vector/hnsw/builder.rs index fbbabcd321d..802d9f7dc8f 100644 --- a/rust/lance-index/src/vector/hnsw/builder.rs +++ b/rust/lance-index/src/vector/hnsw/builder.rs @@ -43,7 +43,7 @@ use crate::vector::graph::{ Visited, beam_search_acorn, beam_search_borrowed, greedy_search, greedy_search_borrowed, }; use crate::vector::storage::{DistCalculator, VectorStore}; -use crate::vector::v3::subindex::IvfSubIndex; +use crate::vector::v3::subindex::{IvfSubIndex, SubIndexBuildAccelerator}; use crate::vector::{ApproxMode, Query, VECTOR_RESULT_SCHEMA}; pub const HNSW_METADATA_KEY: &str = "lance:hnsw"; @@ -264,6 +264,55 @@ impl HNSW { } } + /// Import a fixed-degree neighbor graph as a one-level HNSW graph. + /// + /// CAGRA emits a row-major `num_nodes x graph_degree` graph. Node ids are + /// local to the IVF partition and therefore line up with Lance's storage + /// offsets without remapping. + pub(crate) fn from_neighbor_graph( + params: HnswBuildParams, + neighbors: Vec, + graph_degree: usize, + ) -> Result { + params.validate()?; + if graph_degree == 0 { + return Err(Error::invalid_input( + "CAGRA graph_degree must be greater than 0".to_string(), + )); + } + if !neighbors.len().is_multiple_of(graph_degree) { + return Err(Error::invalid_input(format!( + "CAGRA graph has {} neighbor ids, which is not divisible by graph_degree {}", + neighbors.len(), + graph_degree + ))); + } + let num_nodes = neighbors.len() / graph_degree; + if num_nodes > u32::MAX as usize { + return Err(Error::invalid_input(format!( + "CAGRA graph has {num_nodes} nodes, exceeding the u32 node-id limit" + ))); + } + if let Some(invalid_neighbor) = neighbors + .iter() + .copied() + .find(|neighbor| *neighbor as usize >= num_nodes) + { + return Err(Error::invalid_input(format!( + "CAGRA graph neighbor id {invalid_neighbor} is outside the node range 0..{num_nodes}" + ))); + } + + let nodes = neighbors + .chunks_exact(graph_degree) + .map(|neighbors| { + let neighbors = Arc::new(neighbors.to_vec()); + GraphBuilderNode::from_parts(vec![neighbors.clone()], vec![Vec::new()], neighbors) + }) + .collect(); + Ok(Self::from_parts(params, nodes, vec![num_nodes], 0)) + } + pub fn empty() -> Self { Self { inner: Arc::new(HnswCore { @@ -1447,6 +1496,37 @@ impl IvfSubIndex for HNSW { Ok(builder.finish()) } + fn index_vectors_with_accelerator( + storage: &impl VectorStore, + params: Self::BuildParams, + accelerator: &SubIndexBuildAccelerator, + ) -> Result + where + Self: Sized, + { + match accelerator { + SubIndexBuildAccelerator::Cagra(accelerator) => { + if accelerator.is_disabled() + || !super::cagra::supports_partition(storage.len(), ¶ms) + { + return Self::index_vectors(storage, params); + } + match super::cagra::build(storage, params.clone(), accelerator.library_path()) { + Ok(index) => Ok(index), + Err(error) => { + if accelerator.disable() { + log::warn!( + "cuVS CAGRA HNSW build failed; falling back to CPU: {}", + error + ); + } + Self::index_vectors(storage, params) + } + } + } + } + } + fn remap( &self, _mapping: &RowAddrRemap, // we don't need the mapping here because we rebuild the graph from remapped storage @@ -1540,10 +1620,11 @@ mod tests { use arrow_array::{ ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, UInt8Array, UInt32Array, + UInt64Array, cast::AsArray, types::UInt32Type, }; use arrow_schema::Schema; use lance_arrow::FixedSizeListArrayExt; - use lance_core::{Error, deepsize::DeepSizeOf}; + use lance_core::{Error, ROW_ID, deepsize::DeepSizeOf}; use lance_file::versions::v1::{ reader::FileReader as V1FileReader, writer::{FileWriter as V1FileWriter, FileWriterOptions as V1FileWriterOptions}, @@ -1562,11 +1643,13 @@ mod tests { ImmutableHnswLevelView, MIN_HNSW_M, random_level_with, }; use crate::vector::graph::builder::GraphBuilderNode; + use crate::vector::sq::storage::ScalarQuantizationStorage; use crate::vector::storage::{DistCalculator, VectorStore}; - use crate::vector::v3::subindex::IvfSubIndex; + use crate::vector::v3::subindex::{IvfSubIndex, SubIndexBuildAccelerator}; use crate::vector::{ + SQ_CODE_COLUMN, flat::storage::{FlatBinStorage, FlatFloatStorage}, - graph::{DISTS_FIELD, NEIGHBORS_FIELD, OrderedNode, VisitedGenerator}, + graph::{DISTS_FIELD, NEIGHBORS_COL, NEIGHBORS_FIELD, OrderedNode, VisitedGenerator}, hnsw::{ HNSW, HnswMetadata, VECTOR_ID_FIELD, builder::{HnswBuildParams, HnswQueryParams}, @@ -2074,6 +2157,94 @@ mod tests { ); } + #[test] + fn test_import_cagra_neighbor_graph() { + let params = HnswBuildParams::default(); + let graph = vec![1, 2, 0, 2, 0, 3, 1, 2]; + let hnsw = HNSW::from_neighbor_graph(params, graph, 2).unwrap(); + + assert_eq!(hnsw.len(), 4); + assert_eq!(hnsw.max_level(), 1); + let batch = hnsw.to_batch().unwrap(); + assert_eq!(batch.num_rows(), 4); + let neighbors = batch[NEIGHBORS_COL].as_list::(); + assert_eq!( + neighbors.value(0).as_primitive::().values(), + &[1, 2] + ); + assert_eq!( + neighbors.value(3).as_primitive::().values(), + &[1, 2] + ); + + let loaded = HNSW::load(batch).unwrap(); + assert_eq!(loaded.len(), 4); + assert_eq!(loaded.max_level(), 1); + } + + #[test] + fn test_import_cagra_neighbor_graph_rejects_invalid_node() { + let error = + HNSW::from_neighbor_graph(HnswBuildParams::default(), vec![1, 0, 0, 3], 2).unwrap_err(); + + assert!(matches!(error, Error::InvalidInput { .. })); + assert!(error.to_string().contains("neighbor id 3")); + assert!(error.to_string().contains("0..2")); + } + + #[test] + fn test_cagra_missing_library_falls_back_to_cpu() { + let storage = make_sq_storage(8); + let accelerator = SubIndexBuildAccelerator::cagra("/missing/libcuvs_c.so"); + + let hnsw = HNSW::index_vectors_with_accelerator( + &storage, + HnswBuildParams::default().num_edges(4).ef_construction(10), + &accelerator, + ) + .unwrap(); + + assert_eq!(hnsw.len(), 8); + let SubIndexBuildAccelerator::Cagra(accelerator) = accelerator; + assert!(accelerator.is_disabled()); + } + + #[test] + fn test_cagra_small_partition_does_not_disable_acceleration() { + let storage = make_sq_storage(4); + let accelerator = SubIndexBuildAccelerator::cagra("/missing/libcuvs_c.so"); + + let hnsw = HNSW::index_vectors_with_accelerator( + &storage, + HnswBuildParams::default().num_edges(4).ef_construction(10), + &accelerator, + ) + .unwrap(); + + assert_eq!(hnsw.len(), 4); + let SubIndexBuildAccelerator::Cagra(accelerator) = accelerator; + assert!(!accelerator.is_disabled()); + } + + fn make_sq_storage(num_rows: usize) -> ScalarQuantizationStorage { + let values = (0..num_rows) + .flat_map(|row| { + let value = (row * 255 / num_rows.saturating_sub(1).max(1)) as u8; + [value, value] + }) + .collect::>(); + let codes = FixedSizeListArray::try_new_from_values(UInt8Array::from(values), 2).unwrap(); + let batch = RecordBatch::try_from_iter([ + ( + ROW_ID, + Arc::new(UInt64Array::from((0..num_rows as u64).collect::>())) as ArrayRef, + ), + (SQ_CODE_COLUMN, Arc::new(codes) as ArrayRef), + ]) + .unwrap(); + ScalarQuantizationStorage::try_new(8, DistanceType::L2, 0.0..1.0, [batch], None).unwrap() + } + /// Brute-force top-`k` restricted to mask-passing ids. fn brute_force_topk_masked( store: &FlatFloatStorage, diff --git a/rust/lance-index/src/vector/hnsw/cagra.rs b/rust/lance-index/src/vector/hnsw/cagra.rs new file mode 100644 index 00000000000..8b8b39525ce --- /dev/null +++ b/rust/lance-index/src/vector/hnsw/cagra.rs @@ -0,0 +1,555 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Runtime cuVS CAGRA integration. +//! +//! cuVS remains an optional runtime dependency. The Python layer supplies the +//! exact `libcuvs_c` path, and this module resolves only the stable C symbols +//! needed to build and copy a CAGRA graph. + +use std::ffi::{CStr, CString, c_char, c_int, c_void}; +use std::mem::MaybeUninit; +use std::path::Path; +use std::ptr; + +use lance_core::{Error, Result}; +use lance_linalg::distance::DistanceType; + +use super::builder::{HNSW, HnswBuildParams}; +use crate::vector::sq::storage::ScalarQuantizationStorage; +use crate::vector::storage::VectorStore; + +const CUVS_SUCCESS: c_int = 1; +const DL_CPU: c_int = 1; +const DL_UINT: u8 = 1; +const DL_FLOAT: u8 = 2; +const CAGRA_HNSW_SIMILAR_SEARCH_PERFORMANCE: c_int = 0; + +#[repr(C)] +#[derive(Clone, Copy)] +struct DLDevice { + device_type: c_int, + device_id: i32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct DLDataType { + code: u8, + bits: u8, + lanes: u16, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct DLTensor { + data: *mut c_void, + device: DLDevice, + ndim: i32, + dtype: DLDataType, + shape: *mut i64, + strides: *mut i64, + byte_offset: u64, +} + +#[repr(C)] +struct DLManagedTensor { + dl_tensor: DLTensor, + manager_ctx: *mut c_void, + deleter: Option, +} + +type CuvsStatus = c_int; +type CuvsResources = usize; +type CuvsDataset = *mut c_void; +type CuvsCagraParams = *mut c_void; +type CuvsCagraIndex = *mut c_void; + +type GetLastErrorText = unsafe extern "C" fn() -> *const c_char; +type ResourcesCreate = unsafe extern "C" fn(*mut CuvsResources) -> CuvsStatus; +type ResourcesDestroy = unsafe extern "C" fn(CuvsResources) -> CuvsStatus; +type StreamSync = unsafe extern "C" fn(CuvsResources) -> CuvsStatus; +type MatrixCopy = + unsafe extern "C" fn(CuvsResources, *mut DLManagedTensor, *mut DLManagedTensor) -> CuvsStatus; +type DatasetMakeStandardView = + unsafe extern "C" fn(CuvsResources, *mut DLManagedTensor, *mut CuvsDataset) -> CuvsStatus; +type DatasetDestroy = unsafe extern "C" fn(CuvsDataset) -> CuvsStatus; +type ParamsCreate = unsafe extern "C" fn(*mut CuvsCagraParams) -> CuvsStatus; +type ParamsDestroy = unsafe extern "C" fn(CuvsCagraParams) -> CuvsStatus; +type ParamsFromHnsw = + unsafe extern "C" fn(CuvsCagraParams, i64, i64, c_int, c_int, c_int, c_int) -> CuvsStatus; +type IndexCreate = unsafe extern "C" fn(*mut CuvsCagraIndex) -> CuvsStatus; +type IndexDestroy = unsafe extern "C" fn(CuvsCagraIndex) -> CuvsStatus; +type CagraBuild = + unsafe extern "C" fn(CuvsResources, CuvsCagraParams, CuvsDataset, CuvsCagraIndex) -> CuvsStatus; +type IndexGetGraph = unsafe extern "C" fn(CuvsCagraIndex, *mut DLManagedTensor) -> CuvsStatus; + +#[cfg(unix)] +struct DynamicLibrary(*mut c_void); + +#[cfg(unix)] +impl DynamicLibrary { + fn open(path: &Path) -> Result { + use std::os::unix::ffi::OsStrExt; + + let path_display = path.display().to_string(); + let path = CString::new(path.as_os_str().as_bytes()).map_err(|_| { + Error::invalid_input(format!( + "cuVS library path contains an interior NUL byte: {}", + path_display + )) + })?; + // SAFETY: `path` is a valid NUL-terminated string. The returned handle + // is retained until every resolved function pointer is no longer used. + let handle = unsafe { libc::dlopen(path.as_ptr(), libc::RTLD_NOW | libc::RTLD_LOCAL) }; + if handle.is_null() { + return Err(Error::io(format!( + "failed to load cuVS library {path_display}: {}", + dl_error_message() + ))); + } + Ok(Self(handle)) + } + + fn symbol(&self, name: &'static [u8]) -> Result { + debug_assert_eq!( + name.last(), + Some(&0), + "cuVS dynamic symbol names must be NUL-terminated" + ); + // SAFETY: the symbol name is NUL-terminated and the library handle is + // live. Each caller supplies the signature from the cuVS C header. + let symbol = unsafe { libc::dlsym(self.0, name.as_ptr().cast()) }; + if symbol.is_null() { + let name = match name.strip_suffix(&[0]) { + Some(name) => String::from_utf8_lossy(name), + None => String::from_utf8_lossy(name), + }; + return Err(Error::io(format!( + "cuVS library is missing required symbol {name}: {}", + dl_error_message() + ))); + } + // SAFETY: `symbol` was resolved by name for the exact C function type + // requested at each call site. Function pointers are pointer-sized and + // `T: Copy`, so copying the representation is valid. + Ok(unsafe { std::mem::transmute_copy(&symbol) }) + } +} + +#[cfg(unix)] +impl Drop for DynamicLibrary { + fn drop(&mut self) { + // SAFETY: this is the live handle returned by `dlopen`, closed once. + let status = unsafe { libc::dlclose(self.0) }; + if status != 0 { + log::warn!("failed to unload cuVS library: {}", dl_error_message()); + } + } +} + +#[cfg(unix)] +fn dl_error_message() -> String { + // SAFETY: `dlerror` returns either NULL or a process-owned NUL-terminated + // diagnostic string that remains valid until the next loader call. + let error = unsafe { libc::dlerror() }; + if error.is_null() { + "unknown dynamic loader error".to_string() + } else { + // SAFETY: non-NULL `dlerror` results are valid C strings. + unsafe { CStr::from_ptr(error) } + .to_string_lossy() + .into_owned() + } +} + +#[cfg(unix)] +struct CuvsApi { + _library: DynamicLibrary, + get_last_error_text: GetLastErrorText, + resources_create: ResourcesCreate, + resources_destroy: ResourcesDestroy, + stream_sync: StreamSync, + matrix_copy: MatrixCopy, + dataset_make_standard_view: DatasetMakeStandardView, + dataset_destroy: DatasetDestroy, + params_create: ParamsCreate, + params_destroy: ParamsDestroy, + params_from_hnsw: ParamsFromHnsw, + index_create: IndexCreate, + index_destroy: IndexDestroy, + cagra_build: CagraBuild, + index_get_graph: IndexGetGraph, +} + +#[cfg(unix)] +impl CuvsApi { + fn load(path: &Path) -> Result { + let library = DynamicLibrary::open(path)?; + Ok(Self { + get_last_error_text: library.symbol(b"cuvsGetLastErrorText\0")?, + resources_create: library.symbol(b"cuvsResourcesCreate\0")?, + resources_destroy: library.symbol(b"cuvsResourcesDestroy\0")?, + stream_sync: library.symbol(b"cuvsStreamSync\0")?, + matrix_copy: library.symbol(b"cuvsMatrixCopy\0")?, + dataset_make_standard_view: library.symbol(b"cuvsDatasetMakeStandardView\0")?, + dataset_destroy: library.symbol(b"cuvsDatasetDestroy\0")?, + params_create: library.symbol(b"cuvsCagraIndexParamsCreate\0")?, + params_destroy: library.symbol(b"cuvsCagraIndexParamsDestroy\0")?, + params_from_hnsw: library.symbol(b"cuvsCagraIndexParamsFromHnswParams\0")?, + index_create: library.symbol(b"cuvsCagraIndexCreate\0")?, + index_destroy: library.symbol(b"cuvsCagraIndexDestroy\0")?, + cagra_build: library.symbol(b"cuvsCagraBuild\0")?, + index_get_graph: library.symbol(b"cuvsCagraIndexGetGraph\0")?, + _library: library, + }) + } + + fn check(&self, status: CuvsStatus, operation: &str) -> Result<()> { + if status == CUVS_SUCCESS { + return Ok(()); + } + // SAFETY: the function pointer was resolved with its C header + // signature and returns either NULL or a NUL-terminated error string. + let error = unsafe { (self.get_last_error_text)() }; + let detail = if error.is_null() { + format!("status {status}") + } else { + // SAFETY: non-NULL cuVS error text is a valid C string. + unsafe { CStr::from_ptr(error) } + .to_string_lossy() + .into_owned() + }; + Err(Error::io(format!("cuVS failed to {operation}: {detail}"))) + } +} + +#[cfg(unix)] +struct Resources<'a> { + api: &'a CuvsApi, + handle: CuvsResources, +} + +#[cfg(unix)] +impl<'a> Resources<'a> { + fn create(api: &'a CuvsApi) -> Result { + let mut handle = 0; + // SAFETY: the output pointer is valid and the function signature was + // resolved from the cuVS C API. + api.check( + unsafe { (api.resources_create)(&mut handle) }, + "create resources", + )?; + Ok(Self { api, handle }) + } +} + +#[cfg(unix)] +impl Drop for Resources<'_> { + fn drop(&mut self) { + // SAFETY: this handle was created by the matching API and is dropped once. + let status = unsafe { (self.api.resources_destroy)(self.handle) }; + if let Err(error) = self.api.check(status, "destroy resources") { + log::warn!("{error}"); + } + } +} + +#[cfg(unix)] +struct DatasetView<'a> { + api: &'a CuvsApi, + handle: CuvsDataset, +} + +#[cfg(unix)] +impl Drop for DatasetView<'_> { + fn drop(&mut self) { + // SAFETY: this handle was created by the matching API and is dropped once. + let status = unsafe { (self.api.dataset_destroy)(self.handle) }; + if let Err(error) = self.api.check(status, "destroy the dataset view") { + log::warn!("{error}"); + } + } +} + +#[cfg(unix)] +struct CagraParams<'a> { + api: &'a CuvsApi, + handle: CuvsCagraParams, +} + +#[cfg(unix)] +impl Drop for CagraParams<'_> { + fn drop(&mut self) { + // SAFETY: this handle was created by the matching API and is dropped once. + let status = unsafe { (self.api.params_destroy)(self.handle) }; + if let Err(error) = self.api.check(status, "destroy CAGRA parameters") { + log::warn!("{error}"); + } + } +} + +#[cfg(unix)] +struct CagraIndex<'a> { + api: &'a CuvsApi, + handle: CuvsCagraIndex, +} + +#[cfg(unix)] +impl Drop for CagraIndex<'_> { + fn drop(&mut self) { + // SAFETY: this handle was created by the matching API and is dropped once. + let status = unsafe { (self.api.index_destroy)(self.handle) }; + if let Err(error) = self.api.check(status, "destroy the CAGRA index") { + log::warn!("{error}"); + } + } +} + +struct ManagedGraph(DLManagedTensor); + +impl Drop for ManagedGraph { + fn drop(&mut self) { + if let Some(deleter) = self.0.deleter { + // SAFETY: cuVS installed this deleter for this exact managed tensor; + // it releases only the view metadata, not the index-owned graph. + unsafe { deleter(&mut self.0) }; + } + } +} + +fn host_tensor(values: &mut [T], shape: &mut [i64], dtype: DLDataType) -> DLManagedTensor { + DLManagedTensor { + dl_tensor: DLTensor { + data: values.as_mut_ptr().cast(), + device: DLDevice { + device_type: DL_CPU, + device_id: 0, + }, + ndim: shape.len() as i32, + dtype, + shape: shape.as_mut_ptr(), + strides: ptr::null_mut(), + byte_offset: 0, + }, + manager_ctx: ptr::null_mut(), + deleter: None, + } +} + +fn cuvs_distance_type(distance_type: DistanceType) -> Result { + match distance_type { + DistanceType::L2 => Ok(0), + DistanceType::Cosine => Ok(2), + DistanceType::Dot => Ok(6), + DistanceType::Hamming => Err(Error::invalid_input( + "CAGRA does not support Lance bitwise Hamming vectors".to_string(), + )), + } +} + +/// Whether a partition is large enough for cuVS's HNSW-compatible graph. +/// +/// The similar-search-performance heuristic derives an intermediate graph +/// degree of `M + M * ef_construction / 256`. A graph needs at least one more +/// row than its degree because a node cannot be its own neighbor. Tiny IVF +/// partitions are cheaper and safer to build directly on the CPU. +pub(super) fn supports_partition(num_rows: usize, params: &HnswBuildParams) -> bool { + let intermediate_graph_degree = params.m.saturating_add( + params + .m + .saturating_mul(params.ef_construction) + .saturating_div(256), + ); + num_rows > intermediate_graph_degree +} + +/// Build a one-level Lance HNSW graph with cuVS CAGRA. +pub(super) fn build( + storage: &impl VectorStore, + params: HnswBuildParams, + library_path: &str, +) -> Result { + #[cfg(not(unix))] + { + let _ = (storage, params, library_path); + return Err(Error::io( + "CAGRA HNSW acceleration is currently supported only on Unix".to_string(), + )); + } + + #[cfg(unix)] + { + params.validate()?; + let sq_storage = storage + .as_any() + .downcast_ref::() + .ok_or_else(|| { + Error::invalid_input( + "CAGRA HNSW acceleration currently requires scalar-quantized storage" + .to_string(), + ) + })?; + let (mut vectors, dim) = sq_storage.to_f32_matrix()?; + let num_rows = storage.len(); + if num_rows == 0 { + return Ok(HNSW::empty()); + } + let num_rows_i64 = i64::try_from(num_rows) + .map_err(|_| Error::invalid_input("CAGRA row count exceeds i64::MAX".to_string()))?; + let dim_i64 = i64::try_from(dim) + .map_err(|_| Error::invalid_input("CAGRA dimension exceeds i64::MAX".to_string()))?; + let m = c_int::try_from(params.m) + .map_err(|_| Error::invalid_input("HNSW m exceeds C int range".to_string()))?; + let ef_construction = c_int::try_from(params.ef_construction).map_err(|_| { + Error::invalid_input("HNSW ef_construction exceeds C int range".to_string()) + })?; + + let api = CuvsApi::load(Path::new(library_path))?; + let resources = Resources::create(&api)?; + let mut vector_shape = [num_rows_i64, dim_i64]; + let mut vector_tensor = host_tensor( + &mut vectors, + &mut vector_shape, + DLDataType { + code: DL_FLOAT, + bits: 32, + lanes: 1, + }, + ); + + let mut dataset_handle = ptr::null_mut(); + // SAFETY: all pointers refer to live stack metadata and initialized + // row-major host storage for the duration of the call. + api.check( + unsafe { + (api.dataset_make_standard_view)( + resources.handle, + &mut vector_tensor, + &mut dataset_handle, + ) + }, + "create a dataset view", + )?; + let dataset = DatasetView { + api: &api, + handle: dataset_handle, + }; + + let mut params_handle = ptr::null_mut(); + // SAFETY: the output pointer is valid for the matching create call. + api.check( + unsafe { (api.params_create)(&mut params_handle) }, + "create CAGRA parameters", + )?; + let cagra_params = CagraParams { + api: &api, + handle: params_handle, + }; + // SAFETY: handles are live and scalar arguments were range-checked. + api.check( + unsafe { + (api.params_from_hnsw)( + cagra_params.handle, + num_rows_i64, + dim_i64, + m, + ef_construction, + CAGRA_HNSW_SIMILAR_SEARCH_PERFORMANCE, + cuvs_distance_type(storage.distance_type())?, + ) + }, + "derive CAGRA parameters from HNSW parameters", + )?; + + let mut index_handle = ptr::null_mut(); + // SAFETY: the output pointer is valid for the matching create call. + api.check( + unsafe { (api.index_create)(&mut index_handle) }, + "create a CAGRA index", + )?; + let index = CagraIndex { + api: &api, + handle: index_handle, + }; + // SAFETY: every handle is live and the borrowed dataset outlives the + // index build and graph copy. + api.check( + unsafe { + (api.cagra_build)( + resources.handle, + cagra_params.handle, + dataset.handle, + index.handle, + ) + }, + "build a CAGRA graph", + )?; + + let mut graph_tensor = MaybeUninit::::zeroed(); + // SAFETY: cuVS initializes the complete managed tensor on success. + api.check( + unsafe { (api.index_get_graph)(index.handle, graph_tensor.as_mut_ptr()) }, + "get the CAGRA graph", + )?; + // SAFETY: the preceding successful C call initialized the value. + let mut graph = ManagedGraph(unsafe { graph_tensor.assume_init() }); + let tensor = &graph.0.dl_tensor; + if tensor.ndim != 2 || tensor.shape.is_null() || tensor.data.is_null() { + return Err(Error::io(format!( + "CAGRA returned invalid graph metadata: rank={}, shape_null={}, data_null={}", + tensor.ndim, + tensor.shape.is_null(), + tensor.data.is_null() + ))); + } + if tensor.dtype.code != DL_UINT || tensor.dtype.bits != 32 || tensor.dtype.lanes != 1 { + return Err(Error::io(format!( + "CAGRA returned graph dtype code={}, bits={}, lanes={}; expected uint32", + tensor.dtype.code, tensor.dtype.bits, tensor.dtype.lanes + ))); + } + // SAFETY: a rank-2 successful DLPack result has two shape elements. + let graph_shape = unsafe { std::slice::from_raw_parts(tensor.shape, 2) }; + if graph_shape[0] != num_rows_i64 || graph_shape[1] <= 0 { + return Err(Error::io(format!( + "CAGRA returned graph shape {:?}; expected ({num_rows}, degree)", + graph_shape + ))); + } + let graph_degree = usize::try_from(graph_shape[1]) + .map_err(|_| Error::io("CAGRA graph degree exceeds usize::MAX".to_string()))?; + let graph_len = num_rows + .checked_mul(graph_degree) + .ok_or_else(|| Error::io("CAGRA graph size overflow".to_string()))?; + let mut neighbors = vec![0_u32; graph_len]; + let mut host_shape = [num_rows_i64, graph_shape[1]]; + let mut host_graph = host_tensor( + &mut neighbors, + &mut host_shape, + DLDataType { + code: DL_UINT, + bits: 32, + lanes: 1, + }, + ); + // SAFETY: source and destination tensors have matching shapes and + // dtypes, and both buffers remain live through synchronization. + api.check( + unsafe { (api.matrix_copy)(resources.handle, &mut graph.0, &mut host_graph) }, + "copy the CAGRA graph to host memory", + )?; + // SAFETY: this live resources handle owns the stream used by the copy. + api.check( + unsafe { (api.stream_sync)(resources.handle) }, + "synchronize the CAGRA graph copy", + )?; + + log::info!( + "Built HNSW graph with cuVS CAGRA: num={}, degree={}", + num_rows, + graph_degree + ); + HNSW::from_neighbor_graph(params, neighbors, graph_degree) + } +} diff --git a/rust/lance-index/src/vector/sq/storage.rs b/rust/lance-index/src/vector/sq/storage.rs index 5011fbada46..736ec1d2857 100644 --- a/rust/lance-index/src/vector/sq/storage.rs +++ b/rust/lance-index/src/vector/sq/storage.rs @@ -255,6 +255,38 @@ impl ScalarQuantizationStorage { Ok(new) } } + + /// Reconstruct the quantized vectors as one row-major `f32` matrix. + /// + /// CAGRA builds a neighbor graph from floating-point vectors. Reconstructing + /// from the persisted SQ codes keeps its graph input aligned exactly with + /// the vector ordering used by Lance's SQ storage. + pub(crate) fn to_f32_matrix(&self) -> Result<(Vec, usize)> { + let dim = self.quantizer.metadata.dim; + let value_count = self + .len() + .checked_mul(dim) + .ok_or_else(|| Error::index("SQ vector matrix size overflow".to_string()))?; + let bounds = self.quantizer.bounds(); + let lower_bound = bounds.start as f32; + let value_scale = sq_value_scale(&bounds); + let mut values = Vec::with_capacity(value_count); + for chunk in &self.chunks { + values.extend( + chunk + .sq_codes + .values() + .iter() + .map(|value| lower_bound + f32::from(*value) * value_scale), + ); + } + debug_assert_eq!( + values.len(), + value_count, + "reconstructed SQ matrix length must match rows times dimension" + ); + Ok((values, dim)) + } } #[async_trait] diff --git a/rust/lance-index/src/vector/v3/subindex.rs b/rust/lance-index/src/vector/v3/subindex.rs index cf712c1b246..1f7894d409d 100644 --- a/rust/lance-index/src/vector/v3/subindex.rs +++ b/rust/lance-index/src/vector/v3/subindex.rs @@ -4,7 +4,10 @@ use lance_core::utils::row_addr_remap::RowAddrRemap; use std::collections::BinaryHeap; use std::fmt::Debug; -use std::sync::Arc; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; use arrow_array::{ArrayRef, RecordBatch}; use lance_core::deepsize::DeepSizeOf; @@ -15,6 +18,52 @@ use crate::vector::graph::OrderedNode; use crate::vector::storage::{QueryResidual, QueryScratch, VectorStore}; use crate::vector::{flat, hnsw}; use crate::{prefilter::PreFilter, vector::Query}; + +/// Optional external accelerator for building an IVF sub-index. +/// +/// The accelerator only changes how the in-memory graph is constructed. Lance +/// still owns quantization, persistence, and query execution. +#[derive(Debug, Clone)] +pub enum SubIndexBuildAccelerator { + /// Build a CAGRA graph through the cuVS C API and import it as HNSW. + Cagra(CagraBuildAccelerator), +} + +impl SubIndexBuildAccelerator { + /// Create a cuVS CAGRA accelerator from an absolute `libcuvs_c` path. + pub fn cagra(library_path: impl Into) -> Self { + Self::Cagra(CagraBuildAccelerator { + library_path: library_path.into(), + is_disabled: Arc::new(AtomicBool::new(false)), + }) + } +} + +/// Shared state for one cuVS CAGRA index build. +/// +/// A failing cuVS call disables acceleration for the remaining IVF partitions, +/// avoiding repeated loader or CUDA failures before their CPU fallback. +#[derive(Debug, Clone)] +pub struct CagraBuildAccelerator { + library_path: String, + is_disabled: Arc, +} + +impl CagraBuildAccelerator { + pub(crate) fn library_path(&self) -> &str { + &self.library_path + } + + pub(crate) fn is_disabled(&self) -> bool { + self.is_disabled.load(Ordering::Relaxed) + } + + /// Disable the accelerator and return whether this call changed its state. + pub(crate) fn disable(&self) -> bool { + !self.is_disabled.swap(true, Ordering::Relaxed) + } +} + /// A sub index for IVF index pub trait IvfSubIndex: Send + Sync + Debug + DeepSizeOf { type QueryParams: Send + Sync + for<'a> From<&'a Query>; @@ -116,6 +165,21 @@ pub trait IvfSubIndex: Send + Sync + Debug + DeepSizeOf { where Self: Sized; + /// Build with an optional external accelerator. + /// + /// Sub-index implementations that do not support the requested accelerator + /// retain their normal CPU build behavior. + fn index_vectors_with_accelerator( + storage: &impl VectorStore, + params: Self::BuildParams, + _accelerator: &SubIndexBuildAccelerator, + ) -> Result + where + Self: Sized, + { + Self::index_vectors(storage, params) + } + fn remap(&self, mapping: &RowAddrRemap, store: &impl VectorStore) -> Result where Self: Sized; diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index c4a6376a562..b17df9dc795 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -43,7 +43,7 @@ use lance_arrow::FixedSizeListArrayExt; use lance_index::vector::pq::ProductQuantizer; use lance_index::vector::quantizer::QuantizationType; use lance_index::vector::v3::shuffler::{Shuffler, create_ivf_shuffler}; -use lance_index::vector::v3::subindex::SubIndexType; +use lance_index::vector::v3::subindex::{SubIndexBuildAccelerator, SubIndexType}; use lance_index::vector::{ VectorIndex, hnsw::{ @@ -71,6 +71,12 @@ use crate::{Error, Result, dataset::Dataset, index::pb::vector_index_stage::Stag pub const LANCE_VECTOR_INDEX: &str = "__lance_vector_index"; +/// Transient runtime hint carrying the cuVS C library used for CAGRA builds. +/// +/// This value controls only the current build and is removed before index +/// details are persisted. +pub(crate) const CUVS_LIBRARY_RUNTIME_HINT: &str = "lancedb.cuvs_library"; + /// A materialized snapshot of one logical vector index and all of its segments. #[derive(Debug)] pub struct LogicalVectorIndex { @@ -504,6 +510,13 @@ impl IndexParams for VectorIndexParams { } } +fn cagra_accelerator(params: &VectorIndexParams) -> Option { + params + .runtime_hints + .get(CUVS_LIBRARY_RUNTIME_HINT) + .map(|library_path| SubIndexBuildAccelerator::cagra(library_path.clone())) +} + /// Prepare the shared build inputs used by both direct local builds and /// staged shard builds. /// @@ -893,7 +906,7 @@ pub(crate) async fn build_distributed_vector_index( stages ))); }; - let summary = IvfIndexBuilder::::new( + let mut builder = IvfIndexBuilder::::new( filtered_dataset, column.to_owned(), index_dir.clone(), @@ -903,11 +916,15 @@ pub(crate) async fn build_distributed_vector_index( Some(sq_params.clone()), hnsw_params.clone(), frag_reuse_index, - )? - .with_fragment_filter(fragment_filter) - .with_progress(progress.clone()) - .build() - .await?; + )?; + if let Some(accelerator) = cagra_accelerator(params) { + builder.with_sub_index_accelerator(accelerator); + } + let summary = builder + .with_fragment_filter(fragment_filter) + .with_progress(progress.clone()) + .build() + .await?; return Ok((segment_uuid, summary.files)); } @@ -1262,21 +1279,25 @@ async fn build_vector_index_impl( stages ))); }; - let summary = IvfIndexBuilder::::new( + let mut builder = IvfIndexBuilder::::new( dataset.clone(), column.to_owned(), - dataset.indices_dir().clone().join(uuid.to_string()), + dataset.indices_dir().join(uuid.to_string()), params.metric_type, shuffler, Some(ivf_params), Some(sq_params.clone()), hnsw_params.clone(), frag_reuse_index, - )? - .with_optional_fragment_filter(fragment_ids) - .with_progress(progress.clone()) - .build() - .await?; + )?; + if let Some(accelerator) = cagra_accelerator(params) { + builder.with_sub_index_accelerator(accelerator); + } + let summary = builder + .with_optional_fragment_filter(fragment_ids) + .with_progress(progress.clone()) + .build() + .await?; Ok(summary.files) } _ => Err(Error::index(format!( diff --git a/rust/lance/src/index/vector/builder.rs b/rust/lance/src/index/vector/builder.rs index addc16ecbc3..5984bd0a486 100644 --- a/rust/lance/src/index/vector/builder.rs +++ b/rust/lance/src/index/vector/builder.rs @@ -59,7 +59,7 @@ use lance_index::{ transform::Transformer, v3::{ shuffler::{ShuffleReader, Shuffler}, - subindex::IvfSubIndex, + subindex::{IvfSubIndex, SubIndexBuildAccelerator}, }, }, }; @@ -250,6 +250,9 @@ pub struct IvfIndexBuilder { // whether to transpose codes when building storage transpose_codes: bool, + // Optional accelerator used only while constructing each partition's sub-index. + sub_index_accelerator: Option, + // lance file version for writing index files format_version: LanceFileVersion, @@ -305,6 +308,7 @@ impl IvfIndexBuilder optimize_options: None, merged_num: 0, transpose_codes: true, + sub_index_accelerator: None, format_version, progress: Arc::new(NoopIndexBuildProgress), }) @@ -372,6 +376,7 @@ impl IvfIndexBuilder optimize_options: None, merged_num: 0, transpose_codes: true, + sub_index_accelerator: None, format_version, progress: Arc::new(NoopIndexBuildProgress), }) @@ -509,6 +514,15 @@ impl IvfIndexBuilder self } + /// Use an external accelerator when constructing each partition's sub-index. + pub fn with_sub_index_accelerator( + &mut self, + accelerator: SubIndexBuildAccelerator, + ) -> &mut Self { + self.sub_index_accelerator = Some(accelerator); + self + } + /// Set progress callback for index building pub fn with_progress(&mut self, progress: Arc) -> &mut Self { self.progress = progress; @@ -987,6 +1001,14 @@ impl IvfIndexBuilder let column = self.column.clone(); let frag_reuse_index = self.frag_reuse_index.clone(); let partition_adjustment = Arc::new(partition_adjustment); + let sub_index_accelerator = self.sub_index_accelerator.clone(); + // A cuVS resources handle owns one CUDA stream. Build partitions serially + // so concurrent partitions do not compete for all remaining GPU memory. + let partition_concurrency = if sub_index_accelerator.is_some() { + 1 + } else { + get_num_compute_intensive_cpus() + }; let build_iter = assign_batches .into_iter() @@ -1000,6 +1022,7 @@ impl IvfIndexBuilder let column = column.clone(); let frag_reuse_index = frag_reuse_index.clone(); let partition_adjustment = partition_adjustment.clone(); + let sub_index_accelerator = sub_index_accelerator.clone(); async move { let (is_affected, split_reader) = match partition_adjustment.as_ref() { Some(PartitionAdjustment::Split { @@ -1050,7 +1073,8 @@ impl IvfIndexBuilder loss += extra_loss; } - spawn_cpu(move || { + let is_accelerated = sub_index_accelerator.is_some(); + let build_partition = move || { // Apply assign_batch for join operations (splits no // longer use assign_batches) if let Some((assign_batch, deleted_row_ids)) = assign_batch { @@ -1089,14 +1113,22 @@ impl IvfIndexBuilder batches, column, frag_reuse_index, + sub_index_accelerator.as_ref(), )?; Ok(Some((storage, sub_index, loss))) - }) - .await + }; + + if is_accelerated { + // CAGRA waits on GPU work and must not occupy Lance's + // pure-CPU pool while its CUDA stream is running. + tokio::task::spawn_blocking(build_partition).await? + } else { + spawn_cpu(build_partition).await + } } }); Ok(stream::iter(build_iter) - .buffered(get_num_compute_intensive_cpus()) + .buffered(partition_concurrency) .boxed()) } @@ -1109,10 +1141,16 @@ impl IvfIndexBuilder batches: Vec, column: String, frag_reuse_index: Option>, + sub_index_accelerator: Option<&SubIndexBuildAccelerator>, ) -> Result<(Q::Storage, S)> { let storage = StorageBuilder::new(column, distance_type, quantizer, frag_reuse_index)? .build(batches)?; - let sub_index = S::index_vectors(&storage, sub_index_params)?; + let sub_index = match sub_index_accelerator { + Some(accelerator) => { + S::index_vectors_with_accelerator(&storage, sub_index_params, accelerator)? + } + None => S::index_vectors(&storage, sub_index_params)?, + }; Ok((storage, sub_index)) } diff --git a/rust/lance/src/index/vector/details.rs b/rust/lance/src/index/vector/details.rs index 3ccfd6e0e1e..92a4947abb9 100644 --- a/rust/lance/src/index/vector/details.rs +++ b/rust/lance/src/index/vector/details.rs @@ -32,7 +32,7 @@ use lance_index::vector::ivf::IvfBuildParams; use lance_index::vector::pq::PQBuildParams; use lance_index::vector::sq::builder::SQBuildParams; -use super::{StageParams, VectorIndexParams}; +use super::{CUVS_LIBRARY_RUNTIME_HINT, StageParams, VectorIndexParams}; use crate::dataset::Dataset; use crate::index::open_index_proto; use crate::{Error, Result}; @@ -95,6 +95,7 @@ pub fn vector_index_details(params: &VectorIndexParams) -> prost_types::Any { let mut hnsw_index_config = None; let mut compression = None; let mut runtime_hints: HashMap = params.runtime_hints.clone(); + runtime_hints.remove(CUVS_LIBRARY_RUNTIME_HINT); for stage in ¶ms.stages { match stage { @@ -1114,7 +1115,7 @@ mod tests { use lance_linalg::distance::DistanceType; // Non-default values for IVF and PQ hints - let params = VectorIndexParams::with_ivf_pq_params( + let mut params = VectorIndexParams::with_ivf_pq_params( DistanceType::L2, IvfBuildParams { max_iters: 100, @@ -1132,6 +1133,10 @@ mod tests { ..Default::default() }, ); + params.runtime_hints.insert( + CUVS_LIBRARY_RUNTIME_HINT.to_string(), + "/opt/libcuvs_c.so".to_string(), + ); let any = vector_index_details(¶ms); let details = any.to_msg::().unwrap(); @@ -1195,6 +1200,11 @@ mod tests { details.runtime_hints.get("lance.skip_transpose"), Some(&"false".to_string()) ); + assert!( + !details + .runtime_hints + .contains_key(CUVS_LIBRARY_RUNTIME_HINT) + ); // Roundtrip: apply hints back to a fresh params struct let mut restored = VectorIndexParams::with_ivf_pq_params( From 10649ff5d600e7d1b93d828e8c98b759472c6ea6 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 20:48:42 +0000 Subject: [PATCH 2/4] fix: preserve CAGRA SQ graph compatibility --- python/python/tests/test_vector_index.py | 41 +++ rust/lance-index/src/vector/hnsw/builder.rs | 62 +++- rust/lance-index/src/vector/hnsw/cagra.rs | 376 +++++++++++++++++++- 3 files changed, 470 insertions(+), 9 deletions(-) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 93e8715c6a4..46b79db426b 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -784,6 +784,47 @@ def test_create_index_cagra_accelerator_dispatch(tmp_path, caplog, monkeypatch): ) +@pytest.mark.cuda +@pytest.mark.parametrize("metric", ["l2", "cosine", "dot"]) +def test_create_index_cagra_accelerated_recall(tmp_path, caplog, 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() + + with caplog.at_level(logging.WARNING): + indexed = dataset.create_index( + "vector", + index_type="IVF_HNSW_SQ", + metric=metric, + num_partitions=1, + accelerator="cuda", + ) + assert not any( + "CAGRA HNSW build failed; falling back to CPU" in record.message + for record in caplog.records + ) + + 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( diff --git a/rust/lance-index/src/vector/hnsw/builder.rs b/rust/lance-index/src/vector/hnsw/builder.rs index 802d9f7dc8f..7946b3915c6 100644 --- a/rust/lance-index/src/vector/hnsw/builder.rs +++ b/rust/lance-index/src/vector/hnsw/builder.rs @@ -270,6 +270,7 @@ impl HNSW { /// local to the IVF partition and therefore line up with Lance's storage /// offsets without remapping. pub(crate) fn from_neighbor_graph( + storage: &impl VectorStore, params: HnswBuildParams, neighbors: Vec, graph_degree: usize, @@ -288,6 +289,12 @@ impl HNSW { ))); } let num_nodes = neighbors.len() / graph_degree; + if num_nodes != storage.len() { + return Err(Error::invalid_input(format!( + "CAGRA graph has {num_nodes} nodes, but vector storage has {} rows", + storage.len() + ))); + } if num_nodes > u32::MAX as usize { return Err(Error::invalid_input(format!( "CAGRA graph has {num_nodes} nodes, exceeding the u32 node-id limit" @@ -305,9 +312,21 @@ impl HNSW { let nodes = neighbors .chunks_exact(graph_degree) - .map(|neighbors| { + .enumerate() + .map(|(node_id, neighbors)| { + let dist_calculator = storage.dist_calculator_from_id(node_id as u32); + let neighbors_ranked = neighbors + .iter() + .map(|neighbor| { + OrderedNode::new(*neighbor, dist_calculator.distance(*neighbor).into()) + }) + .collect(); let neighbors = Arc::new(neighbors.to_vec()); - GraphBuilderNode::from_parts(vec![neighbors.clone()], vec![Vec::new()], neighbors) + GraphBuilderNode::from_parts( + vec![neighbors.clone()], + vec![neighbors_ranked], + neighbors, + ) }) .collect(); Ok(Self::from_parts(params, nodes, vec![num_nodes], 0)) @@ -1647,7 +1666,7 @@ mod tests { use crate::vector::storage::{DistCalculator, VectorStore}; use crate::vector::v3::subindex::{IvfSubIndex, SubIndexBuildAccelerator}; use crate::vector::{ - SQ_CODE_COLUMN, + DIST_COL, SQ_CODE_COLUMN, flat::storage::{FlatBinStorage, FlatFloatStorage}, graph::{DISTS_FIELD, NEIGHBORS_COL, NEIGHBORS_FIELD, OrderedNode, VisitedGenerator}, hnsw::{ @@ -2159,9 +2178,10 @@ mod tests { #[test] fn test_import_cagra_neighbor_graph() { + let storage = make_sq_storage(4); let params = HnswBuildParams::default(); let graph = vec![1, 2, 0, 2, 0, 3, 1, 2]; - let hnsw = HNSW::from_neighbor_graph(params, graph, 2).unwrap(); + let hnsw = HNSW::from_neighbor_graph(&storage, params, graph, 2).unwrap(); assert_eq!(hnsw.len(), 4); assert_eq!(hnsw.max_level(), 1); @@ -2176,6 +2196,36 @@ mod tests { neighbors.value(3).as_primitive::().values(), &[1, 2] ); + let distances = batch[DIST_COL].as_list::(); + for node_id in 0..batch.num_rows() { + let node_neighbors = neighbors.value(node_id); + let node_neighbors = node_neighbors.as_primitive::(); + let node_distances = distances.value(node_id); + let node_distances = node_distances.as_primitive::(); + + assert_eq!(node_neighbors.len(), node_distances.len()); + for (neighbor, distance) in node_neighbors.values().iter().zip(node_distances.values()) + { + assert_eq!( + *distance, + storage.dist_between(node_id as u32, *neighbor), + "serialized distance for edge {node_id}->{neighbor}" + ); + } + } + + // Released v3 readers reconstruct adjacency by zipping these two + // lists. Keep this exact compatibility contract covered even though + // the current zero-copy reader no longer materializes edge distances. + let released_reader_edge_count = (0..batch.num_rows()) + .map(|node_id| { + neighbors + .value(node_id) + .len() + .min(distances.value(node_id).len()) + }) + .sum::(); + assert_eq!(released_reader_edge_count, 8); let loaded = HNSW::load(batch).unwrap(); assert_eq!(loaded.len(), 4); @@ -2184,8 +2234,10 @@ mod tests { #[test] fn test_import_cagra_neighbor_graph_rejects_invalid_node() { + let storage = make_sq_storage(2); let error = - HNSW::from_neighbor_graph(HnswBuildParams::default(), vec![1, 0, 0, 3], 2).unwrap_err(); + HNSW::from_neighbor_graph(&storage, HnswBuildParams::default(), vec![1, 0, 0, 3], 2) + .unwrap_err(); assert!(matches!(error, Error::InvalidInput { .. })); assert!(error.to_string().contains("neighbor id 3")); diff --git a/rust/lance-index/src/vector/hnsw/cagra.rs b/rust/lance-index/src/vector/hnsw/cagra.rs index 8b8b39525ce..1120dbebb16 100644 --- a/rust/lance-index/src/vector/hnsw/cagra.rs +++ b/rust/lance-index/src/vector/hnsw/cagra.rs @@ -165,7 +165,7 @@ fn dl_error_message() -> String { #[cfg(unix)] struct CuvsApi { - _library: DynamicLibrary, + _library: Option, get_last_error_text: GetLastErrorText, resources_create: ResourcesCreate, resources_destroy: ResourcesDestroy, @@ -185,6 +185,11 @@ struct CuvsApi { #[cfg(unix)] impl CuvsApi { fn load(path: &Path) -> Result { + #[cfg(test)] + if path == Path::new(tests::FAKE_CUVS_LIBRARY_PATH) { + return Ok(tests::fake_api()); + } + let library = DynamicLibrary::open(path)?; Ok(Self { get_last_error_text: library.symbol(b"cuvsGetLastErrorText\0")?, @@ -201,7 +206,7 @@ impl CuvsApi { index_destroy: library.symbol(b"cuvsCagraIndexDestroy\0")?, cagra_build: library.symbol(b"cuvsCagraBuild\0")?, index_get_graph: library.symbol(b"cuvsCagraIndexGetGraph\0")?, - _library: library, + _library: Some(library), }) } @@ -340,7 +345,10 @@ fn host_tensor(values: &mut [T], shape: &mut [i64], dtype: DLDataType) -> DLM fn cuvs_distance_type(distance_type: DistanceType) -> Result { match distance_type { DistanceType::L2 => Ok(0), - DistanceType::Cosine => Ok(2), + // SQ evaluates cosine queries with scaled squared L2 over normalized + // codes. Reconstructed vectors are not exactly unit length, so asking + // CAGRA for cosine here can produce a different topology ordering. + DistanceType::Cosine => Ok(0), DistanceType::Dot => Ok(6), DistanceType::Hamming => Err(Error::invalid_input( "CAGRA does not support Lance bitwise Hamming vectors".to_string(), @@ -550,6 +558,366 @@ pub(super) fn build( num_rows, graph_degree ); - HNSW::from_neighbor_graph(params, neighbors, graph_degree) + HNSW::from_neighbor_graph(storage, params, neighbors, graph_degree) + } +} + +#[cfg(all(test, unix))] +mod tests { + use std::collections::HashSet; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow_array::{ArrayRef, FixedSizeListArray, Float32Array, RecordBatch, UInt8Array}; + use lance_arrow::FixedSizeListArrayExt; + use lance_core::ROW_ID; + use rstest::rstest; + + use super::*; + use crate::vector::SQ_CODE_COLUMN; + use crate::vector::hnsw::builder::HnswQueryParams; + use crate::vector::storage::DistCalculator; + use crate::vector::v3::subindex::{IvfSubIndex, SubIndexBuildAccelerator}; + + pub(super) const FAKE_CUVS_LIBRARY_PATH: &str = "/fake/libcuvs_c.so"; + + static BUILD_CALLS: AtomicUsize = AtomicUsize::new(0); + + struct FakeParams { + num_rows: usize, + graph_degree: usize, + } + + struct FakeIndex { + graph: Vec, + shape: [i64; 2], + } + + unsafe extern "C" fn fake_get_last_error_text() -> *const c_char { + static ERROR: &[u8] = b"fake cuVS error\0"; + ERROR.as_ptr().cast() + } + + unsafe extern "C" fn fake_resources_create(resources: *mut CuvsResources) -> CuvsStatus { + if resources.is_null() { + return 0; + } + // SAFETY: the fake ABI caller supplies a valid output pointer. + unsafe { resources.write(1) }; + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_resources_destroy(_resources: CuvsResources) -> CuvsStatus { + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_stream_sync(_resources: CuvsResources) -> CuvsStatus { + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_matrix_copy( + _resources: CuvsResources, + source: *mut DLManagedTensor, + destination: *mut DLManagedTensor, + ) -> CuvsStatus { + if source.is_null() || destination.is_null() { + return 0; + } + // SAFETY: both managed tensors are live for this synchronous fake call. + let (source, destination) = unsafe { (&*source, &mut *destination) }; + let source_tensor = &source.dl_tensor; + let destination_tensor = &mut destination.dl_tensor; + if source_tensor.ndim != 2 + || source_tensor.shape.is_null() + || source_tensor.data.is_null() + || destination_tensor.data.is_null() + { + return 0; + } + // SAFETY: rank two was checked above and the fake index owns both shape values. + let shape = unsafe { std::slice::from_raw_parts(source_tensor.shape, 2) }; + let Ok(rows) = usize::try_from(shape[0]) else { + return 0; + }; + let Ok(columns) = usize::try_from(shape[1]) else { + return 0; + }; + let Some(len) = rows.checked_mul(columns) else { + return 0; + }; + // SAFETY: source and destination were allocated for the matching graph shape. + unsafe { + ptr::copy_nonoverlapping( + source_tensor.data.cast::(), + destination_tensor.data.cast::(), + len, + ) + }; + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_dataset_make_standard_view( + _resources: CuvsResources, + tensor: *mut DLManagedTensor, + dataset: *mut CuvsDataset, + ) -> CuvsStatus { + if tensor.is_null() || dataset.is_null() { + return 0; + } + // SAFETY: the tensor outlives the synchronous fake CAGRA build. + unsafe { dataset.write(tensor.cast()) }; + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_dataset_destroy(_dataset: CuvsDataset) -> CuvsStatus { + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_params_create(params: *mut CuvsCagraParams) -> CuvsStatus { + if params.is_null() { + return 0; + } + let params_value = Box::new(FakeParams { + num_rows: 0, + graph_degree: 0, + }); + // SAFETY: ownership transfers to the matching fake destroy function. + unsafe { params.write(Box::into_raw(params_value).cast()) }; + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_params_destroy(params: CuvsCagraParams) -> CuvsStatus { + if !params.is_null() { + // SAFETY: this pointer was allocated by fake_params_create and is dropped once. + unsafe { drop(Box::from_raw(params.cast::())) }; + } + CUVS_SUCCESS + } + + #[allow(clippy::too_many_arguments)] + unsafe extern "C" fn fake_params_from_hnsw( + params: CuvsCagraParams, + num_rows: i64, + _dim: i64, + m: c_int, + _ef_construction: c_int, + _strategy: c_int, + _metric: c_int, + ) -> CuvsStatus { + if params.is_null() { + return 0; + } + let (Ok(num_rows), Ok(graph_degree)) = (usize::try_from(num_rows), usize::try_from(m)) + else { + return 0; + }; + // SAFETY: this pointer was allocated by fake_params_create and remains live. + let params = unsafe { &mut *params.cast::() }; + params.num_rows = num_rows; + params.graph_degree = graph_degree; + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_index_create(index: *mut CuvsCagraIndex) -> CuvsStatus { + if index.is_null() { + return 0; + } + let index_value = Box::new(FakeIndex { + graph: Vec::new(), + shape: [0, 0], + }); + // SAFETY: ownership transfers to the matching fake destroy function. + unsafe { index.write(Box::into_raw(index_value).cast()) }; + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_index_destroy(index: CuvsCagraIndex) -> CuvsStatus { + if !index.is_null() { + // SAFETY: this pointer was allocated by fake_index_create and is dropped once. + unsafe { drop(Box::from_raw(index.cast::())) }; + } + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_cagra_build( + _resources: CuvsResources, + params: CuvsCagraParams, + dataset: CuvsDataset, + index: CuvsCagraIndex, + ) -> CuvsStatus { + if params.is_null() || dataset.is_null() || index.is_null() { + return 0; + } + // SAFETY: the fake create functions allocated both live handles. + let (params, index) = unsafe { + ( + &*params.cast::(), + &mut *index.cast::(), + ) + }; + let graph_degree = params.graph_degree.min(params.num_rows.saturating_sub(1)); + if graph_degree == 0 { + return 0; + } + index.graph = (0..params.num_rows) + .flat_map(|node| { + (1..=graph_degree).map(move |offset| ((node + offset) % params.num_rows) as u32) + }) + .collect(); + index.shape = [params.num_rows as i64, graph_degree as i64]; + BUILD_CALLS.fetch_add(1, Ordering::Relaxed); + CUVS_SUCCESS + } + + unsafe extern "C" fn fake_index_get_graph( + index: CuvsCagraIndex, + graph: *mut DLManagedTensor, + ) -> CuvsStatus { + if index.is_null() || graph.is_null() { + return 0; + } + // SAFETY: the fake index remains live until after the graph copy. + let index = unsafe { &mut *index.cast::() }; + let managed = DLManagedTensor { + dl_tensor: DLTensor { + data: index.graph.as_mut_ptr().cast(), + device: DLDevice { + device_type: DL_CPU, + device_id: 0, + }, + ndim: 2, + dtype: DLDataType { + code: DL_UINT, + bits: 32, + lanes: 1, + }, + shape: index.shape.as_mut_ptr(), + strides: ptr::null_mut(), + byte_offset: 0, + }, + manager_ctx: ptr::null_mut(), + deleter: None, + }; + // SAFETY: the caller provided uninitialized output storage for the full value. + unsafe { graph.write(managed) }; + CUVS_SUCCESS + } + + pub(super) fn fake_api() -> CuvsApi { + CuvsApi { + _library: None, + get_last_error_text: fake_get_last_error_text, + resources_create: fake_resources_create, + resources_destroy: fake_resources_destroy, + stream_sync: fake_stream_sync, + matrix_copy: fake_matrix_copy, + dataset_make_standard_view: fake_dataset_make_standard_view, + dataset_destroy: fake_dataset_destroy, + params_create: fake_params_create, + params_destroy: fake_params_destroy, + params_from_hnsw: fake_params_from_hnsw, + index_create: fake_index_create, + index_destroy: fake_index_destroy, + cagra_build: fake_cagra_build, + index_get_graph: fake_index_get_graph, + } + } + + fn make_sq_storage(distance_type: DistanceType) -> (ScalarQuantizationStorage, ArrayRef) { + const NUM_ROWS: usize = 64; + const DIM: usize = 8; + const QUERY_ROW: usize = 17; + + let sq_codes = (0..NUM_ROWS) + .flat_map(|row| { + (0..DIM).map(move |column| ((row * 37 + column * 53 + row * column) % 256) as u8) + }) + .collect::>(); + let query = Arc::new(Float32Array::from( + sq_codes[QUERY_ROW * DIM..(QUERY_ROW + 1) * DIM] + .iter() + .map(|code| -1.0 + 2.0 * *code as f32 / 255.0) + .collect::>(), + )) as ArrayRef; + let codes = FixedSizeListArray::try_new_from_values(UInt8Array::from(sq_codes), DIM as i32) + .unwrap(); + let batch = RecordBatch::try_from_iter([ + ( + ROW_ID, + Arc::new(arrow_array::UInt64Array::from_iter_values( + 0..NUM_ROWS as u64, + )) as ArrayRef, + ), + (SQ_CODE_COLUMN, Arc::new(codes) as ArrayRef), + ]) + .unwrap(); + let storage = + ScalarQuantizationStorage::try_new(8, distance_type, -1.0..1.0, [batch], None).unwrap(); + (storage, query) + } + + #[rstest] + #[case::l2(DistanceType::L2, 0)] + #[case::cosine(DistanceType::Cosine, 0)] + #[case::dot(DistanceType::Dot, 6)] + fn test_cuvs_distance_matches_sq_contract( + #[case] distance_type: DistanceType, + #[case] expected: c_int, + ) { + assert_eq!(cuvs_distance_type(distance_type).unwrap(), expected); + } + + #[rstest] + #[case::l2(DistanceType::L2)] + #[case::cosine(DistanceType::Cosine)] + #[case::dot(DistanceType::Dot)] + fn test_successful_cagra_backend_recall(#[case] distance_type: DistanceType) { + let (storage, query) = make_sq_storage(distance_type); + let accelerator = SubIndexBuildAccelerator::cagra(FAKE_CUVS_LIBRARY_PATH); + let calls_before = BUILD_CALLS.load(Ordering::Relaxed); + let hnsw = HNSW::index_vectors_with_accelerator( + &storage, + HnswBuildParams::default().num_edges(4).ef_construction(10), + &accelerator, + ) + .unwrap(); + assert!(BUILD_CALLS.load(Ordering::Relaxed) > calls_before); + assert_eq!(hnsw.max_level(), 1); + + // Search the serialized graph so recall also covers the v3 write/read path. + let loaded = HNSW::load(hnsw.to_batch().unwrap()).unwrap(); + let k = 10; + let query_params = HnswQueryParams { + ef: storage.len(), + lower_bound: None, + upper_bound: None, + dist_q_c: 0.0, + use_acorn: false, + }; + let results = loaded + .search_basic(query.clone(), k, &query_params, None, &storage) + .unwrap(); + + let distances = storage + .dist_calculator(query, 0.0) + .distance_all(storage.len()); + let mut truth = distances.into_iter().enumerate().collect::>(); + truth.sort_by(|(left_id, left), (right_id, right)| { + left.total_cmp(right).then_with(|| left_id.cmp(right_id)) + }); + let truth = truth + .into_iter() + .take(k) + .map(|(id, _)| id as u32) + .collect::>(); + let hits = results + .iter() + .filter(|result| truth.contains(&result.id)) + .count(); + let recall = hits as f32 / k as f32; + assert!( + recall >= 0.5, + "{distance_type} CAGRA recall {recall} is below 0.5" + ); } } From b68dc8160020ab4178d8612e3c1a0f5185971254 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 21:49:29 +0000 Subject: [PATCH 3/4] test: require CAGRA in accelerated recall coverage --- python/python/lance/dataset.py | 20 ++++++++ python/python/tests/test_vector_index.py | 37 +++++++++----- python/src/dataset.rs | 7 +++ rust/lance-index/src/vector/hnsw/builder.rs | 53 +++++++++++++++++++-- rust/lance-index/src/vector/v3/subindex.rs | 15 ++++++ rust/lance/src/index/vector.rs | 14 +++++- rust/lance/src/index/vector/details.rs | 13 ++++- 7 files changed, 142 insertions(+), 17 deletions(-) diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index ecce451b6ea..7704be59029 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -3815,6 +3815,19 @@ 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 = {} cagra_library = None @@ -3827,12 +3840,19 @@ def _create_index_impl( 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 diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index 46b79db426b..f26a63e7068 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -784,9 +784,26 @@ def test_create_index_cagra_accelerator_dispatch(tmp_path, caplog, monkeypatch): ) +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="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, caplog, metric): +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( @@ -804,17 +821,13 @@ def test_create_index_cagra_accelerated_recall(tmp_path, caplog, metric): nearest = {"column": "vector", "q": query, "k": 10, "metric": metric} ground_truth = dataset.to_table(columns=["id"], nearest=nearest)["id"].to_numpy() - with caplog.at_level(logging.WARNING): - indexed = dataset.create_index( - "vector", - index_type="IVF_HNSW_SQ", - metric=metric, - num_partitions=1, - accelerator="cuda", - ) - assert not any( - "CAGRA HNSW build failed; falling back to CPU" in record.message - for record in caplog.records + 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() diff --git a/python/src/dataset.rs b/python/src/dataset.rs index 8c6dee34ec7..30439bf527a 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -5162,6 +5162,13 @@ fn prepare_vector_index_params( .runtime_hints .insert("lancedb.cuvs_library".to_string(), library_path.extract()?); } + if let Some(required) = kwargs.get_item("_require_cagra")? + && required.extract::()? + { + params + .runtime_hints + .insert("lancedb.cuvs_required".to_string(), "true".to_string()); + } } Ok(params) } diff --git a/rust/lance-index/src/vector/hnsw/builder.rs b/rust/lance-index/src/vector/hnsw/builder.rs index 7946b3915c6..9d550d04fb5 100644 --- a/rust/lance-index/src/vector/hnsw/builder.rs +++ b/rust/lance-index/src/vector/hnsw/builder.rs @@ -1525,13 +1525,27 @@ impl IvfSubIndex for HNSW { { match accelerator { SubIndexBuildAccelerator::Cagra(accelerator) => { - if accelerator.is_disabled() - || !super::cagra::supports_partition(storage.len(), ¶ms) - { + if accelerator.is_disabled() { + debug_assert!( + !accelerator.is_required(), + "required CAGRA acceleration cannot be disabled" + ); + return Self::index_vectors(storage, params); + } + if !super::cagra::supports_partition(storage.len(), ¶ms) { + if accelerator.is_required() { + return Err(Error::not_supported(format!( + "CAGRA was required, but partition with {} rows is too small for HNSW parameters m={} and ef_construction={}", + storage.len(), + params.m, + params.ef_construction + ))); + } return Self::index_vectors(storage, params); } match super::cagra::build(storage, params.clone(), accelerator.library_path()) { Ok(index) => Ok(index), + Err(error) if accelerator.is_required() => Err(error), Err(error) => { if accelerator.disable() { log::warn!( @@ -2261,6 +2275,39 @@ mod tests { assert!(accelerator.is_disabled()); } + #[test] + fn test_cagra_missing_library_errors_when_required() { + let storage = make_sq_storage(8); + let accelerator = SubIndexBuildAccelerator::cagra_required("/missing/libcuvs_c.so"); + + let error = HNSW::index_vectors_with_accelerator( + &storage, + HnswBuildParams::default().num_edges(4).ef_construction(10), + &accelerator, + ) + .unwrap_err(); + + assert!(matches!(error, Error::IO { .. })); + let SubIndexBuildAccelerator::Cagra(accelerator) = accelerator; + assert!(!accelerator.is_disabled()); + } + + #[test] + fn test_cagra_small_partition_errors_when_required() { + let storage = make_sq_storage(4); + let accelerator = SubIndexBuildAccelerator::cagra_required("/missing/libcuvs_c.so"); + + let error = HNSW::index_vectors_with_accelerator( + &storage, + HnswBuildParams::default().num_edges(4).ef_construction(10), + &accelerator, + ) + .unwrap_err(); + + assert!(matches!(error, Error::NotSupported { .. })); + assert!(error.to_string().contains("partition with 4 rows")); + } + #[test] fn test_cagra_small_partition_does_not_disable_acceleration() { let storage = make_sq_storage(4); diff --git a/rust/lance-index/src/vector/v3/subindex.rs b/rust/lance-index/src/vector/v3/subindex.rs index 1f7894d409d..20e79189213 100644 --- a/rust/lance-index/src/vector/v3/subindex.rs +++ b/rust/lance-index/src/vector/v3/subindex.rs @@ -35,6 +35,16 @@ impl SubIndexBuildAccelerator { Self::Cagra(CagraBuildAccelerator { library_path: library_path.into(), is_disabled: Arc::new(AtomicBool::new(false)), + is_required: false, + }) + } + + /// Create a cuVS CAGRA accelerator that does not fall back to CPU. + pub fn cagra_required(library_path: impl Into) -> Self { + Self::Cagra(CagraBuildAccelerator { + library_path: library_path.into(), + is_disabled: Arc::new(AtomicBool::new(false)), + is_required: true, }) } } @@ -47,6 +57,7 @@ impl SubIndexBuildAccelerator { pub struct CagraBuildAccelerator { library_path: String, is_disabled: Arc, + is_required: bool, } impl CagraBuildAccelerator { @@ -58,6 +69,10 @@ impl CagraBuildAccelerator { self.is_disabled.load(Ordering::Relaxed) } + pub(crate) fn is_required(&self) -> bool { + self.is_required + } + /// Disable the accelerator and return whether this call changed its state. pub(crate) fn disable(&self) -> bool { !self.is_disabled.swap(true, Ordering::Relaxed) diff --git a/rust/lance/src/index/vector.rs b/rust/lance/src/index/vector.rs index b17df9dc795..c5617447861 100644 --- a/rust/lance/src/index/vector.rs +++ b/rust/lance/src/index/vector.rs @@ -77,6 +77,9 @@ pub const LANCE_VECTOR_INDEX: &str = "__lance_vector_index"; /// details are persisted. pub(crate) const CUVS_LIBRARY_RUNTIME_HINT: &str = "lancedb.cuvs_library"; +/// Transient test mode requiring each HNSW partition build to use CAGRA or fail. +pub(crate) const CUVS_REQUIRED_RUNTIME_HINT: &str = "lancedb.cuvs_required"; + /// A materialized snapshot of one logical vector index and all of its segments. #[derive(Debug)] pub struct LogicalVectorIndex { @@ -514,7 +517,16 @@ fn cagra_accelerator(params: &VectorIndexParams) -> Option prost_types::Any { let mut compression = None; let mut runtime_hints: HashMap = params.runtime_hints.clone(); runtime_hints.remove(CUVS_LIBRARY_RUNTIME_HINT); + runtime_hints.remove(CUVS_REQUIRED_RUNTIME_HINT); for stage in ¶ms.stages { match stage { @@ -1137,6 +1140,9 @@ mod tests { CUVS_LIBRARY_RUNTIME_HINT.to_string(), "/opt/libcuvs_c.so".to_string(), ); + params + .runtime_hints + .insert(CUVS_REQUIRED_RUNTIME_HINT.to_string(), "true".to_string()); let any = vector_index_details(¶ms); let details = any.to_msg::().unwrap(); @@ -1205,6 +1211,11 @@ mod tests { .runtime_hints .contains_key(CUVS_LIBRARY_RUNTIME_HINT) ); + assert!( + !details + .runtime_hints + .contains_key(CUVS_REQUIRED_RUNTIME_HINT) + ); // Roundtrip: apply hints back to a fresh params struct let mut restored = VectorIndexParams::with_ivf_pq_params( From eb05f37caf8dbddcf2f3662b12c26f6cc5fdc13d Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 22:29:01 +0000 Subject: [PATCH 4/4] test: accept platform-specific CAGRA failure --- python/python/tests/test_vector_index.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index f26a63e7068..49fb3e9bfe1 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -791,7 +791,13 @@ def test_create_index_cagra_required_disables_fallback(tmp_path, monkeypatch): ) dataset = lance.write_dataset(create_table(), tmp_path) - with pytest.raises(OSError, match="failed to load cuVS library"): + 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",