Add optional PostgreSQL native materializer

.github/workflows/ci.yml

@@ -43,12 +43,12 @@ jobs:
             experimental: true
     continue-on-error: ${{matrix.experimental}}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
       - uses: ruby/setup-ruby@v1
         with:
           ruby-version: ${{ matrix.ruby }}
       - name: Bundler cache
-        uses: actions/cache@v2
+        uses: actions/cache@v4
         with:
           path: vendor/bundle
           key: ${{ runner.os }}-${{ matrix.ruby }}-gems-${{ hashFiles('**/Gemfile.lock') }}
@@ -73,7 +73,7 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
 
       - name: Release Gem
         uses: discourse/publish-rubygems-action@v2

CHANGELOG

@@ -1,3 +1,8 @@
+Unreleased
+
+- FEATURE: add optional `mini_sql-pg_native` companion gem for faster PostgreSQL row materialization
+- FIX: initialize PostgreSQL type map before streaming queries enter single-row mode
+
 2024-08-21 - 1.6.0
 
 - FEATURE: optionaly allow encoding pg arrays efficiently

Gemfile

@@ -5,4 +5,4 @@ source "https://rubygems.org"
 git_source(:github) { |repo_name| "https://github.com/#{repo_name}" }
 
 # Specify your gem's dependencies in mini_sql.gemspec
-gemspec
+gemspec name: "mini_sql"

README.md

