rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-12 14:52:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(python): stop 2-phase reads from aliasing the caller's table handle

Browse Source 
_open_pinned_table() reopened a private handle via
_table_from_pickle_state(_table_to_pickle_state(self._table)) so
checking it out to the block plan's pinned version would never mutate
self._table. That held for local (on-disk) tables, whose pickle state
is just reopen coordinates, but not for the other two table kinds:

- RemoteTable: the pickle-state helpers special-case remote by
  returning the *same* object (correct for their original purpose --
  cross-process pickling, where the unpickled copy necessarily lives
  in a different process -- but wrong for same-process reuse here).
  checkout() is a real in-place mutation on RemoteTable, so every
  2-phase read pinned (and, since Cloud rejects writes through a
  version-pinned handle, effectively made read-only) the caller's own
  live table object.

- memory://: the pickle state *is* the row data (to_arrow()), captured
  fresh on every _open_pinned_table() call instead of once at
  construction, so a write made after construction leaked into every
  later read -- the opposite of what a version pin is supposed to
  guarantee. A table already past version 1 also failed construction
  outright, since the rebuilt table is always a fresh version-1 table.

Fix: deep-copy the RemoteTable instead (round-trips through its own
__getstate__/__setstate__, which carries only reopen coordinates and
clears the live connection handle, so the copy reconnects
independently on first use). For local/memory tables, capture
_pinned_table_state once at construction instead of recomputing it
per call, and skip checkout() for memory tables -- a reopened memory
table already is exactly the pinned snapshot by construction, and
checking out a version number on it would fail since it only has one.

Surfaced in the process: a reopened memory table's to_lance() always
failed (pylance's raw lance.dataset() open can't resolve an isolated
in-memory store), which forced every memory-table block read through
the slower generic query builder with no error -- previously a
gatekeeper-acknowledged non-blocking limitation, since it failed loudly
rather than returning wrong data. _pinned_lance_dataset now recognizes
a memory:// table up front (not only RemoteTable's NotImplementedError)
and skips the native fast path for it, the same way it already does for
RemoteTable, so memory-table block reads now also work correctly
instead of merely failing safely.

