[python][ray] Preserve row kind in TableRead.to_ray()

paimon-python/pypaimon/read/datasource/ray_datasource.py

@@ -127,11 +127,15 @@ def get_read_tasks(self, parallelism: int, **kwargs) -> List:
         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 @@ def _get_read_task(
                 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 @@ def _get_read_task(
             # 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 @@ def _get_read_task(
             schema=schema,
             limit=limit,
             nested_name_paths=nested_name_paths,
+            include_row_kind=include_row_kind,
         )
 
         read_tasks = []

paimon-python/pypaimon/read/datasource/split_provider.py

@@ -67,6 +67,10 @@ def limit(self) -> Optional[int]:
         """
         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 @@ def read_type(self):
     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
 

paimon-python/pypaimon/read/table_read.py

@@ -666,6 +666,8 @@ def to_ray(
 
         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 @@ def to_ray(
                 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(

paimon-python/pypaimon/tests/ray_integration_test.py

@@ -103,6 +103,60 @@ def test_read_paimon_basic(self):
         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

paimon-python/pypaimon/tests/split_provider_test.py

@@ -86,6 +86,7 @@ def test_catalog_provider_resolves_table_and_splits(self):
         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 @@ def test_pre_resolved_provider_returns_inputs(self):
         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()