@@ -8,6 +8,16 @@ Add this line to your application's Gemfile:
 gem 'mini_sql'
 ```
 
+For faster PostgreSQL row materialization, add the optional native companion gem:
+
+```ruby
+gem 'mini_sql'
+gem 'mini_sql-pg_native'
+```
+
+`mini_sql-pg_native` only affects the PostgreSQL adapter and can be disabled with
+`MINI_SQL_PG_NATIVE=0`.
+
 And then execute:
 
     $ bundle

Rakefile

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
-require "bundler/gem_tasks"
+require "bundler/gem_helper"
+Bundler::GemHelper.install_tasks(name: "mini_sql")
 require "rake/testtask"
 
 if RUBY_ENGINE == 'jruby' # Excluding sqlite3 tests

bench/pg_native_materialize_only_perf.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require 'pg'
+require 'benchmark/ips'
+require_relative '../lib/mini_sql'
+
+conn = PG.connect(dbname: ENV.fetch('PGDATABASE', 'discourse_sql_ft'), user: ENV.fetch('PGUSER', 'agent'))
+conn.async_exec <<~SQL
+  DROP TABLE IF EXISTS mini_sql_pg_native_bench;
+  CREATE UNLOGGED TABLE mini_sql_pg_native_bench AS
+  SELECT
+    g AS id,
+    'title ' || g AS title,
+    'slug-' || g AS slug,
+    (g % 1000) AS user_id,
+    (g % 100) AS category_id,
+    (g % 2 = 0) AS visible,
+    now() - (g || ' seconds')::interval AS created_at,
+    now() AS updated_at,
+    repeat('x', 32) AS payload,
+    g::numeric / 10 AS score
+  FROM generate_series(1, 50000) g;
+SQL
+
+mini = MiniSql::Connection.get(conn)
+mode = ENV['MINI_SQL_PG_NATIVE'] == '0' ? 'ruby' : 'pg_native'
+puts "mode=#{mode} pg_native_defined=#{defined?(MiniSql::Postgres::Native::RowMaterializer).inspect} ruby=#{RUBY_VERSION} pg=#{PG.library_version}"
+
+RESULTS = {
+  'mat_1k_2col' => conn.async_exec('select id, title from mini_sql_pg_native_bench order by id limit 1000'),
+  'mat_1k_10col' => conn.async_exec('select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 1000'),
+  'mat_10k_10col' => conn.async_exec('select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 10000'),
+}.freeze
+
+RESULTS.each_value { |r| r.type_map = mini.type_map }
+cache = MiniSql::Postgres::DeserializerCache.new
+RESULTS.each_value { |r| 3.times { cache.materialize(r) } }
+
+Benchmark.ips do |x|
+  x.config(warmup: Float(ENV.fetch('WARMUP', '2')), time: Float(ENV.fetch('TIME', '5')))
+
+  RESULTS.each do |name, result|
+    x.report(name) do |n|
+      while n > 0
+        rows = cache.materialize(result)
+        first = rows[0]
+        raise 'bad row' unless first && first.id
+        n -= 1
+      end
+    end
+  end
+end

bench/pg_native_materialize_perf.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+require 'pg'
+require 'benchmark/ips'
+require_relative '../lib/mini_sql'
+
+conn = PG.connect(dbname: ENV.fetch('PGDATABASE', 'discourse_sql_ft'), user: ENV.fetch('PGUSER', 'agent'))
+conn.async_exec <<~SQL
+  DROP TABLE IF EXISTS mini_sql_pg_native_bench;
+  CREATE UNLOGGED TABLE mini_sql_pg_native_bench AS
+  SELECT
+    g AS id,
+    'title ' || g AS title,
+    'slug-' || g AS slug,
+    (g % 1000) AS user_id,
+    (g % 100) AS category_id,
+    (g % 2 = 0) AS visible,
+    now() - (g || ' seconds')::interval AS created_at,
+    now() AS updated_at,
+    repeat('x', 32) AS payload,
+    g::numeric / 10 AS score
+  FROM generate_series(1, 50000) g;
+  CREATE INDEX ON mini_sql_pg_native_bench(id);
+SQL
+
+mini = MiniSql::Connection.get(conn)
+mode = ENV['MINI_SQL_PG_NATIVE'] == '0' ? 'ruby' : 'pg_native'
+puts "mode=#{mode} pg_native_defined=#{defined?(MiniSql::Postgres::Native::RowMaterializer).inspect} ruby=#{RUBY_VERSION} pg=#{PG.library_version}"
+
+QUERIES = {
+  'query_1k_2col' => 'select id, title from mini_sql_pg_native_bench order by id limit 1000',
+  'query_1k_10col' => 'select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 1000',
+  'query_10k_10col' => 'select id, title, slug, user_id, category_id, visible, created_at, updated_at, payload, score from mini_sql_pg_native_bench order by id limit 10000',
+}.freeze
+
+# Warm caches: PostgreSQL buffer cache, pg type map, mini_sql materializer cache.
+QUERIES.each_value { |sql| 3.times { mini.query(sql) } }
+
+Benchmark.ips do |x|
+  x.config(warmup: Float(ENV.fetch('WARMUP', '2')), time: Float(ENV.fetch('TIME', '5')))
+
+  QUERIES.each do |name, sql|
+    x.report(name) do |n|
+      while n > 0
+        rows = mini.query(sql)
+        # Touch values so dead-code elimination has nowhere cute to hide.
+        first = rows[0]
+        raise 'bad row' unless first && first.id
+        n -= 1
+      end
+    end
+  end
+end

bench/pg_native_simple_materialize_only_perf.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require 'pg'
+require 'benchmark/ips'
+require_relative '../lib/mini_sql'
+
+conn = PG.connect(dbname: ENV.fetch('PGDATABASE', 'discourse_sql_ft'), user: ENV.fetch('PGUSER', 'agent'))
+conn.async_exec <<~SQL
+  DROP TABLE IF EXISTS mini_sql_pg_native_simple_bench;
+  CREATE UNLOGGED TABLE mini_sql_pg_native_simple_bench AS
+  SELECT
+    g AS id,
+    'title ' || g AS title,
+    'slug-' || g AS slug,
+    (g % 1000) AS user_id,
+    (g % 100) AS category_id,
+    repeat('x', 32) AS payload
+  FROM generate_series(1, 50000) g;
+SQL
+
+mini = MiniSql::Connection.get(conn)
+mode = ENV['MINI_SQL_PG_NATIVE'] == '0' ? 'ruby' : 'pg_native'
+puts "mode=#{mode} pg_native_defined=#{defined?(MiniSql::Postgres::Native::RowMaterializer).inspect} ruby=#{RUBY_VERSION} pg=#{PG.library_version}"
+
+RESULTS = {
+  'simple_mat_1k_2col' => conn.async_exec('select id, title from mini_sql_pg_native_simple_bench order by id limit 1000'),
+  'simple_mat_1k_6col' => conn.async_exec('select id, title, slug, user_id, category_id, payload from mini_sql_pg_native_simple_bench order by id limit 1000'),
+  'simple_mat_10k_6col' => conn.async_exec('select id, title, slug, user_id, category_id, payload from mini_sql_pg_native_simple_bench order by id limit 10000'),
+}.freeze
+
+RESULTS.each_value { |r| r.type_map = mini.type_map }
+cache = MiniSql::Postgres::DeserializerCache.new
+RESULTS.each_value { |r| 3.times { cache.materialize(r) } }
+
+Benchmark.ips do |x|
+  x.config(warmup: Float(ENV.fetch('WARMUP', '2')), time: Float(ENV.fetch('TIME', '5')))
+  RESULTS.each do |name, result|
+    x.report(name) do |n|
+      while n > 0
+        rows = cache.materialize(result)
+        raise 'bad row' unless rows[0]&.id
+        n -= 1
+      end
+    end
+  end
+end

ext/mini_sql/pg_native/extconf.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+require "mkmf"
+require "rubygems"
+
+pg_spec = Gem::Specification.find_by_name("pg")
+# rubocop:disable Style/GlobalVars
+$INCFLAGS << " -I#{File.join(pg_spec.full_gem_path, "ext")}"
+# rubocop:enable Style/GlobalVars
+
+create_makefile("mini_sql/pg_native/pg_native")

ext/mini_sql/pg_native/pg_native.c

@@ -0,0 +1,143 @@
+#include "ruby.h"
+#include "pg.h"
+
+typedef struct {
+  VALUE klass;
+  long nfields;
+  ID *ivars;
+} row_materializer_t;
+
+static VALUE cRowMaterializer;
+
+static void rm_mark(void *ptr) {
+  row_materializer_t *rm = (row_materializer_t *)ptr;
+  if (rm) rb_gc_mark_movable(rm->klass);
+}
+
+static void rm_compact(void *ptr) {
+  row_materializer_t *rm = (row_materializer_t *)ptr;
+  if (rm) rm->klass = rb_gc_location(rm->klass);
+}
+
+static void rm_free(void *ptr) {
+  row_materializer_t *rm = (row_materializer_t *)ptr;
+  if (rm) {
+    xfree(rm->ivars);
+    xfree(rm);
+  }
+}
+
+static size_t rm_memsize(const void *ptr) {
+  const row_materializer_t *rm = (const row_materializer_t *)ptr;
+  return rm ? sizeof(row_materializer_t) + (sizeof(ID) * rm->nfields) : 0;
+}
+
+static const rb_data_type_t rm_type = {
+  "MiniSql::Postgres::Native::RowMaterializer",
+  { rm_mark, rm_free, rm_memsize, rm_compact, },
+  0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+};
+
+static VALUE rm_alloc(VALUE klass) {
+  row_materializer_t *rm;
+  VALUE obj = TypedData_Make_Struct(klass, row_materializer_t, &rm_type, rm);
+  rm->klass = Qnil;
+  rm->nfields = 0;
+  rm->ivars = NULL;
+  return obj;
+}
+
+static VALUE rm_initialize(VALUE self, VALUE row_class, VALUE fields) {
+  row_materializer_t *rm;
+  TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm);
+
+  Check_Type(fields, T_ARRAY);
+  rm->klass = row_class;
+  rm->nfields = RARRAY_LEN(fields);
+  xfree(rm->ivars);
+  rm->ivars = ALLOC_N(ID, rm->nfields);
+
+  for (long i = 0; i < rm->nfields; i++) {
+    VALUE field = rb_obj_as_string(rb_ary_entry(fields, i));
+    VALUE ivar_name = rb_str_plus(rb_str_new_cstr("@"), field);
+    rm->ivars[i] = rb_intern_str(ivar_name);
+  }
+
+  return self;
+}
+
+static void check_pgresult(t_pg_result *pgresult) {
+  if (!pgresult->pgresult) {
+    rb_raise(rb_eNoResultError, "PG::Result has been cleared");
+  }
+}
+
+static VALUE rm_materialize(VALUE self, VALUE result, VALUE index) {
+  row_materializer_t *rm;
+  TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm);
+
+  int idx = NUM2INT(index);
+  t_pg_result *pgresult = pgresult_get_this(result);
+  check_pgresult(pgresult);
+
+  int rows = PQntuples(pgresult->pgresult);
+  if (idx < 0 || idx >= rows) {
+    rb_raise(rb_eArgError, "invalid tuple number %d", idx);
+  }
+
+  VALUE row = rb_obj_alloc(rm->klass);
+  long nfields = rm->nfields;
+  ID *ivars = rm->ivars;
+  t_typemap *typemap = pgresult->p_typemap;
+
+  for (long col = 0; col < nfields; col++) {
+    VALUE val = typemap->funcs.typecast_result_value(typemap, result, idx, col);
+    rb_ivar_set(row, ivars[col], val);
+  }
+
+  return row;
+}
+
+static VALUE rm_materialize_all(VALUE self, VALUE result) {
+  row_materializer_t *rm;
+  TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm);
+
+  t_pg_result *pgresult = pgresult_get_this(result);
+  check_pgresult(pgresult);
+
+  int rows = PQntuples(pgresult->pgresult);
+  long nfields = rm->nfields;
+  ID *ivars = rm->ivars;
+  t_typemap *typemap = pgresult->p_typemap;
+  VALUE ary = rb_ary_new_capa(rows);
+
+  for (int idx = 0; idx < rows; idx++) {
+    VALUE row = rb_obj_alloc(rm->klass);
+    for (long col = 0; col < nfields; col++) {
+      VALUE val = typemap->funcs.typecast_result_value(typemap, result, idx, col);
+      rb_ivar_set(row, ivars[col], val);
+    }
+    rb_ary_store(ary, idx, row);
+  }
+
+  return ary;
+}
+
+static VALUE rm_row_class(VALUE self) {
+  row_materializer_t *rm;
+  TypedData_Get_Struct(self, row_materializer_t, &rm_type, rm);
+  return rm->klass;
+}
+
+void Init_pg_native(void) {
+  VALUE mMiniSql = rb_define_module("MiniSql");
+  VALUE mPostgres = rb_define_module_under(mMiniSql, "Postgres");
+  VALUE mNative = rb_define_module_under(mPostgres, "Native");
+
+  cRowMaterializer = rb_define_class_under(mNative, "RowMaterializer", rb_cObject);
+  rb_define_alloc_func(cRowMaterializer, rm_alloc);
+  rb_define_method(cRowMaterializer, "initialize", rm_initialize, 2);
+  rb_define_method(cRowMaterializer, "materialize", rm_materialize, 2);
+  rb_define_method(cRowMaterializer, "materialize_all", rm_materialize_all, 1);
+  rb_define_method(cRowMaterializer, "row_class", rm_row_class, 0);
+}

lib/mini_sql/mysql/connection.rb

@@ -29,8 +29,8 @@ def query_array(sql, *params)
       end
 
       def exec(sql, *params)
-        run(sql, :array, params)
-        raw_connection.affected_rows
+        result = run(sql, :array, params)
+        result ? result.count : raw_connection.affected_rows
       end
 
       def query(sql, *params)

lib/mini_sql/pg_native.rb

@@ -0,0 +1,4 @@
+# frozen_string_literal: true
+
+require "pg"
+require "mini_sql/pg_native/pg_native"