Added support for halfvec type

ankane · ankane · commit 26673ed9298e · 2025-06-03T17:10:24.000-07:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,6 @@
 ## 0.2.0 (unreleased)
 
-- Added support for `sparsevec` type
+- Added support for `halfvec` and `sparsevec` type
 
 ## 0.1.1 (2023-03-09)
 
diff --git a/src/pgvector.lua b/src/pgvector.lua
@@ -6,6 +6,12 @@ local vector_mt = {
   end
 }
 
+local halfvec_mt = {
+  pgmoon_serialize = function(v)
+    return 0, halfvec_serialize(v)
+  end
+}
+
 local sparsevec_mt = {
   pgmoon_serialize = function(v)
     return 0, sparsevec_serialize(v)
@@ -36,6 +42,30 @@ function pgvector.deserialize(v)
   return setmetatable(vec, vector_mt)
 end
 
+function pgvector.halfvec(v)
+  local vec = {}
+  for _, x in ipairs(v) do
+    table.insert(vec, x)
+  end
+  return setmetatable(vec, halfvec_mt)
+end
+
+function halfvec_serialize(v)
+  for _, v in ipairs(v) do
+    assert(type(v) == "number")
+  end
+  return "[" .. table.concat(v, ",") .. "]"
+end
+
+function halfvec_deserialize(v)
+  local vec = {}
+  for x in string.gmatch(string.sub(v, 2, -2), "[^,]+") do
+    table.insert(vec, tonumber(x))
+  end
+  -- pgvector.halfvec without copy
+  return setmetatable(vec, halfvec_mt)
+end
+
 function pgvector.sparsevec(elements, dim)
   for k, v in pairs(elements) do
     assert(type(k) == "number")
@@ -62,10 +92,13 @@ function sparsevec_deserialize(v)
 end
 
 function pgvector.setup_vector(pg)
-  local row = pg:query("SELECT to_regtype('vector')::oid AS vector_oid, to_regtype('sparsevec')::oid AS sparsevec_oid")[1]
+  local row = pg:query("SELECT to_regtype('vector')::oid AS vector_oid, to_regtype('halfvec')::oid AS halfvec_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
+  if row["halfvec_oid"] then
+    pg:set_type_deserializer(row["halfvec_oid"], "halfvec", function(self, v) return halfvec_deserialize(v) end)
+  end
+  -- if row["sparsevec_oid"] then
   --   pg:set_type_deserializer(row["sparsevec_oid"], "sparsevec", function(self, v) return sparsevec_deserialize(v) end)
   -- end
 end
