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 e1d1f60c72 [python] Reuse REST table data token for native reader
initialization (#10436)
e1d1f60c72 is described below
commit e1d1f60c7243bfc4263111000cdfbfc783d71ab6
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Oct 8 21:15:51 2026 +0800
[python] Reuse REST table data token for native reader initialization
(#10436)
---
.../pypaimon/catalog/rest/rest_token_file_io.py | 5 ++
paimon-python/pypaimon/read/native_plan.py | 33 +++++++-
.../pypaimon/tests/native_plan_rest_test.py | 58 +++++++++++++
paimon-python/pypaimon/tests/native_plan_test.py | 98 ++++++++++++++++++++++
.../pypaimon/tests/rest/rest_token_file_io_test.py | 16 ++++
5 files changed, 206 insertions(+), 4 deletions(-)
diff --git a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py
b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py
index 2931712851..861e3d3bb1 100644
--- a/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py
+++ b/paimon-python/pypaimon/catalog/rest/rest_token_file_io.py
@@ -326,5 +326,10 @@ class RESTTokenFileIO(FileIO):
self.try_to_refresh_token()
return self.token
+ def _existing_valid_token(self) -> Optional[RESTToken]:
+ """Return this instance's valid token without locks or refresh
requests."""
+ token = self.token
+ return token if not self._is_token_expired(token) else None
+
def close(self):
pass
diff --git a/paimon-python/pypaimon/read/native_plan.py
b/paimon-python/pypaimon/read/native_plan.py
index bd41fb6185..58a6149b22 100644
--- a/paimon-python/pypaimon/read/native_plan.py
+++ b/paimon-python/pypaimon/read/native_plan.py
@@ -287,7 +287,7 @@ class _NativeRestTableCache:
def __setstate__(self, state):
self.__init__()
- def get(self, response, database, table, options):
+ def get(self, response, database, table, options, token_loader=None):
from pypaimon_rust.datafusion import Table
pid = os.getpid()
@@ -300,12 +300,36 @@ class _NativeRestTableCache:
key = (response, database, table, tuple(sorted(options.items())))
with state.lock:
if state.entry is None or state.entry[0] != key:
- native_table = Table.from_rest_response(
- response, database=database, table=table,
rest_options=options)
+ reuse_token = getattr(Table, 'from_rest_response_with_token',
None)
+ token = token_loader() if callable(reuse_token) and
token_loader else None
+ if token is None:
+ native_table = Table.from_rest_response(
+ response, database=database, table=table,
rest_options=options)
+ else:
+ native_table = reuse_token(
+ response, database=database, table=table,
rest_options=options,
+ data_token=dict(token.token),
+ expires_at_millis=token.expire_at_millis)
state.entry = (key, native_table)
return state.entry[1]
+def _rest_data_token(table):
+ """Reuse only an existing valid token bound to this exact table."""
+ from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO
+ from pypaimon.filesystem.caching_file_io import CachingFileIO
+
+ file_io = table.file_io
+ if type(file_io) is CachingFileIO:
+ file_io = file_io._delegate
+ if (table.current_branch() != 'main'
+ or type(file_io) is not RESTTokenFileIO
+ or file_io.identifier != table.identifier
+ or file_io.path != table.table_path):
+ return None
+ return file_io._existing_valid_token()
+
+
def _native_read_builder(table):
"""Return a fresh builder with the current schema and shared REST
FileIO."""
rest_response = _resolved_rest_table_response(table)
@@ -317,7 +341,8 @@ def _native_read_builder(table):
rest_response,
database=table.identifier.get_database_name(),
table=table.identifier.get_object_name(),
- options=_catalog_options(table))
+ options=_catalog_options(table),
+ token_loader=lambda: _rest_data_token(table))
rt = rt.copy_with_resolved_schema(_resolved_schema_json(table),
branch=table.current_branch())
elif file_io_options is not None:
from pypaimon_rust.datafusion import Table
diff --git a/paimon-python/pypaimon/tests/native_plan_rest_test.py
b/paimon-python/pypaimon/tests/native_plan_rest_test.py
index 320bfd61bc..ef3294c2f3 100644
--- a/paimon-python/pypaimon/tests/native_plan_rest_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_rest_test.py
@@ -149,6 +149,64 @@ def
test_resolved_rest_table_keeps_refreshable_file_io(rest_source, rest_catalog
refresh.assert_called()
[email protected](not native_method_available('Table',
'from_rest_response_with_token'),
+ reason='REST data token reuse binding required')
[email protected]('local_cache', [False, True])
+def test_native_rest_reuses_python_data_token(
+ rest_source, rest_catalog, tmp_path, local_cache):
+ import time
+
+ from pypaimon.catalog.rest.rest_token import RESTToken
+ from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO
+ from pypaimon.filesystem.caching_file_io import CachingFileIO
+ from pypaimon.read.native_plan import _catalog_options,
_resolved_rest_table_response
+ from pypaimon_rust.datafusion import Table as NativeTable
+
+ source, server, _ = rest_source
+ catalog, _ = rest_catalog
+ options = dict(catalog.context.options.to_map())
+ options['data-token.enabled'] = 'true'
+ if local_cache:
+ options.update({'local-cache.enabled': 'true',
+ 'local-cache.dir': str(tmp_path / 'cache')})
+ server.set_table_token(
+ source.identifier, RESTToken({}, int(time.time() * 1000) + 7_200_000))
+ table = CatalogFactory.create(options).get_table(source.identifier)
+ reused_table = CatalogFactory.create(options).get_table(source.identifier)
+ uncached_table =
CatalogFactory.create(options).get_table(source.identifier)
+
+ def token_file_io(table):
+ file_io = table.file_io
+ assert (type(file_io) is CachingFileIO) == local_cache
+ if type(file_io) is CachingFileIO:
+ file_io = file_io._delegate
+ assert type(file_io) is RESTTokenFileIO
+ return file_io
+
+ # The old bridge fetched a Python token and an independent Rust token.
+ with patch.object(RESTTokenFileIO, '_TOKEN_CACHE', {}):
+ with patch.object(server, '_table_token_handle',
+ wraps=server._table_token_handle) as load:
+ token_file_io(table).valid_token()
+ baseline = NativeTable.from_rest_response(
+ _resolved_rest_table_response(table), database='default',
table='t',
+ rest_options=_catalog_options(table))
+ assert baseline.new_read_builder().new_scan().plan().snapshot_id()
== 2
+ assert load.call_count == 2
+ load.reset_mock()
+ RESTTokenFileIO._TOKEN_CACHE.clear()
+
+ token_file_io(reused_table).token = token_file_io(table).token
+
+ for _ in range(2):
+ assert
reused_table.new_read_builder().new_scan().plan().snapshot_id == 2
+ assert load.call_count == 0
+
+ # No instance token: Rust obtains one instead of making Python
refresh.
+ assert
uncached_table.new_read_builder().new_scan().plan().snapshot_id == 2
+ assert load.call_count == 1
+
+
@pytest.mark.parametrize('branch', [None, 'dev'])
def test_rest_dotted_database_and_table_keep_identity(rest_catalog, branch):
from pypaimon.common.identifier import Identifier
diff --git a/paimon-python/pypaimon/tests/native_plan_test.py
b/paimon-python/pypaimon/tests/native_plan_test.py
index 63075bb1b7..19bc268b70 100644
--- a/paimon-python/pypaimon/tests/native_plan_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_test.py
@@ -1033,6 +1033,104 @@ class NativePlanTest(unittest.TestCase):
release.set()
holder.join(5)
+ def test_native_rest_token_lookup_after_fork_avoids_inherited_lock(self):
+ import multiprocessing
+ from threading import Event, Thread
+
+ from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO
+ from pypaimon.common.identifier import Identifier
+ from pypaimon.read.native_plan import _NativeRestTableCache,
_rest_data_token
+
+ if 'fork' not in multiprocessing.get_all_start_methods():
+ self.skipTest('fork required')
+ context = multiprocessing.get_context('fork')
+ identifier = Identifier('default', 't')
+ path = 'oss://bucket/table'
+ file_io = RESTTokenFileIO(identifier, path)
+ table = SimpleNamespace(file_io=file_io, identifier=identifier,
+ table_path=path, current_branch=lambda: 'main')
+ fake_df = ModuleType('pypaimon_rust.datafusion')
+ fake_df.Table = Mock()
+ fake_df.Table.from_rest_response.return_value = object()
+ cache = _NativeRestTableCache()
+ held, release = Event(), Event()
+
+ def hold_token_lock():
+ with RESTTokenFileIO._TOKEN_LOCKS_LOCK:
+ held.set()
+ release.wait()
+
+ def child(connection):
+ try:
+ native_table = cache.get(
+ 'response', 'default', 't', {},
+ token_loader=lambda: _rest_data_token(table))
+ connection.send(('ok', native_table is not None))
+ except BaseException as exc:
+ connection.send(('error', repr(exc)))
+ finally:
+ connection.close()
+
+ with patch.dict(sys.modules, {'pypaimon_rust.datafusion': fake_df}), \
+ patch.object(file_io, '_build_cache_key',
return_value='token-key'):
+ holder = Thread(target=hold_token_lock, daemon=True)
+ holder.start()
+ self.assertTrue(held.wait(5))
+ receiving, sending = context.Pipe(duplex=False)
+ process = context.Process(target=child, args=(sending,))
+ try:
+ process.start()
+ sending.close()
+ self.assertTrue(receiving.poll(3), 'child deadlocked on
inherited token lock')
+ self.assertEqual(receiving.recv(), ('ok', True))
+ process.join(5)
+ self.assertEqual(process.exitcode, 0)
+ finally:
+ if process.is_alive():
+ process.terminate()
+ process.join(5)
+ receiving.close()
+ sending.close()
+ release.set()
+ holder.join(5)
+
+ def test_native_rest_cache_passes_only_existing_token(self):
+ import time
+
+ from pypaimon.catalog.rest.rest_token import RESTToken
+ from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO
+ from pypaimon.common.identifier import Identifier
+ from pypaimon.read.native_plan import _NativeRestTableCache,
_rest_data_token
+
+ identifier = Identifier('default', 't')
+ path = 'oss://bucket/table'
+ file_io = RESTTokenFileIO(identifier, path)
+ token = RESTToken({'key': 'value'}, int(time.time() * 1000) +
7_200_000)
+ file_io.token = token
+ table = SimpleNamespace(file_io=file_io, identifier=identifier,
+ table_path=path, current_branch=lambda: 'main')
+ fake_df = ModuleType('pypaimon_rust.datafusion')
+ fake_df.Table = Mock()
+ fake_df.Table.from_rest_response_with_token.return_value = object()
+
+ with patch.dict(sys.modules, {'pypaimon_rust.datafusion': fake_df}), \
+ patch.object(file_io, 'valid_token',
+ side_effect=AssertionError('must not refresh')):
+ cache = _NativeRestTableCache()
+ cache.get('response', 'default', 't', {},
+ token_loader=lambda: _rest_data_token(table))
+
fake_df.Table.from_rest_response_with_token.assert_called_once_with(
+ 'response', database='default', table='t', rest_options={},
+ data_token=token.token,
expires_at_millis=token.expire_at_millis)
+ fake_df.Table.from_rest_response.assert_not_called()
+
+ file_io.token = None
+ cache = _NativeRestTableCache()
+ cache.get('response', 'default', 't', {},
+ token_loader=lambda: _rest_data_token(table))
+ fake_df.Table.from_rest_response.assert_called_once_with(
+ 'response', database='default', table='t', rest_options={})
+
def test_native_rest_cache_retries_failed_construction(self):
from pypaimon.read.native_plan import _NativeRestTableCache
diff --git a/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py
b/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py
index 4e94b21d72..8255a78ea9 100644
--- a/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py
+++ b/paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py
@@ -49,6 +49,22 @@ class RESTTokenFileIOTest(unittest.TestCase):
if os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
+ def test_existing_valid_token_does_not_refresh_or_read_shared_cache(self):
+ file_io = RESTTokenFileIO(self.identifier, self.warehouse_path)
+ now = 1700000000
+ valid = RESTToken({'key': 'value'}, (now + 7200) * 1000)
+ expiring = RESTToken({'key': 'value'}, (now + 1800) * 1000)
+ with patch('pypaimon.catalog.rest.rest_token_file_io.time.time',
return_value=now), \
+ patch.object(file_io, 'try_to_refresh_token') as refresh, \
+ patch.object(file_io, '_get_cached_token') as shared_cache:
+ self.assertIsNone(file_io._existing_valid_token())
+ file_io.token = valid
+ self.assertIs(file_io._existing_valid_token(), valid)
+ file_io.token = expiring
+ self.assertIsNone(file_io._existing_valid_token())
+ refresh.assert_not_called()
+ shared_cache.assert_not_called()
+
def test_blob_presigned_url_bound_table_root(self):
root = "oss://bucket/table-a"
file_io = RESTTokenFileIO(self.identifier, root, self.catalog_options)
