1
0
Fork 0
mirror of https://github.com/rails/rails.git synced 2022-11-09 12:12:34 -05:00

Use loaded records where possible when preloading

Prior to this commit the preloader would avoid querying if a particular
association was already loaded for all the owner records being
preloaded:

```rb
assert post1.association(:author).loaded?
assert post2.association(:author).loaded?

assert_no_queries do
  ActiveRecord::Associations::Preloader.new(records: [post1, post2], associations: :author).call
end
```

But if only some of owners had the association loaded we'd load it for
all of them:

```rb
assert post1.association(:author).loaded?
assert_not post2.association(:author).loaded?

ActiveRecord::Associations::Preloader.new(records: [post1, post2], associations: :author).call
```

With this commit, we only load the association for owners that don't
already have it loaded:

```rb
assert post1.association(:author).loaded?
assert_not post2.association(:author).loaded?

ActiveRecord::Associations::Preloader.new(records: [post1, post2], associations: :author).call
```

This limits the data returned from the query, and the number of new
records we need to instantiate.

This also makes the `available_records` option useful for more cases.
Prior to this commit we'd use the `available_records` to set the
association target on any owner where it was available, but if available
records weren't provided for ALL of the owners we'd end up fetching them
all. With this commit passing incomplete `available_records` can still
limit the amount of data we query for.

Because of the way we group queries, this also means we can use already
loaded associations for on one type to avoid or limit queries on another
type. (See test_preload_grouped_queries_with_already_loaded_records)

Implementation details
---

This commit gets rid of the special cases for all owners being loaded,
deleting the `already_loaded?` method entirely.

The gathering of already loaded records that used to happen in
`fetch_from_preloaded_records` now happens all the time, within
`LoaderRecords`.

The purpose of the `LoaderRecords` class is to find all the records,
either by looking for already loaded ones, or by loading them. It puts
together a list `already_loaded_records` and of `keys_to_load`, then
passes those off to `LoaderQuery` to do the actual loading.
This commit is contained in:
Daniel Colson 2021-08-31 09:24:55 -04:00
parent ce546c7d3a
commit 605acb9617
No known key found for this signature in database
GPG key ID: 88A364BBE77B1353
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),
@ -774,6 +786,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
@ -787,6 +834,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)
@ -840,6 +906,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