|
| 1 | +local pgmoon = require("pgmoon") |
| 2 | +local pgvector = require("./src/pgvector") |
| 3 | + |
| 4 | +local pg = pgmoon.new({ |
| 5 | + database = "pgvector_lua_test", |
| 6 | + user = os.getenv("USER") |
| 7 | +}) |
| 8 | + |
| 9 | +assert(pg:connect()) |
| 10 | +assert(pg:query("CREATE EXTENSION IF NOT EXISTS vector")) |
| 11 | +pgvector.setup_vector(pg) |
| 12 | + |
| 13 | +local vec = pg:query("SELECT $1::vector::text", pgvector.new({1, 2, 3}))[1]["text"] |
| 14 | +assert(vec == "[1,2,3]") |
| 15 | + |
| 16 | +local vec = pg:query("SELECT '[1,2,3]'::vector")[1]["vector"] |
| 17 | +assert(#vec == 3) |
| 18 | +assert(vec[1] == 1) |
| 19 | +assert(vec[2] == 2) |
| 20 | +assert(vec[3] == 3) |
| 21 | + |
| 22 | +local vec = pg:query("SELECT $1::halfvec::text", pgvector.halfvec({1, 2, 3}))[1]["text"] |
| 23 | +assert(vec == "[1,2,3]") |
| 24 | + |
| 25 | +local vec = pg:query("SELECT '[1,2,3]'::halfvec")[1]["halfvec"] |
| 26 | +assert(#vec == 3) |
| 27 | +assert(vec[1] == 1) |
| 28 | +assert(vec[2] == 2) |
| 29 | +assert(vec[3] == 3) |
| 30 | + |
| 31 | +local vec = pg:query("SELECT $1::sparsevec::text", pgvector.sparsevec({[1] = 1, [3] = 2, [5] = 3}, 6))[1]["text"] |
| 32 | +assert(vec == "{1:1,3:2,5:3}/6") |
| 33 | + |
| 34 | +local vec = pg:query("SELECT '{1:1,3:2,5:3}/6'::sparsevec")[1]["sparsevec"] |
| 35 | +assert(vec["elements"][1] == 1) |
| 36 | +assert(vec["elements"][2] == nil) |
| 37 | +assert(vec["elements"][3] == 2) |
| 38 | +assert(vec["elements"][4] == nil) |
| 39 | +assert(vec["elements"][5] == 3) |
| 40 | +assert(vec["elements"][6] == nil) |
| 41 | +assert(vec["dim"] == 6) |
0 commit comments