Merge pull request #20362 from kddeisz/enumerable_pluck

Allow Enumerable#pluck to take a splat.

4 files changed, 43 insertions, 2 deletions

diff --git a/activerecord/lib/active_record/relation/calculations.rb b/activerecord/lib/active_record/relation/calculations.rb
index 402b317d9c..7a28a98721 100644
--- a/activerecord/lib/active_record/relation/calculations.rb
+++ b/activerecord/lib/active_record/relation/calculations.rb
@@ -161,6 +161,10 @@ module ActiveRecord
         end
       end
 
+      if loaded? && (column_names - @klass.column_names).empty?
+        return @records.pluck(*column_names)
+      end
+
       if has_include?(column_names.first)
         construct_relation_for_association_calculations.pluck(*column_names)
       else
diff --git a/activerecord/test/cases/calculations_test.rb b/activerecord/test/cases/calculations_test.rb
index b246eae5f5..4a9bc7845f 100644
--- a/activerecord/test/cases/calculations_test.rb
+++ b/activerecord/test/cases/calculations_test.rb
@@ -632,6 +632,27 @@ class CalculationsTest < ActiveRecord::TestCase
     assert_equal [part.id], ShipPart.joins(:trinkets).pluck(:id)
   end
 
+  def test_pluck_loaded_relation
+    companies = Company.order(:id).limit(3).load
+    assert_no_queries do
+      assert_equal ['37signals', 'Summit', 'Microsoft'], companies.pluck(:name)
+    end
+  end
+
+  def test_pluck_loaded_relation_multiple_columns
+    companies = Company.order(:id).limit(3).load
+    assert_no_queries do
+      assert_equal [[1, '37signals'], [2, 'Summit'], [3, 'Microsoft']], companies.pluck(:id, :name)
+    end
+  end
+
+  def test_pluck_loaded_relation_sql_fragment
+    companies = Company.order(:id).limit(3).load
+    assert_queries 1 do
+      assert_equal ['37signals', 'Summit', 'Microsoft'], companies.pluck('DISTINCT name')
+    end
+  end
+
   def test_grouped_calculation_with_polymorphic_relation
     part = ShipPart.create!(name: "has trinket")
     part.trinkets.create!
diff --git a/activesupport/lib/active_support/core_ext/enumerable.rb b/activesupport/lib/active_support/core_ext/enumerable.rb
index d28f26260e..fc7531d088 100644
--- a/activesupport/lib/active_support/core_ext/enumerable.rb
+++ b/activesupport/lib/active_support/core_ext/enumerable.rb
@@ -76,8 +76,15 @@ module Enumerable
   #
   #   [{ name: "David" }, { name: "Rafael" }, { name: "Aaron" }].pluck(:name)
   #     => ["David", "Rafael", "Aaron"]
-  def pluck(key)
-    map { |element| element[key] }
+  #
+  #   [{ id: 1, name: "David" }, { id: 2, name: "Rafael" }].pluck(:id, :name)
+  #     => [[1, "David"], [2, "Rafael"]]
+  def pluck(*keys)
+    if keys.many?
+      map { |element| keys.map { |key| element[key] } }
+    else
+      map { |element| element[keys.first] }
+    end
   end
 end
 
diff --git a/activesupport/test/core_ext/enumerable_test.rb b/activesupport/test/core_ext/enumerable_test.rb
index 21743cdea5..f09b7d8850 100644
--- a/activesupport/test/core_ext/enumerable_test.rb
+++ b/activesupport/test/core_ext/enumerable_test.rb
@@ -3,6 +3,8 @@ require 'active_support/core_ext/array'
 require 'active_support/core_ext/enumerable'
 
 Payment = Struct.new(:price)
+ExpandedPayment = Struct.new(:dollars, :cents)
+
 class SummablePayment < Payment
   def +(p) self.class.new(price + p.price) end
 end
@@ -114,5 +116,12 @@ class EnumerableTests < ActiveSupport::TestCase
   def test_pluck
     payments = GenericEnumerable.new([ Payment.new(5), Payment.new(15), Payment.new(10) ])
     assert_equal [5, 15, 10], payments.pluck(:price)
+
+    payments = GenericEnumerable.new([
+      ExpandedPayment.new(5, 99),
+      ExpandedPayment.new(15, 0),
+      ExpandedPayment.new(10, 50)
+    ])
+    assert_equal [[5, 99], [15, 0], [10, 50]], payments.pluck(:dollars, :cents)
   end
 end