[python] Reuse REST table data token for native reader initialization

paimon-python/pypaimon/catalog/rest/rest_token_file_io.py

@@ -326,5 +326,10 @@ def valid_token(self):
         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

paimon-python/pypaimon/read/native_plan.py

@@ -309,7 +309,7 @@ def __getstate__(self):
     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()
@@ -322,12 +322,36 @@ def get(self, response, database, table, options):
         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)
@@ -339,7 +363,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

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()
 
 
+@pytest.mark.skipif(not native_method_available('Table', 'from_rest_response_with_token'),
+                    reason='REST data token reuse binding required')
+@pytest.mark.parametrize('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

paimon-python/pypaimon/tests/native_plan_test.py

@@ -1115,6 +1115,104 @@ def access():
                 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
 

paimon-python/pypaimon/tests/rest/rest_token_file_io_test.py

@@ -49,6 +49,22 @@ def tearDown(self):
         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)