Merge pull request #43137 from composerinteralia/preload-already-loaded

Use loaded records where possible when preloading

Author: eileencodes

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

@@ -23,11 +23,7 @@ def hash
           end
 
           def records_for(loaders)
-            ids = loaders.flat_map(&:owner_keys).uniq
-
-            scope.where(association_key_name => ids).load do |record|
-              loaders.each { |l| l.set_inverse(record) }
-            end
+            LoaderRecords.new(loaders, self).records
           end
 
           def load_records_in_batch(loaders)
@@ -38,6 +34,52 @@ def load_records_in_batch(loaders)
               loader.run
             end
           end
+
+          def load_records_for_keys(keys, &block)
+            scope.where(association_key_name => keys).load(&block)
+          end
+        end
+
+        class LoaderRecords
+          def initialize(loaders, loader_query)
+            @loader_query = loader_query
+            @loaders = loaders
+            @keys_to_load = Set.new
+            @already_loaded_records_by_key = {}
+
+            populate_keys_to_load_and_already_loaded_records
+          end
+
+          def records
+            load_records + already_loaded_records
+          end
+
+          private
+            attr_reader :loader_query, :loaders, :keys_to_load, :already_loaded_records_by_key
+
+            def populate_keys_to_load_and_already_loaded_records
+              loaders.each do |loader|
+                loader.owners_by_key.each do |key, owners|
+                  if loaded_owner = owners.find { |owner| loader.loaded?(owner) }
+                    already_loaded_records_by_key[key] = loader.target_for(loaded_owner)
+                  else
+                    keys_to_load << key
+                  end
+                end
+              end
+
+              @keys_to_load.subtract(already_loaded_records_by_key.keys)
+            end
+
+            def load_records
+              loader_query.load_records_for_keys(keys_to_load) do |record|
+                loaders.each { |l| l.set_inverse(record) }
+              end
+            end
+
+            def already_loaded_records
+              already_loaded_records_by_key.values.flatten
+            end
         end
 
         attr_reader :klass
@@ -57,12 +99,8 @@ def table_name
           @klass.table_name
         end
 
-        def data_available?
-          already_loaded?
-        end
-
         def future_classes
-          if run? || already_loaded?
+          if run?
             []
           else
             [@klass]
@@ -81,11 +119,6 @@ def run
           return self if run?
           @run = true
 
-          if already_loaded?
-            fetch_from_preloaded_records
-            return self
-          end
-
           records = records_by_owner
 
           owners.each do |owner|
@@ -96,25 +129,17 @@ def run
         end
 
         def records_by_owner
-          ensure_loaded unless defined?(@records_by_owner)
+          load_records unless defined?(@records_by_owner)
 
           @records_by_owner
         end
 
         def preloaded_records
-          ensure_loaded unless defined?(@preloaded_records)
+          load_records unless defined?(@preloaded_records)
 
           @preloaded_records
         end
 
-        def ensure_loaded
-          if already_loaded?
-            fetch_from_preloaded_records
-          else
-            load_records
-          end
-        end
-
         # The name of the key on the associated records
         def association_key_name
           reflection.join_primary_key(klass)
@@ -124,8 +149,19 @@ def loader_query
           LoaderQuery.new(scope, association_key_name)
         end
 
-        def owner_keys
-          @owner_keys ||= owners_by_key.keys
+        def owners_by_key
+          @owners_by_key ||= owners.each_with_object({}) do |owner, result|
+            key = convert_key(owner[owner_key_name])
+            (result[key] ||= []) << owner if key
+          end
+        end
+
+        def loaded?(owner)
+          owner.association(reflection.name).loaded?
+        end
+
+        def target_for(owner)
+          Array.wrap(owner.association(reflection.name).target)
         end
 
         def scope
@@ -185,39 +221,23 @@ def associate_records_from_unscoped(unscoped_records)
         private
           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|
-              Array(owner.association(reflection.name).target)
-            end
-
-            @preloaded_records = records_by_owner.flat_map(&:last)
-          end
-
           # The name of the key on the model which declares the association
           def owner_key_name
             reflection.join_foreign_key
           end
 
           def associate_records_to_owner(owner, records)
+            return if loaded?(owner)
+
             association = owner.association(reflection.name)
+
             if reflection.collection?
               association.target = records
             else
               association.target = records.first
             end
           end
 
-          def owners_by_key
-            @owners_by_key ||= owners.each_with_object({}) do |owner, result|
-              key = convert_key(owner[owner_key_name])
-              (result[key] ||= []) << owner if key
-            end
-          end
-
           def key_conversion_required?
             unless defined?(@key_conversion_required)
               @key_conversion_required = (association_key_type != owner_key_type)

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

@@ -16,10 +16,7 @@ def call
 
             loaders.each { |loader| loader.associate_records_from_unscoped(@available_records[loader.klass]) }
 
-            already_loaded = loaders.select(&:data_available?)
-            if already_loaded.any?
-              already_loaded.each(&:run)
-            elsif loaders.any?
+            if loaders.any?
               future_tables = branches.flat_map do |branch|
                 branch.future_classes - branch.runnable_loaders.map(&:klass)
               end.map(&:table_name).uniq

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

@@ -10,10 +10,13 @@ def preloaded_records
 
         def records_by_owner
           return @records_by_owner if defined?(@records_by_owner)
