fix(python): pin table version and drop to_lance() for 2-phase reads

Two independent fixes to 2-phase shuffled reads, found by the same
gatekeeper review pass:

- Packed checkpoint resume: the 2-phase branch of _iter_owned read the
  (empty, in packed mode) row-mode cursor maps instead of
  _pack_consumed, so restoring a packed checkpoint silently replayed
  already-consumed documents. Now mirrors the 1-phase branch's check.

- Version pinning + remote (Cloud) table support: 2-phase previously
  called self._table.to_lance() fresh at iteration time, which (a)
  raises NotImplementedError on RemoteTable, since to_lance() only
  works against a local table, and (b) re-reads whatever version is
  live at that moment rather than the one the block plan was computed
  against, so a write to the table between construction and iteration
  (or between epochs) could shift which rows a block refers to.

  Now pins table.version (and rejects a table checked out to a
  non-default branch, which isn't supported yet) at construction, and
  reads through a private handle reopened via the same
  connection_factory/pickle-state machinery already used for worker
  reconnects, checked out to that pinned version. Reads go through
  Table's backend-agnostic query builder (search/select/take_offsets),
  the same interface RemoteTable implements, instead of pylance's
  to_lance(). take_offsets doesn't guarantee its output order, so
  filtered block reads always additionally select "_rowoffset" and
  sort locally by it to restore the requested order.

Extends MockPermutationServer (utils.py) to answer the offset/limit
scan and _rowoffset-filtered take request shapes 2-phase now sends,
and adds tests for both fixes plus 2-phase over a (mocked) remote
table, unfiltered and filtered+shuffled.

python/python/lancedb/streaming.py

@@ -704,9 +704,7 @@ class StreamingDataset(IterableDataset):
             if window_blocks is not None:
                 raise ValueError("window_blocks requires block_size to be set")
             if max_shuffle_distance is not None:
-                raise ValueError(
-                    "max_shuffle_distance requires block_size to be set"
-                )
+                raise ValueError("max_shuffle_distance requires block_size to be set")
             if reshuffle_blocks_per_epoch is not None:
                 raise ValueError(
                     "reshuffle_blocks_per_epoch requires block_size to be set"
@@ -877,9 +875,31 @@ class StreamingDataset(IterableDataset):
             #
             # block_size counts LIVE rows (rows passing `filter`), not raw
             # rows.
+            #
+            # Pin the exact version/branch the block plan below is computed
+            # against, so a write to the table between construction and a
+            # later iteration (or between epochs) cannot shift which rows a
+            # block refers to.  Read paths in _read_block always reopen a
+            # private handle checked out to this pin (see
+            # _open_pinned_table) rather than reading the table's live,
+            # possibly-since-mutated state.  Branch pinning isn't
+            # implemented yet -- the reopened handle always lands on the
+            # default branch (see _table_from_pickle_state) -- so a table
+            # checked out to a non-default branch is rejected below rather
+            # than silently pinning the wrong branch.
+            self._pinned_branch = table.current_branch()
+            if self._pinned_branch is not None:
+                raise ValueError(
+                    "block_size does not yet support a table checked out to "
+                    f"a non-default branch ({self._pinned_branch!r}); check "
+                    "out the default branch before constructing StreamingDataset"
+                )
+            self._pinned_version = table.version
+
             block_live_offsets: Optional[list[np.ndarray]] = None
+            self._all_columns: Optional[list[str]] = None
             if filter is None:
-                total_rows = table.count_rows()
+                total_rows = self._open_pinned_table().count_rows()
                 # Discard the trailing runt block (the remainder rows left
                 # over when block_size doesn't evenly divide the row count)
                 # rather than reading it as an undersized final block.
@@ -891,10 +911,9 @@ class StreamingDataset(IterableDataset):
                 num_rows = num_blocks * block_size
             else:
                 # One pass, computed once here and never repeated at read
-                # time: a single filtered scan over the whole table asks
-                # Lance to evaluate `filter` natively (pushed down against
-                # each fragment's own data, using its usual page/zone-map
-                # skipping) and hand back only each live row's
+                # time: a filtered scan over the whole table asks Lance to
+                # evaluate `filter` natively (pushed down against each
+                # fragment's own data) and hand back only each live row's
                 # "_rowoffset" -- a system column giving its 0-indexed
                 # dataset-wide position -- without materializing any real
                 # column data.  Blocks are then just contiguous chunks of
@@ -902,11 +921,26 @@ class StreamingDataset(IterableDataset):
                 # each; no more approximating by snapping to fragment
                 # edges), and at read time each block's exact live offsets
                 # are already known, so a sparse take() can fetch just
-                # those rows -- see _read_block in _iter_owned.
+                # those rows -- see _read_block in _iter_owned.  Goes
+                # through the table's own query builder (search/select),
+                # not pylance's to_lance(), so this also works against a
+                # remote (Cloud) table.
+                pinned_for_setup = self._open_pinned_table()
+                # take_offsets doesn't guarantee its output order, so
+                # filtered block reads always additionally select
+                # "_rowoffset" and sort locally by it (see _read_block);
+                # when the caller didn't narrow columns, that means
+                # explicitly listing every column instead of relying on an
+                # implicit "everything", hence caching the full list once
+                # here rather than re-querying the schema per block.
+                self._all_columns: Optional[list[str]] = (
+                    pinned_for_setup.schema.names if columns is None else None
+                )
                 live_offsets = (
-                    table.to_lance()
-                    .scanner(filter=filter, columns=["_rowoffset"])
-                    .to_table()
+                    pinned_for_setup.search()
+                    .where(filter)
+                    .select(["_rowoffset"])
+                    .to_arrow()
                     .column("_rowoffset")
                     .to_numpy()
                 )
@@ -979,7 +1013,7 @@ class StreamingDataset(IterableDataset):
                 if max_shuffle_distance is not None
                 else self._window_blocks * block_size
             )
-        else: #1-phase shuffled read
+        else:  # 1-phase shuffled read
             # Build the permutation table once, deterministically.
             builder = permutation_builder(table)
             if filter is not None:
@@ -1002,31 +1036,59 @@ class StreamingDataset(IterableDataset):
             range(rank_start, rank_start + splits_per_rank)
         )
 
