From 40ff1d52a6df258d9b9768e9606339e0193576ab Mon Sep 17 00:00:00 2001 From: jasonz-lance Date: Thu, 8 Oct 2026 17:49:54 -0700 Subject: [PATCH] fix(python): stop 2-phase reads from aliasing the caller's table handle _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 --- python/python/lancedb/streaming.py | 96 ++++++++++++++++--- .../python/tests/test_elastic_dataloader.py | 58 +++++++++++ 2 files changed, 139 insertions(+), 15 deletions(-) diff --git a/python/python/lancedb/streaming.py b/python/python/lancedb/streaming.py index 530a70d9c..93cd218b9 100644 --- a/python/python/lancedb/streaming.py +++ b/python/python/lancedb/streaming.py @@ -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( diff --git a/python/python/tests/test_elastic_dataloader.py b/python/python/tests/test_elastic_dataloader.py index da1eeb863..3941974a9 100644 --- a/python/python/tests/test_elastic_dataloader.py +++ b/python/python/tests/test_elastic_dataloader.py @@ -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