optimizations when doing pluck / calculations

This commit is contained in:
Dylan Knutson
2025-07-29 14:54:37 +00:00
parent 3704239a6c
commit eba66a85c1
5 changed files with 220 additions and 65 deletions

View File

@@ -5,4 +5,3 @@ config append skip_path gems/sorbet-runtime-
config append skip_path gems/rspec-core-
config append skip_path lib/active_record/transactions.rb
config append skip_path lib/active_support/notifications.rb
config append skip_path has_aux_table/util.rb

View File

@@ -121,47 +121,28 @@ module HasAuxTable
pluck_method = relation_class.instance_method(:pluck)
relation_class.send(:define_method, :pluck) do |column_names|
T.bind(self, ActiveRecord::Relation)
all_predicates =
self.where_clause.instance_eval do
T.cast(predicates, T::Array[Arel::Nodes::Binary])
end
filtered_predicates =
all_predicates.filter do |node|
if node.is_a?(Arel::Nodes::Equality)
if Util.is_same_table?(node.left.relation, aux_config.main.table)
# if it's on the main table, ignore if it if's the primary key or type key
name = node.left.name
next false if aux_config.main.is_primary_key?(name)
next false if aux_config.main.is_type_key?(name)
end
end
true
end
all_on_aux_table =
filtered_predicates.all? do |node|
# if it's a field on the aux table, then it can be plucked
Util.is_same_table?(node.left.relation, aux_config.aux.table)
end
if all_on_aux_table
# the eager load generates a join which creates table alias nodes on attributes instead
# of the original table, so we need to replace those with the original table
filtered_predicates.each do |node|
if (attribute = node.left) && (table_alias = attribute.relation) &&
table_alias.is_a?(Arel::Nodes::TableAlias)
attribute.relation = table_alias.left
end
end
aux_relation = aux_config.aux.klass.where(filtered_predicates)
if (predicates = Util.try_relation_optimization(self, aux_config))
aux_relation = aux_config.aux.klass.where(predicates)
aux_relation.pluck(*column_names)
else
pluck_method.bind(self).call(*column_names)
end
end
calculate_method = relation_class.instance_method(:calculate)
relation_class.send(
:define_method,
:calculate
) do |operation, column_name|
T.bind(self, ActiveRecord::Relation)
if (predicates = Util.try_relation_optimization(self, aux_config))
aux_relation = aux_config.aux.klass.where(predicates)
aux_relation.calculate(operation, column_name)
else
calculate_method.bind(self).call(operation, column_name)
end
end
[
[relation_class, :build_where_clause],
[collection_proxy_class, :where]

View File

@@ -77,5 +77,71 @@ module HasAuxTable
right_table = right.is_a?(Arel::Nodes::TableAlias) ? right.left : right
left_table == right_table
end
sig do
params(
relation: ActiveRecord::Relation,
aux_config: HasAuxTable::AuxTableConfig
).returns(T.nilable(T::Array[Arel::Nodes::Node]))
end
def self.try_relation_optimization(relation, aux_config)
present_clauses = relation.values.keys
# optimize if there are no joins etc on any other tables
unless (present_clauses - %i[where eager_load references]).empty?
return nil
end
# optimize if no other eager loads are present other than the aux association
if relation.eager_load_values.any? &&
relation.eager_load_values != [aux_config.aux_association_name]
return nil
end
# same as eager_load_values but for references
if relation.references_values.any? &&
relation.references_values != [aux_config.aux_association_name.to_s]
return nil
end
all_predicates =
relation.where_clause.instance_eval do
T.cast(predicates, T::Array[Arel::Nodes::Binary])
end
filtered_predicates =
all_predicates.filter do |node|
if node.is_a?(Arel::Nodes::Equality)
if Util.is_same_table?(node.left.relation, aux_config.main.table)
# if it's on the main table, ignore if it if's the primary key or type key
name = node.left.name
next false if aux_config.main.is_primary_key?(name)
next false if aux_config.main.is_type_key?(name)
end
end
true
end
all_on_aux_table =
filtered_predicates.all? do |node|
# if it's a field on the aux table, then it can be plucked
Util.is_same_table?(node.left.relation, aux_config.aux.table)
end
if all_on_aux_table
# the eager load generates a join which creates table alias nodes on attributes instead
# of the original table, so we need to replace those with the original table
filtered_predicates.each do |node|
if (attribute = node.left) && (table_alias = attribute.relation) &&
table_alias.is_a?(Arel::Nodes::TableAlias)
attribute.relation = table_alias.left
end
end
filtered_predicates
else
nil
end
end
end
end

View File

@@ -11,6 +11,12 @@ class Arel::Attributes::Attribute
end
end
class ActiveRecord::Relation
sig { returns(T::Hash[Symbol, T.untyped]) }
def values
end
end
module ActiveRecord::QueryMethods
sig { returns(ActiveRecord::Relation::WhereClause) }
def where_clause
@@ -22,3 +28,18 @@ class ActiveRecord::Relation::WhereClause
def extract_attributes
end
end
class BasicObject
sig do
type_parameters(:T)
.params(
code: String,
filename: String,
lineno: Numeric,
block: T.proc.bind(T.self_type).returns(T.type_parameter(:T))
)
.returns(T.type_parameter(:T))
end
def instance_eval(code = "", filename, lineno, &block)
end
end

View File

@@ -15,41 +15,129 @@ RSpec.describe "loading optimizations" do
)
end
it "queries only the aux table if no main table columns are referenced" do
queries =
SpecHelper.capture_queries do
expect(Car.pluck(:fuel_type)).to eq(%w[gasoline hybrid electric])
end
expect(queries.length).to eq(1)
expect(queries.first).not_to include("JOIN")
shared_examples "queries only the aux table" do
it "queries only the aux table" do
expect(@queries.length).to eq(1)
expect(@queries.first).not_to include("JOIN")
expect(@queries.first).to match(/\bvehicles_car_aux\b/)
expect(@queries.first).not_to match(/\bvehicles\b/)
end
end
it "queries only the aux table if all columns are on the aux table" do
queries =
SpecHelper.capture_queries do
expect(Car.where(engine_size: 1.4..1.9).pluck(:fuel_type)).to eq(
%w[hybrid electric]
)
end
expect(queries.length).to eq(1)
expect(queries.first).not_to include("JOIN")
expect(queries.first).to include("BETWEEN")
expect(queries.first).to match(/\bvehicles_car_aux\b/)
expect(queries.first).not_to match(/\bvehicles\b/)
shared_examples "queries both tables" do
it "queries both tables" do
expect(@queries.length).to eq(1)
expect(@queries.first).to include("JOIN")
expect(@queries.first).to match(/\bvehicles\b/)
expect(@queries.first).to match(/\bvehicles_car_aux\b/)
end
end
it "queries both tables if main table column is referenced" do
queries =
SpecHelper.capture_queries do
rel = Car.where(name: "Toyota Camry")
rel = rel.pluck(:fuel_type)
expect(rel).to eq(%w[gasoline])
describe "pluck" do
context "aux columns are referenced" do
before do
@queries =
SpecHelper.capture_queries do
expect(Car.pluck(:fuel_type)).to eq(%w[gasoline hybrid electric])
end
end
expect(queries.length).to eq(1)
expect(queries.first).to include("JOIN")
it_behaves_like "queries only the aux table"
end
context "aux columns are chained on a where clause" do
before do
@queries =
SpecHelper.capture_queries do
expect(Car.where(engine_size: 1.4..1.9).pluck(:fuel_type)).to eq(
%w[hybrid electric]
)
end
end
it_behaves_like "queries only the aux table"
it "applies the BETWEEN clause" do
expect(@queries.first).to include("BETWEEN")
end
end
context "main table columns are referenced" do
before do
@queries =
SpecHelper.capture_queries do
expect(Car.where(name: "Toyota Camry").pluck(:fuel_type)).to eq(
%w[gasoline]
)
end
end
it_behaves_like "queries both tables"
end
context "main table columns are chained on a where clause" do
before do
@queries =
SpecHelper.capture_queries do
expect(Car.where(name: "Toyota Camry").pluck(:fuel_type)).to eq(
%w[gasoline]
)
end
end
it_behaves_like "queries both tables"
end
end
describe "maximum" do
context "aux columns are referenced" do
before do
@queries =
SpecHelper.capture_queries do
expect(Car.maximum(:engine_size)).to eq(2.0)
end
end
it_behaves_like "queries only the aux table"
end
context "aux columns are chained on a where clause" do
before do
@queries =
SpecHelper.capture_queries do
expect(
Car.where(engine_size: 1.4..1.9).maximum(:engine_size)
).to eq(1.8)
end
end
it_behaves_like "queries only the aux table"
end
context "main table columns are referenced" do
before do
@queries =
SpecHelper.capture_queries do
expect(
Car.where(name: "Toyota Camry").maximum(:engine_size)
).to eq(2.0)
end
end
it_behaves_like "queries both tables"
end
context "main table columns are chained on a where clause" do
before do
@queries =
SpecHelper.capture_queries do
expect(
Car.where(name: "Toyota Camry").maximum(:engine_size)
).to eq(2.0)
end
end
it_behaves_like "queries both tables"
end
end
end
end