optimize pluck

lib/has_aux_table/relation_extensions.rb

@@ -122,50 +122,42 @@ module HasAuxTable
       relation_class.send(:define_method, :pluck) do |column_names|
         T.bind(self, ActiveRecord::Relation)
 
-        filtered_attributes =
-          self.where_clause.extract_attributes.select do |attr|
-            column_name = attr.name
-            if attr.relation == aux_config.main.table
-              # if it's on the main table, ignore if it if's the primary key or type key
-              next false if aux_config.main.is_primary_key?(column_name)
-              next false if aux_config.main.is_type_key?(column_name)
+        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_attributes.all? do |attr|
-            column_name = attr.name
-
+          filtered_predicates.all? do |node|
             # if it's a field on the aux table, then it can be plucked
-            if attr.relation == aux_config.aux.table
-              puts "optimize: #{column_name} is on #{attr.relation.name}"
-              next true
-            end
-
-            # if it's on the main table, ignore if it if's the primary key or type key
-            if attr.relation == aux_config.main.table
-              if aux_config.main.is_primary_key?(column_name)
-                puts "optimize: #{column_name} is primary key on #{aux_config.main.table.name}"
-                next true
-              end
-              if aux_config.main.is_type_key?(column_name)
-                puts "optimize: #{column_name} is type key on #{aux_config.main.table.name}"
-                next true
-              end
-            end
-
-            puts "skip optimization: #{column_name} is on #{attr.relation.name}"
-            false
+            Util.is_same_table?(node.left.relation, aux_config.aux.table)
           end
 
         if all_on_aux_table
-          Kernel.puts "pluck is only for aux columns: #{column_names}"
-          binding.pry
-          aux_relation = aux_config.aux.klass.where(where_clause)
+          # 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
-          Kernel.puts "pluck proxied to original: #{column_names}"
           pluck_method.bind(self).call(*column_names)
         end
       end

lib/has_aux_table/util.rb

@@ -69,5 +69,13 @@ 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
   end
 end

spec/loading_optimizations_spec.rb

@@ -6,12 +6,16 @@ require "spec_helper"
 RSpec.describe "loading optimizations" do
   context "cars table" do
     before do
-      Car.create!(name: "Toyota Camry", fuel_type: "gasoline")
-      Car.create!(name: "Toyota Prius", fuel_type: "hybrid")
-      Car.create!(name: "Toyota Corolla", fuel_type: "electric")
+      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
 
-    it "queries only the aux table if plucking values that are on the aux table" do
+    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])
@@ -21,6 +25,21 @@ RSpec.describe "loading optimizations" do
       expect(queries.first).not_to include("JOIN")
     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/)
+    end
+
     it "queries both tables if main table column is referenced" do
       queries =
         SpecHelper.capture_queries do