optimizations when doing pluck / calculations

.devcontainer/.rdbgrc

@@ -0,0 +1,7 @@
+config append skip_path bin/bundle
+config append skip_path bin/rspec
+config append skip_path /bundler/
+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

.devcontainer/Dockerfile

@@ -31,3 +31,4 @@ RUN BUNDLE_FROZEN=true MAKE="make -j$(nproc)" bundle install --jobs $(nproc)
 RUN echo 'alias rspec="bundle exec rspec"' >> ~/.bashrc
 RUN echo 'alias tapioca="bundle exec tapioca"' >> ~/.bashrc
 RUN echo 'alias srb="bundle exec srb"' >> ~/.bashrc
+COPY ./.devcontainer/.rdbgrc /home/vscode/.rdbgrc

lib/has_aux_table/model_class_helper.rb

@@ -13,6 +13,26 @@ module HasAuxTable
       column_names.include?(name.to_s)
     end
 
+    sig { params(name: String).returns(T::Boolean) }
+    def is_primary_key?(name)
+      primary_keys.include?(name.to_sym)
+    end
+
+    sig { params(name: String).returns(T::Boolean) }
+    def is_type_key?(name)
+      type_key == name.to_s
+    end
+
+    sig { returns(T.nilable(String)) }
+    def type_key
+      self.klass.inheritance_column
+    end
+
+    sig { returns(Arel::Table) }
+    def table
+      self.klass.arel_table
+    end
+
     sig { returns(T::Array[Symbol]) }
     def primary_keys
       @primary_keys ||=

lib/has_aux_table/relation_extensions.rb

@@ -118,6 +118,31 @@ module HasAuxTable
           ActiveRecord::Associations::CollectionProxy
         )
 
+      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))
+          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]

lib/has_aux_table/util.rb

@@ -31,8 +31,12 @@ module HasAuxTable
         )
 
       target.send(define_method, method_name) do |*args, **kwargs, &block|
-        method = is_instance_method ? target_method.bind(self) : target_method
-        T.unsafe(hook_block).call(method, *args, **kwargs, &block)
+        if is_instance_method
+          method = target_method.bind(self)
+          T.unsafe(hook_block).call(method, *args, **kwargs, &block)
+        else
+          T.unsafe(hook_block).call(target_method, *args, **kwargs, &block)
+        end
       end
     end
 
@@ -65,5 +69,79 @@ module HasAuxTable
         Kernel.raise("#{instance.class.name} not a #{klass.name}")
       end
     end
+
+    TableOrAlias = T.type_alias { T.any(Arel::Nodes::Node, Arel::Table) }
+    sig { params(left: TableOrAlias, right: TableOrAlias).returns(T::Boolean) }
+    def self.is_same_table?(left, right)
+      left_table = left.is_a?(Arel::Nodes::TableAlias) ? left.left : left
+      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

sorbet/rbi/shims/arel.rbi

@@ -0,0 +1,30 @@
+# typed: strict
+# frozen_string_literal: true
+
+class Arel::Attributes::Attribute
+  sig { returns(Arel::Table) }
+  def relation
+  end
+
+  sig { returns(String) }
+  def name
+  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
+  end
+end
+
+class ActiveRecord::Relation::WhereClause
+  sig { returns(T::Array[Arel::Attributes::Attribute]) }
+  def extract_attributes
+  end
+end

spec/has_aux_table_spec.rb

@@ -234,13 +234,13 @@ RSpec.describe HasAuxTable do
 
     it "allows saving the model with auxiliary columns" do
       car = Car.create!(name: "Honda Civic")
-      num_queries =
-        SpecHelper.count_queries do
+      queries =
+        SpecHelper.capture_queries do
           car.fuel_type = "hybrid"
           car.engine_size = 1.8
           car.save!
         end
-      expect(num_queries).to eq(1)
+      expect(queries.length).to eq(1)
     end
   end
 
