Feature: Support High Performance CAGRA GPU Vector Index

Motivation

CAGRA is an ANN index developed by NVIDIA RAPIDS cuVS, offering the following advantages:

• Supports GPU training and search, providing high recall and low latency in GPU-accelerated environments.
• Offers superior search performance on large-scale datasets compared to HNSW/IVF_PQ.
• Other vector databases (such as Milvus) already support CAGRA index, significantly improving index creation and search performance.

This feature will integrate cuVS's CAGRA API into Lance's vector index infrastructure, enabling users to create and query CAGRA index via Python SDK.

We present two implementation approaches for community discussion and selection. Preliminary benchmark tests show that, compared to traditional CPU indexes, Lance's integrated CAGRA GPU index achieves higher index creation and search performance, and also achieves an order-of-magnitude breakthrough in QPS.

---

Approach 1: Python-Side Integration

Implementation Principle

Import the cuVS dependency into the Lance Python module and integrate the CAGRA Python API on the Python side to complete index creation and searching.

Sequence Diagram

Index Creation

sequenceDiagram
    participant UserClient as User Client
    participant Lance as Lance Dataset
    participant CuPy as CuPy (GPU)
    participant CAGRA
    participant FS as File System

    UserClient->>Lance: dataset.create_index(...index_type="CAGRA"...)
    activate Lance
    
    Lance->>CuPy: Convert vector column to GPU Array:<br> vectors_cp = cp.asarray(vectors_np)
    activate CuPy
    
    Lance->>CAGRA: Construct Cagra Params: <br>params = cagra.IndexParams(*kwargs)
    activate CAGRA
    CAGRA->>CAGRA: Create Index:<br>index = cagra.build(params, vectors_cp)
    
    CAGRA->>FS: Write Index to _indices/{uuid}/index.idx<br>cagra.save(path, index)
    activate FS
    
    Lance->>Lance: Commit Index Metadata, update manifest:<br> dataset.commit()
    deactivate Lance

search

sequenceDiagram
    participant User Client
    participant Lance as Lance Dataset
    participant CAGRA as CAGRA
 
    User Client->>Lance: Load Index Object: index = load_index() via cagra.load
    User Client->>Lance: to_table(index,nearest={column='emb', q=query, k=10})
    Lance->>CAGRA: Data format Conversion: queries = cp.asarray(q)

    CAGRA->>CAGRA: Construct search_params via cagra.SearchParams()
    CAGRA->>CAGRA: Execute search: cagra.search(search_params,index, queries, k)
    CAGRA-->>Lance: Returns top k result: (neighbors, distances)
    Lance->>Lance: Map neighbors id to original rows via Dataset.take
    Lance-->>User Client: Returns final pa.Table result

Process description

Index Creation:

1. Create Entry: dataset.create_index(...index_type="CAGRA").
2. Data Conversion: Read vector data from Lance dataset and convert to CuPy array.
3. Index Build: In the create_index function of the Lance Python SDK, add a new branch index_type == "CAGRA", which calls cagra.build() to create the index.
4. Index Persistence: Call cagra.save() to save the index to disk.
5. Metadata Commit: Manually create Lance index metadata and submit transaction via dataset.commit().

Search:

1. Pre-load Index: Add a new load_index interface to the Lance Python SDK, which preloads the index by calling cagra.load(path). This aims to decouple index loading from retrieval execution, preventing the loading process from degrading query performance.
2. Search Entry: dataset.to_table(index, nearest={...}), the index param is Optional.
3. Query Vector Conversion: Convert query vectors to CuPy array.
4. Execute Search: In the scanner function of the Lance Python SDK, add a new CAGRA branch, which create a CagraScanner via cagra.search() to get nearest neighbor IDs and distances.
5. Result Mapping: Use take() operation to map CAGRA's internal IDs to Lance's _rowid and assemble into Arrow Table.

---

Approach 2: Rust-Side Integration and Python Adaptation

Implementation Principle

Deeply integrate CAGRA Rust API on the Rust side as one of Lance's native vector index types. Expose it to Python through PyO3.

Architecture Diagram

