Merge pull request rails#41790 from jhawthorn/preloader_smart_batching

"Smart" ActiveRecord Preloader batching

Author: jhawthorn

activerecord/lib/active_record/associations/preloader/association.rb

@@ -40,6 +40,8 @@ def load_records_in_batch(loaders)
           end
         end
 
+        attr_reader :klass
+
         def initialize(klass, owners, reflection, preload_scope, associate_by_default = true)
           @klass         = klass
           @owners        = owners.uniq(&:__id__)
@@ -50,8 +52,20 @@ def initialize(klass, owners, reflection, preload_scope, associate_by_default =
           @run = false
         end
 
-        def already_loaded?
-          @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? }
+        def table_name
+          @klass.table_name
+        end
+
+        def data_available?
+          already_loaded?
+        end
+
+        def future_classes
+          if run? || already_loaded?
+            []
+          else
+            [@klass]
+          end
         end
 
         def runnable_loaders
@@ -149,7 +163,11 @@ def load_records(raw_records = nil)
         end
 
         private
-          attr_reader :owners, :reflection, :preload_scope, :model, :klass
+          attr_reader :owners, :reflection, :preload_scope, :model
+
+          def already_loaded?
+            @already_loaded ||= owners.all? { |o| o.association(reflection.name).loaded? }
+          end
 
           def fetch_from_preloaded_records
             @records_by_owner = owners.index_with do |owner|

activerecord/lib/active_record/associations/preloader/batch.rb

@@ -13,11 +13,20 @@ def call
           until branches.empty?
             loaders = branches.flat_map(&:runnable_loaders)
 
-            already_loaded, loaders = loaders.partition(&:already_loaded?)
-            already_loaded.each(&:run)
-
-            group_and_load_similar(loaders)
-            loaders.each(&:run)
+            already_loaded = loaders.select(&:data_available?)
+            if already_loaded.any?
+              already_loaded.each(&:run)
+            elsif loaders.any?
+              future_tables = branches.flat_map do |branch|
+                branch.future_classes - branch.runnable_loaders.map(&:klass)
+              end.map(&:table_name).uniq
+
+              target_loaders = loaders.reject { |l| future_tables.include?(l.table_name)  }
+              target_loaders = loaders if target_loaders.empty?
+
+              group_and_load_similar(target_loaders)
+              target_loaders.each(&:run)
+            end
 
             finished, in_progress = branches.partition(&:done?)
 

activerecord/lib/active_record/associations/preloader/branch.rb

@@ -15,6 +15,40 @@ def initialize(association:, children:, parent:, associate_by_default:, scope:)
           @associate_by_default = associate_by_default
 
           @children = build_children(children)
+          @loaders = nil
+        end
+
+        def future_classes
+          (immediate_future_classes + children.flat_map(&:future_classes)).uniq
+        end
+
+        def immediate_future_classes
+          if parent.done?
+            loaders.flat_map(&:future_classes).uniq
+          else
+            likely_reflections.reject(&:polymorphic?).flat_map do |reflection|
+              reflection.
+                chain.
+                map(&:klass)
+            end.uniq
+          end
+        end
+
+        def target_classes
+          if done?
+            preloaded_records.map(&:klass).uniq
+          elsif parent.done?
+            loaders.map(&:klass).uniq
+          else
+            likely_reflections.reject(&:polymorphic?).map(&:klass).uniq
+          end
+        end
+
+        def likely_reflections
+          parent_classes = parent.target_classes
+          parent_classes.map do |parent_klass|
+            parent_klass._reflect_on_association(@association)
+          end.compact
         end
 
         def root?
@@ -30,7 +64,7 @@ def preloaded_records
         end
 
         def done?
-          loaders.all?(&:run?)
+          root? || (@loaders && @loaders.all?(&:run?))
         end
 
         def runnable_loaders

activerecord/lib/active_record/associations/preloader/through_association.rb

@@ -35,16 +35,37 @@ def records_by_owner
           end
         end
 
+        def data_available?
+          return true if super()
+          through_preloaders.all?(&:run?) &&
+            source_preloaders.all?(&:run?)
+        end
+
         def runnable_loaders
-          if already_loaded?
+          if data_available?
             [self]
           elsif through_preloaders.all?(&:run?)
-            [self] + source_preloaders.flat_map(&:runnable_loaders)
+            source_preloaders.flat_map(&:runnable_loaders)
           else
             through_preloaders.flat_map(&:runnable_loaders)
           end
         end
 
+        def future_classes
+          if run? || data_available?
+            []
+          elsif through_preloaders.all?(&:run?)
+            source_preloaders.flat_map(&:future_classes).uniq
+          else
+            through_classes = through_preloaders.flat_map(&:future_classes)
+            source_classes = source_reflection.
+              chain.
+              reject { |reflection| reflection.respond_to?(:polymorphic?) && reflection.polymorphic? }.
+              map(&:klass)
+            (through_classes + source_classes).uniq
+          end
+        end
+
         private
           def source_preloaders
             @source_preloaders ||= ActiveRecord::Associations::Preloader.new(records: middle_records, associations: source_reflection.name, scope: scope, associate_by_default: false).loaders

activerecord/test/cases/associations_test.rb

@@ -362,7 +362,7 @@ def test_requires_symbol_argument
 end
 
 class PreloaderTest < ActiveRecord::TestCase
-  fixtures :posts, :comments, :books, :authors
+  fixtures :posts, :comments, :books, :authors, :tags, :taggings
 
   def test_preload_with_scope
     post = posts(:welcome)
@@ -569,6 +569,40 @@ def test_preload_with_grouping_sets_inverse_association
     end
   end
 
+  def test_preload_can_group_separate_levels
+    mary = authors(:mary)
+    bob = authors(:bob)
+
+    AuthorFavorite.create!(author: mary, favorite_author: bob)
+
+    assert_queries(3) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [mary], associations: [:posts, favorite_authors: :posts])
+      preloader.call
+    end
+
+    assert_no_queries do
+      mary.posts
+      mary.favorite_authors.map(&:posts)
+    end
+  end
+
+  def test_preload_can_group_multi_level_ping_pong_through
+    mary = authors(:mary)
+    bob = authors(:bob)
+
+    AuthorFavorite.create!(author: mary, favorite_author: bob)
+
+    assert_queries(9) do
+      preloader = ActiveRecord::Associations::Preloader.new(records: [mary], associations: { similar_posts: :comments, favorite_authors: { similar_posts: :comments } })
+      preloader.call
+    end
+
+    assert_no_queries do
+      mary.similar_posts.map(&:comments).each(&:to_a)
+      mary.favorite_authors.flat_map(&:similar_posts).map(&:comments).each(&:to_a)
+    end
+  end
+
   def test_preload_does_not_group_same_class_different_scope
     post = posts(:welcome)
     postesque = Postesque.create(author: Author.last)