diff --git a/README.md b/README.md index 3119185..c0de5b3 100644 --- a/README.md +++ b/README.md @@ -276,7 +276,7 @@ To install this gem onto your local machine, run `bundle exec rake install`. To export MINI_SQL_MYSQL_HOST=127.0.0.1 export MINI_SQL_MYSQL_PORT=33306 - docker run --name mini-sql-postgres --rm -it -p 55432:5432 -e POSTGRES_DB=test_mini_sql -e POSTGRES_HOST_AUTH_METHOD=trust -d postgres + docker run --name mini-sql-postgres --rm -it -p 55432:5432 -e POSTGRES_DB=test_mini_sql -e POSTGRES_HOST_AUTH_METHOD=trust -d postgres # or ankane/pgvector for testing vector type decoder export MINI_SQL_PG_USER=postgres export MINI_SQL_PG_HOST=127.0.0.1 export MINI_SQL_PG_PORT=55432 diff --git a/lib/mini_sql/postgres/connection.rb b/lib/mini_sql/postgres/connection.rb index 9f557d0..c7b2a75 100644 --- a/lib/mini_sql/postgres/connection.rb +++ b/lib/mini_sql/postgres/connection.rb @@ -26,6 +26,20 @@ def self.type_map(conn) else map.add_coder(MiniSql::Postgres::Coders::TimestampUtc.new(name: "timestamp", oid: 1114, format: 0)) end + + if defined? Pgvector::PG + vector_oid = + PG::BasicTypeRegistry::CoderMapsBundle + .new(conn) + .typenames_by_oid + .find { |k, v| v == "vector" } + &.first + + if !vector_oid.nil? + map.add_coder(Pgvector::PG::TextDecoder::Vector.new(name: "vector", oid: vector_oid, format: 0)) + map.add_coder(Pgvector::PG::BinaryDecoder::Vector.new(name: "vector", oid: vector_oid, format: 1)) + end + end map end end diff --git a/mini_sql.gemspec b/mini_sql.gemspec index 6c22fb3..f15242f 100644 --- a/mini_sql.gemspec +++ b/mini_sql.gemspec @@ -49,5 +49,8 @@ Gem::Specification.new do |spec| spec.add_development_dependency "mysql2" spec.add_development_dependency "sqlite3", "~> 1.4.4" spec.add_development_dependency "activerecord", "~> 7.0.0" + if RUBY_VERSION >= "3.0" + spec.add_development_dependency "pgvector", "~> 0.2.1" + end end end diff --git a/test/mini_sql/postgres/connection_test.rb b/test/mini_sql/postgres/connection_test.rb index a8dce3e..0986039 100644 --- a/test/mini_sql/postgres/connection_test.rb +++ b/test/mini_sql/postgres/connection_test.rb @@ -1,6 +1,7 @@ # frozen_string_literal: true require "test_helper" +require "pgvector" if RUBY_ENGINE != "jruby" && RUBY_VERSION >= "3.0" class MiniSql::Postgres::TestConnection < Minitest::Test def setup @@ -48,6 +49,15 @@ def test_cidr assert_equal(IPAddr.new("1.2.3.0/24"), ip) end + def test_vector + skip if @connection.query("SELECT 1 FROM pg_available_extensions WHERE name = 'vector';").empty? + + vector = [0.1, -0.2, 0.3] + @connection.exec("SET client_min_messages TO WARNING; CREATE EXTENSION IF NOT EXISTS vector;") + result = @connection.query_single("SELECT '[:vector]'::vector", vector: vector) + assert_equal(vector, result.first) + end + def test_bool b = @connection.query_single("select true").first assert_equal(true, b)