fix(python): reopen native tables in forked workers

python/python/lancedb/permutation.py

@@ -591,6 +591,13 @@ class Permutation:
         then the first split will be used.
         """
         assert base_table is not None, "base_table is required"
+        # A PyTorch fork worker may construct its Permutation lazily from a
+        # table opened in the parent process. Reopen that table before the
+        # Rust reader clones its object-store clients and connection pools.
+        if hasattr(base_table, "_ensure_open"):
+            base_table._ensure_open()
+        if permutation_table is not None and hasattr(permutation_table, "_ensure_open"):
+            permutation_table._ensure_open()
         if split is not None:
             if permutation_table is None:
                 raise ValueError(

python/python/lancedb/table.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import asyncio
 import inspect
 import deprecation
+import os
 import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
@@ -2139,6 +2140,8 @@ class LanceTable(Table):
             namespace_path = []
         self._conn = connection
         self._namespace_path = namespace_path
+        self._storage_options = storage_options
+        self._index_cache_size = index_cache_size
         self._location = location  # Store location for use in _dataset_path
         self._namespace_client = namespace_client
         self._pushdown_operations = pushdown_operations or set()
@@ -2164,10 +2167,65 @@ class LanceTable(Table):
                     managed_versioning=managed_versioning,
                 )
             )
+        self._initialize_reopen_state(name)
+
+    def _initialize_reopen_state(self, name: str) -> None:
+        """Capture the state needed to replace inherited native handles."""
+        self._name = name
+        self._pid = os.getpid()
+        self._branch = self._table.current_branch()
+        self._checkout_version: Optional[int] = None
+
+        # A native table owns object-store clients and connection pools. Those
+        # handles must not be used after fork, so retain a process-independent
+        # connection description while it is still safe to inspect the parent
+        # connection. Some tables constructed from a bare Rust handle cannot
+        # be serialized; those retain the historical best-effort behavior.
+        try:
+            self._connection_state: Optional[str] = self._conn.serialize()
+            self._can_reopen_after_fork = not self._conn.uri.startswith("memory://")
+        except Exception:
+            self._connection_state = None
+            self._can_reopen_after_fork = False
+
+    def _ensure_open(self) -> None:
+        """Reopen native table handles inherited from another process."""
+        pid = os.getpid()
+        if self._pid == pid:
+            return
+
+        if not self._can_reopen_after_fork or self._connection_state is None:
+            # In-memory and opaque Rust-only connections cannot be recreated
+            # from connection metadata. Their local handles retain the prior
+            # best-effort fork behavior.
+            self._pid = pid
+            return
+
+        from lancedb import deserialize_conn
+
+        connection = deserialize_conn(self._connection_state, for_worker=True)
+        reopened = connection.open_table(
+            self._name,
+            namespace_path=self._namespace_path or None,
+            storage_options=self._storage_options,
+            index_cache_size=self._index_cache_size,
+            branch=self._branch,
+            version=self._checkout_version,
+        )
+
+        # Keep this Python object stable because user datasets commonly retain
+        # it across fork. Replace every process-bound component with the fresh
+        # child's equivalent.
+        self._conn = reopened._conn
+        self._table = reopened._table
+        self._namespace_client = reopened._namespace_client
+        self._pushdown_operations = reopened._pushdown_operations
+        self._route_pushdown_to_rust = reopened._route_pushdown_to_rust
+        self._pid = pid
 
     @property
     def name(self) -> str:
-        return self._table.name
+        return self._name
 
     @property
     def namespace(self) -> List[str]:
@@ -2379,18 +2437,20 @@ class LanceTable(Table):
     def _wrap_branch_handle(
         self, async_table: "AsyncTable", version: Optional[int] = None
     ) -> "LanceTable":
-        # version is unused locally: the pin already lives on async_table and a
-        # local handle is not reopened via a serialized connection.
-        return LanceTable(
+        table = LanceTable(
             self._conn,
             async_table.name,
             namespace_path=self._namespace_path,
+            storage_options=self._storage_options,
+            index_cache_size=self._index_cache_size,
             namespace_client=self._namespace_client,
             pushdown_operations=self._pushdown_operations,
             route_pushdown_to_rust=self._route_pushdown_to_rust,
             location=self._location,
             _async=async_table,
         )
+        table._checkout_version = version
+        return table
 
     def checkout(self, version: Union[int, str]):
         """Checkout a version of the table. This is an in-place operation.
@@ -2429,6 +2489,9 @@ class LanceTable(Table):
         0  [1.1, 0.9]  vector
         """
         LOOP.run(self._table.checkout(version))
+        # Resolve tags to their numeric version so a forked child can reopen
+        # the same pinned view through ``open_table(version=...)``.
+        self._checkout_version = self.version
 
     def checkout_latest(self):
         """Checkout the latest version of the table. This is an in-place operation.
