This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new e8a88867bc [python] Fix Arrow offset overflow in merge condition
filtering (#9678)
e8a88867bc is described below
commit e8a88867bc6a29913dc775c9f94c8c5d3302d1d9
Author: XiaoHongbo <[email protected]>
AuthorDate: Tue Sep 8 16:39:26 2026 +0800
[python] Fix Arrow offset overflow in merge condition filtering (#9678)
---
paimon-python/pypaimon/ray/merge_condition.py | 21 +++++++--
.../tests/ray_data_evolution_merge_into_test.py | 54 ++++++++++++++++++++++
2 files changed, 72 insertions(+), 3 deletions(-)
diff --git a/paimon-python/pypaimon/ray/merge_condition.py
b/paimon-python/pypaimon/ray/merge_condition.py
index 8bd437707a..081d05b6de 100644
--- a/paimon-python/pypaimon/ray/merge_condition.py
+++ b/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(
diff --git a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
index ee53ad2298..4bee261eb1 100644
--- a/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
+++ b/paimon-python/pypaimon/tests/ray_data_evolution_merge_into_test.py
@@ -4722,6 +4722,60 @@ class MergeConditionUnitTest(unittest.TestCase):
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()
