fix(python): preserve native table reopen state

This commit is contained in:
Gatefixer
2026-08-06 08:14:08 +00:00
parent feccabd739
commit 36e44ab7a9
6 changed files with 232 additions and 22 deletions
+1 -1
View File
@@ -804,7 +804,7 @@ class LanceDBConnection(DBConnection):
"manifest_enabled": self._manifest_enabled,
"namespace_client_properties": self._namespace_client_properties,
"read_consistency_interval_seconds": (
rci.total_seconds() if rci else None
rci.total_seconds() if rci is not None else None
),
}
)
+7 -1
View File
@@ -493,7 +493,7 @@ class LanceNamespaceDBConnection(DBConnection):
"storage_options": self.storage_options or None,
"read_consistency_interval_seconds": (
self.read_consistency_interval.total_seconds()
if self.read_consistency_interval
if self.read_consistency_interval is not None
else None
),
}
@@ -569,6 +569,7 @@ class LanceNamespaceDBConnection(DBConnection):
self,
name,
namespace_path=namespace_path,
storage_options=storage_options,
namespace_client=self._namespace_client,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
@@ -607,6 +608,8 @@ class LanceNamespaceDBConnection(DBConnection):
self,
name,
namespace_path=namespace_path,
storage_options=storage_options,
index_cache_size=index_cache_size,
namespace_client=self._namespace_client,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
@@ -899,10 +902,13 @@ class LanceNamespaceDBConnection(DBConnection):
self,
name,
namespace_path=namespace_path,
storage_options=storage_options,
index_cache_size=index_cache_size,
location=table_uri,
namespace_client=namespace_client,
managed_versioning=managed_versioning,
pushdown_operations=self._namespace_client_pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
_async=async_table,
)
+147 -19
View File
@@ -164,7 +164,7 @@ def _maybe_add_fts_error_note(
if TYPE_CHECKING:
from .db import LanceDBConnection
from .db import DBConnection, LanceDBConnection
from ._lancedb import (
Table as LanceDBTable,
OptimizeStats,
@@ -2106,6 +2106,22 @@ class Table(ABC):
"""
@dataclass
class _LanceTableReopenState:
"""Process-independent coordinates for reopening a native table."""
connection_state: Optional[str]
can_reopen_after_fork: bool
name: str
namespace_path: List[str]
storage_options: Optional[Dict[str, str]]
index_cache_size: Optional[int]
location: Optional[str]
managed_versioning: Optional[bool]
branch: Optional[str]
checkout_version: Optional[int]
class LanceTable(Table):
"""
A table in a LanceDB database.
@@ -2143,6 +2159,7 @@ class LanceTable(Table):
self._storage_options = storage_options
self._index_cache_size = index_cache_size
self._location = location # Store location for use in _dataset_path
self._managed_versioning = managed_versioning
self._namespace_client = namespace_client
self._pushdown_operations = pushdown_operations or set()
# When the connection built the namespace client natively (e.g. an
@@ -2173,8 +2190,6 @@ class LanceTable(Table):
"""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
@@ -2182,19 +2197,126 @@ class LanceTable(Table):
# 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://")
connection_state: Optional[str] = self._conn.serialize()
can_reopen_after_fork = not self._conn.uri.startswith("memory://")
except Exception:
self._connection_state = None
self._can_reopen_after_fork = False
connection_state = None
can_reopen_after_fork = False
self._reopen_state = _LanceTableReopenState(
connection_state=connection_state,
can_reopen_after_fork=can_reopen_after_fork,
name=name,
namespace_path=list(self._namespace_path),
storage_options=(
dict(self._storage_options)
if self._storage_options is not None
else None
),
index_cache_size=self._index_cache_size,
location=self._location,
managed_versioning=self._managed_versioning,
branch=self._table.current_branch(),
checkout_version=None,
)
@property
def _connection_state(self) -> Optional[str]:
"""Serialized connection retained for worker reconstruction."""
return self._reopen_state.connection_state
@property
def _can_reopen_after_fork(self) -> bool:
return self._reopen_state.can_reopen_after_fork
@property
def _branch(self) -> Optional[str]:
state = getattr(self, "_reopen_state", None)
if state is not None:
return state.branch
return getattr(self, "_legacy_branch", None)
@_branch.setter
def _branch(self, value: Optional[str]) -> None:
state = getattr(self, "_reopen_state", None)
if state is not None:
state.branch = value
else:
self._legacy_branch = value
@property
def _checkout_version(self) -> Optional[int]:
state = getattr(self, "_reopen_state", None)
if state is not None:
return state.checkout_version
return getattr(self, "_legacy_checkout_version", None)
@_checkout_version.setter
def _checkout_version(self, value: Optional[int]) -> None:
state = getattr(self, "_reopen_state", None)
if state is not None:
state.checkout_version = value
else:
self._legacy_checkout_version = value
@classmethod
def _open_from_reopen_state(
cls,
connection: "DBConnection",
state: "_LanceTableReopenState",
) -> "LanceTable":
"""Open a table from its complete process-independent descriptor."""
async_connection = getattr(connection, "_conn", None)
if async_connection is None:
async_connection = connection._inner
namespace_client = getattr(connection, "_namespace_client", None)
async_table = LOOP.run(
async_connection.open_table(
state.name,
namespace_path=state.namespace_path,
storage_options=state.storage_options,
index_cache_size=state.index_cache_size,
location=state.location,
namespace_client=namespace_client,
managed_versioning=state.managed_versioning,
)
)
table = cls(
connection,
state.name,
namespace_path=state.namespace_path,
storage_options=state.storage_options,
index_cache_size=state.index_cache_size,
location=state.location,
namespace_client=namespace_client,
managed_versioning=state.managed_versioning,
pushdown_operations=getattr(
connection, "_namespace_client_pushdown_operations", None
),
route_pushdown_to_rust=getattr(
connection, "_route_pushdown_to_rust", False
),
_async=async_table,
)
if state.branch is not None:
table = table.branches.checkout(state.branch, state.checkout_version)
elif state.checkout_version is not None:
table.checkout(state.checkout_version)
return table
def _ensure_open(self) -> None:
"""Reopen native table handles inherited from another process."""
pid = os.getpid()
if self._pid == pid:
if getattr(self, "_pid", pid) == pid:
return
if not self._can_reopen_after_fork or self._connection_state is None:
state = getattr(self, "_reopen_state", None)
if (
state is None
or not state.can_reopen_after_fork
or state.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.
@@ -2203,14 +2325,10 @@ class LanceTable(Table):
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,
connection = deserialize_conn(state.connection_state, for_worker=True)
reopened = self._open_from_reopen_state(
connection,
state,
)
# Keep this Python object stable because user datasets commonly retain
@@ -2221,11 +2339,16 @@ class LanceTable(Table):
self._namespace_client = reopened._namespace_client
self._pushdown_operations = reopened._pushdown_operations
self._route_pushdown_to_rust = reopened._route_pushdown_to_rust
self._reopen_state = reopened._reopen_state
self._pid = pid
@property
def name(self) -> str:
return self._name
if hasattr(self, "_name"):
return self._name
# Preserve compatibility with lightweight / legacy instances that
# were constructed without running ``LanceTable.__init__``.
return self._table.name
@property
def namespace(self) -> List[str]:
@@ -2451,6 +2574,7 @@ class LanceTable(Table):
pushdown_operations=self._pushdown_operations,
route_pushdown_to_rust=self._route_pushdown_to_rust,
location=self._location,
managed_versioning=self._managed_versioning,
_async=async_table,
)
table._checkout_version = version
@@ -2495,7 +2619,7 @@ class LanceTable(Table):
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
self._checkout_version = version if isinstance(version, int) else self.version
def checkout_latest(self):
"""Checkout the latest version of the table. This is an in-place operation.
@@ -2552,6 +2676,9 @@ class LanceTable(Table):
"""
if version is not None:
LOOP.run(self._table.checkout(version))
self._checkout_version = (
version if isinstance(version, int) else self.version
)
LOOP.run(self._table.restore())
self._checkout_version = None
@@ -3666,6 +3793,7 @@ class LanceTable(Table):
self._namespace_path = namespace_path
self._index_cache_size = None
self._location = location
self._managed_versioning = None
self._namespace_client = namespace_client
self._pushdown_operations = pushdown_operations or set()
self._route_pushdown_to_rust = route_pushdown_to_rust
+12
View File
@@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import json
import re
import sys
from datetime import timedelta
@@ -94,6 +95,17 @@ def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkey
assert db_from_inner.read_consistency_interval == consistency_interval
def test_serialize_preserves_zero_read_consistency_interval(tmp_path):
db = lancedb.connect(tmp_path, read_consistency_interval=timedelta(0))
table = db.create_table("items", pa.table({"x": [1]}))
encoded = json.loads(table._connection_state)
assert encoded["read_consistency_interval_seconds"] == 0.0
restored = lancedb.deserialize_conn(table._connection_state)
assert restored.read_consistency_interval == timedelta(0)
def test_ingest_pd(tmp_path):
db = lancedb.connect(tmp_path)
+26 -1
View File
@@ -6,6 +6,7 @@
import tempfile
import shutil
import importlib
from datetime import timedelta
import pytest
import pyarrow as pa
import lancedb
@@ -419,7 +420,31 @@ class TestNamespaceConnection:
pa.field("vector", pa.list_(pa.float32(), 2)),
]
)
db.create_table("test_table", schema=schema, storage_options=table_opts)
created = db.create_table(
"test_table", schema=schema, storage_options=table_opts
)
assert created._storage_options == table_opts
opened = db.open_table(
"test_table",
storage_options={"allow_http": "true"},
index_cache_size=17,
)
assert opened._storage_options == {"allow_http": "true"}
assert opened._index_cache_size == 17
opened._pid = -1
opened._ensure_open()
assert opened.count_rows() == 0
def test_serialize_preserves_zero_read_consistency_interval(self):
db = lancedb.connect_namespace(
"dir",
{"root": self.temp_dir},
read_consistency_interval=timedelta(0),
)
restored = lancedb.deserialize_conn(db.serialize())
assert restored.read_consistency_interval == timedelta(0)
def test_namespace_operations(self):
"""Test namespace management operations."""
+39
View File
@@ -2059,6 +2059,45 @@ def test_restore(mem_db: DBConnection):
table.restore(0)
def test_restore_tracks_checkout_when_restore_fails():
class FailingRestore:
def __init__(self):
self.live_version = None
async def checkout(self, version):
self.live_version = version
async def restore(self):
raise RuntimeError("injected restore failure")
inner = FailingRestore()
table = LanceTable.__new__(LanceTable)
table._table = inner
table._checkout_version = None
with pytest.raises(RuntimeError, match="injected restore failure"):
table.restore(7)
assert table._checkout_version == inner.live_version
def test_reopen_preserves_explicit_table_location(tmp_path):
db = lancedb.connect(tmp_path / "db")
location = str(tmp_path / "physical-table")
table = LanceTable.create(
db,
"items",
pa.table({"x": [1]}),
location=location,
)
table._pid = -1
table._ensure_open()
assert table.count_rows() == 1
assert table._location == location
def test_restore_with_tags(mem_db: DBConnection):
table = mem_db.create_table(
"my_table",