@@ -423,32 +423,32 @@ RSpec.describe HasAuxTable do
         # All query methods now use single queries with proper LEFT OUTER JOINs
 
         it "loads single model with auxiliary data in one query using find" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               car = Car.find(@car1.id)
               # Access auxiliary attributes to ensure they're loaded
               car.fuel_type
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "loads single model with auxiliary data in one query using find_by" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               car = Car.find_by(name: "Toyota Prius")
               # Access auxiliary attributes to ensure they're loaded
               car.fuel_type
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "loads multiple models with auxiliary data in one query using where" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               cars = Car.where(fuel_type: %w[hybrid electric])
               # Access auxiliary attributes for all cars
               cars.each do |car|
@@ -457,7 +457,7 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "avoids N+1 queries when loading multiple models" do
@@ -472,8 +472,8 @@ RSpec.describe HasAuxTable do
           end
 
           cars = nil
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               cars = Car.where(fuel_type: "gasoline")
               # Access auxiliary attributes for all cars - should not trigger additional queries
               cars.each do |car|
@@ -483,13 +483,13 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1) # Single query regardless of how many cars are loaded
+          expect(queries.length).to eq(1) # Single query regardless of how many cars are loaded
           expect(cars.length).to be >= 1 # At least the original Honda Civic plus new cars
         end
 
         it "uses single query when ordering by auxiliary columns" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               cars = Car.where(engine_size: 1.0..3.0).order(:engine_size)
               # Access all attributes
               cars.each do |car|
@@ -499,12 +499,12 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "uses single query for complex auxiliary column queries" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               cars =
                 Car.where(fuel_type: "hybrid").or(Car.where(engine_size: 0.0))
               # Access all attributes
@@ -515,12 +515,12 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "uses single query when finding by auxiliary columns" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               car = Car.find_by(fuel_type: "hybrid", name: "Toyota Prius")
               # Access all attributes
               car.name
@@ -528,7 +528,7 @@ RSpec.describe HasAuxTable do
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "doesn't trigger additional queries when accessing auxiliary attributes after load" do
@@ -536,8 +536,8 @@ RSpec.describe HasAuxTable do
           car = Car.find(@car1.id)
 
           # Now count queries when accessing auxiliary attributes
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               car.fuel_type
               car.engine_size
               car.fuel_type? # presence check
@@ -545,12 +545,12 @@ RSpec.describe HasAuxTable do
             end
 
           # Currently this should be 0 since auxiliary record is already loaded
-          expect(query_count).to eq(0) # No additional queries should be triggered
+          expect(queries.length).to eq(0) # No additional queries should be triggered
         end
 
         it "handles mixed queries with main and auxiliary columns in single query" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               cars = Car.where(name: "Toyota Prius", fuel_type: "hybrid")
               cars.each do |car|
                 car.name
@@ -559,12 +559,12 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "uses single query for range queries on auxiliary columns" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               cars = Car.where(engine_size: 0.0..1.9)
               cars.each do |car|
                 car.name
@@ -573,19 +573,19 @@ RSpec.describe HasAuxTable do
               end
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
 
         it "maintains single query performance with limit and offset" do
-          query_count =
-            SpecHelper.count_queries do
+          queries =
+            SpecHelper.capture_queries do
               car = Car.where(fuel_type: %w[hybrid electric]).limit(1).first
               # Access auxiliary attributes
               car.fuel_type
               car.engine_size
             end
 
-          expect(query_count).to eq(1)
+          expect(queries.length).to eq(1)
         end
       end
 
@@ -983,8 +983,8 @@ RSpec.describe HasAuxTable do
       end
 
       it "reloads with one query" do
-        num_queries = SpecHelper.count_queries { @car.reload }
-        expect(num_queries).to eq(1)
+        queries = SpecHelper.capture_queries { @car.reload }
+        expect(queries.length).to eq(1)
       end
 
       it "reloads associations" do
@@ -1027,12 +1027,16 @@ RSpec.describe HasAuxTable do
     expect(Car.count).to eq(1)
     expect(Boat.count).to eq(1)
 
