Added package

ankane · ankane · commit d44c22df5ba1 · 2025-06-07T19:12:35.000-07:00
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -16,4 +16,5 @@ jobs:
           make
           sudo make install
       - run: julia --project=. -e "using Pkg; Pkg.instantiate()"
+      - run: julia --project=. -e "using Pkg; Pkg.test()"
       - run: julia --project=. LibPQ/example.jl
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -0,0 +1,3 @@
+## 0.1.0 (unreleased)
+
+- First release
diff --git a/LibPQ/example.jl b/LibPQ/example.jl
@@ -1,28 +1,24 @@
-using LibPQ, Tables
+using LibPQ, Pgvector, Tables
 
 conn = LibPQ.Connection("dbname=pgvector_julia_test host=localhost")
 
 execute(conn, "CREATE EXTENSION IF NOT EXISTS vector")
 execute(conn, "DROP TABLE IF EXISTS items")
 execute(conn, "CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3))")
 
-module Pgvector
-    convert(vec::AbstractVector{T}) where T<:Number = string("[", join(vec, ","), "]")
-
-    parse(str::String) = map(x -> Base.parse(Float32, x), split(str[2:end-1], ","))
-end
+Pgvector.register!(conn)
 
 embeddings = [[1, 1, 1], [2, 2, 2], [1, 1, 2]]
 LibPQ.load!(
-    (embedding = map(Pgvector.convert, embeddings),),
+    (embedding = map(Pgvector.Vector, embeddings),),
     conn,
     "INSERT INTO items (embedding) VALUES (\$1)",
 )
 
-embedding = Pgvector.convert([1, 1, 1])
-result = execute(conn, "SELECT * FROM items ORDER BY embedding <-> \$1 LIMIT 5", [embedding])
+embedding = [1, 1, 1]
+result = execute(conn, "SELECT * FROM items ORDER BY embedding <-> \$1 LIMIT 5", [Pgvector.Vector(embedding)])
 data = columntable(result)
-println(map(Pgvector.parse, data.embedding))
+println(data)
 
 execute(conn, "CREATE INDEX ON items USING hnsw (embedding vector_l2_ops)")
 
diff --git a/Project.toml b/Project.toml
@@ -1,3 +1,18 @@
+name = "Pgvector"
+uuid = "ab7f0314-681e-439b-b2a8-1689834ece97"
+authors = ["Andrew Kane <andrew@ankane.org>"]
+version = "0.1.0"
+
+[compat]
+julia = "1.10"
+
 [deps]
 LibPQ = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+
+[extras]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test"]
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # pgvector-julia
 