@@ -2437,6 +2500,7 @@ class LanceTable(Table):
         version of the table.
         """
         LOOP.run(self._table.checkout_latest())
+        self._checkout_version = None
 
     def restore(self, version: Optional[Union[int, str]] = None):
         """Restore a version of the table. This is an in-place operation.
@@ -2485,6 +2549,7 @@ class LanceTable(Table):
         if version is not None:
             LOOP.run(self._table.checkout(version))
         LOOP.run(self._table.restore())
+        self._checkout_version = None
 
     def count_rows(self, filter: Optional[str] = None) -> int:
         return LOOP.run(self._table.count_rows(filter))
@@ -3595,6 +3660,7 @@ class LanceTable(Table):
         self = cls.__new__(cls)
         self._conn = db
         self._namespace_path = namespace_path
+        self._index_cache_size = None
         self._location = location
         self._namespace_client = namespace_client
         self._pushdown_operations = pushdown_operations or set()
@@ -3623,6 +3689,7 @@ class LanceTable(Table):
                 enable_v2_manifest_paths
             )
 
+        self._storage_options = storage_options
         self._table = LOOP.run(
             self._conn._conn.create_table(
                 name,
@@ -3639,6 +3706,7 @@ class LanceTable(Table):
                 namespace_client=namespace_client,
             )
         )
+        self._initialize_reopen_state(name)
         return self
 
     def delete(self, where: Union[str, Expr]) -> DeleteResult:

python/python/tests/test_torch.py

@@ -342,6 +342,42 @@ def _multiworker_dataloader_target(db_uri: str, result_queue):
     result_queue.put(count)
 
 
+class _LazyPermutationDataset(torch.utils.data.Dataset):
+    """Match applications that create their Permutation inside a fork worker."""
+
+    def __init__(self, table):
+        self._table = table
+        self._permutation = None
+        self._length = table.count_rows()
+
+    def __len__(self):
+        return self._length
+
+    def __getitems__(self, indices):
+        if self._permutation is None:
+            inherited_connection = self._table._conn
+            self._permutation = Permutation.identity(self._table)
+            if self._table._conn is inherited_connection:
+                raise RuntimeError("Permutation reused a connection inherited by fork")
+        return self._permutation.__getitems__(indices)
+
+
+def _lazy_multiworker_dataloader_target(db_uri: str, result_queue):
+    table = lancedb.connect(db_uri).open_table("test_table")
+    dataset = _LazyPermutationDataset(table)
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=10,
+        num_workers=2,
+        multiprocessing_context="fork",
+    )
+    count = 0
+    for batch in dataloader:
+        assert batch["a"].size(0) == 10
+        count += 1
+    result_queue.put(count)
+
+
 def _remote_multiworker_dataloader_target(port: int, result_queue):
     import lancedb
     from lancedb.permutation import Permutation
@@ -410,6 +446,46 @@ def test_permutation_dataloader_fork_workers(tmp_path):
     assert queue.get() == 100
 
 
+@pytest.mark.skipif(
+    sys.platform != "linux",
+    reason=(
+        "fork() is unavailable on Windows and unsafe on macOS "
+        "(Apple frameworks/TLS are not fork-safe)"
+    ),
+)
+def test_lazy_permutation_reopens_inherited_table_in_fork_worker(tmp_path):
+    """A lazily built Permutation must not reuse an inherited table client.
+
+    Object-store table handles contain HTTP connection pools that are unsafe
+    after fork. The local table makes the handle replacement deterministic
+    without requiring an S3 service in the unit-test environment.
+    """
+    db_uri = str(tmp_path / "db")
+    db = lancedb.connect(db_uri)
+    db.create_table("test_table", pa.table({"a": list(range(1000))}))
+
+    ctx = mp.get_context("spawn")
+    queue = ctx.Queue()
+    proc = ctx.Process(
+        target=_lazy_multiworker_dataloader_target,
+        args=(db_uri, queue),
+    )
+    proc.start()
+    proc.join(timeout=30)
+
+    if proc.is_alive():
+        proc.terminate()
+        proc.join(timeout=5)
+        if proc.is_alive():
+            proc.kill()
+            proc.join()
+        pytest.fail("Lazy Permutation hung in a fork-based DataLoader worker")
+
+    assert proc.exitcode == 0, f"child exited with code {proc.exitcode}"
+    assert not queue.empty(), "child produced no batches"
+    assert queue.get() == 100
+
+
 @pytest.mark.skipif(
     sys.platform != "linux",
     reason=(