optimizations when doing pluck / calculations

.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

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)
-
+        if (predicates = Util.try_relation_optimization(self, aux_config))
-        all_predicates =
+          aux_relation = aux_config.aux.klass.where(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)
           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]

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

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

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
+    shared_examples "queries only the aux table" do
-      queries =
+      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
+
+    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
+
+    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
-
-      expect(queries.length).to eq(1)
-      expect(queries.first).not_to include("JOIN")
         end
 
-    it "queries only the aux table if all columns are on the aux table" do
+        it_behaves_like "queries only the aux table"
-      queries =
+      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
-
-      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/)
         end
 
-    it "queries both tables if main table column is referenced" do
+        it_behaves_like "queries only the aux table"
-      queries =
+        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
-          rel = Car.where(name: "Toyota Camry")
+              expect(Car.where(name: "Toyota Camry").pluck(:fuel_type)).to eq(
-          rel = rel.pluck(:fuel_type)
+                %w[gasoline]
-          expect(rel).to eq(%w[gasoline])
+              )
+            end
         end
 
-      expect(queries.length).to eq(1)
+        it_behaves_like "queries both tables"
-      expect(queries.first).to include("JOIN")
+      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