From 605acb961777d9d3ac8712ecd4841c6dcd46d1cc Mon Sep 17 00:00:00 2001 From: Daniel Colson Date: Tue, 31 Aug 2021 09:24:55 -0400 Subject: [PATCH] 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. --- .../associations/preloader/association.rb | 112 +++++++++++------- .../associations/preloader/batch.rb | 5 +- .../preloader/through_association.rb | 28 +++-- activerecord/test/cases/associations_test.rb | 80 +++++++++++++ 4 files changed, 166 insertions(+), 59 deletions(-) diff --git a/activerecord/lib/active_record/associations/preloader/association.rb b/activerecord/lib/active_record/associations/preloader/association.rb index 4785c838ef..e741b04e42 100644 --- a/activerecord/lib/active_record/associations/preloader/association.rb +++ b/activerecord/lib/active_record/associations/preloader/association.rb @@ -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) diff --git a/activerecord/lib/active_record/associations/preloader/batch.rb b/activerecord/lib/active_record/associations/preloader/batch.rb index 97304625d4..17cf01f828 100644 --- a/activerecord/lib/active_record/associations/preloader/batch.rb +++ b/activerecord/lib/active_record/associations/preloader/batch.rb @@ -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 diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb index 212435a6a0..67f6bc4293 100644 --- a/activerecord/lib/active_record/associations/preloader/through_association.rb +++ b/activerecord/lib/active_record/associations/preloader/through_association.rb @@ -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 diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb index 06050fd129..289540a5de 100644 --- a/activerecord/test/cases/associations_test.rb +++ b/activerecord/test/cases/associations_test.rb @@ -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