+    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
+        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.
+        """
+        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)
+        return pinned
+
     def _read_block(
-        self, lance_ds, block_id: int, columns: Optional[list[str]]
+        self, pinned_table, block_id: int, columns: Optional[list[str]]
     ) -> pa.Table:
         """Read one 2-phase block's live rows, in natural (ascending) order."""
         if self._filter is None:
-            # Plain offset/limit: unambiguous raw-position scan that Lance
-            # can seek to directly.
+            # offset/limit is a positional, sequential read, so its output
+            # order is exactly storage order -- no restoration needed.
             row_start, row_end = self._block_ranges[block_id]
-            return lance_ds.scanner(
-                offset=row_start,
-                limit=row_end - row_start,
-                columns=columns,
-            ).to_table()
+            q = pinned_table.search().offset(row_start).limit(row_end - row_start)
+            if columns is not None:
+                q = q.select(columns)
+            return q.to_arrow()
 
         # Filtered: this block's exact live-row dataset offsets were
         # already computed once, up front (from the "_rowoffset" system
-        # column -- see __init__), so a sparse take() fetches just the
-        # live rows themselves, skipping any dead-row gaps entirely.
+        # column -- see __init__), so take_offsets fetches just the live
+        # rows themselves, skipping any dead-row gaps entirely.  Unlike
+        # offset/limit, take_offsets makes no ordering guarantee for an
+        # arbitrary index list, so "_rowoffset" is always also selected
+        # and used to restore the request's ascending order locally.
         offsets = self._block_live_offsets[block_id]
-        return lance_ds.take(offsets.tolist(), columns=columns)
+        select_cols = (columns if columns is not None else self._all_columns) + [
+            "_rowoffset"
+        ]
+        block = (
+            pinned_table.take_offsets(offsets.tolist()).select(select_cols).to_arrow()
+        )
+        order = pc.sort_indices(block.column("_rowoffset"))
+        return block.take(order).drop_columns(["_rowoffset"])
 
     def _make_two_phase_reader(
         self,
         split_idx: int,
-        lance_ds,
+        pinned_table,
         columns: Optional[list[str]],
         skip: int = 0,
     ) -> "_TwoPhaseSplitReader":
