From 9aa59f9d4aa88478c03f323ae7730059fd1d0ff5 Mon Sep 17 00:00:00 2001
From: Michael Fowler <mkrfowler@gmail.com>
Date: Sun, 29 Dec 2019 16:40:16 +1300
Subject: [PATCH] Avoid extraneous preloading when loading across `has_one`
 associations

The Preloader relies on other objects to bind the retrieved records to their
parents. When executed across a hash, it assumes that the results of
`preloaded_records` is the appropriate set of records to pass in to the next
layer.

Filtering based on the reflection properties in `preloaded_records` allows us to
avoid excessive preloading in the instance where we are loading across a
`has_one` association distinguished by an order (e.g. "last comment" or
similar), by dropping these records before they are returned to the
Preloader. In this situation, we avoid potentially very long key lists in
generated queries and the consequential AR object instantiations.

This is mostly relevant if the underlying linked set has relatively many
records, because this is effectively a multiplier on the number of records
returned on the far side of the preload. Unfortunately, avoiding the
over-retrieval of the `has_one` association seems to require substantial changes
to the preloader design, and probably adaptor-specific logic -- it is a
top-by-group problem.
---
 .../associations/preloader/association.rb     | 19 ++++++++-
 .../cascaded_eager_loading_test.rb            | 42 +++++++++++++++++++
 activerecord/test/models/author.rb            |  3 ++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb
index 4c7b0e6f07..5c68e2e3be 100644
--- a/activerecord/lib/active_record/associations/preloader/association.rb
+++ b/activerecord/lib/active_record/associations/preloader/association.rb
@@ -38,7 +38,24 @@ module ActiveRecord
 
         def preloaded_records
           return @preloaded_records if defined?(@preloaded_records)
-          @preloaded_records = owner_keys.empty? ? [] : records_for(owner_keys)
+
+          raw_records = owner_keys.empty? ? [] : records_for(owner_keys)
+          seen_records_by_owner = {}.compare_by_identity
+
+          @preloaded_records = raw_records.select do |record|
+            assignments = []
+
+            owners_by_key[convert_key(record[association_key_name])].each do |owner|
+              entries = (seen_records_by_owner[owner] ||= [])
+
+              if reflection.collection? || entries.empty?
+                entries << record
+                assignments << record
+              end
+            end
+
+            !assignments.empty?
+          end
         end
 
         private
diff --git a/activerecord/test/cases/associations/cascaded_eager_loading_test.rb b/activerecord/test/cases/associations/cascaded_eager_loading_test.rb
index 055a62b2bc..002962c3f9 100644
--- a/activerecord/test/cases/associations/cascaded_eager_loading_test.rb
+++ b/activerecord/test/cases/associations/cascaded_eager_loading_test.rb
@@ -200,4 +200,46 @@ class CascadedEagerLoadingTest < ActiveRecord::TestCase
 
     assert_equal expected, actual
   end
+
+  def test_preloading_across_has_one_constrains_loaded_records
+    author = authors(:david)
+
+    old_post = author.posts.create!(title: "first post", body: "test")
+    old_post.comments.create!(author: authors(:mary), body: "a response")
+
+    recent_post = author.posts.create!(title: "first post", body: "test")
+    last_comment = recent_post.comments.create!(author: authors(:bob), body: "a response")
+
+    authors = Author.where(id: author.id)
+    retrieved_comments = []
+
+    reset_callbacks(Comment, :initialize) do
+      Comment.after_initialize { |record| retrieved_comments << record }
+      authors.preload(recent_post: :comments).load
+    end
+
+    assert_equal 1, retrieved_comments.size
+    assert_equal [last_comment], retrieved_comments
+  end
+
+  def test_preloading_across_has_one_through_constrains_loaded_records
+    author = authors(:david)
+
+    old_post = author.posts.create!(title: "first post", body: "test")
+    old_post.comments.create!(author: authors(:mary), body: "a response")
+
+    recent_post = author.posts.create!(title: "first post", body: "test")
+    recent_post.comments.create!(author: authors(:bob), body: "a response")
+
+    authors = Author.where(id: author.id)
+    retrieved_authors = []
+
+    reset_callbacks(Author, :initialize) do
+      Author.after_initialize { |record| retrieved_authors << record }
+      authors.preload(recent_response: :author).load
+    end
+
+    assert_equal 2, retrieved_authors.size
+    assert_equal [author, authors(:bob)], retrieved_authors
+  end
 end
diff --git a/activerecord/test/models/author.rb b/activerecord/test/models/author.rb
index fef3671253..a8ddc86bdb 100644
--- a/activerecord/test/models/author.rb
+++ b/activerecord/test/models/author.rb
@@ -160,6 +160,9 @@ class Author < ActiveRecord::Base
   has_many :posts_with_signature, ->(record) { where("posts.title LIKE ?", "%by #{record.name.downcase}%") }, class_name: "Post"
   has_many :posts_mentioning_author, ->(record = nil) { where("posts.body LIKE ?", "%#{record&.name&.downcase}%") }, class_name: "Post"
 
+  has_one :recent_post, -> { order(id: :desc) }, class_name: "Post"
+  has_one :recent_response, through: :recent_post, source: :comments
+
   has_many :posts_with_extension, -> { order(:title) }, class_name: "Post" do
     def extension_method; end
   end