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 eaf5958150 [python][ray] Fix count metadata for row-level
authorization (#9681)
eaf5958150 is described below
commit eaf5958150a46e6094c9a2f38a980f529d1aa0f4
Author: QuakeWang <[email protected]>
AuthorDate: Tue Sep 8 22:10:48 2026 +0800
[python][ray] Fix count metadata for row-level authorization (#9681)
---
.../pypaimon/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 7a2f726143..f8d5c7fafe 100644
--- a/paimon-python/pypaimon/read/datasource/ray_datasource.py
+++ b/paimon-python/pypaimon/read/datasource/ray_datasource.py
@@ -197,15 +197,16 @@ class RayDatasource(Datasource):
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 @@ class RayDatasource(Datasource):
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 875a9a3c9e..df8c52ae81 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 pyarrow as pa
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 @@ class RayIntegrationTest(unittest.TestCase):
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.