graph TD
    A[Python API Layer]
    A --> A1[dataset.create_index]
    A --> A2[scanner.nearest]
    Z--> C[Rust Layer: lance]

    C --> D[ExecutionPlan ]
    A1 --> Z[PyO3 Layer]
    A2 --> Z[PyO3 Layer]
    
    B --> B2[VectorIndex trait implementation: CagraVectorIndex]
  
    B[Rust Layer: lance-index] --> B3[VectorIndexParams trait implementation：CagraBuildParams]
    
    C --> C1[CreateIndexBuilder]
    C1 --> C2[Create Cagra Index]
    C1 --> C3[Index Persistence]
    C1 --> C4[Commit Transaction]
    
    D --> D1[CagraSearchExec]
    D --> D2[Execute search]
    D2 --> B
    C2 --> B
    style B2 fill:#add8e6,stroke:#333
    style B3 fill:#add8e6,stroke:#333
    style C2 fill:#add8e6,stroke:#333
    style D1 fill:#add8e6,stroke:#333
   

Process description

Index Creation:

1. Create Entry: dataset.create_index(...index_type="CAGRA").
2. PyO3 Parameter Parsing: In the prepare_vector_index_params function, add a process to parse the cagra IndexParams params, mapping Python parameters to Rust's VectorIndexParams。
3. PyO3 CreateIndexBuilder initing: Call CreateIndexBuilder.into_future() to trigger index building and transaction commit.
4. Index Creating: Rust layer Call cuvs::cagra::Index::build to create cagra index.
5. Index Persistencee: Serialize CAGRA index and write to _indices/{uuid}/index.idx.
6. Commit Transaction: Update dataset manifest file.

Search:

1. Search Entry: dataset.to_table(nearest={...}).
2. Rust Plan Creation: Add a new ExecutionPlan node named CagraSearchExec and construct a scan plan containing this node using LanceScanner::create_plan().
3. Plan Execution: DataFusion executes the physical plan, calling cuvs::cagra::Index::search on the CagraSearchExec node to perform the query.
4. Result Packaging: The CAGRA search results (neighbors, distances) will be encapsulated as a RecordBatchStream streaming interface and connected to the Python layer via PyArrow Reader, ultimately being converted into a complete pa.Table format.

---

Comparison

Aspect	Approach 1	Approach 2
Intrusiveness	✅ Low Only modifies Python code	❌High The Rust needs to be modified.
Development Cost	✅ Low Directly call to cuVS Python API	❌High Need to implement Rust index type, execution node, etc.
Dependency Management	✅ Flexible The compilation and runtime depend on cuVS Python package: cuvs-cu12, cupy-cuda12, etc. No need to compile the cuVS rust and underlying C++ dynamic link library yourself.	❌ Heavy The cuVS C++ dynamic linking library becomes a dependency for both Rust and Python compilation and runtime. Building the Wheel package involves complex dynamic linking library packages.
Performance	❌ Low Data must be stored across disk, CPU memory, and GPU. For large-scale data processing, GPU memory consumption is extremely high.	✅ High IPreliminary tests indicate that GPU memory usage is slightly lower than in approach 1, while index creation performance is slightly higher than in approach 1.
Evolution Capability	❌ Low If LanceDB needs to natively integrate with CAGRA in the future, it will result in a dual implementation of core business logic.	✅ High You only need to maintain the core business logic on the Lance Rust side, while LanceDB is mainly responsible for interface adaptation.
Ecosystem Fit	❌ Low It primarily relies on cuVS CAGRA's native index capabilities, is independent of Lance's native index framework implementation, and is not deeply integrated with DataFusion's concepts at the execution level.	✅ High It is highly compatible with Lance's native index architecture and seamlessly integrates into DataFusion's query execution logic.

First review the two proposed approaches and resolve which integration path is selected. For the Rust path, inspect prepare_vector_index_params, CreateIndexBuilder, the lance-index VectorIndex traits, and the proposed CagraSearchExec entry point; for the Python path, inspect dataset.create_index, load_index, and scanner. Done means CAGRA creation, persistence, search, and result packaging work through the selected API.