-    expect(SpecHelper.count_queries { car = Vehicle.find(car.id) }).to eq(1)
+    expect(
+      SpecHelper.capture_queries { car = Vehicle.find(car.id) }.length
+    ).to eq(1)
     expect(car.fuel_type).to eq("gasoline")
     expect(car.engine_size).to eq(2.0)
     expect(car.name).to eq("Honda Civic")
 
-    expect(SpecHelper.count_queries { boat = Vehicle.find(boat.id) }).to eq(1)
+    expect(
+      SpecHelper.capture_queries { boat = Vehicle.find(boat.id) }.length
+    ).to eq(1)
     expect(boat.only_freshwater).to eq(true)
   end
 
@@ -1235,4 +1239,66 @@ RSpec.describe HasAuxTable do
       expect(specific_b.on_aux).to eq("2b_aux")
     end
   end
+
+  describe "range queries" do
+    ActiveRecord::Schema.define do
+      create_base_table :test_range_models do |t|
+        t.integer :base_field
+        t.create_aux :specific do |t|
+          t.integer :aux_field
+        end
+      end
+    end
+
+    class TestRangeModel < ActiveRecord::Base
+      include HasAuxTable
+    end
+
+    class TestRangeModelSpecific < TestRangeModel
+      aux_table :specific
+    end
+
+    before do
+      @bases = (0..5).map { |i| TestRangeModel.create!(base_field: i) }
+      @specifics =
+        (0..5).map do |i|
+          TestRangeModelSpecific.create!(base_field: i, aux_field: i)
+        end
+    end
+
+    it "works with a from..to range" do
+      expect(TestRangeModel.where(base_field: 1..5)).to eq(
+        @bases[1..5] + @specifics[1..5]
+      )
+      expect(TestRangeModelSpecific.where(base_field: 2..3)).to eq(
+        @specifics[2..3]
+      )
+      expect(TestRangeModelSpecific.where(aux_field: 2..3)).to eq(
+        @specifics[2..3]
+      )
+      expect(TestRangeModelSpecific.where(aux_field: 4..7)).to eq(
+        @specifics[4..5]
+      )
+    end
+
+    it "works with a from.. range" do
+      expect(TestRangeModel.where(base_field: 1..)).to eq(
+        @bases[1..5] + @specifics[1..5]
+      )
+
+      expect(TestRangeModelSpecific.where(aux_field: 1..)).to eq(
+        @specifics[1..5]
+      )
+    end
+
+    it "works with a ..to range" do
+      expect(TestRangeModel.where(base_field: ..4)).to eq(
+        @bases[0..4] + @specifics[0..4]
+      )
+
+      expect(TestRangeModelSpecific.where(aux_field: ..4)).to eq(
+        @specifics[0..4]
+      )
+    end
+  end
 end

spec/loading_optimizations_spec.rb

@@ -0,0 +1,143 @@
+# typed: false
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe "loading optimizations" do
+  context "cars table" do
+    before do
+      Car.create!(name: "Toyota Camry", fuel_type: "gasoline", engine_size: 2.0)
+      Car.create!(name: "Toyota Prius", fuel_type: "hybrid", engine_size: 1.5)
+      Car.create!(
+        name: "Toyota Corolla",
+        fuel_type: "electric",
+        engine_size: 1.8
+      )
+    end
+
+    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
+
+    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
+        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).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

spec/spec_helper.rb

@@ -62,11 +62,13 @@ module SpecHelper
   LOG_QUERIES = T.let(false, T::Boolean)
 
   # Helper method to count queries
-  sig { params(block: T.proc.void).returns(Integer) }
-  def self.count_queries(&block)
-    query_count = 0
+  sig { params(block: T.proc.void).returns(T::Array[String]) }
+  def self.capture_queries(&block)
+    queries = T.let([], T::Array[String])
     query_callback =
-      lambda { |name, start, finish, message_id, values| query_count += 1 }
+      lambda do |name, start, finish, message_id, values|
+        queries << values[:sql]
+      end
 
     ActiveSupport::Notifications.subscribed(
       query_callback,
@@ -81,6 +83,6 @@ module SpecHelper
       ActiveRecord::Base.logger = old_logger if LOG_QUERIES
     end
 
-    query_count
+    queries
   end
 end