rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2025-12-23 05:19:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: use GPU for index creation. (#540)

Browse Source 
Bump lance to 0.8.3 to include GPU training

---------

Co-authored-by: Rob Meng <rob.xu.meng@gmail.com>
This commit is contained in:
Lei Xu
2023-10-05 20:49:00 -07:00
committed by GitHub
parent 88d8d7249e
commit a26c8f3316
8 changed files with 104 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
Cargo.toml
View File

@@ -5,9 +5,9 @@ exclude = ["python"]
resolver = "2" resolver = "2"
[workspace.dependencies] [workspace.dependencies]
lance = { "version" = "=0.8.1", "features" = ["dynamodb"] } lance = { "version" = "=0.8.3", "features" = ["dynamodb"] }
lance-linalg = { "version" = "=0.8.1" } lance-linalg = { "version" = "=0.8.3" }
lance-testing = { "version" = "=0.8.1" } lance-testing = { "version" = "=0.8.3" }
# Note that this one does not include pyarrow # Note that this one does not include pyarrow
arrow = { version = "43.0.0", optional = false } arrow = { version = "43.0.0", optional = false }
arrow-array = "43.0" arrow-array = "43.0"

158
README.md
View File

@@ -1,78 +1,80 @@
<div align="center"> <div align="center">
<p align="center"> <p align="center">
<img width="275" alt="LanceDB Logo" src="https://user-images.githubusercontent.com/917119/226205734-6063d87a-1ecc-45fe-85be-1dea6383a3d8.png"> <img width="275" alt="LanceDB Logo" src="https://user-images.githubusercontent.com/917119/226205734-6063d87a-1ecc-45fe-85be-1dea6383a3d8.png">
**Developer-friendly, serverless vector database for AI applications** **Developer-friendly, serverless vector database for AI applications**
<a href="https://lancedb.github.io/lancedb/">Documentation</a> • <a href="https://lancedb.github.io/lancedb/">Documentation</a> •
<a href="https://blog.lancedb.com/">Blog</a> • <a href="https://blog.lancedb.com/">Blog</a> •
<a href="https://discord.gg/zMM32dvNtd">Discord</a> • <a href="https://discord.gg/zMM32dvNtd">Discord</a> •
<a href="https://twitter.com/lancedb">Twitter</a> <a href="https://twitter.com/lancedb">Twitter</a>
</p> </p>
<img max-width="750px" alt="LanceDB Multimodal Search" src="https://github.com/lancedb/lancedb/assets/917119/09c5afc5-7816-4687-bae4-f2ca194426ec"> <img max-width="750px" alt="LanceDB Multimodal Search" src="https://github.com/lancedb/lancedb/assets/917119/09c5afc5-7816-4687-bae4-f2ca194426ec">
</p> </p>
</div> </div>
<hr /> <hr />
LanceDB is an open-source database for vector-search built with persistent storage, which greatly simplifies retrevial, filtering and management of embeddings. LanceDB is an open-source database for vector-search built with persistent storage, which greatly simplifies retrevial, filtering and management of embeddings.
The key features of LanceDB include: The key features of LanceDB include:
* Production-scale vector search with no servers to manage. * Production-scale vector search with no servers to manage.
* Store, query and filter vectors, metadata and multi-modal data (text, images, videos, point clouds, and more). * Store, query and filter vectors, metadata and multi-modal data (text, images, videos, point clouds, and more).
* Support for vector similarity search, full-text search and SQL. * Support for vector similarity search, full-text search and SQL.
* Native Python and Javascript/Typescript support. * Native Python and Javascript/Typescript support.
* Zero-copy, automatic versioning, manage versions of your data without needing extra infrastructure. * Zero-copy, automatic versioning, manage versions of your data without needing extra infrastructure.
* Ecosystem integrations with [LangChain 🦜️🔗](https://python.langchain.com/en/latest/modules/indexes/vectorstores/examples/lanecdb.html), [LlamaIndex 🦙](https://gpt-index.readthedocs.io/en/latest/examples/vector_stores/LanceDBIndexDemo.html), Apache-Arrow, Pandas, Polars, DuckDB and more on the way. * GPU support in building vector index(*).
LanceDB's core is written in Rust 🦀 and is built using <a href="https://github.com/lancedb/lance">Lance</a>, an open-source columnar format designed for performant ML workloads. * Ecosystem integrations with [LangChain 🦜️🔗](https://python.langchain.com/en/latest/modules/indexes/vectorstores/examples/lanecdb.html), [LlamaIndex 🦙](https://gpt-index.readthedocs.io/en/latest/examples/vector_stores/LanceDBIndexDemo.html), Apache-Arrow, Pandas, Polars, DuckDB and more on the way.
## Quick Start LanceDB's core is written in Rust 🦀 and is built using <a href="https://github.com/lancedb/lance">Lance</a>, an open-source columnar format designed for performant ML workloads.
**Javascript** ## Quick Start
```shell
npm install vectordb **Javascript**
``` ```shell
npm install vectordb
```javascript ```
const lancedb = require('vectordb');
const db = await lancedb.connect('data/sample-lancedb'); ```javascript
const lancedb = require('vectordb');
const table = await db.createTable('vectors', const db = await lancedb.connect('data/sample-lancedb');
[{ id: 1, vector: [0.1, 0.2], item: "foo", price: 10 },
{ id: 2, vector: [1.1, 1.2], item: "bar", price: 50 }]) const table = await db.createTable('vectors',
[{ id: 1, vector: [0.1, 0.2], item: "foo", price: 10 },
const query = table.search([0.1, 0.3]); { id: 2, vector: [1.1, 1.2], item: "bar", price: 50 }])
query.limit = 20;
const results = await query.execute(); const query = table.search([0.1, 0.3]);
``` query.limit = 20;
const results = await query.execute();
**Python** ```
```shell
pip install lancedb **Python**
``` ```shell
pip install lancedb
```python ```
import lancedb
```python
uri = "data/sample-lancedb" import lancedb
db = lancedb.connect(uri)
table = db.create_table("my_table", uri = "data/sample-lancedb"
data=[{"vector": [3.1, 4.1], "item": "foo", "price": 10.0}, db = lancedb.connect(uri)
{"vector": [5.9, 26.5], "item": "bar", "price": 20.0}]) table = db.create_table("my_table",
result = table.search([100, 100]).limit(2).to_df() data=[{"vector": [3.1, 4.1], "item": "foo", "price": 10.0},
``` {"vector": [5.9, 26.5], "item": "bar", "price": 20.0}])
result = table.search([100, 100]).limit(2).to_df()
## Blogs, Tutorials & Videos ```
* 📈 <a href="https://blog.eto.ai/benchmarking-random-access-in-lance-ed690757a826">2000x better performance with Lance over Parquet</a>
* 🤖 <a href="https://github.com/lancedb/lancedb/blob/main/docs/src/notebooks/youtube_transcript_search.ipynb">Build a question and answer bot with LanceDB</a> ## Blogs, Tutorials & Videos
* 📈 <a href="https://blog.eto.ai/benchmarking-random-access-in-lance-ed690757a826">2000x better performance with Lance over Parquet</a>
* 🤖 <a href="https://github.com/lancedb/lancedb/blob/main/docs/src/notebooks/youtube_transcript_search.ipynb">Build a question and answer bot with LanceDB</a>

6
docs/src/ann_indexes.md
View File

@@ -68,6 +68,12 @@ a single PQ code.
<figcaption>IVF_PQ index with <code>num_partitions=2, num_sub_vectors=4</code></figcaption> <figcaption>IVF_PQ index with <code>num_partitions=2, num_sub_vectors=4</code></figcaption>
</figure> </figure>
### Use GPU to build vector index
Lance Python SDK has experimental GPU support for creating IVF index.
You can specify the GPU device to train IVF partitions via
- **accelerator**: Specify to `"cuda"`` to enable GPU training.
## Querying an ANN Index ## Querying an ANN Index

7
docs/test/requirements.txt
View File

@@ -1,5 +1,8 @@
lancedb @ git+https://github.com/lancedb/lancedb.git#egg=subdir&subdirectory=python -e ../../python
numpy numpy
pandas pandas
pylance pylance
duckdb duckdb
--extra-index-url https://download.pytorch.org/whl/cpu
torch

3
python/lancedb/remote/table.py
View File

@@ -13,7 +13,7 @@
import uuid import uuid
from functools import cached_property from functools import cached_property
from typing import Union from typing import Optional, Union
import pyarrow as pa import pyarrow as pa
from lance import json_to_schema from lance import json_to_schema
@@ -62,6 +62,7 @@ class RemoteTable(Table):
num_sub_vectors=96, num_sub_vectors=96,
vector_column_name: str = VECTOR_COLUMN_NAME, vector_column_name: str = VECTOR_COLUMN_NAME,
replace: bool = True, replace: bool = True,
accelerator: Optional[str] = None,
): ):
raise NotImplementedError raise NotImplementedError

6
python/lancedb/table.py
View File

@@ -180,6 +180,7 @@ class Table(ABC):
num_sub_vectors=96, num_sub_vectors=96,
vector_column_name: str = VECTOR_COLUMN_NAME, vector_column_name: str = VECTOR_COLUMN_NAME,
replace: bool = True, replace: bool = True,
accelerator: Optional[str] = None,
): ):
"""Create an index on the table. """Create an index on the table.
@@ -200,6 +201,9 @@ class Table(ABC):
replace: bool, default True replace: bool, default True
If True, replace the existing index if it exists. If True, replace the existing index if it exists.
If False, raise an error if duplicate index exists. If False, raise an error if duplicate index exists.
accelerator: str, default None
If set, use the given accelerator to create the index.
Only support "cuda" for now.
""" """
raise NotImplementedError raise NotImplementedError
@@ -479,6 +483,7 @@ class LanceTable(Table):
num_sub_vectors=96, num_sub_vectors=96,
vector_column_name=VECTOR_COLUMN_NAME, vector_column_name=VECTOR_COLUMN_NAME,
replace: bool = True, replace: bool = True,
accelerator: Optional[str] = None,
): ):
"""Create an index on the table.""" """Create an index on the table."""
self._dataset.create_index( self._dataset.create_index(
@@ -488,6 +493,7 @@ class LanceTable(Table):
num_partitions=num_partitions, num_partitions=num_partitions,
num_sub_vectors=num_sub_vectors, num_sub_vectors=num_sub_vectors,
replace=replace, replace=replace,
accelerator=accelerator,
) )
self._reset_dataset() self._reset_dataset()

2
python/pyproject.toml
View File

@@ -2,7 +2,7 @@
name = "lancedb" name = "lancedb"
version = "0.3.0" version = "0.3.0"
dependencies = [ dependencies = [
"pylance==0.8.1", "pylance==0.8.3",
"ratelimiter~=1.0", "ratelimiter~=1.0",
"retry>=0.9.2", "retry>=0.9.2",
"tqdm>=4.1.0", "tqdm>=4.1.0",

1
python/tests/test_table.py
View File

@@ -223,6 +223,7 @@ def test_create_index_method():
num_partitions=256, num_partitions=256,
num_sub_vectors=96, num_sub_vectors=96,
replace=True, replace=True,
accelerator=None,
) )