updated code for benchmarks

bin/bm_in_find_by_sql.rb

@@ -0,0 +1,59 @@
+require 'rubygems'
+require 'active_record'
+require 'active_record/migration'
+require_relative "../lib/treeify"
+require 'benchmark'
+include Benchmark   
+
+ActiveRecord::Base.establish_connection(
+  :adapter => 'postgresql',
+  :database => 'tree_test',
+  :username => 'postgres'
+)
+
+class Node < ActiveRecord::Base
+  include Treeify
+
+  validates_uniqueness_of :name
+  validates_uniqueness_of :parent_id, :scope=> :id
+end
+
+class NodeSetup < ActiveRecord::Migration
+  class << self
+    def up
+      create_table :nodes do |t|
+        t.text :name
+        t.integer :parent_id
+        t.references :parent
+      end
+      add_index :nodes, [:parent_id, :id], :unique => true
+   end
+   def down
+     drop_table :nodes
+   end
+  end
+end
+NodeSetup.up
+nodes = []
+
+puts "Creating records..."
+10.times do |i|
+  nodes[i] = []
+  parent = Node.create(name: "root_#{i}")
+  10.times do |j|
+    node = Node.new(name: "node_#{i}_#{j}")
+    _parent = nodes[i][rand(nodes[i].size)] || parent
+    node.parent_id = _parent.id
+    node.save
+    nodes[i] << node
+  end
+end
+puts "Created #{Node.count} records"
+
+puts "Running benchmark..."
+Benchmark.benchmark(CAPTION, 7, FORMAT, ">total:", ">avg:") do |x|
+  fb   = x.report("find_by_sql:") { Node.roots.first.descendents * 1000 }
+  w_in = x.report("where in:")   { Node.roots.first.descendents2 * 1000 }
+  [fb+w_in, (fb+w_in)/2]
+end
+NodeSetup.down

lib/treeify.rb

@@ -1,10 +1,15 @@
 require 'active_record'
+require 'pp'
 require "active_support/concern"
 require "active_support/core_ext/module/attribute_accessors"
 require "active_support/core_ext/class/attribute"
 
 module Treeify
-  extend ActiveSupport::Concern 
+  extend ActiveSupport::Concern
+
+  cattr_writer :cols do 
+    []
+  end
 
   included do
     has_many :children,
@@ -17,27 +22,34 @@ module Treeify
     class_attribute :cols
     scope :roots, -> { where(parent_id: nil) }
     scope :tree_for, ->(instance) { self.find_by_sql self.tree_sql_for(instance) }
+    scope :tree_for2, ->(instance) { where("#{table_name}.id IN (#{tree_sql_for2(instance)})").order("#{table_name}.id") }
     scope :tree_for_ancestors, ->(instance) { self.find_by_sql self.tree_sql_for_ancestors(instance) }
   end
 
   module ClassMethods
 
     def tree_config(hash = {})
-      self.cols = hash[:cols]
+      self.cols = !hash[:cols].nil? == true ? hash[:cols] : []
     end
 
     def columns_joined(char=",")
+      self.cols ||= []
       self.cols.join(char)
     end
 
     def columns_with_table_name
+      self.cols ||= []
       self.cols.map{|c| 
         c = "#{table_name}.#{c}" 
       }.join(",")
     end
 
     def has_config_defined_cols?
-      !self.cols.empty?
+      #return true if self.respond_to?("cols") && !self.cols.nil? 
+      if self.respond_to?("cols")
+        return !self.cols.empty? if !self.cols.nil?
+      end
+      false
     end
 
     # sort of hacky, but check to see if we have any columns defined in the config
@@ -67,12 +79,32 @@ def tree_sql(instance)
        )"
     end
 
+    def tree_sql2(instance)
+    "WITH RECURSIVE cte (id, path)  AS (
+     SELECT  id,
+             array[id] AS path
+     FROM    #{table_name}
+     WHERE   id = #{instance.id}
+     UNION ALL
+     SELECT  #{table_name}.id,
+             cte.path || #{table_name}.id
+     FROM    #{table_name}
+     JOIN cte ON #{table_name}.parent_id = cte.id
+     )"
+    end
+
     def tree_sql_for(instance)
       "#{tree_sql(instance)}
        SELECT * FROM cte
        ORDER BY path"
     end
 
+    def tree_sql_for2(instance)
+      "#{tree_sql2(instance)}
+       SELECT id FROM cte
+       ORDER BY path"
+    end
+
     def tree_sql_for_ancestors(instance)
       "#{tree_sql(instance)}
       SELECT * FROM cte 
@@ -84,6 +116,10 @@ def descendents
     self_and_descendents - [self]
   end
 
+  def descendents2
+    self_and_descendents2 - [self]
+  end
+
   def ancestors
     self.class.tree_for_ancestors(self)
   end
@@ -92,6 +128,10 @@ def self_and_descendents
     self.class.tree_for(self)
   end
 
+  def self_and_descendents2
+    self.class.tree_for2(self)
+  end
+
   def is_root?
     self.parent_id != nil
   end
@@ -125,4 +165,6 @@ def build_tree(data)
       nested_hash.has_key? item['parent_id']
     }.values
   end
+
+
 end