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:
parent
d30dad4bc1
commit
c6c0b2e8af
5 changed files with 170 additions and 44 deletions
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Reference in a new issue