From 6c08f819e6154bf0d1fb4776f7b4f08a0f7d42d3 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Wed, 9 Sep 2026 00:11:15 +0800 Subject: [PATCH] [python][ray] Preserve row kind in TableRead.to_ray() TableRead.to_ray() dropped include_row_kind when rebuilding worker readers, causing Ray output to omit _row_kind. Forward the option through split providers and keep read-task and empty-result schemas consistent with Arrow output. Signed-off-by: QuakeWang --- .../read/datasource/ray_datasource.py | 9 +++- .../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 f8d5c7fafe81..4b7e9459da43 100644 --- a/paimon-python/pypaimon/read/datasource/ray_datasource.py +++ b/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 = [] diff --git a/paimon-python/pypaimon/read/datasource/split_provider.py b/paimon-python/pypaimon/read/datasource/split_provider.py index 6060ccce80aa..224a74fe54cc 100644 --- a/paimon-python/pypaimon/read/datasource/split_provider.py +++ b/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 diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 3472666ae5ef..51465486ed7d 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/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( diff --git a/paimon-python/pypaimon/tests/ray_integration_test.py b/paimon-python/pypaimon/tests/ray_integration_test.py index df8c52ae81a5..05aff4bc63f6 100644 --- a/paimon-python/pypaimon/tests/ray_integration_test.py +++ b/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 diff --git a/paimon-python/pypaimon/tests/split_provider_test.py b/paimon-python/pypaimon/tests/split_provider_test.py index 61e14005f5fc..7602a9978e27 100644 --- a/paimon-python/pypaimon/tests/split_provider_test.py +++ b/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()