diff --git a/paimon-python/pypaimon/ray/merge_condition.py b/paimon-python/pypaimon/ray/merge_condition.py index 8bd437707a1e..081d05b6def4 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 ee53ad229836..4bee261eb16b 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 @@ 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()