[python] Fix Arrow offset overflow in merge condition filtering

paimon-python/pypaimon/ray/merge_condition.py

@@ -86,12 +86,27 @@ def filter_batch(
         return batch
     datafusion = _load_datafusion()
     rewritten = condition if _pre_rewritten else rewrite_condition(condition)
-    ctx = datafusion.SessionContext()
-    ctx.register_record_batches("_batch", [batch.to_batches()])
+    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.
+    ctx.register_record_batches(
+        "_batch", [[record_batch] for record_batch in input_batches]
+    )
     result = ctx.sql(
         f'SELECT * FROM _batch WHERE {rewritten}'
     )
-    return result.to_arrow_table()
+    output_batches = [
+        record_batch
+        for partition in result.collect_partitioned()
+        for record_batch in partition
+    ]
+    if not output_batches:
+        return batch.schema.empty_table()
+    return pa.Table.from_batches(output_batches)
 
 
 def apply_condition(

paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py

@@ -4722,6 +4722,60 @@ 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
+
+        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)
+            ]),
+        })
+
+        result = filter_batch(
+            source, '"t.id" >= 0', _pre_rewritten=True,
+        )
+
+        self.assertEqual(result.column('t.id').to_pylist(), expected)
+
+    @unittest.skipIf(_SKIP_CONDITION, _SKIP_REASON)
+    def test_filter_batch_preserves_large_offset_chunks(self):
+        from pypaimon.ray.merge_condition import filter_batch
+
+        child_count = 1_100_000_000
+
+        def large_list():
+            return pa.ListArray.from_arrays(
+                pa.array([0, child_count], type=pa.int32()),
+                pa.nulls(child_count),
+            )
+
+        batch = pa.table({
+            't._ROW_ID': pa.chunked_array([
+                pa.array([0], type=pa.int64()),
+                pa.array([1], type=pa.int64()),
+            ]),
+            't.payload': pa.chunked_array([large_list(), large_list()]),
+        })
+
+        result = filter_batch(
+            batch, '"t._ROW_ID" >= 0', _pre_rewritten=True,
+        )
+
+        self.assertEqual(result.column('t._ROW_ID').to_pylist(), [0, 1])
+        payload = result.column('t.payload')
+        self.assertEqual(payload.num_chunks, 2)
+        self.assertEqual(
+            [len(chunk.values) for chunk in payload.chunks],
+            [child_count, child_count],
+        )
+
 
 if __name__ == '__main__':
     unittest.main()