Fix loaded relation batching for models with composite primary keys

Author: fatkodima

activerecord/lib/active_record/relation/batches.rb

@@ -338,23 +338,17 @@ def get_the_order_of_primary_key(order)
 
       def batch_on_loaded_relation(relation:, start:, finish:, order:, batch_limit:)
         records = relation.to_a
+        order = build_batch_orders(order).map(&:second)
 
         if start || finish
           records = records.filter do |record|
-            id = record.id
-
-            if order == :asc
-              (start.nil? || id >= start) && (finish.nil? || id <= finish)
-            else
-              (start.nil? || id <= start) && (finish.nil? || id >= finish)
-            end
+            (start.nil? || compare_values_for_order(record.id, start, order) >= 0) &&
+              (finish.nil? || compare_values_for_order(record.id, finish, order) <= 0)
           end
         end
 
-        records.sort_by!(&:id)
-
-        if order == :desc
-          records.reverse!
+        records.sort! do |record1, record2|
+          compare_values_for_order(record1.id, record2.id, order)
         end
 
         records.each_slice(batch_limit) do |subrecords|
@@ -367,6 +361,28 @@ def batch_on_loaded_relation(relation:, start:, finish:, order:, batch_limit:)
         nil
       end
 
+      # This is a custom implementation of `<=>` operator,
+      # which also takes into account how the collection will be ordered.
+      def compare_values_for_order(value1, value2, order)
+        # Multiple column values.
+        if value1.is_a?(Array)
+          value1.each_with_index do |element1, index|
+            element2 = value2[index]
+            direction = order[index]
+            comparison = element1 <=> element2
+            comparison = -comparison if direction == :desc
+            return comparison if comparison != 0
+          end
+
+          0
+        # Single column values.
+        elsif order.first == :asc
+          value1 <=> value2
+        else
+          value2 <=> value1
+        end
+      end
+
       def batch_on_unloaded_relation(relation:, start:, finish:, load:, order:, use_ranges:, remaining:, batch_limit:)
         batch_orders = build_batch_orders(order)
         relation = relation.reorder(batch_orders.to_h).limit(batch_limit)

activerecord/test/cases/batches_test.rb

@@ -538,6 +538,25 @@ def test_in_batches_when_loaded_can_return_an_enum
     assert_equal posts.size, batch_count
   end
 
+  def test_in_batches_when_loaded_runs_no_queries_when_batching_over_cpk_model
+    incorrectly_sorted_orders = Cpk::Order.order(shop_id: :asc, id: :desc)
+    incorrectly_sorted_orders.load
+
+    correctly_sorted_orders = Cpk::Order.order(shop_id: :desc, id: :asc).to_a
+    expected_orders = correctly_sorted_orders[1..-2]
+    start_id = expected_orders.first.id
+    finish_id = expected_orders.last.id
+    orders = []
+
+    assert_no_queries do
+      incorrectly_sorted_orders.find_each(batch_size: 1, start: start_id, finish: finish_id, order: [:desc, :asc]) do |order|
+        orders << order
+      end
+    end
+
+    assert_equal expected_orders, orders
+  end
+
   def test_in_batches_should_return_relations
     assert_queries_count(@total + 1) do
       Post.in_batches(of: 1) do |relation|