activerecord 3.1.10 → 4.2.11

data/lib/active_record/relation/batches.rb

@@ -1,80 +1,131 @@
-require 'active_support/core_ext/object/blank'
-
 module ActiveRecord
   module Batches
-    # Yields each record that was found by the find +options+. The find is
-    # performed by find_in_batches with a batch size of 1000 (or as
-    # specified by the <tt>:batch_size</tt> option).
+    # Looping through a collection of records from the database
+    # (using the +all+ method, for example) is very inefficient
+    # since it will try to instantiate all the objects at once.
+    #
+    # In that case, batch processing methods allow you to work
+    # with the records in batches, thereby greatly reducing memory consumption.
     #
-    # Example:
+    # The #find_each method uses #find_in_batches with a batch size of 1000 (or as
+    # specified by the +:batch_size+ option).
+    #
+    #   Person.find_each do |person|
+    #     person.do_awesome_stuff
+    #   end
     #
     #   Person.where("age > 21").find_each do |person|
     #     person.party_all_night!
     #   end
     #
-    # Note: This method is only intended to use for batch processing of
-    # large amounts of records that wouldn't fit in memory all at once. If
-    # you just need to loop over less than 1000 records, it's probably
-    # better just to use the regular find methods.
+    # If you do not provide a block to #find_each, it will return an Enumerator
+    # for chaining with other methods:
+    #
+    #   Person.find_each.with_index do |person, index|
+    #     person.award_trophy(index + 1)
+    #   end
+    #
+    # ==== Options
+    # * <tt>:batch_size</tt> - Specifies the size of the batch. Default to 1000.
+    # * <tt>:start</tt> - Specifies the starting point for the batch processing.
+    # This is especially useful if you want multiple workers dealing with
+    # the same processing queue. You can make worker 1 handle all the records
+    # between id 0 and 10,000 and worker 2 handle from 10,000 and beyond
+    # (by setting the +:start+ option on that worker).
+    #
+    #   # Let's process for a batch of 2000 records, skipping the first 2000 rows
+    #   Person.find_each(start: 2000, batch_size: 2000) do |person|
+    #     person.party_all_night!
+    #   end
+    #
+    # NOTE: It's not possible to set the order. That is automatically set to
+    # ascending on the primary key ("id ASC") to make the batch ordering
+    # work. This also means that this method only works with integer-based
+    # primary keys.
+    #
+    # NOTE: You can't set the limit either, that's used to control
+    # the batch sizes.
     def find_each(options = {})
-      find_in_batches(options) do |records|
-        records.each { |record| yield record }
+      if block_given?
+        find_in_batches(options) do |records|
+          records.each { |record| yield record }
+        end
+      else
+        enum_for :find_each, options do
+          options[:start] ? where(table[primary_key].gteq(options[:start])).size : size
+        end
       end
     end
 
     # Yields each batch of records that was found by the find +options+ as
-    # an array. The size of each batch is set by the <tt>:batch_size</tt>
-    # option; the default is 1000.
+    # an array.
     #
-    # You can control the starting point for the batch processing by
-    # supplying the <tt>:start</tt> option. This is especially useful if you
-    # want multiple workers dealing with the same processing queue. You can
-    # make worker 1 handle all the records between id 0 and 10,000 and
-    # worker 2 handle from 10,000 and beyond (by setting the <tt>:start</tt>
-    # option on that worker).
+    #   Person.where("age > 21").find_in_batches do |group|
+    #     sleep(50) # Make sure it doesn't get too crowded in there!
+    #     group.each { |person| person.party_all_night! }
+    #   end
     #
-    # It's not possible to set the order. That is automatically set to
-    # ascending on the primary key ("id ASC") to make the batch ordering
-    # work. This also mean that this method only works with integer-based
-    # primary keys. You can't set the limit either, that's used to control
-    # the batch sizes.
+    # If you do not provide a block to #find_in_batches, it will return an Enumerator
+    # for chaining with other methods:
     #
-    # Example:
+    #   Person.find_in_batches.with_index do |group, batch|
+    #     puts "Processing group ##{batch}"
+    #     group.each(&:recover_from_last_night!)
+    #   end
     #
-    #   Person.where("age > 21").find_in_batches do |group|
-    #     sleep(50) # Make sure it doesn't get too crowded in there!
+    # To be yielded each record one by one, use #find_each instead.
+    #
+    # ==== Options
+    # * <tt>:batch_size</tt> - Specifies the size of the batch. Default to 1000.
+    # * <tt>:start</tt> - Specifies the starting point for the batch processing.
+    # This is especially useful if you want multiple workers dealing with
+    # the same processing queue. You can make worker 1 handle all the records
+    # between id 0 and 10,000 and worker 2 handle from 10,000 and beyond
+    # (by setting the +:start+ option on that worker).
+    #
+    #   # Let's process the next 2000 records
+    #   Person.find_in_batches(start: 2000, batch_size: 2000) do |group|
     #     group.each { |person| person.party_all_night! }
     #   end
+    #
+    # NOTE: It's not possible to set the order. That is automatically set to
+    # ascending on the primary key ("id ASC") to make the batch ordering
+    # work. This also means that this method only works with integer-based
+    # primary keys.
+    #
+    # NOTE: You can't set the limit either, that's used to control
+    # the batch sizes.
     def find_in_batches(options = {})
+      options.assert_valid_keys(:start, :batch_size)
+
       relation = self
+      start = options[:start]
+      batch_size = options[:batch_size] || 1000
 
-      unless arel.orders.blank? && arel.taken.blank?
-        ActiveRecord::Base.logger.warn("Scoped order and limit are ignored, it's forced to be batch order and batch size")
+      unless block_given?
+        return to_enum(:find_in_batches, options) do
+          total = start ? where(table[primary_key].gteq(start)).size : size
+          (total - 1).div(batch_size) + 1
+        end
       end
 
-      if (finder_options = options.except(:start, :batch_size)).present?
-        raise "You can't specify an order, it's forced to be #{batch_order}" if options[:order].present?
-        raise "You can't specify a limit, it's forced to be the batch_size"  if options[:limit].present?
-
-        relation = apply_finder_options(finder_options)
+      if logger && (arel.orders.present? || arel.taken.present?)
+        logger.warn("Scoped order and limit are ignored, it's forced to be batch order and batch size")
       end
 
-      start = options.delete(:start).to_i
-      batch_size = options.delete(:batch_size) || 1000
-
       relation = relation.reorder(batch_order).limit(batch_size)
-      records = relation.where(table[primary_key].gteq(start)).all
+      records = start ? relation.where(table[primary_key].gteq(start)).to_a : relation.to_a
 
       while records.any?
+        records_size = records.size
+        primary_key_offset = records.last.id
+        raise "Primary key not included in the custom select clause" unless primary_key_offset
+
         yield records
 
-        break if records.size < batch_size
+        break if records_size < batch_size
 
-        if primary_key_offset = records.last.id
-          records = relation.where(table[primary_key].gt(primary_key_offset)).to_a
-        else
-          raise "Primary key not included in the custom select clause"
-        end
+        records = relation.where(table[primary_key].gt(primary_key_offset)).to_a
       end
     end