Added 3 regression tests: the caller's RemoteTable handle must stay
unpinned after iterating a 2-phase dataset over a mocked remote table;
a memory table's pin must survive a write made after construction; and
constructing over a memory table already past version 1 must not
raise. Full suite: 210 passed.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
jasonz-lance
2026-10-08 17:49:54 -07:00
co-authored by Claude Sonnet 5
parent 046987bfa1
commit 40ff1d52a6
2 changed files with 139 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/streaming.py
+81 -15
View File
@@ -914,6 +914,22 @@ class StreamingDataset(IterableDataset):
"first, or construct StreamingDataset without block_size"
)
self._pinned_version = table.version
# Captured once, now, rather than recomputed on every
# _open_pinned_table() call: for a memory table, the "pickle
# state" below IS the row data (to_arrow()), not just reopen
# coordinates, so recomputing it later would silently capture
# whatever the live table looks like *then* instead of the
# snapshot this pin is supposed to mean. Not used for a
# RemoteTable -- see _open_pinned_table, which deep-copies it
# directly instead.
from .remote.table import RemoteTable as _RemoteTable
self._pinned_table_state: Optional[dict] = (
_table_to_pickle_state(table)
if self._connection_factory is None
and not isinstance(table, _RemoteTable)
else None
)
block_live_offsets: Optional[list[np.ndarray]] = None
self._all_columns: Optional[list[str]] = None
@@ -1058,37 +1074,87 @@ class StreamingDataset(IterableDataset):
def _open_pinned_table(self):
"""Open a private table handle checked out to this dataset's pin.
A fresh handle (via ``connection_factory`` if one was supplied,
otherwise the same reopen-by-name state ``__setstate__`` uses for
worker reconnects) so checking it out to the pinned version never
A fresh handle so checking it out to the pinned version never
mutates ``self._table``, which the caller may still be using
elsewhere. Goes through ``Table.checkout``/``take_offsets``/
``search`` -- the same backend-agnostic interface ``RemoteTable``
implements -- instead of pylance's ``to_lance()``, which only
works against a local table.
Three cases, each independent of ``self._table`` by construction:
- ``connection_factory`` set: call it fresh every time (the
documented way to get an unrelated handle; e.g. a worker
reconnect).
- A ``RemoteTable``: ``_table_to_pickle_state``/
``_table_from_pickle_state`` (used below for the other two cases)
special-case "remote" by returning ``self._table`` itself --
correct for their original purpose (cross-process pickling, where
the unpickled copy necessarily lives in a different process), but
wrong here, since this runs in the *same* process: calling
``checkout()`` on that result would pin (and, since LanceDB Cloud
rejects writes through a version-pinned handle, effectively make
read-only) the caller's own table object. ``deepcopy`` instead
round-trips through ``RemoteTable.__getstate__``/``__setstate__``,
which carries only reopen coordinates (connection state, branch,
checkout version) and clears the live connection handle, so the
copy reconnects independently on first real use.
- Local or memory: ``self._pinned_table_state`` (captured once at
construction -- see ``__init__``) is replayed here instead of
recomputing ``_table_to_pickle_state(self._table)`` fresh each
call. For a local (on-disk) table this would barely matter (the
state is just reopen coordinates, not data), but for a memory
table the "pickle state" the helper returns *is* the row data
(``to_arrow()``); recomputing it on every call would silently
capture whatever the live table looks like at that later moment
instead of the snapshot construction actually pinned. A reopened
memory table is always a fresh version-1 table, so there is
nothing meaningful to additionally ``checkout()`` -- it already
is exactly the pinned snapshot by construction.
"""
if self._connection_factory is not None:
pinned = self._connection_factory(self._table.name)
else:
pinned = _table_from_pickle_state(_table_to_pickle_state(self._table))
pinned.checkout(self._pinned_version)
pinned.checkout(self._pinned_version)
return pinned
if self._pinned_table_state is None:
# Only a RemoteTable skips capturing _pinned_table_state (see
# __init__) -- anything else reaching here is a construction
# bug, not a runtime condition callers can hit.
pinned = deepcopy(self._table)
pinned.checkout(self._pinned_version)
return pinned
pinned = _table_from_pickle_state(self._pinned_table_state)
if self._pinned_table_state["kind"] != "memory":
pinned.checkout(self._pinned_version)
return pinned
def _pinned_lance_dataset(self, pinned_table):
"""Native pylance handle for ``pinned_table``, if it's a local table.
"""Native pylance handle for ``pinned_table``, if it's an on-disk table.
``to_lance()`` only works against a local table (it raises
``NotImplementedError`` on ``RemoteTable``); cached per dataset
instance (keyed by nothing -- there's only ever one pin per
``to_lance()`` raises ``NotImplementedError`` on ``RemoteTable``, and
-- for a different reason -- also fails on a reopened ``memory://``
table: ``_open_pinned_table`` rebuilds it as an isolated store (see
its docstring), which pylance's raw ``lance.dataset()`` open (what
``to_lance()`` uses) cannot resolve; only the table's own connection
handle can see its data. Both cases are excluded here up front
rather than only via the ``except``, since the memory failure is a
``ValueError`` indistinguishable from a real error. Cached per
dataset instance (keyed by nothing -- there's only ever one pin per
instance) so the conversion runs at most once rather than once per
block read. ``None`` means ``pinned_table`` is remote and callers
must use the backend-agnostic query builder instead.
block read. ``None`` means callers must use the backend-agnostic
query builder instead.
"""
if not hasattr(self, "_pinned_lance_ds_cache"):
try:
self._pinned_lance_ds_cache = pinned_table.to_lance()
except NotImplementedError:
conn_uri = getattr(getattr(pinned_table, "_conn", None), "uri", "")
if conn_uri.startswith("memory://"):
self._pinned_lance_ds_cache = None
else:
try:
self._pinned_lance_ds_cache = pinned_table.to_lance()
except NotImplementedError:
self._pinned_lance_ds_cache = None
return self._pinned_lance_ds_cache
def _read_block(
python/python/tests/test_elastic_dataloader.py
+58
View File
@@ -3233,6 +3233,29 @@ def test_two_phase_over_remote_table_filtered_and_shuffled():
assert sorted(ids) == live_ids, "only live rows, each exactly once"
def test_two_phase_over_remote_table_does_not_pin_callers_handle():
"""_open_pinned_table() must not mutate the caller's own RemoteTable.
Without a connection_factory, the reopen helpers used for local/memory
tables return the *same* RemoteTable object for the remote case (correct
for their original cross-process pickling purpose, wrong for same-process
reuse here). Checking that object out to the block-mode pin would
therefore pin -- and, since checkout_version != None marks a handle
version-pinned -- effectively make read-only the caller's own live table,
not an unrelated copy.
"""
server = MockPermutationServer(num_rows=8)
with mock_remote_table(server) as table:
assert table._checkout_version is None, "caller's table starts unpinned"
ds = StreamingDataset(table, num_splits=1, block_size=4, shuffle=False)
list(ds) # drive at least one _open_pinned_table() call
assert table._checkout_version is None, (
"iterating a 2-phase dataset must not pin the caller's own "
"table handle to a fixed version"
)
# ---------------------------------------------------------------------------
# 2-phase shuffled reads (block_size)
# ---------------------------------------------------------------------------
@@ -3418,6 +3441,41 @@ def test_two_phase_pinned_version_survives_concurrent_write(tmp_path):
)
def test_two_phase_pinned_memory_table_survives_concurrent_write():
"""Same guarantee as the on-disk test above, but for a memory:// table.
A memory table's "pickle state" is its row data (to_arrow()), not just
reopen coordinates, so it must be captured once at construction -- not
recomputed from the live table on every _open_pinned_table() call --
or a write made after construction would leak into every later read.
"""
db = lancedb.connect("memory://")
table = db.create_table("t", pa.table({"id": list(range(8))}))
ds = StreamingDataset(table, num_splits=1, block_size=4, shuffle=False)
table.delete("id = 0")
table.add(pa.table({"id": [8]}))
ids = [s["id"] for s in ds]
assert ids == list(range(8)), (
"must read the snapshot captured at construction, not the live, "
"since-mutated one"
)
def test_two_phase_memory_table_already_at_a_later_version():
"""Constructing over a memory table already past version 1 must not
fail: the pin is the snapshot's data, not a version number to replay
against a freshly rebuilt (always version-1) in-memory table."""
db = lancedb.connect("memory://")
table = db.create_table("t", pa.table({"id": list(range(4))}))
table.add(pa.table({"id": [4, 5, 6, 7]})) # now at version 2, 8 live rows
ds = StreamingDataset(table, num_splits=1, block_size=4, shuffle=False)
ids = [s["id"] for s in ds]
assert sorted(ids) == list(range(8))
def test_two_phase_rejects_lsm_write_spec_table(tmp_path):
"""An LSM write spec means some acknowledged writes live only in the
MemWAL, invisible to a version-pinned read; block_size must refuse the