Optimize Dataset#all and #with_sql_all

This allocates an array with the expected size up front, and avoids
a C -> Ruby call per row.

Instead of checking whether Sequel::Model::ClassMethods is already
defined during load, check it inside optimize_model_load. This
allows for optimization in the case where sequel/core is required,
then the database is setup, and then sequel is required.

However, this approach also results in unnecessary slowdown of
Dataset#each when models are not used. An alternative approach will
be used in a future commit to avoid this.

Author: jeremyevans

CHANGELOG

@@ -1,5 +1,7 @@
 === master
 
+* Optimize Dataset#all and #with_sql_all (jeremyevans)
+
 * Fix runtime warnings when using Dataset#as_hash and #to_hash_groups with invalid columns (jeremyevans)
 
 * Fix Dataset#map return value when the null_dataset extension is used (jeremyevans)

ext/sequel_pg/sequel_pg.c

@@ -77,6 +77,8 @@
 #define SPG_YIELD_COLUMNS_SET 19
 #define SPG_YIELD_FIRST_SET 20
 #define SPG_YIELD_ARRAY_SET 21
+#define SPG_YIELD_ALL 22
+#define SPG_YIELD_ALL_MODEL 23
 
 /* External functions defined by ruby-pg */
 PGconn* pg_get_pgconn(VALUE);
@@ -114,6 +116,8 @@ static VALUE spg_sym_array_set;
 static VALUE spg_sym_hash;
 static VALUE spg_sym_hash_groups;
 static VALUE spg_sym_model;
+static VALUE spg_sym_all;
+static VALUE spg_sym_all_model;
 static VALUE spg_sym__sequel_pg_type;
 static VALUE spg_sym__sequel_pg_value;
 
@@ -1470,6 +1474,10 @@ static VALUE spg_yield_hash_rows_internal(VALUE self, PGresult *res, int enc_ind
         }
       } else if (pg_type == spg_sym_model && rb_type(pg_value) == T_CLASS) {
         type = SPG_YIELD_MODEL;
+      } else if (pg_type == spg_sym_all_model && rb_type(pg_value) == T_CLASS) {
+        type = SPG_YIELD_ALL_MODEL;
+      } else if (pg_type == spg_sym_all) {
+        type = SPG_YIELD_ALL;
       }
     }
   }
@@ -1760,6 +1768,35 @@ static VALUE spg_yield_hash_rows_internal(VALUE self, PGresult *res, int enc_ind
         rb_yield(pg_type);
       }
       break;
+    case SPG_YIELD_ALL_MODEL:
+      {
+        VALUE ary = rb_ary_new2(ntuples);
+        VALUE obj;
+        for(i=0; i<ntuples; i++) {
+          h = rb_hash_new_capa(nfields);
+          for(j=0; j<nfields; j++) {
+            rb_hash_aset(h, colsyms[j], spg__col_value(self, res, i, j, colconvert, enc_index));
+          }
+          obj = rb_obj_alloc(pg_value);
+          rb_ivar_set(obj, spg_id_values, h);
+          rb_ary_store(ary, i, obj);
+        }
+        rb_yield(ary);
+      }
+      break;
+    case SPG_YIELD_ALL:
+      {
+        VALUE ary = rb_ary_new2(ntuples);
+        for(i=0; i<ntuples; i++) {
+          h = rb_hash_new_capa(nfields);
+          for(j=0; j<nfields; j++) {
+            rb_hash_aset(h, colsyms[j], spg__col_value(self, res, i, j, colconvert, enc_index));
+          }
+          rb_ary_store(ary, i, h);
+        }
+        rb_yield(ary);
+      }
+      break;
   }
   return self;
 }
@@ -2083,6 +2120,8 @@ void Init_sequel_pg(void) {
   spg_sym_hash = ID2SYM(rb_intern("hash"));
   spg_sym_hash_groups = ID2SYM(rb_intern("hash_groups"));
   spg_sym_model = ID2SYM(rb_intern("model"));
+  spg_sym_all = ID2SYM(rb_intern("all"));
+  spg_sym_all_model = ID2SYM(rb_intern("all_model"));
   spg_sym__sequel_pg_type = ID2SYM(rb_intern("_sequel_pg_type"));
   spg_sym__sequel_pg_value = ID2SYM(rb_intern("_sequel_pg_value"));
 

lib/sequel_pg/sequel_pg.rb

@@ -91,16 +91,35 @@ def to_hash_groups(key_column, value_column = nil, opts = Sequel::OPTS)
     end
   end
 
-  if defined?(Sequel::Model::ClassMethods)
-    # If model loads are being optimized and this is a model load, use the optimized
-    # version.
-    def each(&block)
-      if optimize_model_load?
-        clone(:_sequel_pg_type=>:model, :_sequel_pg_value=>row_proc).fetch_rows(sql, &block)
-      else
-        super
+  # If model loads are being optimized and this is a model load, use the optimized
+  # version.
+  def each(&block)
+    if optimize_model_load?
+      clone(:_sequel_pg_type=>:model, :_sequel_pg_value=>row_proc).fetch_rows(sql, &block)
+    else
+      super
+    end
+  end
+
+  # Delegate to with_sql_all using the default SQL
+  def all(&block)
+    with_sql_all(sql, &block)
+  end
+
+  # Always use optimized version
+  def with_sql_all(sql, &block)
+    return super if opts[:graph] || opts[:cursor]
+    type = optimize_model_load? ? :all_model : :all
+    rp = row_proc
+    clone(:_sequel_pg_type=>type, :_sequel_pg_value=>rp).fetch_rows(sql) do |array|
+      if rp && type == :all
+        array.map!{|h| rp.call(h)}
       end
+      post_load(array)
+      array.each(&block) if block
+      return array
     end
+    []
   end
 
   protected
@@ -153,18 +172,17 @@ def _select_set_single
 
   private
 
-  if defined?(Sequel::Model::ClassMethods)
-    # The model load can only be optimized if it's for a model and it's not a graphed dataset
-    # or using a cursor.
-    def optimize_model_load?
+  # The model load can only be optimized if it's for a model and it's not a graphed dataset
+  # or using a cursor.
+  def optimize_model_load?
+    defined?(Sequel::Model::ClassMethods) &&
       (rp = row_proc) &&
-        rp.is_a?(Class) &&
-        rp < Sequel::Model &&
-        rp.method(:call).owner == Sequel::Model::ClassMethods &&
-        opts[:optimize_model_load] != false &&
-        !opts[:use_cursor] &&
-        !opts[:graph]
-    end
+      rp.is_a?(Class) &&
+      rp < Sequel::Model &&
+      rp.method(:call).owner == Sequel::Model::ClassMethods &&
+      opts[:optimize_model_load] != false &&
+      !opts[:use_cursor] &&
+      !opts[:graph]
   end
 end