diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index eeae8bf50..fa0f91fea 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -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 ), } ) diff --git a/python/python/lancedb/namespace.py b/python/python/lancedb/namespace.py index b151395cc..e7821c763 100644 --- a/python/python/lancedb/namespace.py +++ b/python/python/lancedb/namespace.py @@ -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, ) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 5efe5c8b7..faafe641a 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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 diff --git a/python/python/tests/test_db.py b/python/python/tests/test_db.py index 93b791650..91e5f5c6e 100644 --- a/python/python/tests/test_db.py +++ b/python/python/tests/test_db.py @@ -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) diff --git a/python/python/tests/test_namespace.py b/python/python/tests/test_namespace.py index f8cbfe92c..a5ace932f 100644 --- a/python/python/tests/test_namespace.py +++ b/python/python/tests/test_namespace.py @@ -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.""" diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 069527b21..14f703170 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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",