require_relative "test_helper"

class PgTest < Minitest::Test

def setup

skip if RUBY_ENGINE == "jruby"

conn.exec("DELETE FROM pg_items")

end

def test_vector_text

embedding = [1.5, 2, 3]

conn.exec_params("INSERT INTO pg_items (embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec("SELECT * FROM pg_items ORDER BY id").to_a

assert_equal embedding, res[0]["embedding"]

assert_nil res[1]["embedding"]

end

def test_vector_binary

embedding = [1.5, 2, 3]

conn.exec_params("INSERT INTO pg_items (embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec_params("SELECT * FROM pg_items ORDER BY id", [], 1).to_a

assert_equal embedding, res[0]["embedding"]

assert_nil res[1]["embedding"]

end

def test_halfvec_text

embedding = [1.5, 2, 3]

conn.exec_params("INSERT INTO pg_items (half_embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec("SELECT * FROM pg_items ORDER BY id").to_a

assert_equal embedding, res[0]["half_embedding"]

assert_nil res[1]["half_embedding"]

end

def test_bit_text

embedding = "010100001"

conn.exec_params("INSERT INTO pg_items (binary_embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec("SELECT * FROM pg_items ORDER BY id").to_a

assert_equal embedding, res[0]["binary_embedding"]

assert_nil res[1]["binary_embedding"]

end

def test_bit_binary

embedding = "010100001"

conn.exec_params("INSERT INTO pg_items (binary_embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec_params("SELECT * FROM pg_items ORDER BY id", [], 1).to_a

assert_equal embedding, res[0]["binary_embedding"]

assert_nil res[1]["binary_embedding"]

assert_equal "010100001", conn.exec_params("SELECT '010100001'::bit(9)", [], 1).first["bit"]

end

def test_sparsevec_text

embedding = Pgvector::SparseVector.new([1.5, 2, 3])

conn.exec_params("INSERT INTO pg_items (sparse_embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec("SELECT * FROM pg_items ORDER BY id").to_a

assert_equal [1.5, 2, 3], res[0]["sparse_embedding"].to_a

assert_nil res[1]["sparse_embedding"]

end

def test_sparsevec_binary

embedding = Pgvector::SparseVector.new([1.5, 2, 3])

conn.exec_params("INSERT INTO pg_items (sparse_embedding) VALUES ($1), (NULL)", [embedding])

res = conn.exec_params("SELECT * FROM pg_items ORDER BY id", [], 1).to_a

assert_equal [1.5, 2, 3], res[0]["sparse_embedding"].to_a

assert_nil res[1]["sparse_embedding"]

end

def test_type_map_none

coder = PG::TextEncoder::CopyRow.new

assert_equal "[1.0,2.0,3.0]\n", coder.encode([Pgvector::Vector.new([1, 2, 3])])

assert_equal "[1.0,2.0,3.0]\n", coder.encode([Pgvector::HalfVector.new([1, 2, 3])])

assert_equal "101\n", coder.encode([Pgvector::Bit.new([true, false, true])])

assert_equal "{1:1.0,2:2.0,3:3.0}/3\n", coder.encode([Pgvector::SparseVector.new([1, 2, 3])])

end

def test_type_map_text

coder = PG::TextEncoder::CopyRow.new(type_map: Pgvector::PG::TextEncoder.type_map)

assert_equal "[1.0,2.0,3.0]\n", coder.encode([Pgvector::Vector.new([1, 2, 3])])

assert_equal "[1.0,2.0,3.0]\n", coder.encode([Pgvector::HalfVector.new([1, 2, 3])])

assert_equal "101\n", coder.encode([Pgvector::Bit.new([true, false, true])])

assert_equal "{1:1.0,2:2.0,3:3.0}/3\n", coder.encode([Pgvector::SparseVector.new([1, 2, 3])])

end

def test_type_map_binary

vec = Pgvector::Vector.new([1, 2, 3])

binary_vec = Pgvector::Bit.new([true, false, true])

sparse_vec = Pgvector::SparseVector.new([1, 2, 3])

coder = PG::BinaryEncoder::CopyRow.new(type_map: Pgvector::PG::BinaryEncoder.type_map)

assert_match vec.to_binary, coder.encode([vec])

assert_match binary_vec.to_binary, coder.encode([binary_vec])

assert_match sparse_vec.to_binary, coder.encode([sparse_vec])

end

def test_copy_text

embedding = Pgvector::Vector.new([1, 2, 3])

half_embedding = Pgvector::HalfVector.new([1, 2, 3])

binary_embedding = Pgvector::Bit.new([false, true, false, true, false, false, false, false, true])

sparse_embedding = Pgvector::SparseVector.new([1, 2, 3])

coder = PG::TextEncoder::CopyRow.new

conn.copy_data("COPY pg_items (embedding, half_embedding, binary_embedding, sparse_embedding) FROM STDIN", coder) do

conn.put_copy_data([embedding, half_embedding, binary_embedding, sparse_embedding])

end

res = conn.exec("SELECT * FROM pg_items").first

assert_equal [1, 2, 3], res["embedding"]

assert_equal [1, 2, 3], res["half_embedding"].to_a

assert_equal "010100001", res["binary_embedding"]

assert_equal [1, 2, 3], res["sparse_embedding"].to_a

end

def test_copy_binary

embedding = Pgvector::Vector.new([1, 2, 3])

binary_embedding = Pgvector::Bit.new([false, true, false, true, false, false, false, false, true])

sparse_embedding = Pgvector::SparseVector.new([1, 2, 3])

coder = PG::BinaryEncoder::CopyRow.new

conn.copy_data("COPY pg_items (embedding, binary_embedding, sparse_embedding) FROM STDIN WITH (FORMAT BINARY)", coder) do

conn.put_copy_data([embedding.to_binary, binary_embedding.to_binary, sparse_embedding.to_binary])

end

res = conn.exec("SELECT * FROM pg_items").first

assert_equal [1, 2, 3], res["embedding"]

assert_equal "010100001", res["binary_embedding"]

assert_equal [1, 2, 3], res["sparse_embedding"].to_a

end

def conn

@@conn ||= begin

conn = PG.connect(dbname: "pgvector_ruby_test")

unless conn.exec("SELECT 1 FROM pg_extension WHERE extname = 'vector'").any?

conn.exec("CREATE EXTENSION IF NOT EXISTS vector")

end

conn.exec("DROP TABLE IF EXISTS pg_items")

conn.exec("CREATE TABLE pg_items (id bigserial PRIMARY KEY, embedding vector(3), half_embedding halfvec(3), binary_embedding bit(9), sparse_embedding sparsevec(3))")

registry = PG::BasicTypeRegistry.new.define_default_types

Pgvector::PG.register_vector(registry)

conn.type_map_for_queries = PG::BasicTypeMapForQueries.new(conn, registry: registry)

conn.type_map_for_results = PG::BasicTypeMapForResults.new(conn, registry: registry)

conn

end

end

end