-          source_records_by_owner = source_preloaders.map(&:records_by_owner).reduce(:merge)
-          through_records_by_owner = through_preloaders.map(&:records_by_owner).reduce(:merge)
 
           @records_by_owner = owners.each_with_object({}) do |owner, result|
+            if loaded?(owner)
+              result[owner] = target_for(owner)
+              next
+            end
+
             through_records = through_records_by_owner[owner] || []
 
             if owners.first.association(through_reflection.name).loaded?
@@ -35,12 +38,6 @@ 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 data_available?
             [self]
@@ -52,7 +49,7 @@ def runnable_loaders
         end
 
         def future_classes
-          if run? || data_available?
+          if run?
             []
           elsif through_preloaders.all?(&:run?)
             source_preloaders.flat_map(&:future_classes).uniq
@@ -67,6 +64,11 @@ def future_classes
         end
 
         private
+          def data_available?
+            owners.all? { |owner| loaded?(owner) } ||
+              through_preloaders.all?(&:run?) && source_preloaders.all?(&:run?)
+          end
+
           def source_preloaders
             @source_preloaders ||= ActiveRecord::Associations::Preloader.new(records: middle_records, associations: source_reflection.name, scope: scope, associate_by_default: false).loaders
           end
@@ -87,6 +89,14 @@ def source_reflection
             reflection.source_reflection
           end
 
+          def source_records_by_owner
+            @source_records_by_owner ||= source_preloaders.map(&:records_by_owner).reduce(:merge)
+          end
+
+          def through_records_by_owner
+            @through_records_by_owner ||= through_preloaders.map(&:records_by_owner).reduce(:merge)
+          end
+
           def preload_index
             @preload_index ||= preloaded_records.each_with_object({}).with_index do |(record, result), index|
               result[record] = index

activerecord/test/cases/associations_test.rb

@@ -443,6 +443,18 @@ def test_preload_groups_queries_with_same_scope
     end
   end
 
+  def test_preload_grouped_queries_with_already_loaded_records
+    book = books(:awdr)
+    post = posts(:welcome)
+    book.author
+
+    assert_no_queries do
+      ActiveRecord::Associations::Preloader.new(records: [book, post], associations: :author).call
+      book.author
+      post.author
+    end
+  end
+
   def test_preload_grouped_queries_of_middle_records
     comments = [
       comments(:eager_sti_on_associations_s_comment1),
@@ -792,6 +804,41 @@ def test_preload_with_available_records
     end
   end
 
+  def test_preload_with_only_some_records_available
+    bob_post = posts(:misc_by_bob)
+    mary_post = posts(:misc_by_mary)
+    bob = authors(:bob)
+    mary = authors(:mary)
+
+    assert_queries(1) do
+      ActiveRecord::Associations::Preloader.new(records: [bob_post, mary_post], associations: :author, available_records: [bob]).call
+    end
+
+    assert_no_queries do
+      assert_same bob, bob_post.author
+      assert_equal mary, mary_post.author
+    end
+  end
+
+  def test_preload_with_some_records_already_loaded
+    bob_post = posts(:misc_by_bob)
+    mary_post = posts(:misc_by_mary)
+    bob = bob_post.author
+    mary = authors(:mary)
+
+    assert bob_post.association(:author).loaded?
+    assert_not mary_post.association(:author).loaded?
+
+    assert_queries(1) do
+      ActiveRecord::Associations::Preloader.new(records: [bob_post, mary_post], associations: :author).call
+    end
+
+    assert_no_queries do
+      assert_same bob, bob_post.author
+      assert_equal mary, mary_post.author
+    end
+  end
+
   def test_preload_with_available_records_with_through_association
     author = authors(:david)
     categories = Category.all.to_a
@@ -805,6 +852,25 @@ def test_preload_with_available_records_with_through_association
     assert categories.map(&:object_id).include?(author.essay_category.object_id)
   end
 
+  def test_preload_with_only_some_records_available_with_through_associations
+    mary = authors(:mary)
+    mary_essay = essays(:mary_stay_home)
+    mary_category = categories(:technology)
+    mary_essay.update!(category: mary_category)
+
+    dave = authors(:david)
+    dave_category = categories(:general)
+
+    assert_queries(2) do
+      ActiveRecord::Associations::Preloader.new(records: [mary, dave], associations: :essay_category, available_records: [mary_category]).call
+    end
+
+    assert_no_queries do
+      assert_same mary_category, mary.essay_category
+      assert_equal dave_category, dave.essay_category
+    end
+  end
+
   def test_preload_with_available_records_with_multiple_classes
     essay = essays(:david_modest_proposal)
     general = categories(:general)
@@ -858,6 +924,20 @@ def test_preload_with_available_records_queries_when_incomplete
       assert_equal david, post.author
     end
   end
+
+  def test_preload_with_unpersisted_records_no_ops
+    author = Author.new
+    new_post_with_author = Post.new(author: author)
+    new_post_without_author = Post.new
+    posts = [new_post_with_author, new_post_without_author]
+
+    assert_no_queries do
+      ActiveRecord::Associations::Preloader.new(records: posts, associations: :author).call
+
+      assert_same author, new_post_with_author.author
+      assert_nil new_post_without_author.author
+    end
+  end
 end
 
 class GeneratedMethodsTest < ActiveRecord::TestCase