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 94a58cb6ef [python][ray] Preserve row kind in TableRead.to_ray()
(#9693)
94a58cb6ef is described below
commit 94a58cb6efdbe72a38e87391a37368dbbf908832
Author: QuakeWang <[email protected]>
AuthorDate: Fri Sep 11 14:02:16 2026 +0800
[python][ray] Preserve row kind in TableRead.to_ray() (#9693)
---
.../pypaimon/read/datasource/ray_datasource.py | 9 +++-
.../pypaimon/read/datasource/split_provider.py | 11 ++++-
paimon-python/pypaimon/read/table_read.py | 3 ++
.../pypaimon/tests/ray_integration_test.py | 54 ++++++++++++++++++++++
.../pypaimon/tests/split_provider_test.py | 6 +++
5 files changed, 81 insertions(+), 2 deletions(-)
diff --git a/paimon-python/pypaimon/read/datasource/ray_datasource.py
b/paimon-python/pypaimon/read/datasource/ray_datasource.py
index f8d5c7fafe..4b7e9459da 100644
--- a/paimon-python/pypaimon/read/datasource/ray_datasource.py
+++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py
@@ -127,11 +127,15 @@ class RayDatasource(Datasource):
nested_name_paths = self._split_provider.nested_name_paths()
splits = self._split_provider.splits()
limit = self._split_provider.limit()
+ include_row_kind = self._split_provider.include_row_kind()
if not splits:
return []
if self._schema is None:
self._schema = PyarrowFieldParser.from_paimon_schema(read_type)
+ if include_row_kind:
+ from pypaimon.read.table_read import TableRead
+ self._schema = TableRead._add_row_kind_to_schema(self._schema)
schema = self._schema
if parallelism > len(splits):
@@ -150,6 +154,7 @@ class RayDatasource(Datasource):
schema=schema,
limit=limit,
nested_name_paths=nested_name_paths,
+ include_row_kind=include_row_kind,
) -> Iterable[pyarrow.Table]:
"""Read function that will be executed by Ray workers."""
from pypaimon.read.table_read import TableRead
@@ -159,7 +164,8 @@ class RayDatasource(Datasource):
# columns and reads every projected leaf as NULL.
worker_table_read = TableRead(
table, predicate, read_type, limit=limit,
- nested_name_paths=nested_name_paths)
+ nested_name_paths=nested_name_paths,
+ include_row_kind=include_row_kind)
batch_reader = worker_table_read.to_arrow_batch_reader(splits)
has_data = False
@@ -187,6 +193,7 @@ class RayDatasource(Datasource):
schema=schema,
limit=limit,
nested_name_paths=nested_name_paths,
+ include_row_kind=include_row_kind,
)
read_tasks = []
diff --git a/paimon-python/pypaimon/read/datasource/split_provider.py
b/paimon-python/pypaimon/read/datasource/split_provider.py
index 6060ccce80..224a74fe54 100644
--- a/paimon-python/pypaimon/read/datasource/split_provider.py
+++ b/paimon-python/pypaimon/read/datasource/split_provider.py
@@ -67,6 +67,10 @@ class SplitProvider(ABC):
"""
return None
+ def include_row_kind(self) -> bool:
+ """Whether Arrow output should include the row kind column."""
+ return False
+
def nested_name_paths(self) -> Optional[List[List[str]]]:
"""Parallel name paths for a nested-leaf projection, or ``None``.
@@ -212,13 +216,15 @@ class PreResolvedSplitProvider(SplitProvider):
"""
def __init__(self, table, splits: List[Split], read_type, predicate=None,
- limit: Optional[int] = None, nested_name_paths=None):
+ limit: Optional[int] = None, nested_name_paths=None,
+ include_row_kind: bool = False):
self._table = table
self._splits = splits
self._read_type = read_type
self._predicate = predicate
self._limit = limit
self._nested_name_paths = nested_name_paths
+ self._include_row_kind = include_row_kind
def table(self):
return self._table
@@ -232,6 +238,9 @@ class PreResolvedSplitProvider(SplitProvider):
def nested_name_paths(self) -> Optional[List[List[str]]]:
return self._nested_name_paths
+ def include_row_kind(self) -> bool:
+ return self._include_row_kind
+
def predicate(self):
return self._predicate
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 3472666ae5..51465486ed 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -666,6 +666,8 @@ class TableRead:
if not splits:
schema = PyarrowFieldParser.from_paimon_schema(self.read_type)
+ if self.include_row_kind:
+ schema = self._add_row_kind_to_schema(schema)
empty_table = pyarrow.Table.from_arrays(
[pyarrow.array([], type=field.type) for field in schema],
schema=schema
@@ -686,6 +688,7 @@ class TableRead:
predicate=self.predicate,
limit=self.limit,
nested_name_paths=self.nested_name_paths,
+ include_row_kind=self.include_row_kind,
)
)
ds = ray.data.read_datasource(
diff --git a/paimon-python/pypaimon/tests/ray_integration_test.py
b/paimon-python/pypaimon/tests/ray_integration_test.py
index df8c52ae81..05aff4bc63 100644
--- a/paimon-python/pypaimon/tests/ray_integration_test.py
+++ b/paimon-python/pypaimon/tests/ray_integration_test.py
@@ -103,6 +103,60 @@ class RayIntegrationTest(unittest.TestCase):
self.assertEqual(list(df['id']), [1, 2, 3])
self.assertEqual(list(df['name']), ['a', 'b', 'c'])
+ def test_to_ray_preserves_row_kind(self):
+ from pypaimon.read.table_read import TableRead
+
+ pa_schema = pa.schema([('id', pa.int32()), ('name', pa.string())])
+ identifier = self._create_and_populate_table(
+ 'test_to_ray_row_kind', pa_schema,
+ {'id': [1, 2], 'name': ['a', 'b']},
+ )
+ table =
CatalogFactory.create(self.catalog_options).get_table(identifier)
+ rb = table.new_read_builder()
+ splits = rb.new_scan().plan().splits()
+ self.assertTrue(splits)
+ predicate = rb.new_predicate_builder().equal('id', 999)
+ cases = [('data', splits, None), ('no_splits', [], None),
+ ('filtered_empty', splits, predicate)]
+
+ for include_row_kind in (False, True):
+ for name, task_splits, task_predicate in cases:
+ with self.subTest(include_row_kind=include_row_kind,
case=name):
+ read = TableRead(
+ table, task_predicate, rb.read_type(),
+ include_row_kind=include_row_kind)
+ arrow = read.to_arrow(task_splits)
+ expected_schema = pa_schema
+ if include_row_kind:
+ expected_schema = pa.schema(
+ [pa.field('_row_kind', pa.string())] +
list(pa_schema))
+ self.assertEqual(arrow.schema, expected_schema)
+ if name == 'data' and include_row_kind:
+
self.assertEqual(arrow.column('_row_kind').to_pylist(), ['+I', '+I'])
+
+ with patch.object(ray.data, 'read_datasource',
+ wraps=ray.data.read_datasource) as
read_datasource:
+ ds = read.to_ray(task_splits, override_num_blocks=1)
+ self.assertEqual(ds.schema().base_schema, expected_schema)
+ materialized = ds.materialize()
+ self.assertEqual(materialized.schema().base_schema,
expected_schema)
+ self.assertEqual(materialized.take_all(),
arrow.to_pylist())
+
+ if task_splits:
+ datasource = read_datasource.call_args[0][0]
+ tasks = datasource.get_read_tasks(1)
+ self.assertEqual(len(tasks), 1)
+ task = tasks[0]
+ task_schema = (task.schema if hasattr(task, 'schema')
+ else task.metadata.schema)
+ self.assertEqual(task_schema, expected_schema)
+ task =
ray.cloudpickle.loads(ray.cloudpickle.dumps(task))
+ blocks = list(task())
+ self.assertTrue(blocks)
+ for block in blocks:
+ self.assertEqual(block.schema, expected_schema)
+ self.assertEqual(pa.concat_tables(blocks), arrow)
+
def test_read_paimon_with_projection(self):
"""read_paimon() respects column projection."""
from pypaimon.ray import read_paimon
diff --git a/paimon-python/pypaimon/tests/split_provider_test.py
b/paimon-python/pypaimon/tests/split_provider_test.py
index 61e14005f5..7602a9978e 100644
--- a/paimon-python/pypaimon/tests/split_provider_test.py
+++ b/paimon-python/pypaimon/tests/split_provider_test.py
@@ -86,6 +86,7 @@ class SplitProviderTest(unittest.TestCase):
self.assertIs(provider.splits(), splits) # cached
self.assertIsNotNone(provider.read_type())
self.assertIsNone(provider.predicate())
+ self.assertFalse(provider.include_row_kind())
def test_catalog_provider_propagates_projection(self):
"""``projection`` reaches ``ReadBuilder.with_projection`` (visible via
read_type)."""
@@ -282,6 +283,11 @@ class SplitProviderTest(unittest.TestCase):
self.assertIs(provider.read_type(), read_type)
self.assertIsNone(provider.predicate())
+ self.assertFalse(provider.include_row_kind())
+ provider = PreResolvedSplitProvider(
+ table, splits, read_type, include_row_kind=True)
+ self.assertTrue(provider.include_row_kind())
+
if __name__ == '__main__':
unittest.main()
