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

Batch Preloader::Association queries having the same scope

When preloading an association on differing
models, we can execute less queries by batching
queries having the same scope. For example:

`Preloader.new(records: [book, post], associations: :author).call`

Previously this would execute two identical
queries against the authors table. By identifying
that the two queries will be the same, we can
fetch them in one go. We do this by building all
the loaders, then grouping them by scope, then
overriding the `load_records` method to avoid
double fetching. Records are associated with
their owners using existing logic in
`Preloader::Association#run`.

Co-authored-by: John Hawthorn <john@hawthorn.email>
Co-Authored-By: Eileen M. Uchitelle <eileencodes@users.noreply.github.com>
This commit is contained in:
Dinah Shi 2021-01-20 15:27:11 -05:00
parent d30dad4bc1
commit c6c0b2e8af
5 changed files with 170 additions and 44 deletions

View file

@ -94,13 +94,20 @@ module ActiveRecord
@scope = kwargs[:scope]
@associate_by_default = associate_by_default
@polymorphic_parent = polymorphic_parent
@child_preloaders = []
end
end
def call
return [] if associations.nil? || records.length == 0
build_preloaders
loaders = build_preloaders
group_and_load_similar(loaders)
loaders.map(&:run)
child_preloaders.each { |reflection, child, parents| build_child_preloader(reflection, child, parents) }
loaders
end
def preload(records, associations, preload_scope = nil)
@ -110,6 +117,8 @@ module ActiveRecord
end
private
attr_accessor :child_preloaders
def build_preloaders
Array.wrap(associations).flat_map { |association|
Array(association).flat_map { |parent, child|
@ -117,7 +126,7 @@ module ActiveRecord
loaders = preloaders_for_reflection(reflection, reflection_records)
if child
loaders.concat build_child_preloader(reflection, child, loaders)
child_preloaders << [reflection, child, loaders]
end
loaders
@ -135,7 +144,7 @@ module ActiveRecord
def preloaders_for_reflection(reflection, reflection_records)
reflection_records.group_by { |record| record.association(reflection.name).klass }.map do |rhs_klass, rs|
preloader_for(reflection).new(rhs_klass, rs, reflection, scope, associate_by_default).run
preloader_for(reflection).new(rhs_klass, rs, reflection, scope, associate_by_default)
end
end
@ -161,6 +170,15 @@ module ActiveRecord
Association
end
end
def group_and_load_similar(loaders)
loaders.grep_v(ThroughAssociation).group_by(&:grouping_key).each do |(_, _, association_key_name), similar_loaders|
next if similar_loaders.all? { |l| l.already_loaded? }
scope = similar_loaders.first.scope
Association.load_records_in_batch(scope, association_key_name, similar_loaders)
end
end
end
end
end

View file

@ -4,6 +4,16 @@ module ActiveRecord
module Associations
class Preloader
class Association #:nodoc:
def self.load_records_in_batch(scope, association_key_name, loaders)
ids = loaders.flat_map(&:owner_keys).uniq
raw_records = scope.where(association_key_name => ids).load do |record|
loaders.each { |l| l.set_inverse(record) }
end
loaders.each { |l| l.load_records(raw_records) }
end
def initialize(klass, owners, reflection, preload_scope, associate_by_default = true)
@klass = klass
@owners = owners.uniq(&:__id__)
@ -15,6 +25,10 @@ module ActiveRecord
@already_loaded = owners.all? { |o| o.association(reflection.name).loaded? }
end
def already_loaded?
@already_loaded
end
def run
if @already_loaded
fetch_from_preloaded_records
@ -42,6 +56,54 @@ module ActiveRecord
@preloaded_records
end
# The name of the key on the associated records
def association_key_name
reflection.join_primary_key(klass)
end
def grouping_key
[scope.to_sql, scope.preload_values + scope.includes_values, association_key_name]
end
def owner_keys
@owner_keys ||= owners_by_key.keys
end
def scope
@scope ||= build_scope
end
def set_inverse(record)
if owners = owners_by_key[convert_key(record[association_key_name])]
# Processing only the first owner
# because the record is modified but not an owner
association = owners.first.association(reflection.name)
association.set_inverse_instance(record)
end
end
def load_records(raw_records = nil)
# owners can be duplicated when a relation has a collection association join
# #compare_by_identity makes such owners different hash keys
@records_by_owner = {}.compare_by_identity
raw_records ||= owner_keys.empty? ? [] : records_for(owner_keys)
@preloaded_records = raw_records.select do |record|
assignments = false
owners_by_key[convert_key(record[association_key_name])]&.each do |owner|
entries = (@records_by_owner[owner] ||= [])
if reflection.collection? || entries.empty?
entries << record
assignments = true
end
end
assignments
end
end
private
attr_reader :owners, :reflection, :preload_scope, :model, :klass
@ -53,33 +115,6 @@ module ActiveRecord
@preloaded_records = records_by_owner.flat_map(&:last)
end
def load_records
# owners can be duplicated when a relation has a collection association join
# #compare_by_identity makes such owners different hash keys
@records_by_owner = {}.compare_by_identity
raw_records = owner_keys.empty? ? [] : records_for(owner_keys)
@preloaded_records = raw_records.select do |record|
assignments = false
owners_by_key[convert_key(record[association_key_name])].each do |owner|
entries = (@records_by_owner[owner] ||= [])
if reflection.collection? || entries.empty?
entries << record
assignments = true
end
end
assignments
end
end
# The name of the key on the associated records
def association_key_name
reflection.join_primary_key(klass)
end
# The name of the key on the model which declares the association
def owner_key_name
reflection.join_foreign_key
@ -94,10 +129,6 @@ module ActiveRecord
end
end
def owner_keys
@owner_keys ||= owners_by_key.keys
end
def owners_by_key
@owners_by_key ||= owners.each_with_object({}) do |owner, result|
key = convert_key(owner[owner_key_name])
@ -131,18 +162,10 @@ module ActiveRecord
def records_for(ids)
scope.where(association_key_name => ids).load do |record|
# Processing only the first owner
# because the record is modified but not an owner
owner = owners_by_key[convert_key(record[association_key_name])].first
association = owner.association(reflection.name)
association.set_inverse_instance(record)
set_inverse(record)
end
end
def scope
@scope ||= build_scope
end
def reflection_scope
@reflection_scope ||= begin
reflection.join_scopes(klass.arel_table, klass.predicate_builder, klass).inject(&:merge!) || klass.unscoped

View file

@ -9,6 +9,7 @@ require "models/categorization"
require "models/category"
require "models/post"
require "models/author"
require "models/book"
require "models/comment"
require "models/tag"
require "models/tagging"
@ -356,7 +357,7 @@ class OverridingAssociationsTest < ActiveRecord::TestCase
end
class PreloaderTest < ActiveRecord::TestCase
fixtures :posts, :comments
fixtures :posts, :comments, :books, :authors
def test_preload_with_scope
post = posts(:welcome)
@ -398,6 +399,78 @@ class PreloaderTest < ActiveRecord::TestCase
preloader.call
end
end
def test_preload_groups_queries_with_same_scope
book = books(:awdr)
post = posts(:welcome)
assert_queries(1) do
preloader = ActiveRecord::Associations::Preloader.new(records: [book, post], associations: :author)
preloader.call
book.author
post.author
end
end
def test_preload_with_grouping_sets_inverse_association
mary = authors(:mary)
bob = authors(:bob)
AuthorFavorite.create!(author: mary, favorite_author: bob)
favorites = AuthorFavorite.all.load
assert_queries(1) do
preloader = ActiveRecord::Associations::Preloader.new(records: favorites, associations: [:author, :favorite_author])
preloader.call
favorites.first.author
favorites.first.favorite_author
end
end
def test_preload_does_not_group_same_class_different_scope
post = posts(:welcome)
postesque = Postesque.create(author: Author.last)
postesque.reload
# When the scopes differ in the generated SQL:
# SELECT "authors".* FROM "authors" WHERE (name LIKE '%a%') AND "authors"."id" = ?
# SELECT "authors".* FROM "authors" WHERE "authors"."id" = ?.
assert_queries(2) do
preloader = ActiveRecord::Associations::Preloader.new(records: [post, postesque], associations: :author_with_the_letter_a)
preloader.call
post.author_with_the_letter_a
postesque.author_with_the_letter_a
end
post.reload
postesque.reload
# When the generated SQL is identical, but one scope has preload values.
assert_queries(3) do
preloader = ActiveRecord::Associations::Preloader.new(records: [post, postesque], associations: :author_with_address)
preloader.call
post.author_with_address
postesque.author_with_address
end
end
def test_preload_does_not_group_same_scope_different_key_name
post = posts(:welcome)
postesque = Postesque.create(author: Author.last)
postesque.reload
assert_queries(2) do
preloader = ActiveRecord::Associations::Preloader.new(records: [post, postesque], associations: :author)
preloader.call
post.author
postesque.author
end
end
end
class GeneratedMethodsTest < ActiveRecord::TestCase

View file

@ -40,6 +40,7 @@ class Post < ActiveRecord::Base
belongs_to :author_with_posts, -> { includes(:posts) }, class_name: "Author", foreign_key: :author_id
belongs_to :author_with_address, -> { includes(:author_address) }, class_name: "Author", foreign_key: :author_id
belongs_to :author_with_select, -> { select(:id) }, class_name: "Author", foreign_key: :author_id
belongs_to :author_with_the_letter_a, -> { where("name LIKE '%a%'") }, class_name: "Author", foreign_key: :author_id
def first_comment
super.body
@ -368,3 +369,9 @@ class FakeKlass
inherited self
end
class Postesque < ActiveRecord::Base
belongs_to :author, class_name: "Author", foreign_key: :author_name, primary_key: :name
belongs_to :author_with_address, class_name: "Author", foreign_key: :author_id
belongs_to :author_with_the_letter_a, class_name: "Author", foreign_key: :author_id
end

View file

@ -805,6 +805,11 @@ ActiveRecord::Schema.define do
t.integer :tags_with_nullify_count, default: 0
end
create_table :postesques, force: true do |t|
t.string :author_name
t.string :author_id
end
create_table :serialized_posts, force: true do |t|
t.integer :author_id
t.string :title, null: false