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 179277d014 [python][ray] Fix read-task limits and close batch readers
(#10698)
179277d014 is described below
commit 179277d014821a1d303c5025044e76a206b260e2
Author: Zaynt <[email protected]>
AuthorDate: Mon Oct 26 12:15:10 2026 +0800
[python][ray] Fix read-task limits and close batch readers (#10698)
---
.../pypaimon/read/datasource/ray_datasource.py | 4 +
.../pypaimon/tests/ray_datasource_test.py | 132 +++++++++++++++++++++
.../pypaimon/tests/ray_integration_test.py | 26 ++--
3 files changed, 144 insertions(+), 18 deletions(-)
diff --git a/paimon-python/pypaimon/read/datasource/ray_datasource.py
b/paimon-python/pypaimon/read/datasource/ray_datasource.py
index a8cb978e03..b37292f75c 100644
--- a/paimon-python/pypaimon/read/datasource/ray_datasource.py
+++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py
@@ -128,6 +128,8 @@ class RayDatasource(Datasource):
output_projection = self._split_provider.output_projection()
splits = self._split_provider.splits()
limit = self._split_provider.limit()
+ if per_task_row_limit is not None:
+ limit = per_task_row_limit if limit is None else min(limit,
per_task_row_limit)
include_row_kind = self._split_provider.include_row_kind()
if not splits:
return []
@@ -239,6 +241,8 @@ class RayDatasource(Datasource):
num_rows = None # Can't estimate with predicate filtering
else:
num_rows = total_rows
+ if num_rows is not None and limit is not None:
+ num_rows = min(num_rows, limit)
size_bytes = total_size if total_size > 0 else None
metadata_kwargs = {
diff --git a/paimon-python/pypaimon/tests/ray_datasource_test.py
b/paimon-python/pypaimon/tests/ray_datasource_test.py
new file mode 100644
index 0000000000..c127fe271a
--- /dev/null
+++ b/paimon-python/pypaimon/tests/ray_datasource_test.py
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from unittest.mock import Mock, patch
+
+import pyarrow as pa
+import pytest
+import ray.cloudpickle
+from ray.data import DataContext
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.read.datasource.ray_datasource import RayDatasource
+from pypaimon.read.datasource.split_provider import CatalogSplitProvider
+from pypaimon.read.table_read import TableRead
+
+
[email protected]
+def catalog_options(tmp_path):
+ options = {"warehouse": str(tmp_path)}
+ catalog = CatalogFactory.create(options)
+ catalog.create_database("default", True)
+ data = pa.table({"id": list(range(5))})
+ catalog.create_table("default.data",
Schema.from_pyarrow_schema(data.schema), False)
+ table = catalog.get_table("default.data")
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write()
+ commit = builder.new_commit()
+ try:
+ writer.write_arrow(data)
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ return options
+
+
[email protected]("limit,task_limit,expected", [
+ (None, None, 5), (2, None, 2), (None, 2, 2),
+ (4, 2, 2), (2, 4, 2), (0, None, 0), (None, 0, 0),
+ (10, None, 5),
+])
+def test_task_metadata_matches_limited_rows(catalog_options, limit,
task_limit, expected):
+ source = RayDatasource(CatalogSplitProvider(
+ "default.data", catalog_options, limit=limit))
+ tasks = source.get_read_tasks(
+ 1, per_task_row_limit=task_limit,
data_context=DataContext.get_current())
+ # The native planner can return no splits for a zero scan limit.
+ rows = 0
+ for task in tasks:
+ task = ray.cloudpickle.loads(ray.cloudpickle.dumps(task))
+ task_rows = sum(block.num_rows for block in task())
+ assert task.metadata.num_rows == task_rows
+ rows += task_rows
+
+ assert rows == expected
+
+
+def test_filter_keeps_unknown_row_count_with_limit(catalog_options):
+ table = CatalogFactory.create(catalog_options).get_table("default.data")
+ predicate =
table.new_read_builder().new_predicate_builder().greater_than("id", 1)
+ source = RayDatasource(CatalogSplitProvider(
+ "default.data", catalog_options, predicate=predicate, limit=2))
+ task = source.get_read_tasks(1, per_task_row_limit=1)[0]
+
+ assert task.metadata.num_rows is None
+ assert sum(block.num_rows for block in task()) == 1
+
+
+def test_ray_limit_is_applied_inside_worker_reader(catalog_options):
+ source = RayDatasource(CatalogSplitProvider("default.data",
catalog_options, limit=4))
+ task = source.get_read_tasks(1, per_task_row_limit=2)[0]
+ original = TableRead.to_arrow_batch_reader
+ limits = []
+
+ def track_reader(read, splits):
+ limits.append(read.limit)
+ return original(read, splits)
+
+ with patch.object(TableRead, "to_arrow_batch_reader", track_reader):
+ assert sum(block.num_rows for block in task()) == 2
+ assert limits == [2]
+
+
[email protected]("finish", ["exhausted", "cancelled", "failed",
"empty"])
[email protected]("fail_close", [False, True])
+def test_worker_closes_batch_reader(catalog_options, finish, fail_close):
+ source = RayDatasource(CatalogSplitProvider("default.data",
catalog_options))
+ task = source.get_read_tasks(1)[0]
+ batch = pa.record_batch([pa.array([0, 1])], names=["id"])
+ read_error = RuntimeError("read failed")
+ close_error = OSError("close failed") if fail_close else None
+ batches = [batch, StopIteration()]
+ if finish == "failed":
+ batches = [batch, read_error]
+ elif finish == "empty":
+ batches = [StopIteration()]
+ reader = Mock()
+ reader.read_next_batch.side_effect = batches
+ reader.close.side_effect = close_error
+
+ with patch.object(TableRead, "to_arrow_batch_reader", return_value=reader):
+ blocks = task.read_fn()
+ if finish == "cancelled":
+ assert next(blocks).num_rows == 2
+ consume = blocks.close
+ else:
+ consume = lambda: list(blocks)
+
+ expected_error = read_error if finish == "failed" else close_error
+ if expected_error is not None:
+ with pytest.raises(type(expected_error)) as exc_info:
+ consume()
+ assert exc_info.value is expected_error
+ else:
+ result = consume()
+ if finish != "cancelled":
+ assert sum(block.num_rows for block in result) == (0 if finish
== "empty" else 2)
+ reader.close.assert_called_once_with()
diff --git a/paimon-python/pypaimon/tests/ray_integration_test.py
b/paimon-python/pypaimon/tests/ray_integration_test.py
index d95d3855a9..0096dd4c95 100644
--- a/paimon-python/pypaimon/tests/ray_integration_test.py
+++ b/paimon-python/pypaimon/tests/ray_integration_test.py
@@ -358,18 +358,7 @@ class RayIntegrationTest(unittest.TestCase):
self.assertEqual(set(df['category'].tolist()), {'A'})
def test_read_paimon_with_limit(self):
- """``read_paimon(limit=N)`` propagates the limit into the scan plan.
-
- Writes 10 rows across two partitions (5 + 5) so the scan produces two
- raw-convertible splits. ``limit=3`` causes ``FileScanner`` to drop the
- second split once the first already covers the limit, so the Ray
- Dataset contains strictly fewer than the full 10 rows.
-
- We assert ``< 10`` (not ``== N``) because Paimon's scan-time limit is
- a per-split cap — whole-split granularity at this layer — not a
- row-exact hard limit. Row-exact short-circuiting in the reader is a
- separate follow-up.
- """
+ """Limits remain exact before and after materialization across
tasks."""
from pypaimon.ray import read_paimon
pa_schema = pa.schema([
@@ -391,12 +380,13 @@ class RayIntegrationTest(unittest.TestCase):
ds_full = read_paimon(identifier, self.catalog_options)
self.assertEqual(ds_full.count(), 10)
- # With limit=3, the scan plan drops the second partition's split
- # once the first split's row count already covers the limit.
- ds = read_paimon(identifier, self.catalog_options, limit=3)
- limited_count = ds.count()
- self.assertGreater(limited_count, 0)
- self.assertLess(limited_count, 10)
+ for limit in (0, 3, 7, 20):
+ with self.subTest(limit=limit):
+ ds = read_paimon(
+ identifier, self.catalog_options, limit=limit,
+ override_num_blocks=2)
+ self.assertEqual(ds.count(), min(limit, 10))
+ self.assertEqual(ds.materialize().count(), min(limit, 10))
def test_read_paimon_empty_table(self):
"""read_paimon() on an empty table preserves the table schema."""
