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

Use loaded records where possible when preloading
This commit is contained in:
Eileen M. Uchitelle 2021-10-12 11:54:12 -04:00 committed by GitHub
commit dedefdd297
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 166 additions and 59 deletions

View File

@ -23,11 +23,7 @@ module ActiveRecord
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 @@ module ActiveRecord
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 @@ module ActiveRecord
@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 @@ module ActiveRecord
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 @@ module ActiveRecord
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 @@ module ActiveRecord
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,25 +221,16 @@ module ActiveRecord
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
@ -211,13 +238,6 @@ module ActiveRecord
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)

View File

@ -16,10 +16,7 @@ module ActiveRecord
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

View File

@ -10,10 +10,13 @@ module ActiveRecord
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 @@ module ActiveRecord
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 @@ module ActiveRecord
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 @@ module ActiveRecord
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 @@ module ActiveRecord
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

View File

@ -443,6 +443,18 @@ class PreloaderTest < ActiveRecord::TestCase
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 @@ class PreloaderTest < ActiveRecord::TestCase
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 @@ class PreloaderTest < ActiveRecord::TestCase
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 @@ class PreloaderTest < ActiveRecord::TestCase
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