Started support for sparsevec type

Author: ankane

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.2.0 (unreleased)
+
+- Added support for `sparsevec` type
+
 ## 0.1.1 (2023-03-09)
 
 - Added `new` method

examples/sparse/example.lua

@@ -9,6 +9,7 @@ local cjson = require("cjson")
 local http = require("socket.http")
 local ltn12 = require("ltn12")
 local pgmoon = require("pgmoon")
+local pgvector = require("./src/pgvector")
 
 local pg = pgmoon.new({
   database = "pgvector_example",
@@ -20,14 +21,6 @@ assert(pg:query("CREATE EXTENSION IF NOT EXISTS vector"))
 assert(pg:query("DROP TABLE IF EXISTS documents"))
 assert(pg:query("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding sparsevec(30522))"))
 
-function sparsevec(elements, dim)
-  local e = {}
-  for k, v in pairs(elements) do
-    table.insert(e, k .. ":" .. v)
-  end
-  return "{" .. table.concat(e, ",") .. "}/" .. dim
-end
-
 function embed(inputs)
   local url = "http://localhost:3000/embed_sparse"
   local data = {
@@ -67,12 +60,12 @@ local documents = {
 local embeddings = embed(documents)
 for i, content in ipairs(documents) do
   local embedding = embeddings[i]
-  assert(pg:query("INSERT INTO documents (content, embedding) VALUES ($1, $2::text::sparsevec)", content, sparsevec(embedding, 30522)))
+  assert(pg:query("INSERT INTO documents (content, embedding) VALUES ($1, $2)", content, pgvector.sparsevec(embedding, 30522)))
 end
 
 local query = "forest"
 local embedding = embed({query})[1]
-local res = assert(pg:query("SELECT content FROM documents ORDER BY embedding <#> $1::text::sparsevec LIMIT 5", sparsevec(embedding, 30522)))
+local res = assert(pg:query("SELECT content FROM documents ORDER BY embedding <#> $1 LIMIT 5", pgvector.sparsevec(embedding, 30522)))
 for i, row in ipairs(res) do
   print(row["content"])
 end

src/pgvector.lua

@@ -6,6 +6,12 @@ local vector_mt = {
   end
 }
 
+local sparsevec_mt = {
+  pgmoon_serialize = function(v)
+    return 0, sparsevec_serialize(v)
+  end
+}
+
 function pgvector.new(v)
   local vec = {}
   for _, x in ipairs(v) do
@@ -30,10 +36,38 @@ function pgvector.deserialize(v)
   return setmetatable(vec, vector_mt)
 end
 
+function pgvector.sparsevec(elements, dim)
+  for k, v in pairs(elements) do
+    assert(type(k) == "number")
+    assert(type(v) == "number")
+  end
+  assert(type(dim) == "number")
+
+  local vec = {}
+  vec["elements"] = elements
+  vec["dim"] = dim
+  return setmetatable(vec, sparsevec_mt)
+end
+
+function sparsevec_serialize(vec)
+  local elements = {}
+  for i, v in pairs(vec["elements"]) do
+    table.insert(elements, tonumber(i) .. ":" .. tonumber(v))
+  end
+  return "{" .. table.concat(elements, ",") .. "}/" .. tonumber(vec["dim"])
+end
+
+function sparsevec_deserialize(v)
+  -- TODO
+end
+
 function pgvector.setup_vector(pg)
-  local row = pg:query("SELECT to_regtype('vector')::oid AS vector_oid")[1]
+  local row = pg:query("SELECT to_regtype('vector')::oid AS vector_oid, to_regtype('sparsevec')::oid AS sparsevec_oid")[1]
   assert(row["vector_oid"], "vector type not found in the database")
   pg:set_type_deserializer(row["vector_oid"], "vector", function(self, v) return pgvector.deserialize(v) end)
+  -- if row["sparsevec_oid"] do
+  --   pg:set_type_deserializer(row["sparsevec_oid"], "sparsevec", function(self, v) return sparsevec_deserialize(v) end)
+  -- end
 end
 
 return pgvector