Preserve order within individual batches for batch iteration

Author: fatkodima

activerecord/lib/active_record/relation/batches.rb

@@ -435,20 +435,20 @@ def batch_on_unloaded_relation(relation:, start:, finish:, load:, cursor:, order
           if load
             records = batch_relation.records
             values = records.pluck(*cursor)
-            yielded_relation = where(cursor => values)
+            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)
 
             finish = values.last
             if finish
               yielded_relation = apply_finish_limit(batch_relation, cursor, finish, batch_orders)
-              yielded_relation = yielded_relation.except(:limit, :order)
+              yielded_relation = yielded_relation.except(:limit).reorder(batch_orders.to_h)
               yielded_relation.skip_query_cache!(false)
             end
           else
             values = batch_relation.pluck(*cursor)
-            yielded_relation = where(cursor => values)
+            yielded_relation = where(cursor => values).order(batch_orders.to_h)
           end
 
           break if values.empty?

activerecord/test/cases/batches_test.rb

@@ -455,6 +455,36 @@ def test_in_batches_should_be_loaded
     end
   end
 
+  def test_in_loaded_batches_preserves_order_within_batches
+    expected_posts = Post.order(id: :desc).to_a
+    posts = []
+
+    Post.in_batches(of: 2, load: true, order: :desc) do |relation|
+      posts.concat(relation.where("1 = 1"))
+    end
+    assert_equal expected_posts, posts
+  end
+
+  def test_in_range_batches_preserves_order_within_batches
+    expected_posts = Post.order(id: :desc).to_a
+    posts = []
+
+    Post.in_batches(of: 2, order: :desc, use_ranges: true) do |relation|
+      posts.concat(relation)
+    end
+    assert_equal expected_posts, posts
+  end
+
+  def test_in_scoped_batches_preserves_order_within_batches
+    expected_posts = Post.order(id: :desc).to_a
+    posts = []
+
+    Post.where("id > 0").in_batches(of: 2, order: :desc) do |relation|
+      posts.concat(relation)
+    end
+    assert_equal expected_posts, posts
+  end
+
   def test_in_batches_if_not_loaded_executes_more_queries
     assert_queries_count(@total + 1) do
       Post.in_batches(of: 1, load: false) do |relation|
@@ -638,13 +668,6 @@ def test_in_batches_executes_range_queries_when_constrained_and_opted_in_into_ra
     end
   end
 
-  def test_in_batches_no_subqueries_for_whole_tables_batching
-    quoted_posts_id = Regexp.escape(quote_table_name("posts.id"))
-    assert_queries_match(/DELETE FROM #{Regexp.escape(quote_table_name("posts"))} WHERE #{quoted_posts_id} > .+ AND #{quoted_posts_id} <=/i) do
-      Post.in_batches(of: 2).delete_all
-    end
-  end
-
   def test_in_batches_shouldnt_execute_query_unless_needed
     assert_queries_count(2) do
       Post.in_batches(of: @total) { |relation| assert_kind_of ActiveRecord::Relation, relation }