From cbc7d0e134dfbf361a397487f0255e3cbcfd02ab Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Tue, 8 Sep 2026 10:49:15 +0800 Subject: [PATCH] [python][ray] Fix count metadata for row-level authorization Signed-off-by: QuakeWang --- .../read/datasource/ray_datasource.py | 15 +++-- .../pypaimon/tests/ray_integration_test.py | 66 +++++++++++++++++++ 2 files changed, 74 insertions(+), 7 deletions(-) diff --git a/paimon-python/pypaimon/read/datasource/ray_datasource.py b/paimon-python/pypaimon/read/datasource/ray_datasource.py index 7a2f72614312..f8d5c7fafe81 100644 --- a/paimon-python/pypaimon/read/datasource/ray_datasource.py +++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py @@ -197,15 +197,16 @@ def _get_read_task( continue # Calculate metadata for this chunk - total_rows = 0 + total_rows: Optional[int] = 0 total_size = 0 for split in chunk_splits: - if predicate is None: - # Only estimate rows if no predicate (predicate filtering changes row count) - merged = split.merged_row_count() - row_count = merged if merged is not None else split.row_count - if row_count > 0: + if predicate is None and total_rows is not None: + row_count = split.merged_row_count() + if row_count is None: + # Physical rows cannot replace unknown counts after filtering or merging. + total_rows = None + else: total_rows += row_count if hasattr(split, 'file_size') and split.file_size > 0: total_size += split.file_size @@ -222,7 +223,7 @@ def _get_read_task( elif predicate is not None: num_rows = None # Can't estimate with predicate filtering else: - num_rows = total_rows if total_rows > 0 else None + num_rows = total_rows size_bytes = total_size if total_size > 0 else None metadata_kwargs = { diff --git a/paimon-python/pypaimon/tests/ray_integration_test.py b/paimon-python/pypaimon/tests/ray_integration_test.py index 875a9a3c9e04..df8c52ae81a5 100644 --- a/paimon-python/pypaimon/tests/ray_integration_test.py +++ b/paimon-python/pypaimon/tests/ray_integration_test.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import json import os import shutil import tempfile @@ -25,6 +26,8 @@ import ray from pypaimon import CatalogFactory, Schema +from pypaimon.catalog.catalog_environment import CatalogEnvironment +from pypaimon.catalog.table_query_auth import TableQueryAuthResult class RayIntegrationTest(unittest.TestCase): @@ -119,6 +122,69 @@ def test_read_paimon_with_projection(self): self.assertEqual(set(df.columns), {'id', 'name'}) self.assertEqual(len(df), 2) + def test_read_paimon_count_with_query_auth(self): + from pypaimon.ray import read_paimon + + identifier = self._create_and_populate_table( + 'test_count_query_auth', + pa.schema([('id', pa.int64()), ('dept', pa.string()), ('part', pa.string())]), + {'id': list(range(100)), + 'dept': ['eng' if i % 5 == 0 else 'sales' for i in range(100)], + 'part': ['a'] * 50 + ['b'] * 50}, + partition_keys=['part'], + ) + for dept, expected_ids in [('eng', list(range(0, 100, 5))), ('missing', [])]: + with self.subTest(dept=dept): + auth = TableQueryAuthResult([json.dumps({ + 'kind': 'LEAF', + 'transform': {'name': 'FIELD_REF', 'fieldRef': {'name': 'dept'}}, + 'function': 'EQUAL', + 'literals': [dept], + })], None) + # Replace only the catalog authorization response; filtering runs in Ray workers. + with patch.object(CatalogEnvironment, 'table_query_auth', + return_value=lambda select: auth): + ds = read_paimon(identifier, self.catalog_options, override_num_blocks=2) + # Count before materialization exercises the datasource metadata path. + self.assertEqual(ds.count(), len(expected_ids)) + self.assertEqual(sorted(row['id'] for row in ds.take_all()), expected_ids) + + def test_ray_datasource_logical_row_count_metadata(self): + from pypaimon.read.datasource.ray_datasource import RayDatasource + from pypaimon.read.datasource.split_provider import PreResolvedSplitProvider + from pypaimon.read.query_auth_split import QueryAuthSplit + from pypaimon.read.split import DataSplit + + identifier = self._create_and_populate_table( + 'test_count_metadata', pa.schema([('id', pa.int64()), ('part', pa.string())]), + {'id': [1, 2, 3, 4], 'part': ['a', 'a', 'b', 'b']}, partition_keys=['part'], + ) + table = CatalogFactory.create(self.catalog_options).get_table(identifier) + rb = table.new_read_builder() + splits = rb.new_scan().plan().splits() + self.assertEqual(len(splits), 2) + auth = TableQueryAuthResult([json.dumps({ + 'kind': 'LEAF', + 'transform': {'name': 'FIELD_REF', 'fieldRef': {'name': 'id'}}, + 'function': 'EQUAL', 'literals': [1], + })], None) + masked = TableQueryAuthResult(None, {'id': json.dumps({'name': 'NULL'})}) + empty = DataSplit([], splits[0].partition, splits[0].bucket, raw_convertible=True) + cases = [ + ('known', splits, 4), + ('masked', [QueryAuthSplit(s, masked) for s in splits], 4), + ('first_unknown', [QueryAuthSplit(splits[0], auth), splits[1]], None), + ('second_unknown', [splits[0], QueryAuthSplit(splits[1], auth)], None), + ('zero', [empty], 0), + ] + for name, task_splits, expected in cases: + with self.subTest(name=name): + datasource = RayDatasource(PreResolvedSplitProvider( + table, task_splits, rb.read_type())) + tasks = datasource.get_read_tasks(1) + self.assertEqual(len(tasks), 1) + self.assertEqual(tasks[0].metadata.num_rows, expected) + def test_read_paimon_with_nested_projection(self): """read_paimon() respects a nested-leaf projection.