-[pgvector](https://github.com/pgvector/pgvector) examples for Julia
+[pgvector](https://github.com/pgvector/pgvector) support for Julia
 
 Supports [LibPQ.jl](https://github.com/iamed2/LibPQ.jl)
 
@@ -85,7 +85,7 @@ git clone https://github.com/pgvector/pgvector-julia.git
 cd pgvector-julia
 createdb pgvector_julia_test
 julia --project=. -e "using Pkg; Pkg.instantiate()"
-julia --project=. LibPQ/example.jl
+julia --project=. -e "using Pkg; Pkg.test()"
 ```
 
 To run an example:
diff --git a/examples/hybrid/Project.toml b/examples/hybrid/Project.toml
@@ -2,4 +2,8 @@
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LibPQ = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1"
+Pgvector = "ab7f0314-681e-439b-b2a8-1689834ece97"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+
+[sources]
+Pgvector = {path = "../.."}
diff --git a/examples/hybrid/example.jl b/examples/hybrid/example.jl
@@ -1,4 +1,4 @@
-using HTTP, JSON, LibPQ, Tables
+using HTTP, JSON, LibPQ, Pgvector, Tables
 
 conn = LibPQ.Connection("dbname=pgvector_example host=localhost")
 
@@ -7,10 +7,6 @@ execute(conn, "DROP TABLE IF EXISTS documents")
 execute(conn, "CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(768))")
 execute(conn, "CREATE INDEX ON documents USING GIN (to_tsvector('english', content))")
 
-module Pgvector
-    convert(v::AbstractVector{T}) where T<:Real = string("[", join(v, ","), "]")
-end
-
 function embed(input, task)
     # nomic-embed-text uses a task prefix
     # https://huggingface.co/nomic-ai/nomic-embed-text-v1.5
@@ -35,7 +31,7 @@ input = [
 ]
 embeddings = embed(input, "search_document")
 LibPQ.load!(
-    (content = input, embedding = map(Pgvector.convert, embeddings)),
+    (content = input, embedding = map(Pgvector.Vector, embeddings)),
     conn,
     "INSERT INTO documents (content, embedding) VALUES (\$1, \$2)",
 )
@@ -66,7 +62,7 @@ LIMIT 5
 query = "growling bear"
 embedding = embed([query], "search_query")[1]
 k = 60
-result = execute(conn, sql, [query, Pgvector.convert(embedding), k])
+result = execute(conn, sql, [query, Pgvector.Vector(embedding), k])
 rows = Tables.rows(columntable(result))
 for row in rows
     id = Tables.getcolumn(row, 1)
diff --git a/examples/openai/Project.toml b/examples/openai/Project.toml
@@ -2,4 +2,8 @@
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LibPQ = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1"
+Pgvector = "ab7f0314-681e-439b-b2a8-1689834ece97"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+
+[sources]
+Pgvector = {path = "../.."}
diff --git a/examples/openai/example.jl b/examples/openai/example.jl
@@ -1,15 +1,11 @@
-using HTTP, JSON, LibPQ, Tables
+using HTTP, JSON, LibPQ, Pgvector, Tables
 
 conn = LibPQ.Connection("dbname=pgvector_example host=localhost")
 
 execute(conn, "CREATE EXTENSION IF NOT EXISTS vector")
 execute(conn, "DROP TABLE IF EXISTS documents")
 execute(conn, "CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(1536))")
 
-module Pgvector
-    convert(v::AbstractVector{T}) where T<:Real = string("[", join(v, ","), "]")
-end
-
 function embed(input)
     url = "https://api.openai.com/v1/embeddings"
     data = Dict(
@@ -31,14 +27,14 @@ input = [
 ]
 embeddings = embed(input)
 LibPQ.load!(
-    (content = input, embedding = map(Pgvector.convert, embeddings)),
+    (content = input, embedding = map(Pgvector.Vector, embeddings)),
     conn,
     "INSERT INTO documents (content, embedding) VALUES (\$1, \$2)",
 )
 
 query = "forest"
 embedding = embed([query])[1]
-result = execute(conn, "SELECT content FROM documents ORDER BY embedding <=> \$1 LIMIT 5", [Pgvector.convert(embedding)])
+result = execute(conn, "SELECT content FROM documents ORDER BY embedding <=> \$1 LIMIT 5", [Pgvector.Vector(embedding)])
 rows = Tables.rows(columntable(result))
 for row in rows
     println(Tables.getcolumn(row, 1))
diff --git a/examples/sparse/Project.toml b/examples/sparse/Project.toml
@@ -2,4 +2,8 @@
 HTTP = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LibPQ = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1"
+Pgvector = "ab7f0314-681e-439b-b2a8-1689834ece97"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+
+[sources]
+Pgvector = {path = "../.."}
diff --git a/examples/sparse/example.jl b/examples/sparse/example.jl
@@ -1,20 +1,11 @@
-using HTTP, JSON, LibPQ, SparseArrays, Tables
+using HTTP, JSON, LibPQ, Pgvector, SparseArrays, Tables
 
 conn = LibPQ.Connection("dbname=pgvector_example host=localhost")
 
 execute(conn, "CREATE EXTENSION IF NOT EXISTS vector")
 execute(conn, "DROP TABLE IF EXISTS documents")
 execute(conn, "CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding sparsevec(30522))")
 
-module Pgvector
-    using SparseArrays
-
-    function convert(vec::AbstractSparseVector{Tv,Ti}) where {Tv<:Real,Ti<:Integer}
-        elements = [string(i, ":", v) for (i, v) in zip(vec.nzind, vec.nzval)]
-        string("{", join(elements, ","), "}/", vec.n)
-    end
-end
-
 function embed(inputs)
     url = "http://localhost:3000/embed_sparse"
     data = Dict(
@@ -34,14 +25,14 @@ input = [
 ]
 embeddings = embed(input)
 LibPQ.load!(
-    (content = input, embedding = map(Pgvector.convert, embeddings)),
+    (content = input, embedding = map(Pgvector.SparseVector, embeddings)),
     conn,
     "INSERT INTO documents (content, embedding) VALUES (\$1, \$2)",
 )
 
 query = "forest"
 embedding = embed([query])[1]
-result = execute(conn, "SELECT content FROM documents ORDER BY embedding <#> \$1 LIMIT 5", [Pgvector.convert(embedding)])
+result = execute(conn, "SELECT content FROM documents ORDER BY embedding <#> \$1 LIMIT 5", [Pgvector.SparseVector(embedding)])
 rows = Tables.rows(columntable(result))
 for row in rows
     println(Tables.getcolumn(row, 1))
diff --git a/src/Pgvector.jl b/src/Pgvector.jl
@@ -0,0 +1,8 @@
+module Pgvector
+
+include("halfvec.jl")
+include("register.jl")
+include("sparsevec.jl")
+include("vector.jl")
+
+end
diff --git a/src/halfvec.jl b/src/halfvec.jl
@@ -0,0 +1,11 @@
+struct HalfVector
+    vec::Base.Vector{Float16}
+end
+
+function Base.parse(::Type{HalfVector}, s::String)
+    HalfVector(map(x -> Base.parse(Float16, x), split(s[2:end-1], ",")))
+end
+
+function Base.show(io::IO, vec::HalfVector)
+    print(io, string("[", join(vec.vec, ","), "]"))
+end
diff --git a/src/register.jl b/src/register.jl
@@ -0,0 +1,24 @@
+using LibPQ, Tables
+
+function register!(conn::LibPQ.Connection)
+    result = execute(conn, "SELECT to_regtype('vector')::oid AS vector_oid, to_regtype('halfvec')::oid AS halfvec_oid, to_regtype('sparsevec')::oid AS sparsevec_oid")
+    data = columntable(result)
+
+    vector_oid = data.vector_oid[1]
+    if ismissing(vector_oid)
+        error("vector type not found in the database")
+    end
+    conn.type_map[vector_oid] = Vector
+
+    halfvec_oid = data.halfvec_oid[1]
+    if !ismissing(halfvec_oid)
+        conn.type_map[halfvec_oid] = HalfVector
+    end
+
+    sparsevec_oid = data.sparsevec_oid[1]
+    if !ismissing(sparsevec_oid)
+        conn.type_map[sparsevec_oid] = SparseVector
+    end
+
+    nothing
+end
diff --git a/src/sparsevec.jl b/src/sparsevec.jl
@@ -0,0 +1,18 @@
+using SparseArrays
+
+struct SparseVector
+    vec::SparseArrays.SparseVector{Float32,Int32}
+end
+
+function Base.parse(::Type{SparseVector}, s::String)
+    elements, dim = rsplit(s, "/"; limit=2)
+    elements = map(e -> split(e, ":"; limit=2), split(elements[2:end-1], ","))
+    indices = map(v -> Base.parse(Int32, v[1]), elements)
+    values = map(v -> Base.parse(Float32, v[2]), elements)
+    SparseVector(sparsevec(indices, values, Base.parse(Int32, dim)))
+end
+
+function Base.show(io::IO, vec::SparseVector)
+    elements = [string(i, ":", v) for (i, v) in zip(vec.vec.nzind, vec.vec.nzval)]
+    print(io, string("{", join(elements, ","), "}/", vec.vec.n))
+end
diff --git a/src/vector.jl b/src/vector.jl
@@ -0,0 +1,11 @@
+struct Vector
+    vec::Base.Vector{Float32}
+end
+
+function Base.parse(::Type{Vector}, s::String)
+    Vector(map(x -> Base.parse(Float32, x), split(s[2:end-1], ",")))
+end
+
+function Base.show(io::IO, vec::Vector)
+    print(io, string("[", join(vec.vec, ","), "]"))
+end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -0,0 +1,29 @@
+using LibPQ, Pgvector, SparseArrays, Tables, Test
+
+conn = LibPQ.Connection("dbname=pgvector_julia_test host=localhost")
+execute(conn, "CREATE EXTENSION IF NOT EXISTS vector")
+Pgvector.register!(conn)
+
+@testset "Vector" begin
+    vec = Pgvector.Vector([1, 2, 3])
+    @test string(vec) == "[1.0,2.0,3.0]"
+
+    vec = columntable(execute(conn, "SELECT '[1,2,3]'::vector"))[1][1]
+    @test vec.vec == [1, 2, 3]
+end
+
+@testset "HalfVector" begin
+    vec = Pgvector.HalfVector([1, 2, 3])
+    @test string(vec) == "[1.0,2.0,3.0]"
+
+    vec = columntable(execute(conn, "SELECT '[1,2,3]'::halfvec"))[1][1]
+    @test vec.vec == [1, 2, 3]
+end
+
+@testset "SparseVector" begin
+    vec = Pgvector.SparseVector(sparsevec([1, 0, 2, 0, 3, 0]))
+    @test string(vec) == "{1:1.0,3:2.0,5:3.0}/6"
+
+    vec = columntable(execute(conn, "SELECT '{1:1,3:2,5:3}/6'::sparsevec"))[1][1]
+    @test vec.vec == sparsevec([1, 0, 2, 0, 3, 0])
+end

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+## 0.1.0 (unreleased)`
	`2`	`+`
	`3`	`+- First release`