Optimize Active Record batching further when using ranges

When `use_ranges: true` option is used, we do not need to return the
whole list of values for each range. This wastes network resources and
takes longer than necessary, as only the last value from the range would
be needed.

We can instead use OFFSET to peek at the last value of the next batch,
instead of returning the whole list of values from the range.

There is a trade-off, however, as we need an additional query to confirm
the size of the very last batch, and get the last value. Unless we only
have a handful of small batches, this strategy is a winner overall, as
it reduces the time and network resources spent generating batches.

For example, `[10, 11, 12, ..., 20]` was previously loaded to generate a
relation in the form of `WHERE id > 10 AND id <= 20` for a batch. Since
`10` would already be known from the previous iteration, we only need to
retrieve the last value in the range (`20`) using a `LIMIT 1 OFFSET ω`
construct (ω = batch size - 1), thus avoiding the unnecessary loading
and discarding of all other values in the range.

Author: maximerety

activerecord/CHANGELOG.md

@@ -1,3 +1,11 @@
+*   Optimize Active Record batching further when using ranges.
+
+    Tested on a PostgreSQL table with 10M records and batches of 10k records, the generation
+    of relations for the 1000 batches was `4.8x` faster (`6.8s` vs. `1.4s`), used `900x`
+    less bandwidth (`180MB` vs. `0.2MB`) and allocated `45x` less memory (`490MB` vs. `11MB`).
+
+    *Maxime Réty*, *fatkodima*
+
 *   Include current character length in error messages for index and table name length validations.
 
     *Joshua Young*

activerecord/lib/active_record/relation/batches.rb

