From b4fd61428079e764f6f7c2e69e2246747f8c5d9b Mon Sep 17 00:00:00 2001 From: Dylan Knutson Date: Tue, 29 Jul 2025 14:54:37 +0000 Subject: [PATCH] optimizations when doing pluck / calculations --- .devcontainer/.rdbgrc | 1 - lib/has_aux_table/relation_extensions.rb | 51 +++----- lib/has_aux_table/util.rb | 67 +++++++++++ sorbet/rbi/shims/arel.rbi | 6 + spec/loading_optimizations_spec.rb | 146 ++++++++++++++++++----- 5 files changed, 206 insertions(+), 65 deletions(-) diff --git a/.devcontainer/.rdbgrc b/.devcontainer/.rdbgrc index 8ca08d9..747bcef 100644 --- a/.devcontainer/.rdbgrc +++ b/.devcontainer/.rdbgrc @@ -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 diff --git a/lib/has_aux_table/relation_extensions.rb b/lib/has_aux_table/relation_extensions.rb index a109ce2..df70816 100644 --- a/lib/has_aux_table/relation_extensions.rb +++ b/lib/has_aux_table/relation_extensions.rb @@ -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] diff --git a/lib/has_aux_table/util.rb b/lib/has_aux_table/util.rb index fa1e86a..8a70d0e 100644 --- a/lib/has_aux_table/util.rb +++ b/lib/has_aux_table/util.rb @@ -77,5 +77,72 @@ 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 = + T.cast( + relation.where_clause.instance_variable_get(:@predicates), + T::Array[Arel::Nodes::Binary] + ) + + 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 diff --git a/sorbet/rbi/shims/arel.rbi b/sorbet/rbi/shims/arel.rbi index 68e0125..4e08eb1 100644 --- a/sorbet/rbi/shims/arel.rbi +++ b/sorbet/rbi/shims/arel.rbi @@ -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 diff --git a/spec/loading_optimizations_spec.rb b/spec/loading_optimizations_spec.rb index 8ce3c73..6b40f92 100644 --- a/spec/loading_optimizations_spec.rb +++ b/spec/loading_optimizations_spec.rb @@ -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