[python] Disable round-robin merge condition repartition

Author: XiaoHongbo-Hope

paimon-python/pypaimon/ray/merge_condition.py

@@ -86,7 +86,10 @@ def filter_batch(
         return batch
     datafusion = _load_datafusion()
     rewritten = condition if _pre_rewritten else rewrite_condition(condition)
-    ctx = datafusion.SessionContext()
+    config = datafusion.SessionConfig().set(
+        "datafusion.optimizer.enable_round_robin_repartition", "false"
+    )
+    ctx = datafusion.SessionContext(config)
     input_batches = batch.to_batches()
     # Use one batch per partition and rebuild from partitioned batches so
     # DataFusion neither concatenates 32-bit offsets nor reorders the input.

paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py

@@ -4722,28 +4722,27 @@ def test_filter_batch(self):
         result = filter_batch(batch, 's.age > t.age')
         self.assertEqual(result.column('s.id').to_pylist(), [2, 3])
 
+    @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
     def test_filter_batch_preserves_partition_order(self):
         from pypaimon.ray.merge_condition import filter_batch
 
-        first = pa.record_batch({'t.id': pa.array([0])})
-        second = pa.record_batch({'t.id': pa.array([1])})
-        dataframe = Mock()
-        dataframe.collect.return_value = [second, first]
-        dataframe.collect_partitioned.return_value = [[first], [second]]
-        context = Mock()
-        context.sql.return_value = dataframe
-        fake_datafusion = Mock()
-        fake_datafusion.SessionContext.return_value = context
-        source = pa.Table.from_batches([first, second])
+        batch_size = 20_000
+        expected = list(range(4 * batch_size))
+        source = pa.table({
+            't.id': pa.chunked_array([
+                pa.array(
+                    range(i * batch_size, (i + 1) * batch_size),
+                    type=pa.int64(),
+                )
+                for i in range(4)
+            ]),
+        })
 
-        with patch(
-            'pypaimon.ray.merge_condition._load_datafusion',
-            return_value=fake_datafusion,
-        ):
-            result = filter_batch(source, 'TRUE', _pre_rewritten=True)
+        result = filter_batch(
+            source, '"t.id" >= 0', _pre_rewritten=True,
+        )
 
-        self.assertEqual(result.column('t.id').to_pylist(), [0, 1])
-        dataframe.collect.assert_not_called()
+        self.assertEqual(result.column('t.id').to_pylist(), expected)
 
     @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
     def test_filter_batch_preserves_large_offset_chunks(self):