@@ -435,35 +435,49 @@ def batch_on_unloaded_relation(relation:, start:, finish:, load:, cursor:, order
           if load
             records = batch_relation.records
             values = records.pluck(*cursor)
+            values_size = values.size
+            values_last = values.last
             yielded_relation = where(cursor => values).order(batch_orders.to_h)
             yielded_relation.load_records(records)
           elsif (empty_scope && use_ranges != false) || use_ranges
-            values = batch_relation.pluck(*cursor)
+            # Efficiently peak at the last value for the next batch using offset and limit.
+            values_size = batch_limit
+            values_last = batch_relation.offset(batch_limit - 1).pick(*cursor)
+
+            # If the last value is not found using offset, there is at most one more batch of size < batch_limit.
+            # Retry by getting the whole list of remaining values so that we have the exact size and last value.
+            unless values_last
+              values = batch_relation.pluck(*cursor)
+              values_size = values.size
+              values_last = values.last
+            end
 
-            finish = values.last
-            if finish
-              yielded_relation = apply_finish_limit(batch_relation, cursor, finish, batch_orders)
+            # Finally, build the yielded relation if at least one value found.
+            if values_last
+              yielded_relation = apply_finish_limit(batch_relation, cursor, values_last, batch_orders)
               yielded_relation = yielded_relation.except(:limit).reorder(batch_orders.to_h)
               yielded_relation.skip_query_cache!(false)
             end
           else
             values = batch_relation.pluck(*cursor)
+            values_size = values.size
+            values_last = values.last
             yielded_relation = where(cursor => values).order(batch_orders.to_h)
           end
 
-          break if values.empty?
+          break if values_size == 0
 
-          if values.flatten.any?(nil)
+          if [values_last].flatten.any?(nil)
             raise ArgumentError, "Not all of the batch cursor columns were included in the custom select clause "\
                                   "or some columns contain nil."
           end
 
           yield yielded_relation
 
-          break if values.length < batch_limit
+          break if values_size < batch_limit
 
           if limit_value
-            remaining -= values.length
+            remaining -= values_size
 
             if remaining == 0
               # Saves a useless iteration when the limit is a multiple of the
@@ -481,7 +495,7 @@ def batch_on_unloaded_relation(relation:, start:, finish:, load:, cursor:, order
           end
           operators << (last_order == :desc ? :lt : :gt)
 
-          cursor_value = values.last
+          cursor_value = values_last
           batch_relation = batch_condition(relation, cursor, cursor_value, operators)
         end
 

activerecord/test/cases/batches_test.rb

@@ -335,15 +335,15 @@ def test_in_batches_has_attribute_readers
   end
 
   def test_in_batches_should_yield_relation_if_block_given
-    assert_queries_count(6) do
+    assert_queries_count(7) do
       Post.in_batches(of: 2) do |relation|
         assert_kind_of ActiveRecord::Relation, relation
       end
     end
   end
 
   def test_in_batches_should_be_enumerable_if_no_block_given
-    assert_queries_count(6) do
+    assert_queries_count(7) do
       Post.in_batches(of: 2).each do |relation|
         assert_kind_of ActiveRecord::Relation, relation
       end
@@ -378,10 +378,10 @@ def test_in_batches_each_record_should_be_ordered_by_id
   end
 
   def test_in_batches_update_all_affect_all_records
-    assert_queries_count(6 + 6) do # 6 selects, 6 updates
+    assert_queries_count(7 + 6) do # 7 selects, 6 updates
       Post.in_batches(of: 2).update_all(title: "updated-title")
     end
-    assert_equal Post.all.pluck(:title), ["updated-title"] * Post.count
+    assert_equal ["updated-title"] * Post.count, Post.all.pluck(:title)
   end
 
   def test_in_batches_update_all_returns_rows_affected
@@ -394,7 +394,7 @@ def test_in_batches_update_all_returns_zero_when_no_batches
 
   def test_in_batches_touch_all_affect_all_records
     time = Time.new(2000, 1, 1, 0, 0, 0)
-    assert_queries_count(6 + 6) do # 6 selects, 6 updates
+    assert_queries_count(7 + 6) do # 7 selects, 6 updates
       Developer.in_batches(of: 2).touch_all(time: time)
     end
     assert_equal [time] * Developer.count, Developer.all.pluck(:updated_at)
@@ -486,7 +486,7 @@ def test_in_scoped_batches_preserves_order_within_batches
   end
 
   def test_in_batches_if_not_loaded_executes_more_queries
-    assert_queries_count(@total + 1) do
+    assert_queries_count(@total + 2) do
       Post.in_batches(of: 1, load: false) do |relation|
         assert_not_predicate relation, :loaded?
       end
@@ -617,7 +617,7 @@ def test_in_batches_when_loaded_iterates_using_custom_column
   end
 
   def test_in_batches_should_return_relations
-    assert_queries_count(@total + 1) do
+    assert_queries_count(@total + 2) do
       Post.in_batches(of: 1) do |relation|
         assert_kind_of ActiveRecord::Relation, relation
       end
@@ -634,46 +634,70 @@ def test_in_batches_should_start_from_the_start_option
 
   def test_in_batches_should_end_at_the_finish_option
     post = Post.order("id DESC").where("id <= ?", 5).first
-    assert_queries_count(7) do
+    assert_queries_count(8) do
       relation = Post.in_batches(of: 1, finish: 5, load: true).reverse_each.first
       assert_equal post, relation.last
     end
   end
 
   def test_in_batches_executes_range_queries_when_unconstrained
     quoted_posts_id = Regexp.escape(quote_table_name("posts.id"))
+
+    relations = assert_queries_match(/ORDER BY #{quoted_posts_id} ASC LIMIT \S+ OFFSET \S+\z/i, count: 6) do
+      assert_queries_match(/ORDER BY #{quoted_posts_id} ASC LIMIT \S+\z/i, count: 1) do
+        Post.in_batches(of: 2).to_a
+      end
+    end
+
     assert_queries_match(/WHERE #{quoted_posts_id} > .+ AND #{quoted_posts_id} <= .+/i) do
-      Post.in_batches(of: 2) { |relation| assert_kind_of Post, relation.first }
+      relations.each { |relation| assert_kind_of Post, relation.first }
     end
   end
 
   def test_in_batches_executes_in_queries_when_unconstrained_and_opted_out_of_ranges
     quoted_posts_id = Regexp.escape(quote_table_name("posts.id"))
+
+    relations = assert_queries_match(/ORDER BY #{quoted_posts_id} ASC LIMIT \S+\z/i, count: 6) do
+      Post.in_batches(of: 2, use_ranges: false).to_a
+    end
+
     assert_queries_match(/#{quoted_posts_id} IN \(.+\)/i) do
-      Post.in_batches(of: 2, use_ranges: false) { |relation| assert_kind_of Post, relation.first }
+      relations.each { |relation| assert_kind_of Post, relation.first }
     end
   end
 
   def test_in_batches_executes_in_queries_when_constrained
     quoted_posts_id = Regexp.escape(quote_table_name("posts.id"))
+
+    relations = assert_queries_match(/ORDER BY #{quoted_posts_id} ASC LIMIT \S+\z/i, count: 3) do
+      Post.where("id < ?", 5).in_batches(of: 2).to_a
+    end
+
     assert_queries_match(/#{quoted_posts_id} IN \(.+\)/i) do
-      Post.where("id < ?", 5).in_batches(of: 2) { |relation| assert_kind_of Post, relation.first }
+      relations.each { |relation| assert_kind_of Post, relation.first }
     end
   end
 
   def test_in_batches_executes_range_queries_when_constrained_and_opted_in_into_ranges
     quoted_posts_id = Regexp.escape(quote_table_name("posts.id"))
+
+    relations = assert_queries_match(/ORDER BY #{quoted_posts_id} ASC LIMIT \S+ OFFSET \S+\z/i, count: 3) do
+      assert_queries_match(/ORDER BY #{quoted_posts_id} ASC LIMIT \S+\z/i, count: 1) do
+        Post.where("id < ?", 5).in_batches(of: 2, use_ranges: true).to_a
+      end
+    end
+
     assert_queries_match(/#{quoted_posts_id} > .+ AND #{quoted_posts_id} <= .+/i) do
-      Post.where("id < ?", 5).in_batches(of: 2, use_ranges: true) { |relation| assert_kind_of Post, relation.first }
+      relations.each { |relation| assert_kind_of Post, relation.first }
     end
   end
 
   def test_in_batches_shouldnt_execute_query_unless_needed
-    assert_queries_count(2) do
+    assert_queries_count(3) do
       Post.in_batches(of: @total) { |relation| assert_kind_of ActiveRecord::Relation, relation }
     end
 
-    assert_queries_count(1) do
+    assert_queries_count(2) do
       Post.in_batches(of: @total + 1) { |relation| assert_kind_of ActiveRecord::Relation, relation }
     end
   end
@@ -997,7 +1021,7 @@ def test_find_in_batches_should_return_a_sized_enumerator
 
   test ".in_batches bypasses the query cache for its own queries" do
     Post.cache do
-      assert_queries_count(2) do
+      assert_queries_count(4) do
         Post.in_batches { }
         Post.in_batches { }
       end