@@ -1057,7 +1119,7 @@ class StreamingDataset(IterableDataset):
             max_shuffle_distance=self._max_shuffle_distance,
             seed=sigma_seed,
             read_block_fn=lambda block_id: self._read_block(
-                lance_ds, block_id, columns
+                pinned_table, block_id, columns
             ),
             columns=columns,
             skip=skip,
@@ -1086,7 +1148,9 @@ class StreamingDataset(IterableDataset):
             "pass an explicit value for exact epoch sizing",
         )
 
-        lance_ds = self._table.to_lance() if self._block_size is not None else None
+        pinned_table = (
+            self._open_pinned_table() if self._block_size is not None else None
+        )
         for split in range(self._num_splits):
             if self._block_size is not None:
                 # Sample through the same 2-phase reader used at read time,
@@ -1094,7 +1158,7 @@ class StreamingDataset(IterableDataset):
                 # assignment and shuffle bound rather than a separate,
                 # unrelated approximation.
                 permutation = self._make_two_phase_reader(
-                    split, lance_ds, [token_column]
+                    split, pinned_table, [token_column]
                 )
             else:
                 permutation = Permutation.from_tables(
@@ -1259,7 +1323,7 @@ class StreamingDataset(IterableDataset):
             # 2-phase flows through the exact same stage-1 I/O machinery
             # (_fill_io/_submit_io/_drain_io, io_queue_depth prefetch) as
             # 1-phase.  No 2-phase-specific code exists past this point.
-            lance_ds = self._table.to_lance()
+            pinned_table = self._open_pinned_table()
 
             permutations: list[_TwoPhaseSplitReader] = []
             initial_samples = []
@@ -1269,12 +1333,20 @@ class StreamingDataset(IterableDataset):
                 # loaded checkpoint's positions_consumed_per_split (falling
                 # back to the sample count, then to the global offset) is
                 # this split's saved absolute position in its own
-                # deterministic read order.
+                # deterministic read order.  Packing stores its own cursor
+                # separately (it also checkpoints partial blocks), so it
+                # must be consulted here too, exactly as the 1-phase branch
+                # above does -- otherwise a packed resume silently ignores
+                # _pack_consumed and replays already-consumed documents.
                 sample_count = self._resume_samples.get(split_idx, self._resume_offset)
-                start_pos = self._resume_positions.get(split_idx, sample_count)
+                start_pos = (
+                    self._pack_consumed[split_idx]
+                    if self._pack_sequences is not None
+                    else self._resume_positions.get(split_idx, sample_count)
+                )
                 permutations.append(
                     self._make_two_phase_reader(
-                        split_idx, lance_ds, self._columns, skip=start_pos
+                        split_idx, pinned_table, self._columns, skip=start_pos
                     )
                 )
                 initial_samples.append(sample_count)

python/python/tests/test_elastic_dataloader.py

@@ -3197,6 +3197,42 @@ def test_streaming_dataset_over_remote_table():
     assert_server_safe_row_id_requests(server)
 
 
+def test_two_phase_over_remote_table():
+    """2-phase (block_size) reads a remote table via offset/limit scans,
+    not pylance's to_lance() (which only works locally)."""
+    server = MockPermutationServer(num_rows=16)
+
+    with mock_remote_table(server) as table:
+        ds = StreamingDataset(table, num_splits=2, block_size=4, shuffle=False)
+        ids = [row["id"] for row in ds]
+
+    assert sorted(ids) == list(range(server.num_rows)), (
+        "Every row of the remote table must be yielded exactly once"
+    )
+
+
+def test_two_phase_over_remote_table_filtered_and_shuffled():
+    """2-phase's filtered live-offset precompute and take_offsets-based
+    block reads both work against a remote table, including restoring
+    take_offsets' unspecified order via the _rowoffset it also selects."""
+    server = MockPermutationServer(num_rows=20)
+    live_ids = list(range(16))  # 16 live rows -> 4 blocks, divides num_splits=2
+
+    with mock_remote_table(server) as table:
+        ds = StreamingDataset(
+            table,
+            num_splits=2,
+            block_size=4,
+            shuffle=True,
+            shuffle_seed=SHUFFLE_SEED,
+            filter="id IN (" + ",".join(str(i) for i in live_ids) + ")",
+            columns=["id"],
+        )
+        ids = [row["id"] for row in ds]
+
+    assert sorted(ids) == live_ids, "only live rows, each exactly once"
+
+
 # ---------------------------------------------------------------------------
 # 2-phase shuffled reads (block_size)
 # ---------------------------------------------------------------------------
@@ -3366,6 +3402,22 @@ def test_two_phase_resume_checkpoint_round_trip(lance_table):
     )
 
 
