Merge pull request rails#51243 from maximerety/optim-in-batches-using-offset

Optimize Active Record batching further when using ranges

Author: fatkodima

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