Vendor dependencies

This commit is contained in:
2026-08-01 16:11:49 +03:00
parent 7f139a0241
commit 6b5e7f0f8b
29706 changed files with 9575646 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
{"$comment":"This file only protects against accidental modifications. It is not a security mechanism and does not protect against malicious changes.","files":{".cargo_vcs_info.json":"ad508b258663235a1d41423fe8ddbc1bb27e88bb5e022ef41969c9e65c65ce8c",".github/workflows/build.yml":"fa8f7314c42b812c3e4cfd268324701cefb2facd7d7a295c1c123530b75f159f","CHANGELOG.md":"53850a160088ed9757f236a5e2b409e196c43db957f348f3a0ae830ead39eaeb","Cargo.lock":"5dd8c6e89a381ef5cbb8aae6e96583dcb580355e6af36d5a63c6a8d056e3af02","Cargo.toml":"6bf954b6b37dd456437ca4cad8c4341d1dee6c4130302cd0bfad59e5a10abf9f","Cargo.toml.orig":"e5c545915d14aa5951c85e5cb80f3b3290c3e89323c60540eb2743532f24d9a4","LICENSE-APACHE":"cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30","LICENSE-MIT":"8a631592536e9967bae49aaa121c1312c1571eab9544ab288ff96b08c21d5893","README.md":"31d3b88883760909a760e5dea890b54da8daa50957e7439e6cf62adefbaa4e68","src/bit.rs":"ff4fb132318e728a912ca7d455c003084df8cdca96710482413a8099cf747a55","src/diesel_ext/bit.rs":"f6c77fb666953998cfb0196232602a4490e10e9c314a1fcb72e3fed514f6bb55","src/diesel_ext/expression_methods.rs":"dfaab944a6e403c16abfba5a78cacb041ddef8af1d2dac50a4901e13a306693d","src/diesel_ext/halfvec.rs":"0f43644aa14606fa618efa738e295340ee9075066395f471feb991b31eab04d7","src/diesel_ext/mod.rs":"1ad00de749f4841b32d69621e7bf4d7b5dd50fdfac777985747f55415349e668","src/diesel_ext/sparsevec.rs":"39410e563a929134a22c50f47e19fa847a07899ea5a34a547b003f2086a36ff6","src/diesel_ext/vector.rs":"e83d667237af8c94c005a475680e082c2d6aa26b39aaefe1380577ce39fb650a","src/halfvec.rs":"c885553cf45cd914fd91e3fa0c50c44ee6006c2a25944b500719b4fff3e41a2e","src/lib.rs":"95b3be360bff7bf1b7a8159a6096388033f060a840371c3e76317e0349784469","src/postgres_ext/bit.rs":"ddfb83a41457ea7e60c78992c922c8d267384fd841f8f8f926f44c00944dbbb1","src/postgres_ext/halfvec.rs":"517de19f7a467a15607a402f5f1e408048c31ba82dd411af96df4c4984954f69","src/postgres_ext/mod.rs":"2c9994fe215cd9a4c864dbbd3d8b770e1056456c5fb009400dce502168c8142c","src/postgres_ext/sparsevec.rs":"9f04dbfe46146d5dfc8cd241945c2ccd3e6bb025b3254dfae4c7148a6515f260","src/postgres_ext/vector.rs":"fa79b9e7a699a00b82085f22af9ca330f2e04a27fb915d4b44a7cb50e8830f5a","src/sparsevec.rs":"2067c7271189e60752a5533137334bcc0770d63fe01c6e9c9ae439e08425b7a9","src/sqlx_ext/bit.rs":"a9379149c6bc6133c9c3c12899f017dcd06418e3d1f8f2f4e9d99bdfcb4f4ffa","src/sqlx_ext/halfvec.rs":"00eeb89934a44b2d9250226311ca6f52d713f1e0e1e0ea53db2a5a4ee0091030","src/sqlx_ext/mod.rs":"2c9994fe215cd9a4c864dbbd3d8b770e1056456c5fb009400dce502168c8142c","src/sqlx_ext/sparsevec.rs":"5f1130e306e1e3571602b754e128dc047d3ad8bca133947051bdab3e3dc7bee0","src/sqlx_ext/vector.rs":"38ff25c6653bb102a3446f2fbddd1abf806e41fb5feb0449823f47ca02973a30","src/vector.rs":"708f6e528407b572855729e8034091dd9847f5713dcb3bb5db3034c5ca2975ff"},"package":"3673cba5b9a124916096a423b806a9f29620972c6c97b08db5f2053e9428b481"}
+6
View File
@@ -0,0 +1,6 @@
{
"git": {
"sha1": "71506c34dfe22577f4b28750bd1dcbf12ba6e0e8"
},
"path_in_vcs": ""
}
@@ -0,0 +1,27 @@
name: build
on: [push, pull_request]
jobs:
build:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v6
- uses: ankane/setup-postgres@v1
with:
database: pgvector_rust_test
dev-files: true
- run: |
cd /tmp
git clone --branch v0.8.2 https://github.com/pgvector/pgvector.git
cd pgvector
make
sudo make install
- run: psql -d pgvector_rust_test -c "CREATE EXTENSION vector"
# test features individually
- run: cargo test --features postgres
- run: cargo test --features sqlx
- run: cargo test --features diesel
- run: cargo test --features serde
- run: cargo test --features postgres,halfvec
- run: cargo test --features sqlx,halfvec
- run: cargo test --features diesel,halfvec
+73
View File
@@ -0,0 +1,73 @@
## 0.4.2 (2026-05-22)
- Added support for SQLx 0.9
## 0.4.1 (2025-05-20)
- Added `from_f32_slice` function to `HalfVector`
## 0.4.0 (2024-07-28)
- Added support for SQLx 0.8
- Dropped support for SQLx < 0.8
## 0.3.4 (2024-07-17)
- Added `Eq` trait to `Bit`
## 0.3.3 (2024-06-25)
- Added support for `halfvec`, `bit`, and `sparsevec` types to Rust-Postgres
- Added support for `halfvec`, `bit`, and `sparsevec` type to SQLx
- Added support for `halfvec`, `bit`, and `sparsevec` type to Diesel
- Added `l1_distance`, `hamming_distance`, and `jaccard_distance` functions for Diesel
## 0.3.2 (2023-10-30)
- Fixed error with Diesel without `with-deprecated` feature
## 0.3.1 (2023-10-19)
- Added `as_slice` method
## 0.3.0 (2023-10-17)
- Added `serde` feature
- Removed `postgres` from default features
- Reduced dependencies
- Updated Rust edition to 2021
## 0.2.2 (2023-06-02)
- Added `Clone` trait to `Vector`
- Fixed deprecation warning with Diesel 2.1
## 0.2.1 (2023-05-23)
- Added support for `vector[]` type with SQLx
## 0.2.0 (2022-09-05)
- Added support for Diesel 2
- Dropped support for Diesel 1
## 0.1.4 (2022-01-12)
- Added `into`
## 0.1.3 (2021-06-22)
- Fixed SQL type of Diesel operators
## 0.1.2 (2021-06-18)
- Added support for SQLx
- Added operators for Diesel
## 0.1.1 (2021-06-17)
- Added support for Diesel
## 0.1.0 (2021-06-09)
- First release
+2267
View File
File diff suppressed because it is too large Load Diff
+101
View File
@@ -0,0 +1,101 @@
# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO
#
# When uploading crates to the registry Cargo will automatically
# "normalize" Cargo.toml files for maximal compatibility
# with all versions of Cargo and also rewrite `path` dependencies
# to registry (e.g., crates.io) dependencies.
#
# If you are reading this file be aware that the original Cargo.toml
# will likely look very different (and much more reasonable).
# See Cargo.toml.orig for the original contents.
[package]
edition = "2021"
rust-version = "1.60.0"
name = "pgvector"
version = "0.4.2"
authors = ["Andrew Kane <andrew@ankane.org>"]
build = false
autolib = false
autobins = false
autoexamples = false
autotests = false
autobenches = false
description = "pgvector support for Rust"
readme = "README.md"
license = "MIT OR Apache-2.0"
repository = "https://github.com/pgvector/pgvector-rust"
[package.metadata.docs.rs]
features = ["halfvec"]
[features]
halfvec = ["dep:half"]
postgres = [
"dep:postgres-types",
"dep:bytes",
]
[lib]
name = "pgvector"
path = "src/lib.rs"
doctest = false
[dependencies.bytes]
version = "1"
optional = true
[dependencies.diesel]
version = "2"
features = ["postgres"]
optional = true
default-features = false
[dependencies.half]
version = "2"
optional = true
default-features = false
[dependencies.postgres-types]
version = "0.2"
optional = true
default-features = false
[dependencies.serde]
version = "1"
features = ["derive"]
optional = true
[dependencies.sqlx]
version = ">= 0.8, < 0.10"
features = ["postgres"]
optional = true
default-features = false
[dev-dependencies.diesel]
version = "2"
features = ["32-column-tables"]
default-features = false
[dev-dependencies.postgres]
version = "0.19"
default-features = false
[dev-dependencies.serde_json]
version = "1"
[dev-dependencies.sqlx]
version = "0"
features = [
"runtime-tokio",
"tls-native-tls",
]
default-features = false
[dev-dependencies.tokio]
version = "1"
features = ["full"]
[dev-dependencies.tokio-postgres]
version = "0.7"
default-features = false
+36
View File
@@ -0,0 +1,36 @@
[package]
name = "pgvector"
version = "0.4.2"
description = "pgvector support for Rust"
repository = "https://github.com/pgvector/pgvector-rust"
license = "MIT OR Apache-2.0"
authors = ["Andrew Kane <andrew@ankane.org>"]
edition = "2021"
rust-version = "1.60.0"
readme = "README.md"
[lib]
doctest = false
[dependencies]
bytes = { version = "1", optional = true }
postgres-types = { version = "0.2", default-features = false, optional = true }
diesel = { version = "2", default-features = false, features = ["postgres"], optional = true }
sqlx = { version = ">= 0.8, < 0.10", default-features = false, features = ["postgres"], optional = true }
serde = { version = "1", features = ["derive"], optional = true }
half = { version = "2", default-features = false, optional = true }
[dev-dependencies]
postgres = { version = "0.19", default-features = false }
tokio-postgres = { version = "0.7", default-features = false }
diesel = { version = "2", default-features = false, features = ["32-column-tables"] }
sqlx = { version = "0", default-features = false, features = ["runtime-tokio", "tls-native-tls"] }
tokio = { version = "1", features = ["full"] }
serde_json = "1"
[features]
postgres = ["dep:postgres-types", "dep:bytes"]
halfvec = ["dep:half"]
[package.metadata.docs.rs]
features = ["halfvec"]
+202
View File
@@ -0,0 +1,202 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
+21
View File
@@ -0,0 +1,21 @@
The MIT License (MIT)
Copyright (c) 2021-2026 Andrew Kane
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
+416
View File
@@ -0,0 +1,416 @@
# pgvector-rust
[pgvector](https://github.com/pgvector/pgvector) support for Rust
Supports [Rust-Postgres](https://github.com/sfackler/rust-postgres), [SQLx](https://github.com/launchbadge/sqlx), and [Diesel](https://github.com/diesel-rs/diesel)
[![Build Status](https://github.com/pgvector/pgvector-rust/actions/workflows/build.yml/badge.svg)](https://github.com/pgvector/pgvector-rust/actions)
## Getting Started
Follow the instructions for your database library:
- [Rust-Postgres](#rust-postgres)
- [SQLx](#sqlx)
- [Diesel](#diesel)
Or check out some examples:
- [Embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/openai/src/main.rs) with OpenAI
- [Binary embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/cohere/src/main.rs) with Cohere
- [Sentence embeddings](https://github.com/pgvector/pgvector-rust/blob/master/examples/candle/src/main.rs) with Candle
- [Hybrid search](https://github.com/pgvector/pgvector-rust/blob/master/examples/hybrid_search/src/main.rs) with Candle (Reciprocal Rank Fusion)
- [Recommendations](https://github.com/pgvector/pgvector-rust/blob/master/examples/disco/src/main.rs) with Disco
- [Horizontal scaling](https://github.com/pgvector/pgvector-rust/blob/master/examples/citus/src/main.rs) with Citus
- [Bulk loading](https://github.com/pgvector/pgvector-rust/blob/master/examples/loading/src/main.rs) with `COPY`
## Rust-Postgres
Add this line to your applications `Cargo.toml` under `[dependencies]`:
```toml
pgvector = { version = "0.4", features = ["postgres"] }
```
Enable the extension
```rust
client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
```
Create a table
```rust
client.execute("CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))", &[])?;
```
Create a vector from a `Vec<f32>`
```rust
use pgvector::Vector;
let embedding = Vector::from(vec![1.0, 2.0, 3.0]);
```
Insert a vector
```rust
client.execute("INSERT INTO items (embedding) VALUES ($1)", &[&embedding])?;
```
Get the nearest neighbor
```rust
let row = client.query_one(
"SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 1",
&[&embedding],
)?;
```
Retrieve a vector
```rust
let row = client.query_one("SELECT embedding FROM items LIMIT 1", &[])?;
let embedding: Vector = row.get(0);
```
Use `Option` if the value could be `NULL`
```rust
let embedding: Option<Vector> = row.get(0);
```
## SQLx
Add this line to your applications `Cargo.toml` under `[dependencies]`:
```toml
pgvector = { version = "0.4", features = ["sqlx"] }
```
For SQLx < 0.8, use `version = "0.3"` and [this readme](https://github.com/pgvector/pgvector-rust/blob/v0.3.4/README.md).
Enable the extension
```rust
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(&pool)
.await?;
```
Create a table
```rust
sqlx::query("CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))")
.execute(&pool)
.await?;
```
Create a vector from a `Vec<f32>`
```rust
use pgvector::Vector;
let embedding = Vector::from(vec![1.0, 2.0, 3.0]);
```
Insert a vector
```rust
sqlx::query("INSERT INTO items (embedding) VALUES ($1)")
.bind(embedding)
.execute(&pool)
.await?;
```
Get the nearest neighbors
```rust
let rows = sqlx::query("SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 1")
.bind(embedding)
.fetch_all(&pool)
.await?;
```
Retrieve a vector
```rust
let row = sqlx::query("SELECT embedding FROM items LIMIT 1").fetch_one(&pool).await?;
let embedding: Vector = row.try_get("embedding")?;
```
## Diesel
Add this line to your applications `Cargo.toml` under `[dependencies]`:
```toml
pgvector = { version = "0.4", features = ["diesel"] }
```
And update your applications `diesel.toml` under `[print_schema]`:
```toml
import_types = ["diesel::sql_types::*", "pgvector::sql_types::*"]
generate_missing_sql_type_definitions = false
```
Create a migration
```sh
diesel migration generate create_vector_extension
```
with `up.sql`:
```sql
CREATE EXTENSION vector
```
and `down.sql`:
```sql
DROP EXTENSION vector
```
Run the migration
```sql
diesel migration run
```
You can now use the `vector` type in future migrations
```sql
CREATE TABLE items (
id SERIAL PRIMARY KEY,
embedding VECTOR(3)
)
```
For models, use:
```rust
use pgvector::Vector;
#[derive(Queryable)]
#[diesel(table_name = items)]
pub struct Item {
pub id: i32,
pub embedding: Option<Vector>,
}
#[derive(Insertable)]
#[diesel(table_name = items)]
pub struct NewItem {
pub embedding: Option<Vector>,
}
```
Create a vector from a `Vec<f32>`
```rust
let embedding = Vector::from(vec![1.0, 2.0, 3.0]);
```
Insert a vector
```rust
let new_item = NewItem {
embedding: Some(embedding)
};
diesel::insert_into(items::table)
.values(&new_item)
.get_result::<Item>(&mut conn)?;
```
Get the nearest neighbors
```rust
use pgvector::VectorExpressionMethods;
let neighbors = items::table
.order(items::embedding.l2_distance(embedding))
.limit(5)
.load::<Item>(&mut conn)?;
```
Also supports `max_inner_product`, `cosine_distance`, `l1_distance`, `hamming_distance`, and `jaccard_distance`
Get the distances
```rust
let distances = items::table
.select(items::embedding.l2_distance(embedding))
.load::<Option<f64>>(&mut conn)?;
```
Add an approximate index in a migration
```sql
CREATE INDEX my_index ON items USING hnsw (embedding vector_l2_ops)
-- or
CREATE INDEX my_index ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100)
```
Use `vector_ip_ops` for inner product and `vector_cosine_ops` for cosine distance
## Serialization
Use the `serde` feature to enable serialization
## Reference
### Vectors
Create a vector
```rust
use pgvector::Vector;
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
```
Convert to a `Vec<f32>`
```rust
let f32_vec: Vec<f32> = vec.into();
```
Get a slice
```rust
let slice = vec.as_slice();
```
### Half Vectors
Note: Use the `halfvec` feature to enable half vectors
Create a half vector from a `Vec<f16>`
```rust
use half::f16;
use pgvector::HalfVector;
let vec = HalfVector::from(vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]);
```
Or a `f32` slice
```rust
let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
```
Convert to a `Vec<f16>`
```rust
let f16_vec: Vec<f16> = vec.into();
```
Get a slice
```rust
let slice = vec.as_slice();
```
### Binary Vectors
Create a binary vector from a slice of bits
```rust
use pgvector::Bit;
let vec = Bit::new(&[true, false, true]);
```
Or a slice of bytes
```rust
let vec = Bit::from_bytes(&[0b00000000, 0b11111111]);
```
Get the number of bits
```rust
let len = vec.len();
```
Get a slice of bytes
```rust
let bytes = vec.as_bytes();
```
### Sparse Vectors
Create a sparse vector from a dense vector
```rust
use pgvector::SparseVector;
let vec = SparseVector::from_dense(vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0]);
```
Or a map of non-zero elements
```rust
let map = HashMap::from([(0, 1.0), (2, 2.0), (4, 3.0)]);
let vec = SparseVector::from_map(&map, 6);
```
Note: Indices start at 0
Get the number of dimensions
```rust
let dim = vec.dimensions();
```
Get the indices of non-zero elements
```rust
let indices = vec.indices();
```
Get the values of non-zero elements
```rust
let values = vec.values();
```
Get a dense vector
```rust
let f32_vec = vec.to_vec();
```
## History
View the [changelog](https://github.com/pgvector/pgvector-rust/blob/master/CHANGELOG.md)
## Contributing
Everyone is encouraged to help improve this project. Here are a few ways you can help:
- [Report bugs](https://github.com/pgvector/pgvector-rust/issues)
- Fix bugs and [submit pull requests](https://github.com/pgvector/pgvector-rust/pulls)
- Write, clarify, or fix documentation
- Suggest or add new features
To get started with development:
```sh
git clone https://github.com/pgvector/pgvector-rust.git
cd pgvector-rust
createdb pgvector_rust_test
cargo test --all-features
```
To run an example:
```sh
cd examples/loading
createdb pgvector_example
cargo run
```
+83
View File
@@ -0,0 +1,83 @@
#[cfg(feature = "diesel")]
use crate::diesel_ext::bit::BitType;
#[cfg(feature = "diesel")]
use diesel::{deserialize::FromSqlRow, expression::AsExpression};
/// A bit string.
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "diesel", derive(FromSqlRow, AsExpression))]
#[cfg_attr(feature = "diesel", diesel(sql_type = BitType))]
pub struct Bit {
pub(crate) len: usize,
pub(crate) data: Vec<u8>,
}
impl Bit {
/// Creates a bit string from a slice of bits.
pub fn new(data: &[bool]) -> Bit {
let len = data.len();
let mut bytes = vec![0; (len + 7) / 8];
for (i, v) in data.iter().enumerate() {
bytes[i / 8] |= u8::from(*v) << (7 - (i % 8));
}
Bit { len, data: bytes }
}
/// Creates a bit string from a slice of bytes.
pub fn from_bytes(data: &[u8]) -> Bit {
Bit {
len: data.len().checked_mul(8).unwrap(),
data: data.to_vec(),
}
}
/// Returns the number of bits in the bit string.
pub fn len(&self) -> usize {
self.len
}
/// Returns whether the bit string is empty.
pub fn is_empty(&self) -> bool {
self.len == 0
}
/// Returns the bit string as a slice of bytes.
pub fn as_bytes(&self) -> &[u8] {
self.data.as_slice()
}
#[cfg(any(feature = "postgres", feature = "sqlx", feature = "diesel"))]
pub(crate) fn from_sql(buf: &[u8]) -> Result<Bit, Box<dyn std::error::Error + Sync + Send>> {
let len = i32::from_be_bytes(buf[0..4].try_into()?).try_into()?;
let data = buf[4..4 + (len + 7) / 8].to_vec();
Ok(Bit { len, data })
}
}
#[cfg(test)]
mod tests {
use crate::Bit;
#[test]
fn test_from_bytes() {
let vec = Bit::from_bytes(&[0b00000000, 0b11111111]);
assert_eq!(16, vec.len());
assert_eq!(&[0b00000000, 0b11111111], vec.as_bytes());
}
#[test]
fn test_as_bytes() {
let vec = Bit::new(&[true, false, true]);
assert_eq!(3, vec.len());
assert_eq!(&[0b10100000], vec.as_bytes());
}
#[test]
fn test_is_empty() {
let vec = Bit::new(&[]);
assert_eq!(0, vec.len());
assert!(vec.is_empty());
}
}
+133
View File
@@ -0,0 +1,133 @@
use diesel::deserialize::{self, FromSql};
use diesel::pg::{Pg, PgValue};
use diesel::query_builder::QueryId;
use diesel::serialize::{self, IsNull, Output, ToSql};
use diesel::sql_types::SqlType;
use std::convert::TryFrom;
use std::io::Write;
use crate::Bit;
#[derive(SqlType, QueryId)]
#[diesel(postgres_type(name = "bit"))]
pub struct BitType;
impl ToSql<BitType, Pg> for Bit {
fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result {
let len = self.len;
out.write_all(&i32::try_from(len)?.to_be_bytes())?;
out.write_all(&self.data)?;
Ok(IsNull::No)
}
}
impl FromSql<BitType, Pg> for Bit {
fn from_sql(value: PgValue<'_>) -> deserialize::Result<Self> {
Bit::from_sql(value.as_bytes())
}
}
#[cfg(test)]
mod tests {
use crate::{Bit, VectorExpressionMethods};
use diesel::prelude::*;
table! {
use diesel::sql_types::*;
diesel_bit_items (id) {
id -> Int4,
embedding -> Nullable<crate::sql_types::Bit>,
}
}
use diesel_bit_items as items;
#[derive(Queryable)]
#[diesel(table_name = items)]
struct Item {
pub id: i32,
pub embedding: Option<Bit>,
}
#[derive(Insertable)]
#[diesel(table_name = items)]
struct NewItem {
pub embedding: Option<Bit>,
}
#[test]
fn it_works() -> Result<(), diesel::result::Error> {
let mut conn = PgConnection::establish("postgres://localhost/pgvector_rust_test").unwrap();
diesel::sql_query("CREATE EXTENSION IF NOT EXISTS vector").execute(&mut conn)?;
diesel::sql_query("DROP TABLE IF EXISTS diesel_bit_items").execute(&mut conn)?;
diesel::sql_query(
"CREATE TABLE diesel_bit_items (id serial PRIMARY KEY, embedding bit(9))",
)
.execute(&mut conn)?;
let new_items = vec![
NewItem {
embedding: Some(Bit::new(&[
false, false, false, false, false, false, false, false, true,
])),
},
NewItem {
embedding: Some(Bit::new(&[
false, true, false, true, false, false, false, false, true,
])),
},
NewItem {
embedding: Some(Bit::new(&[
false, true, true, true, false, false, false, false, true,
])),
},
NewItem { embedding: None },
];
diesel::insert_into(items::table)
.values(&new_items)
.get_results::<Item>(&mut conn)?;
let all = items::table.load::<Item>(&mut conn)?;
assert_eq!(4, all.len());
let neighbors = items::table
.order(items::embedding.hamming_distance(Bit::new(&[
false, true, false, true, false, false, false, false, true,
])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![2, 3, 1, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
assert_eq!(
Some(Bit::new(&[
false, true, false, true, false, false, false, false, true
])),
neighbors.first().unwrap().embedding
);
let neighbors = items::table
.order(items::embedding.jaccard_distance(Bit::new(&[
false, true, false, true, false, false, false, false, true,
])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![2, 3, 1, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let distances = items::table
.select(items::embedding.hamming_distance(Bit::new(&[
false, true, false, true, false, false, false, false, true,
])))
.order(items::id)
.load::<Option<f64>>(&mut conn)?;
assert_eq!(vec![Some(2.0), Some(0.0), Some(1.0), None], distances);
Ok(())
}
}
@@ -0,0 +1,62 @@
use diesel::expression::{AsExpression, Expression};
use diesel::pg::Pg;
use diesel::sql_types::{Double, SqlType};
diesel::infix_operator!(L2Distance, " <-> ", Double, backend: Pg);
diesel::infix_operator!(MaxInnerProduct, " <#> ", Double, backend: Pg);
diesel::infix_operator!(CosineDistance, " <=> ", Double, backend: Pg);
diesel::infix_operator!(L1Distance, " <+> ", Double, backend: Pg);
diesel::infix_operator!(HammingDistance, " <~> ", Double, backend: Pg);
diesel::infix_operator!(JaccardDistance, " <%> ", Double, backend: Pg);
pub trait VectorExpressionMethods: Expression + Sized {
fn l2_distance<T>(self, other: T) -> L2Distance<Self, T::Expression>
where
Self::SqlType: SqlType,
T: AsExpression<Self::SqlType>,
{
L2Distance::new(self, other.as_expression())
}
fn max_inner_product<T>(self, other: T) -> MaxInnerProduct<Self, T::Expression>
where
Self::SqlType: SqlType,
T: AsExpression<Self::SqlType>,
{
MaxInnerProduct::new(self, other.as_expression())
}
fn cosine_distance<T>(self, other: T) -> CosineDistance<Self, T::Expression>
where
Self::SqlType: SqlType,
T: AsExpression<Self::SqlType>,
{
CosineDistance::new(self, other.as_expression())
}
fn l1_distance<T>(self, other: T) -> L1Distance<Self, T::Expression>
where
Self::SqlType: SqlType,
T: AsExpression<Self::SqlType>,
{
L1Distance::new(self, other.as_expression())
}
fn hamming_distance<T>(self, other: T) -> HammingDistance<Self, T::Expression>
where
Self::SqlType: SqlType,
T: AsExpression<Self::SqlType>,
{
HammingDistance::new(self, other.as_expression())
}
fn jaccard_distance<T>(self, other: T) -> JaccardDistance<Self, T::Expression>
where
Self::SqlType: SqlType,
T: AsExpression<Self::SqlType>,
{
JaccardDistance::new(self, other.as_expression())
}
}
impl<T: Expression> VectorExpressionMethods for T {}
+144
View File
@@ -0,0 +1,144 @@
use diesel::deserialize::{self, FromSql};
use diesel::pg::{Pg, PgValue};
use diesel::query_builder::QueryId;
use diesel::serialize::{self, IsNull, Output, ToSql};
use diesel::sql_types::SqlType;
use std::convert::TryFrom;
use std::io::Write;
use crate::HalfVector;
#[derive(SqlType, QueryId)]
#[diesel(postgres_type(name = "halfvec"))]
pub struct HalfVectorType;
impl ToSql<HalfVectorType, Pg> for HalfVector {
fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result {
let dim = self.0.len();
out.write_all(&u16::try_from(dim)?.to_be_bytes())?;
out.write_all(&0_u16.to_be_bytes())?;
for v in &self.0 {
out.write_all(&v.to_be_bytes())?;
}
Ok(IsNull::No)
}
}
impl FromSql<HalfVectorType, Pg> for HalfVector {
fn from_sql(value: PgValue<'_>) -> deserialize::Result<Self> {
HalfVector::from_sql(value.as_bytes())
}
}
#[cfg(test)]
mod tests {
use crate::{HalfVector, VectorExpressionMethods};
use diesel::prelude::*;
table! {
use diesel::sql_types::*;
diesel_half_items (id) {
id -> Int4,
embedding -> Nullable<crate::sql_types::HalfVector>,
}
}
use diesel_half_items as items;
#[derive(Queryable)]
#[diesel(table_name = items)]
struct Item {
pub id: i32,
pub embedding: Option<HalfVector>,
}
#[derive(Insertable)]
#[diesel(table_name = items)]
struct NewItem {
pub embedding: Option<HalfVector>,
}
#[test]
fn it_works() -> Result<(), diesel::result::Error> {
let mut conn = PgConnection::establish("postgres://localhost/pgvector_rust_test").unwrap();
diesel::sql_query("CREATE EXTENSION IF NOT EXISTS vector").execute(&mut conn)?;
diesel::sql_query("DROP TABLE IF EXISTS diesel_half_items").execute(&mut conn)?;
diesel::sql_query(
"CREATE TABLE diesel_half_items (id serial PRIMARY KEY, embedding halfvec(3))",
)
.execute(&mut conn)?;
let new_items = vec![
NewItem {
embedding: Some(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])),
},
NewItem {
embedding: Some(HalfVector::from_f32_slice(&[2.0, 2.0, 2.0])),
},
NewItem {
embedding: Some(HalfVector::from_f32_slice(&[1.0, 1.0, 2.0])),
},
NewItem { embedding: None },
];
diesel::insert_into(items::table)
.values(&new_items)
.get_results::<Item>(&mut conn)?;
let all = items::table.load::<Item>(&mut conn)?;
assert_eq!(4, all.len());
let neighbors = items::table
.order(items::embedding.l2_distance(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 3, 2, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
assert_eq!(
Some(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])),
neighbors.first().unwrap().embedding
);
let neighbors = items::table
.order(items::embedding.max_inner_product(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![2, 3, 1, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let neighbors = items::table
.order(items::embedding.cosine_distance(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 2, 3, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let neighbors = items::table
.order(items::embedding.l1_distance(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 3, 2, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let distances = items::table
.select(
items::embedding.max_inner_product(HalfVector::from_f32_slice(&[1.0, 1.0, 1.0])),
)
.order(items::id)
.load::<Option<f64>>(&mut conn)?;
assert_eq!(vec![Some(-3.0), Some(-6.0), Some(-4.0), None], distances);
Ok(())
}
}
+7
View File
@@ -0,0 +1,7 @@
pub(crate) mod bit;
pub(crate) mod expression_methods;
pub(crate) mod sparsevec;
pub(crate) mod vector;
#[cfg(feature = "halfvec")]
pub(crate) mod halfvec;
@@ -0,0 +1,148 @@
use diesel::deserialize::{self, FromSql};
use diesel::pg::{Pg, PgValue};
use diesel::query_builder::QueryId;
use diesel::serialize::{self, IsNull, Output, ToSql};
use diesel::sql_types::SqlType;
use std::convert::TryFrom;
use std::io::Write;
use crate::SparseVector;
#[derive(SqlType, QueryId)]
#[diesel(postgres_type(name = "sparsevec"))]
pub struct SparseVectorType;
impl ToSql<SparseVectorType, Pg> for SparseVector {
fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result {
let dim = self.dim;
let nnz = self.indices.len();
out.write_all(&dim.to_be_bytes())?;
out.write_all(&i32::try_from(nnz)?.to_be_bytes())?;
out.write_all(&0_i32.to_be_bytes())?;
for v in &self.indices {
out.write_all(&v.to_be_bytes())?;
}
for v in &self.values {
out.write_all(&v.to_be_bytes())?;
}
Ok(IsNull::No)
}
}
impl FromSql<SparseVectorType, Pg> for SparseVector {
fn from_sql(value: PgValue<'_>) -> deserialize::Result<Self> {
SparseVector::from_sql(value.as_bytes())
}
}
#[cfg(test)]
mod tests {
use crate::{SparseVector, VectorExpressionMethods};
use diesel::prelude::*;
table! {
use diesel::sql_types::*;
diesel_sparse_items (id) {
id -> Int4,
embedding -> Nullable<crate::sql_types::SparseVector>,
}
}
use diesel_sparse_items as items;
#[derive(Queryable)]
#[diesel(table_name = items)]
struct Item {
pub id: i32,
pub embedding: Option<SparseVector>,
}
#[derive(Insertable)]
#[diesel(table_name = items)]
struct NewItem {
pub embedding: Option<SparseVector>,
}
#[test]
fn it_works() -> Result<(), diesel::result::Error> {
let mut conn = PgConnection::establish("postgres://localhost/pgvector_rust_test").unwrap();
diesel::sql_query("CREATE EXTENSION IF NOT EXISTS vector").execute(&mut conn)?;
diesel::sql_query("DROP TABLE IF EXISTS diesel_sparse_items").execute(&mut conn)?;
diesel::sql_query(
"CREATE TABLE diesel_sparse_items (id serial PRIMARY KEY, embedding sparsevec(3))",
)
.execute(&mut conn)?;
let new_items = vec![
NewItem {
embedding: Some(SparseVector::from_dense(&[1.0, 1.0, 1.0])),
},
NewItem {
embedding: Some(SparseVector::from_dense(&[2.0, 2.0, 2.0])),
},
NewItem {
embedding: Some(SparseVector::from_dense(&[1.0, 1.0, 2.0])),
},
NewItem { embedding: None },
];
diesel::insert_into(items::table)
.values(&new_items)
.get_results::<Item>(&mut conn)?;
let all = items::table.load::<Item>(&mut conn)?;
assert_eq!(4, all.len());
let neighbors = items::table
.order(items::embedding.l2_distance(SparseVector::from_dense(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 3, 2, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
assert_eq!(
Some(SparseVector::from_dense(&[1.0, 1.0, 1.0])),
neighbors.first().unwrap().embedding
);
let neighbors = items::table
.order(items::embedding.max_inner_product(SparseVector::from_dense(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![2, 3, 1, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let neighbors = items::table
.order(items::embedding.cosine_distance(SparseVector::from_dense(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 2, 3, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let neighbors = items::table
.order(items::embedding.l1_distance(SparseVector::from_dense(&[1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 3, 2, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let distances = items::table
.select(items::embedding.max_inner_product(SparseVector::from_dense(&[1.0, 1.0, 1.0])))
.order(items::id)
.load::<Option<f64>>(&mut conn)?;
assert_eq!(vec![Some(-3.0), Some(-6.0), Some(-4.0), None], distances);
Ok(())
}
}
+140
View File
@@ -0,0 +1,140 @@
use diesel::deserialize::{self, FromSql};
use diesel::pg::{Pg, PgValue};
use diesel::query_builder::QueryId;
use diesel::serialize::{self, IsNull, Output, ToSql};
use diesel::sql_types::SqlType;
use std::convert::TryFrom;
use std::io::Write;
use crate::Vector;
#[derive(SqlType, QueryId)]
#[diesel(postgres_type(name = "vector"))]
pub struct VectorType;
impl ToSql<VectorType, Pg> for Vector {
fn to_sql<'b>(&'b self, out: &mut Output<'b, '_, Pg>) -> serialize::Result {
let dim = self.0.len();
out.write_all(&u16::try_from(dim)?.to_be_bytes())?;
out.write_all(&0_u16.to_be_bytes())?;
for v in &self.0 {
out.write_all(&v.to_be_bytes())?;
}
Ok(IsNull::No)
}
}
impl FromSql<VectorType, Pg> for Vector {
fn from_sql(value: PgValue<'_>) -> deserialize::Result<Self> {
Vector::from_sql(value.as_bytes())
}
}
#[cfg(test)]
mod tests {
use crate::{Vector, VectorExpressionMethods};
use diesel::prelude::*;
table! {
use diesel::sql_types::*;
diesel_items (id) {
id -> Int4,
embedding -> Nullable<crate::sql_types::Vector>,
}
}
use diesel_items as items;
#[derive(Queryable)]
#[diesel(table_name = items)]
struct Item {
pub id: i32,
pub embedding: Option<Vector>,
}
#[derive(Insertable)]
#[diesel(table_name = items)]
struct NewItem {
pub embedding: Option<Vector>,
}
#[test]
fn it_works() -> Result<(), diesel::result::Error> {
let mut conn = PgConnection::establish("postgres://localhost/pgvector_rust_test").unwrap();
diesel::sql_query("CREATE EXTENSION IF NOT EXISTS vector").execute(&mut conn)?;
diesel::sql_query("DROP TABLE IF EXISTS diesel_items").execute(&mut conn)?;
diesel::sql_query("CREATE TABLE diesel_items (id serial PRIMARY KEY, embedding vector(3))")
.execute(&mut conn)?;
let new_items = vec![
NewItem {
embedding: Some(Vector::from(vec![1.0, 1.0, 1.0])),
},
NewItem {
embedding: Some(Vector::from(vec![2.0, 2.0, 2.0])),
},
NewItem {
embedding: Some(Vector::from(vec![1.0, 1.0, 2.0])),
},
NewItem { embedding: None },
];
diesel::insert_into(items::table)
.values(&new_items)
.get_results::<Item>(&mut conn)?;
let all = items::table.load::<Item>(&mut conn)?;
assert_eq!(4, all.len());
let neighbors = items::table
.order(items::embedding.l2_distance(Vector::from(vec![1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 3, 2, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
assert_eq!(
Some(Vector::from(vec![1.0, 1.0, 1.0])),
neighbors.first().unwrap().embedding
);
let neighbors = items::table
.order(items::embedding.max_inner_product(Vector::from(vec![1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![2, 3, 1, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let neighbors = items::table
.order(items::embedding.cosine_distance(Vector::from(vec![1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 2, 3, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let neighbors = items::table
.order(items::embedding.l1_distance(Vector::from(vec![1.0, 1.0, 1.0])))
.limit(5)
.load::<Item>(&mut conn)?;
assert_eq!(
vec![1, 3, 2, 4],
neighbors.iter().map(|v| v.id).collect::<Vec<i32>>()
);
let distances = items::table
.select(items::embedding.max_inner_product(Vector::from(vec![1.0, 1.0, 1.0])))
.order(items::id)
.load::<Option<f64>>(&mut conn)?;
assert_eq!(vec![Some(-3.0), Some(-6.0), Some(-4.0), None], distances);
Ok(())
}
}
+99
View File
@@ -0,0 +1,99 @@
use half::f16;
#[cfg(feature = "diesel")]
use crate::diesel_ext::halfvec::HalfVectorType;
#[cfg(feature = "diesel")]
use diesel::{deserialize::FromSqlRow, expression::AsExpression};
/// A half vector.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "diesel", derive(FromSqlRow, AsExpression))]
#[cfg_attr(feature = "diesel", diesel(sql_type = HalfVectorType))]
pub struct HalfVector(pub(crate) Vec<f16>);
impl From<Vec<f16>> for HalfVector {
fn from(v: Vec<f16>) -> Self {
HalfVector(v)
}
}
impl From<HalfVector> for Vec<f16> {
fn from(val: HalfVector) -> Self {
val.0
}
}
impl HalfVector {
/// Creates a half vector from a `f32` slice.
pub fn from_f32_slice(slice: &[f32]) -> HalfVector {
HalfVector(slice.iter().map(|v| f16::from_f32(*v)).collect())
}
/// Returns a copy of the half vector as a `Vec<f16>`.
pub fn to_vec(&self) -> Vec<f16> {
self.0.clone()
}
/// Returns the half vector as a slice.
pub fn as_slice(&self) -> &[f16] {
self.0.as_slice()
}
#[cfg(any(feature = "postgres", feature = "sqlx", feature = "diesel"))]
pub(crate) fn from_sql(
buf: &[u8],
) -> Result<HalfVector, Box<dyn std::error::Error + Sync + Send>> {
let dim = u16::from_be_bytes(buf[0..2].try_into()?).into();
let unused = u16::from_be_bytes(buf[2..4].try_into()?);
if unused != 0 {
return Err("expected unused to be 0".into());
}
let mut vec = Vec::with_capacity(dim);
for i in 0..dim {
let s = 4 + 2 * i;
vec.push(f16::from_be_bytes(buf[s..s + 2].try_into()?));
}
Ok(HalfVector(vec))
}
}
#[cfg(test)]
mod tests {
use crate::HalfVector;
use half::f16;
#[test]
fn test_into() {
let vec = HalfVector::from(vec![
f16::from_f32(1.0),
f16::from_f32(2.0),
f16::from_f32(3.0),
]);
let f16_vec: Vec<f16> = vec.into();
assert_eq!(
f16_vec,
vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]
);
}
#[test]
fn test_to_vec() {
let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
assert_eq!(
vec.to_vec(),
vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]
);
}
#[test]
fn test_as_slice() {
let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
assert_eq!(
vec.as_slice(),
&[f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)]
);
}
}
+37
View File
@@ -0,0 +1,37 @@
#![doc = include_str!("../README.md")]
mod bit;
mod sparsevec;
mod vector;
pub use bit::Bit;
pub use sparsevec::SparseVector;
pub use vector::Vector;
#[cfg(feature = "halfvec")]
mod halfvec;
#[cfg(feature = "halfvec")]
pub use halfvec::HalfVector;
#[cfg(feature = "postgres")]
mod postgres_ext;
#[cfg(feature = "sqlx")]
mod sqlx_ext;
#[cfg(feature = "diesel")]
mod diesel_ext;
#[cfg(feature = "diesel")]
pub mod sql_types {
pub use super::diesel_ext::bit::BitType as Bit;
pub use super::diesel_ext::sparsevec::SparseVectorType as SparseVector;
pub use super::diesel_ext::vector::VectorType as Vector;
#[cfg(feature = "halfvec")]
pub use super::diesel_ext::halfvec::HalfVectorType as HalfVector;
}
#[cfg(feature = "diesel")]
pub use diesel_ext::expression_methods::VectorExpressionMethods;
+103
View File
@@ -0,0 +1,103 @@
use bytes::{BufMut, BytesMut};
use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
use std::convert::TryInto;
use std::error::Error;
use crate::Bit;
impl<'a> FromSql<'a> for Bit {
fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<Bit, Box<dyn Error + Sync + Send>> {
Bit::from_sql(raw)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "bit"
}
}
impl ToSql for Bit {
fn to_sql(&self, _ty: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
let len = self.len;
w.put_i32(len.try_into()?);
w.put(&self.data[..]);
Ok(IsNull::No)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "bit"
}
to_sql_checked!();
}
#[cfg(test)]
mod tests {
use crate::Bit;
use postgres::binary_copy::BinaryCopyInWriter;
use postgres::types::Type;
use postgres::{Client, NoTls};
#[test]
fn it_works() -> Result<(), postgres::Error> {
let user = std::env::var("USER").unwrap();
let mut client = Client::configure()
.host("localhost")
.dbname("pgvector_rust_test")
.user(user.as_str())
.connect(NoTls)?;
client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
client.execute("DROP TABLE IF EXISTS postgres_bit_items", &[])?;
client.execute(
"CREATE TABLE postgres_bit_items (id bigserial PRIMARY KEY, embedding bit(9))",
&[],
)?;
let vec = Bit::new(&[false, true, false, true, false, false, false, false, true]);
let vec2 = Bit::new(&[false, false, true, false, false, false, false, false, true]);
client.execute(
"INSERT INTO postgres_bit_items (embedding) VALUES ($1), ($2), (NULL)",
&[&vec, &vec2],
)?;
let query_vec = Bit::new(&[false, true, false, true, false, false, false, false, true]);
let row = client.query_one(
"SELECT embedding FROM postgres_bit_items ORDER BY embedding <~> $1 LIMIT 1",
&[&query_vec],
)?;
let res_vec: Bit = row.get(0);
assert_eq!(vec, res_vec);
assert_eq!(9, res_vec.len());
assert_eq!(&[0b01010000, 0b10000000], res_vec.as_bytes());
let null_row = client.query_one(
"SELECT embedding FROM postgres_bit_items WHERE embedding IS NULL LIMIT 1",
&[],
)?;
let null_res: Option<Bit> = null_row.get(0);
assert!(null_res.is_none());
// ensures binary format is correct
let text_row = client.query_one(
"SELECT embedding::text FROM postgres_bit_items ORDER BY id LIMIT 1",
&[],
)?;
let text_res: String = text_row.get(0);
assert_eq!("010100001", text_res);
// copy
let bit_type = Type::BIT;
let writer = client
.copy_in("COPY postgres_bit_items (embedding) FROM STDIN WITH (FORMAT BINARY)")?;
let mut writer = BinaryCopyInWriter::new(writer, &[bit_type]);
writer.write(&[&Bit::new(&[
false, true, false, true, false, false, false, false, true,
])])?;
writer.write(&[&Bit::new(&[
false, false, true, false, false, false, false, false, true,
])])?;
writer.finish()?;
Ok(())
}
}
@@ -0,0 +1,136 @@
use bytes::{BufMut, BytesMut};
use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
use std::convert::TryInto;
use std::error::Error;
use crate::HalfVector;
impl<'a> FromSql<'a> for HalfVector {
fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<HalfVector, Box<dyn Error + Sync + Send>> {
HalfVector::from_sql(raw)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "halfvec"
}
}
impl ToSql for HalfVector {
fn to_sql(&self, _ty: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
let dim = self.0.len();
w.put_u16(dim.try_into()?);
w.put_u16(0);
for v in &self.0 {
w.put(&v.to_be_bytes()[..]);
}
Ok(IsNull::No)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "halfvec"
}
to_sql_checked!();
}
#[cfg(test)]
mod tests {
use crate::HalfVector;
use half::f16;
use postgres::binary_copy::BinaryCopyInWriter;
use postgres::types::{Kind, Type};
use postgres::{Client, NoTls};
#[test]
fn it_works() -> Result<(), postgres::Error> {
let user = std::env::var("USER").unwrap();
let mut client = Client::configure()
.host("localhost")
.dbname("pgvector_rust_test")
.user(user.as_str())
.connect(NoTls)?;
client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
client.execute("DROP TABLE IF EXISTS postgres_half_items", &[])?;
client.execute(
"CREATE TABLE postgres_half_items (id bigserial PRIMARY KEY, embedding halfvec(3))",
&[],
)?;
let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
let vec2 = HalfVector::from_f32_slice(&[4.0, 5.0, 6.0]);
client.execute(
"INSERT INTO postgres_half_items (embedding) VALUES ($1), ($2), (NULL)",
&[&vec, &vec2],
)?;
let query_vec = HalfVector::from_f32_slice(&[3.0, 1.0, 2.0]);
let row = client.query_one(
"SELECT embedding FROM postgres_half_items ORDER BY embedding <-> $1 LIMIT 1",
&[&query_vec],
)?;
let res_vec: HalfVector = row.get(0);
assert_eq!(vec, res_vec);
assert_eq!(
vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)],
res_vec.to_vec()
);
let empty_vec = HalfVector::from(vec![]);
let empty_res = client.execute(
"INSERT INTO postgres_half_items (embedding) VALUES ($1)",
&[&empty_vec],
);
assert!(empty_res.is_err());
assert_eq!(
empty_res.unwrap_err().as_db_error().unwrap().message(),
"halfvec must have at least 1 dimension"
);
let null_row = client.query_one(
"SELECT embedding FROM postgres_half_items WHERE embedding IS NULL LIMIT 1",
&[],
)?;
let null_res: Option<HalfVector> = null_row.get(0);
assert!(null_res.is_none());
// ensures binary format is correct
let text_row = client.query_one(
"SELECT embedding::text FROM postgres_half_items ORDER BY id LIMIT 1",
&[],
)?;
let text_res: String = text_row.get(0);
assert_eq!("[1,2,3]", text_res);
// copy
let halfvec_type = get_type(&mut client, "halfvec")?;
let writer = client
.copy_in("COPY postgres_half_items (embedding) FROM STDIN WITH (FORMAT BINARY)")?;
let mut writer = BinaryCopyInWriter::new(writer, &[halfvec_type]);
writer.write(&[&HalfVector::from(vec![
f16::from_f32(1.0),
f16::from_f32(2.0),
f16::from_f32(3.0),
])])?;
writer.write(&[&HalfVector::from(vec![
f16::from_f32(4.0),
f16::from_f32(5.0),
f16::from_f32(6.0),
])])?;
writer.finish()?;
Ok(())
}
fn get_type(client: &mut Client, name: &str) -> Result<Type, postgres::Error> {
let row = client.query_one("SELECT pg_type.oid, nspname AS schema FROM pg_type INNER JOIN pg_namespace ON pg_namespace.oid = pg_type.typnamespace WHERE typname = $1", &[&name])?;
Ok(Type::new(
name.into(),
row.get("oid"),
Kind::Simple,
row.get("schema"),
))
}
}
+6
View File
@@ -0,0 +1,6 @@
mod bit;
mod sparsevec;
mod vector;
#[cfg(feature = "halfvec")]
mod halfvec;
@@ -0,0 +1,118 @@
use bytes::{BufMut, BytesMut};
use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
use std::convert::TryInto;
use std::error::Error;
use crate::SparseVector;
impl<'a> FromSql<'a> for SparseVector {
fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<SparseVector, Box<dyn Error + Sync + Send>> {
SparseVector::from_sql(raw)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "sparsevec"
}
}
impl ToSql for SparseVector {
fn to_sql(&self, _ty: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
let dim = self.dim;
let nnz = self.indices.len();
w.put_i32(dim);
w.put_i32(nnz.try_into()?);
w.put_i32(0);
for v in &self.indices {
w.put_i32(*v);
}
for v in &self.values {
w.put_f32(*v);
}
Ok(IsNull::No)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "sparsevec"
}
to_sql_checked!();
}
#[cfg(test)]
mod tests {
use crate::SparseVector;
use postgres::binary_copy::BinaryCopyInWriter;
use postgres::types::{Kind, Type};
use postgres::{Client, NoTls};
#[test]
fn it_works() -> Result<(), postgres::Error> {
let user = std::env::var("USER").unwrap();
let mut client = Client::configure()
.host("localhost")
.dbname("pgvector_rust_test")
.user(user.as_str())
.connect(NoTls)?;
client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
client.execute("DROP TABLE IF EXISTS postgres_sparse_items", &[])?;
client.execute(
"CREATE TABLE postgres_sparse_items (id bigserial PRIMARY KEY, embedding sparsevec(3))",
&[],
)?;
let vec = SparseVector::from_dense(&[1.0, 2.0, 3.0]);
let vec2 = SparseVector::from_dense(&[4.0, 5.0, 6.0]);
client.execute(
"INSERT INTO postgres_sparse_items (embedding) VALUES ($1), ($2), (NULL)",
&[&vec, &vec2],
)?;
let query_vec = SparseVector::from_dense(&[3.0, 1.0, 2.0]);
let row = client.query_one(
"SELECT embedding FROM postgres_sparse_items ORDER BY embedding <-> $1 LIMIT 1",
&[&query_vec],
)?;
let res_vec: SparseVector = row.get(0);
assert_eq!(vec, res_vec);
let null_row = client.query_one(
"SELECT embedding FROM postgres_sparse_items WHERE embedding IS NULL LIMIT 1",
&[],
)?;
let null_res: Option<SparseVector> = null_row.get(0);
assert!(null_res.is_none());
// ensures binary format is correct
let text_row = client.query_one(
"SELECT embedding::text FROM postgres_sparse_items ORDER BY id LIMIT 1",
&[],
)?;
let text_res: String = text_row.get(0);
assert_eq!("{1:1,2:2,3:3}/3", text_res);
// copy
let sparsevec_type = get_type(&mut client, "sparsevec")?;
let writer = client
.copy_in("COPY postgres_sparse_items (embedding) FROM STDIN WITH (FORMAT BINARY)")?;
let mut writer = BinaryCopyInWriter::new(writer, &[sparsevec_type]);
writer.write(&[&SparseVector::from_dense(&[1.0, 2.0, 3.0])])?;
writer.write(&[&SparseVector::from_dense(&[4.0, 5.0, 6.0])])?;
writer.finish()?;
Ok(())
}
fn get_type(client: &mut Client, name: &str) -> Result<Type, postgres::Error> {
let row = client.query_one("SELECT pg_type.oid, nspname AS schema FROM pg_type INNER JOIN pg_namespace ON pg_namespace.oid = pg_type.typnamespace WHERE typname = $1", &[&name])?;
Ok(Type::new(
name.into(),
row.get("oid"),
Kind::Simple,
row.get("schema"),
))
}
}
+174
View File
@@ -0,0 +1,174 @@
use bytes::{BufMut, BytesMut};
use postgres_types::{to_sql_checked, FromSql, IsNull, ToSql, Type};
use std::convert::TryInto;
use std::error::Error;
use crate::Vector;
impl<'a> FromSql<'a> for Vector {
fn from_sql(_ty: &Type, raw: &'a [u8]) -> Result<Vector, Box<dyn Error + Sync + Send>> {
Vector::from_sql(raw)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "vector"
}
}
impl ToSql for Vector {
fn to_sql(&self, _ty: &Type, w: &mut BytesMut) -> Result<IsNull, Box<dyn Error + Sync + Send>> {
let dim = self.0.len();
w.put_u16(dim.try_into()?);
w.put_u16(0);
for v in &self.0 {
w.put_f32(*v);
}
Ok(IsNull::No)
}
fn accepts(ty: &Type) -> bool {
ty.name() == "vector"
}
to_sql_checked!();
}
#[cfg(test)]
mod tests {
use crate::Vector;
use postgres::binary_copy::BinaryCopyInWriter;
use postgres::types::{Kind, Type};
use postgres::{Client, NoTls};
#[test]
fn it_works() -> Result<(), postgres::Error> {
let user = std::env::var("USER").unwrap();
let mut client = Client::configure()
.host("localhost")
.dbname("pgvector_rust_test")
.user(user.as_str())
.connect(NoTls)?;
client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
client.execute("DROP TABLE IF EXISTS postgres_items", &[])?;
client.execute(
"CREATE TABLE postgres_items (id bigserial PRIMARY KEY, embedding vector(3))",
&[],
)?;
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
let vec2 = Vector::from(vec![4.0, 5.0, 6.0]);
client.execute(
"INSERT INTO postgres_items (embedding) VALUES ($1), ($2), (NULL)",
&[&vec, &vec2],
)?;
let query_vec = Vector::from(vec![3.0, 1.0, 2.0]);
let row = client.query_one(
"SELECT embedding FROM postgres_items ORDER BY embedding <-> $1 LIMIT 1",
&[&query_vec],
)?;
let res_vec: Vector = row.get(0);
assert_eq!(vec, res_vec);
assert_eq!(vec![1.0, 2.0, 3.0], res_vec.to_vec());
let empty_vec = Vector::from(vec![]);
let empty_res = client.execute(
"INSERT INTO postgres_items (embedding) VALUES ($1)",
&[&empty_vec],
);
assert!(empty_res.is_err());
assert_eq!(
empty_res.unwrap_err().as_db_error().unwrap().message(),
"vector must have at least 1 dimension"
);
let null_row = client.query_one(
"SELECT embedding FROM postgres_items WHERE embedding IS NULL LIMIT 1",
&[],
)?;
let null_res: Option<Vector> = null_row.get(0);
assert!(null_res.is_none());
// ensures binary format is correct
let text_row = client.query_one(
"SELECT embedding::text FROM postgres_items ORDER BY id LIMIT 1",
&[],
)?;
let text_res: String = text_row.get(0);
assert_eq!("[1,2,3]", text_res);
// copy
let vector_type = get_type(&mut client, "vector")?;
let writer =
client.copy_in("COPY postgres_items (embedding) FROM STDIN WITH (FORMAT BINARY)")?;
let mut writer = BinaryCopyInWriter::new(writer, &[vector_type]);
writer.write(&[&Vector::from(vec![1.0, 2.0, 3.0])])?;
writer.write(&[&Vector::from(vec![4.0, 5.0, 6.0])])?;
writer.finish()?;
Ok(())
}
fn get_type(client: &mut Client, name: &str) -> Result<Type, postgres::Error> {
let row = client.query_one("SELECT pg_type.oid, nspname AS schema FROM pg_type INNER JOIN pg_namespace ON pg_namespace.oid = pg_type.typnamespace WHERE typname = $1", &[&name])?;
Ok(Type::new(
name.into(),
row.get("oid"),
Kind::Simple,
row.get("schema"),
))
}
#[tokio::test]
async fn tokio_works() -> Result<(), tokio_postgres::Error> {
let (client, connection) = tokio_postgres::connect(
"host=localhost dbname=pgvector_rust_test",
tokio_postgres::NoTls,
)
.await?;
tokio::spawn(async move {
if let Err(e) = connection.await {
eprintln!("connection error: {}", e);
}
});
client
.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])
.await?;
client
.execute("DROP TABLE IF EXISTS tokio_postgres_items", &[])
.await?;
client
.execute(
"CREATE TABLE tokio_postgres_items (id bigserial PRIMARY KEY, embedding vector(3))",
&[],
)
.await?;
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
let vec2 = Vector::from(vec![4.0, 5.0, 6.0]);
client
.execute(
"INSERT INTO tokio_postgres_items (embedding) VALUES ($1), ($2), (NULL)",
&[&vec, &vec2],
)
.await?;
let query_vec = Vector::from(vec![3.0, 1.0, 2.0]);
let row = client
.query_one(
"SELECT embedding FROM tokio_postgres_items ORDER BY embedding <-> $1 LIMIT 1",
&[&query_vec],
)
.await?;
let res_vec: Vector = row.get(0);
assert_eq!(vec, res_vec);
assert_eq!(vec![1.0, 2.0, 3.0], res_vec.to_vec());
Ok(())
}
}
+154
View File
@@ -0,0 +1,154 @@
#[cfg(feature = "diesel")]
use crate::diesel_ext::sparsevec::SparseVectorType;
#[cfg(feature = "diesel")]
use diesel::{deserialize::FromSqlRow, expression::AsExpression};
/// A sparse vector.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "diesel", derive(FromSqlRow, AsExpression))]
#[cfg_attr(feature = "diesel", diesel(sql_type = SparseVectorType))]
pub struct SparseVector {
pub(crate) dim: i32,
pub(crate) indices: Vec<i32>,
pub(crate) values: Vec<f32>,
}
impl SparseVector {
/// Creates a sparse vector from a dense vector.
pub fn from_dense(vec: &[f32]) -> SparseVector {
let dim: i32 = vec.len().try_into().unwrap();
let mut indices = Vec::new();
let mut values = Vec::new();
for (i, v) in vec.iter().enumerate() {
if *v != 0.0 {
indices.push(i.try_into().unwrap());
values.push(*v);
}
}
SparseVector {
dim,
indices,
values,
}
}
/// Creates a sparse vector from a map of non-zero elements.
pub fn from_map<'a, I: IntoIterator<Item = (&'a i32, &'a f32)>>(
map: I,
dim: i32,
) -> SparseVector {
let mut elements: Vec<(&i32, &f32)> = map.into_iter().filter(|v| *v.1 != 0.0).collect();
elements.sort_by_key(|v| *v.0);
let indices: Vec<i32> = elements.iter().map(|v| *v.0).collect();
let values: Vec<f32> = elements.iter().map(|v| *v.1).collect();
SparseVector {
dim,
indices,
values,
}
}
/// Returns the number of dimensions.
pub fn dimensions(&self) -> i32 {
self.dim
}
/// Returns the non-zero indices.
pub fn indices(&self) -> &[i32] {
&self.indices
}
/// Returns the non-zero values.
pub fn values(&self) -> &[f32] {
&self.values
}
/// Returns the sparse vector as a `Vec<f32>`.
pub fn to_vec(&self) -> Vec<f32> {
let mut vec = vec![0.0; self.dim.try_into().unwrap()];
for (i, v) in self.indices.iter().zip(&self.values) {
vec[usize::try_from(*i).unwrap()] = *v;
}
vec
}
#[cfg(any(feature = "postgres", feature = "sqlx", feature = "diesel"))]
pub(crate) fn from_sql(
buf: &[u8],
) -> Result<SparseVector, Box<dyn std::error::Error + Sync + Send>> {
let dim = i32::from_be_bytes(buf[0..4].try_into()?);
let nnz = i32::from_be_bytes(buf[4..8].try_into()?).try_into()?;
let unused = i32::from_be_bytes(buf[8..12].try_into()?);
if unused != 0 {
return Err("expected unused to be 0".into());
}
let mut indices = Vec::with_capacity(nnz);
for i in 0..nnz {
let s = 12 + 4 * i;
indices.push(i32::from_be_bytes(buf[s..s + 4].try_into()?));
}
let mut values = Vec::with_capacity(nnz);
for i in 0..nnz {
let s = 12 + 4 * nnz + 4 * i;
values.push(f32::from_be_bytes(buf[s..s + 4].try_into()?));
}
Ok(SparseVector {
dim,
indices,
values,
})
}
}
#[cfg(test)]
mod tests {
use crate::SparseVector;
use std::collections::{BTreeMap, HashMap};
#[test]
fn test_from_dense() {
let vec = SparseVector::from_dense(&[1.0, 0.0, 2.0, 0.0, 3.0, 0.0]);
assert_eq!(vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0], vec.to_vec());
assert_eq!(6, vec.dimensions());
assert_eq!(&[0, 2, 4], vec.indices());
assert_eq!(&[1.0, 2.0, 3.0], vec.values());
}
#[test]
fn test_from_hash_map() {
let map = HashMap::from([(0, 1.0), (2, 2.0), (4, 3.0)]);
let vec = SparseVector::from_map(&map, 6);
assert_eq!(vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0], vec.to_vec());
assert_eq!(6, vec.dimensions());
assert_eq!(&[0, 2, 4], vec.indices());
assert_eq!(&[1.0, 2.0, 3.0], vec.values());
}
#[test]
fn test_from_btree_map() {
let map = BTreeMap::from([(0, 1.0), (2, 2.0), (4, 3.0)]);
let vec = SparseVector::from_map(&map, 6);
assert_eq!(vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0], vec.to_vec());
assert_eq!(6, vec.dimensions());
assert_eq!(&[0, 2, 4], vec.indices());
assert_eq!(&[1.0, 2.0, 3.0], vec.values());
}
#[test]
fn test_from_vec_map() {
let vec = vec![(0, 1.0), (2, 2.0), (4, 3.0)];
let map = vec.iter().map(|v| (&v.0, &v.1));
let vec = SparseVector::from_map(map, 6);
assert_eq!(vec![1.0, 0.0, 2.0, 0.0, 3.0, 0.0], vec.to_vec());
assert_eq!(6, vec.dimensions());
assert_eq!(&[0, 2, 4], vec.indices());
assert_eq!(&[1.0, 2.0, 3.0], vec.values());
}
}
+105
View File
@@ -0,0 +1,105 @@
use sqlx::encode::IsNull;
use sqlx::error::BoxDynError;
use sqlx::postgres::{PgArgumentBuffer, PgHasArrayType, PgTypeInfo, PgValueRef};
use sqlx::{Decode, Encode, Postgres, Type};
use std::convert::TryFrom;
use crate::Bit;
impl Type<Postgres> for Bit {
fn type_info() -> PgTypeInfo {
PgTypeInfo::with_name("bit")
}
}
impl Encode<'_, Postgres> for Bit {
fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
let len = self.len;
buf.extend(&i32::try_from(len)?.to_be_bytes());
buf.extend(&self.data);
Ok(IsNull::No)
}
}
impl Decode<'_, Postgres> for Bit {
fn decode(value: PgValueRef<'_>) -> Result<Self, BoxDynError> {
let buf = <&[u8] as Decode<Postgres>>::decode(value)?;
Bit::from_sql(buf)
}
}
impl PgHasArrayType for Bit {
fn array_type_info() -> PgTypeInfo {
PgTypeInfo::with_name("_bit")
}
}
#[cfg(test)]
mod tests {
use crate::Bit;
use sqlx::postgres::PgPoolOptions;
use sqlx::Row;
#[tokio::test]
async fn it_works() -> Result<(), sqlx::Error> {
let pool = PgPoolOptions::new()
.max_connections(1)
.connect("postgres://localhost/pgvector_rust_test")
.await?;
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(&pool)
.await?;
sqlx::query("DROP TABLE IF EXISTS sqlx_bit_items")
.execute(&pool)
.await?;
sqlx::query("CREATE TABLE sqlx_bit_items (id bigserial PRIMARY KEY, embedding bit(9))")
.execute(&pool)
.await?;
let vec = Bit::new(&[false, true, false, true, false, false, false, false, true]);
let vec2 = Bit::new(&[false, false, true, false, false, false, false, false, true]);
sqlx::query("INSERT INTO sqlx_bit_items (embedding) VALUES ($1), ($2), (NULL)")
.bind(&vec)
.bind(&vec2)
.execute(&pool)
.await?;
let query_vec = Bit::new(&[false, true, false, true, false, false, false, false, true]);
let row =
sqlx::query("SELECT embedding FROM sqlx_bit_items ORDER BY embedding <~> $1 LIMIT 1")
.bind(query_vec)
.fetch_one(&pool)
.await?;
let res_vec: Bit = row.try_get("embedding").unwrap();
assert_eq!(vec, res_vec);
assert_eq!(&[0b01010000, 0b10000000], res_vec.as_bytes());
let null_row =
sqlx::query("SELECT embedding FROM sqlx_bit_items WHERE embedding IS NULL LIMIT 1")
.fetch_one(&pool)
.await?;
let null_res: Option<Bit> = null_row.try_get("embedding").unwrap();
assert!(null_res.is_none());
// ensures binary format is correct
let text_row =
sqlx::query("SELECT embedding::text FROM sqlx_bit_items ORDER BY id LIMIT 1")
.fetch_one(&pool)
.await?;
let text_res: String = text_row.try_get("embedding").unwrap();
assert_eq!("010100001", text_res);
sqlx::query("ALTER TABLE sqlx_bit_items ADD COLUMN factors bit(9)[]")
.execute(&pool)
.await?;
let vecs = &[vec, vec2];
sqlx::query("INSERT INTO sqlx_bit_items (factors) VALUES ($1)")
.bind(vecs)
.execute(&pool)
.await?;
Ok(())
}
}
+127
View File
@@ -0,0 +1,127 @@
use sqlx::encode::IsNull;
use sqlx::error::BoxDynError;
use sqlx::postgres::{PgArgumentBuffer, PgHasArrayType, PgTypeInfo, PgValueRef};
use sqlx::{Decode, Encode, Postgres, Type};
use std::convert::TryFrom;
use crate::HalfVector;
impl Type<Postgres> for HalfVector {
fn type_info() -> PgTypeInfo {
PgTypeInfo::with_name("halfvec")
}
}
impl Encode<'_, Postgres> for HalfVector {
fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
let dim = self.0.len();
buf.extend(&u16::try_from(dim)?.to_be_bytes());
buf.extend(&0_u16.to_be_bytes());
for v in &self.0 {
buf.extend(&v.to_be_bytes());
}
Ok(IsNull::No)
}
}
impl Decode<'_, Postgres> for HalfVector {
fn decode(value: PgValueRef<'_>) -> Result<Self, BoxDynError> {
let buf = <&[u8] as Decode<Postgres>>::decode(value)?;
HalfVector::from_sql(buf)
}
}
impl PgHasArrayType for HalfVector {
fn array_type_info() -> PgTypeInfo {
PgTypeInfo::with_name("_halfvec")
}
}
#[cfg(test)]
mod tests {
use crate::HalfVector;
use half::f16;
use sqlx::postgres::PgPoolOptions;
use sqlx::Row;
#[tokio::test]
async fn it_works() -> Result<(), sqlx::Error> {
let pool = PgPoolOptions::new()
.max_connections(1)
.connect("postgres://localhost/pgvector_rust_test")
.await?;
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(&pool)
.await?;
sqlx::query("DROP TABLE IF EXISTS sqlx_half_items")
.execute(&pool)
.await?;
sqlx::query(
"CREATE TABLE sqlx_half_items (id bigserial PRIMARY KEY, embedding halfvec(3))",
)
.execute(&pool)
.await?;
let vec = HalfVector::from_f32_slice(&[1.0, 2.0, 3.0]);
let vec2 = HalfVector::from_f32_slice(&[4.0, 5.0, 6.0]);
sqlx::query("INSERT INTO sqlx_half_items (embedding) VALUES ($1), ($2), (NULL)")
.bind(&vec)
.bind(&vec2)
.execute(&pool)
.await?;
let query_vec = HalfVector::from_f32_slice(&[3.0, 1.0, 2.0]);
let row =
sqlx::query("SELECT embedding FROM sqlx_half_items ORDER BY embedding <-> $1 LIMIT 1")
.bind(query_vec)
.fetch_one(&pool)
.await?;
let res_vec: HalfVector = row.try_get("embedding").unwrap();
assert_eq!(vec, res_vec);
assert_eq!(
vec![f16::from_f32(1.0), f16::from_f32(2.0), f16::from_f32(3.0)],
res_vec.to_vec()
);
let empty_vec = HalfVector::from(vec![]);
let empty_res = sqlx::query("INSERT INTO sqlx_half_items (embedding) VALUES ($1)")
.bind(&empty_vec)
.execute(&pool)
.await;
assert!(empty_res.is_err());
assert!(empty_res
.unwrap_err()
.to_string()
.contains("halfvec must have at least 1 dimension"));
let null_row =
sqlx::query("SELECT embedding FROM sqlx_half_items WHERE embedding IS NULL LIMIT 1")
.fetch_one(&pool)
.await?;
let null_res: Option<HalfVector> = null_row.try_get("embedding").unwrap();
assert!(null_res.is_none());
// ensures binary format is correct
let text_row =
sqlx::query("SELECT embedding::text FROM sqlx_half_items ORDER BY id LIMIT 1")
.fetch_one(&pool)
.await?;
let text_res: String = text_row.try_get("embedding").unwrap();
assert_eq!("[1,2,3]", text_res);
sqlx::query("ALTER TABLE sqlx_half_items ADD COLUMN factors halfvec[]")
.execute(&pool)
.await?;
let vecs = &[vec, vec2];
sqlx::query("INSERT INTO sqlx_half_items (factors) VALUES ($1)")
.bind(vecs)
.execute(&pool)
.await?;
Ok(())
}
}
+6
View File
@@ -0,0 +1,6 @@
mod bit;
mod sparsevec;
mod vector;
#[cfg(feature = "halfvec")]
mod halfvec;
+130
View File
@@ -0,0 +1,130 @@
use sqlx::encode::IsNull;
use sqlx::error::BoxDynError;
use sqlx::postgres::{PgArgumentBuffer, PgHasArrayType, PgTypeInfo, PgValueRef};
use sqlx::{Decode, Encode, Postgres, Type};
use std::convert::TryFrom;
use crate::SparseVector;
impl Type<Postgres> for SparseVector {
fn type_info() -> PgTypeInfo {
PgTypeInfo::with_name("sparsevec")
}
}
impl Encode<'_, Postgres> for SparseVector {
fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
let dim = self.dim;
let nnz = self.indices.len();
buf.extend(dim.to_be_bytes());
buf.extend(&i32::try_from(nnz)?.to_be_bytes());
buf.extend(&0_i32.to_be_bytes());
for v in &self.indices {
buf.extend(&v.to_be_bytes());
}
for v in &self.values {
buf.extend(&v.to_be_bytes());
}
Ok(IsNull::No)
}
}
impl Decode<'_, Postgres> for SparseVector {
fn decode(value: PgValueRef<'_>) -> Result<Self, BoxDynError> {
let buf = <&[u8] as Decode<Postgres>>::decode(value)?;
SparseVector::from_sql(buf)
}
}
impl PgHasArrayType for SparseVector {
fn array_type_info() -> PgTypeInfo {
PgTypeInfo::with_name("_sparsevec")
}
}
#[cfg(test)]
mod tests {
use crate::SparseVector;
use sqlx::postgres::PgPoolOptions;
use sqlx::Row;
#[tokio::test]
async fn it_works() -> Result<(), sqlx::Error> {
let pool = PgPoolOptions::new()
.max_connections(1)
.connect("postgres://localhost/pgvector_rust_test")
.await?;
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(&pool)
.await?;
sqlx::query("DROP TABLE IF EXISTS sqlx_sparse_items")
.execute(&pool)
.await?;
sqlx::query(
"CREATE TABLE sqlx_sparse_items (id bigserial PRIMARY KEY, embedding sparsevec(3))",
)
.execute(&pool)
.await?;
let vec = SparseVector::from_dense(&[1.0, 2.0, 3.0]);
let vec2 = SparseVector::from_dense(&[4.0, 5.0, 6.0]);
sqlx::query("INSERT INTO sqlx_sparse_items (embedding) VALUES ($1), ($2), (NULL)")
.bind(&vec)
.bind(&vec2)
.execute(&pool)
.await?;
let query_vec = SparseVector::from_dense(&[3.0, 1.0, 2.0]);
let row = sqlx::query(
"SELECT embedding FROM sqlx_sparse_items ORDER BY embedding <-> $1 LIMIT 1",
)
.bind(query_vec)
.fetch_one(&pool)
.await?;
let res_vec: SparseVector = row.try_get("embedding").unwrap();
assert_eq!(vec, res_vec);
assert_eq!(vec![1.0, 2.0, 3.0], res_vec.to_vec());
let empty_vec = SparseVector::from_dense(&[]);
let empty_res = sqlx::query("INSERT INTO sqlx_sparse_items (embedding) VALUES ($1)")
.bind(&empty_vec)
.execute(&pool)
.await;
assert!(empty_res.is_err());
assert!(empty_res
.unwrap_err()
.to_string()
.contains("sparsevec must have at least 1 dimension"));
let null_row =
sqlx::query("SELECT embedding FROM sqlx_sparse_items WHERE embedding IS NULL LIMIT 1")
.fetch_one(&pool)
.await?;
let null_res: Option<SparseVector> = null_row.try_get("embedding").unwrap();
assert!(null_res.is_none());
// ensures binary format is correct
let text_row =
sqlx::query("SELECT embedding::text FROM sqlx_sparse_items ORDER BY id LIMIT 1")
.fetch_one(&pool)
.await?;
let text_res: String = text_row.try_get("embedding").unwrap();
assert_eq!("{1:1,2:2,3:3}/3", text_res);
sqlx::query("ALTER TABLE sqlx_sparse_items ADD COLUMN factors sparsevec[]")
.execute(&pool)
.await?;
let vecs = &[vec, vec2];
sqlx::query("INSERT INTO sqlx_sparse_items (factors) VALUES ($1)")
.bind(vecs)
.execute(&pool)
.await?;
Ok(())
}
}
+119
View File
@@ -0,0 +1,119 @@
use sqlx::encode::IsNull;
use sqlx::error::BoxDynError;
use sqlx::postgres::{PgArgumentBuffer, PgHasArrayType, PgTypeInfo, PgValueRef};
use sqlx::{Decode, Encode, Postgres, Type};
use std::convert::TryFrom;
use crate::Vector;
impl Type<Postgres> for Vector {
fn type_info() -> PgTypeInfo {
PgTypeInfo::with_name("vector")
}
}
impl Encode<'_, Postgres> for Vector {
fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
let dim = self.0.len();
buf.extend(&u16::try_from(dim)?.to_be_bytes());
buf.extend(&0_u16.to_be_bytes());
for v in &self.0 {
buf.extend(&v.to_be_bytes());
}
Ok(IsNull::No)
}
}
impl Decode<'_, Postgres> for Vector {
fn decode(value: PgValueRef<'_>) -> Result<Self, BoxDynError> {
let buf = <&[u8] as Decode<Postgres>>::decode(value)?;
Vector::from_sql(buf)
}
}
impl PgHasArrayType for Vector {
fn array_type_info() -> PgTypeInfo {
PgTypeInfo::with_name("_vector")
}
}
#[cfg(test)]
mod tests {
use crate::Vector;
use sqlx::postgres::PgPoolOptions;
use sqlx::Row;
#[tokio::test]
async fn it_works() -> Result<(), sqlx::Error> {
let pool = PgPoolOptions::new()
.max_connections(1)
.connect("postgres://localhost/pgvector_rust_test")
.await?;
sqlx::query("CREATE EXTENSION IF NOT EXISTS vector")
.execute(&pool)
.await?;
sqlx::query("DROP TABLE IF EXISTS sqlx_items")
.execute(&pool)
.await?;
sqlx::query("CREATE TABLE sqlx_items (id bigserial PRIMARY KEY, embedding vector(3))")
.execute(&pool)
.await?;
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
let vec2 = Vector::from(vec![4.0, 5.0, 6.0]);
sqlx::query("INSERT INTO sqlx_items (embedding) VALUES ($1), ($2), (NULL)")
.bind(&vec)
.bind(&vec2)
.execute(&pool)
.await?;
let query_vec = Vector::from(vec![3.0, 1.0, 2.0]);
let row = sqlx::query("SELECT embedding FROM sqlx_items ORDER BY embedding <-> $1 LIMIT 1")
.bind(query_vec)
.fetch_one(&pool)
.await?;
let res_vec: Vector = row.try_get("embedding").unwrap();
assert_eq!(vec, res_vec);
assert_eq!(vec![1.0, 2.0, 3.0], res_vec.to_vec());
let empty_vec = Vector::from(vec![]);
let empty_res = sqlx::query("INSERT INTO sqlx_items (embedding) VALUES ($1)")
.bind(&empty_vec)
.execute(&pool)
.await;
assert!(empty_res.is_err());
assert!(empty_res
.unwrap_err()
.to_string()
.contains("vector must have at least 1 dimension"));
let null_row =
sqlx::query("SELECT embedding FROM sqlx_items WHERE embedding IS NULL LIMIT 1")
.fetch_one(&pool)
.await?;
let null_res: Option<Vector> = null_row.try_get("embedding").unwrap();
assert!(null_res.is_none());
// ensures binary format is correct
let text_row = sqlx::query("SELECT embedding::text FROM sqlx_items ORDER BY id LIMIT 1")
.fetch_one(&pool)
.await?;
let text_res: String = text_row.try_get("embedding").unwrap();
assert_eq!("[1,2,3]", text_res);
sqlx::query("ALTER TABLE sqlx_items ADD COLUMN factors vector[]")
.execute(&pool)
.await?;
let vecs = &[vec, vec2];
sqlx::query("INSERT INTO sqlx_items (factors) VALUES ($1)")
.bind(vecs)
.execute(&pool)
.await?;
Ok(())
}
}
+99
View File
@@ -0,0 +1,99 @@
#[cfg(any(feature = "postgres", feature = "sqlx", feature = "diesel"))]
use std::convert::TryInto;
#[cfg(feature = "diesel")]
use crate::diesel_ext::vector::VectorType;
#[cfg(feature = "diesel")]
use diesel::{deserialize::FromSqlRow, expression::AsExpression};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
/// A vector.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "diesel", derive(FromSqlRow, AsExpression))]
#[cfg_attr(feature = "diesel", diesel(sql_type = VectorType))]
#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
pub struct Vector(pub(crate) Vec<f32>);
impl From<Vec<f32>> for Vector {
fn from(v: Vec<f32>) -> Self {
Vector(v)
}
}
impl From<Vector> for Vec<f32> {
fn from(val: Vector) -> Self {
val.0
}
}
impl Vector {
/// Returns a copy of the vector as a `Vec<f32>`.
pub fn to_vec(&self) -> Vec<f32> {
self.0.clone()
}
/// Returns the vector as a slice.
pub fn as_slice(&self) -> &[f32] {
self.0.as_slice()
}
#[cfg(any(feature = "postgres", feature = "sqlx", feature = "diesel"))]
pub(crate) fn from_sql(buf: &[u8]) -> Result<Vector, Box<dyn std::error::Error + Sync + Send>> {
let dim = u16::from_be_bytes(buf[0..2].try_into()?).into();
let unused = u16::from_be_bytes(buf[2..4].try_into()?);
if unused != 0 {
return Err("expected unused to be 0".into());
}
let mut vec = Vec::with_capacity(dim);
for i in 0..dim {
let s = 4 + 4 * i;
vec.push(f32::from_be_bytes(buf[s..s + 4].try_into()?));
}
Ok(Vector(vec))
}
}
#[cfg(test)]
mod tests {
use crate::Vector;
#[test]
fn test_into() {
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
let f32_vec: Vec<f32> = vec.into();
assert_eq!(f32_vec, vec![1.0, 2.0, 3.0]);
}
#[test]
fn test_to_vec() {
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
assert_eq!(vec.to_vec(), vec![1.0, 2.0, 3.0]);
}
#[test]
fn test_as_slice() {
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
assert_eq!(vec.as_slice(), &[1.0, 2.0, 3.0]);
}
#[cfg(feature = "serde")]
#[test]
fn test_serialize() {
let vec = Vector::from(vec![1.0, 2.0, 3.0]);
let json = serde_json::to_string(&vec).unwrap();
assert_eq!(json, "[1.0,2.0,3.0]");
}
#[cfg(feature = "serde")]
#[test]
fn test_deserialize() {
let json = "[1.0,2.0,3.0]";
let vec: Vector = serde_json::from_str(json).unwrap();
assert_eq!(vec, Vector::from(vec![1.0, 2.0, 3.0]));
}
}