+def test_two_phase_pinned_version_survives_concurrent_write(tmp_path):
+    """The block plan is pinned to the table version at construction; a
+    write between construction and iteration must not change what an
+    in-progress epoch reads."""
+    db = lancedb.connect(tmp_path)
+    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")  # mutate after construction, before iteration
+
+    ids = [s["id"] for s in ds]
+    assert ids == list(range(8)), (
+        "must read the version pinned at construction, not the live, since-mutated one"
+    )
+
+
 def test_two_phase_reshuffle_blocks_per_epoch_toggle(lance_table):
     """reshuffle_blocks_per_epoch=False keeps each split's blocks fixed
     across epochs; the default (True) reassigns them every epoch."""

python/python/tests/utils.py

@@ -36,6 +36,15 @@ def is_row_id_take(body) -> bool:
     return "_rowid" in (body.get("filter") or "")
 
 
+def is_row_offset_take(body) -> bool:
+    """True when a query body fetches specific rows by ``_rowoffset``.
+
+    This is the request shape ``Table.take_offsets`` sends on the remote
+    backend, distinct from ``is_row_id_take``'s ``_rowid IN (...)``.
+    """
+    return "_rowoffset" in (body.get("filter") or "")
+
+
 def arrow_file_bytes(table: pa.Table) -> bytes:
     """Serialize to the Arrow IPC *file* framing the /query/ route answers with."""
     sink = pa.BufferOutputStream()
@@ -149,6 +158,49 @@ class MockPermutationServer:
                 pa.table({"_rowid": pa.array(row_ids, pa.uint64())}),
             )
 
+        if is_row_offset_take(body):
+            # A _rowoffset take (2-phase's filtered block reads): this
+            # table's rows are never deleted/reordered, so id == _rowoffset
+            # for every row, same invariant as id == _rowid above.
+            row_offsets = sorted(parse_in_list(body["filter"]))
+            columns = body.get("columns") or ["id", "_rowoffset"]
+            values = {"id": row_offsets, "_rowoffset": row_offsets}
+            return self._arrow(
+                request,
+                pa.table(
+                    {
+                        col: pa.array(
+                            values[col],
+                            pa.uint64() if col == "_rowoffset" else pa.int64(),
+                        )
+                        for col in columns
+                    }
+                ),
+            )
+
+        if body.get("columns") == ["_rowoffset"] and "IN" in (body.get("filter") or ""):
+            # 2-phase's live-offset precompute scan: a user filter over
+            # "_rowoffset" alone.  This mock only understands `<col> IN
+            # (...)` filters (not arbitrary SQL), and since id == _rowoffset
+            # for every row here, reuses parse_in_list regardless of which
+            # column the filter actually names.
+            row_offsets = sorted(parse_in_list(body["filter"]))
+            return self._arrow(
+                request, pa.table({"_rowoffset": pa.array(row_offsets, pa.uint64())})
+            )
+
+        if "offset" in body and not body.get("filter"):
+            # A plain offset/limit scan (2-phase's unfiltered contiguous
+            # block reads): a positional, sequential slice of storage order.
+            start = body["offset"]
+            end = start + body["k"]
+            row_ids = list(range(self.num_rows))[start:end]
+            columns = body.get("columns") or ["id"]
+            return self._arrow(
+                request,
+                pa.table({col: pa.array(row_ids, pa.int64()) for col in columns}),
+            )
+
         # The schema probe: filtered to nothing, so it carries schema and no rows.
         return self._arrow(request, pa.table({"id": pa.array([], pa.int64())}))