feat: support remote tables in the data loader (#3981)

`StreamingDataset`, `PermutationBuilder`, and `Permutation` now work
against a `RemoteTable` (LanceDB Cloud and Enterprise), which unblocks
benchmarking the loader against the enterprise cluster cache.

```python
db = lancedb.connect("db://my-db", api_key=..., host_override=...)
ds = StreamingDataset(db.open_table("training"), world_size=8, rank=r)
```

Rows are addressed by `_rowid` exactly as before —
`PermutationReader::load_batch` already built the same `_rowid IN (...)`
filter that `Table::take_row_ids` sends, so the loader's fetch was
always the take path. It just was never allowed to run.

### The guard

`PermutationBuilder.__init__` rejected anything without `_inner`, so a
`RemoteTable` raised `TypeError` before reaching the PyO3 layer — which
already unwraps one via `_table._inner`.

### A bounded schema lookup

`PermutationReader::output_schema` reads the schema off a query plan,
and building a plan on a remote table *executes* the query
(`create_plan` → `execute_query`). With no limit that is `k =
isize::MAX`, so asking a remote table for its output schema pulled the
whole table over HTTP and threw it away — once per assigned split, on
every epoch, since `StreamingDataset.__iter__` constructs a
`Permutation` per split.

One row rather than zero, deliberately: lance gates its limit node on
`self.limit.unwrap_or(0) > 0`, so `Some(0)` means *no limit*.

### Tables with an LSM write spec are refused

A permutation references rows by row id, and rows that have not been
flushed to the base table do not have one yet. The loader could read
around them, but they would then be missing from training with nothing
said about it, so the build refuses such a table up front instead of
half supporting it.

### Fallible identity construction

`PermutationReader::identity` resolved `inner_new` with `unwrap`. That
was near total against a local dataset, but construction counts the base
table — an HTTP round trip for a remote one — so a transient network or
auth failure became a panic across the PyO3 boundary.

### Tests

End-to-end `permutation_builder` and `StreamingDataset` runs against a
mock server, the former torch-free so it runs wherever the suite does,
plus a test that a build succeeds without an LSM write spec and is
refused once one is installed.
This commit is contained in:
Jack Ye
2026-08-21 13:45:39 -07:00
committed by GitHub
parent 1baada89ef
commit f39a7a4dd9
9 changed files with 441 additions and 27 deletions
@@ -37,6 +37,11 @@ from unittest.mock import patch
import lancedb
import pyarrow as pa
import pytest
from utils import (
MockPermutationServer,
assert_server_safe_row_id_requests,
mock_remote_table,
)
torch = pytest.importorskip("torch")
streaming = pytest.importorskip("lancedb.streaming")
@@ -2118,3 +2123,27 @@ def test_doc_example_checkpoint(lance_table):
assert sorted(consumed + remaining_original) == list(range(NUM_ROWS)), (
"Consumed + remaining must cover every row exactly once"
)
# ---------------------------------------------------------------------------
# Remote tables (LanceDB Cloud / Enterprise)
# ---------------------------------------------------------------------------
def test_streaming_dataset_over_remote_table():
"""StreamingDataset reads a remote table, with server-safe requests.
Builds a permutation over a remote table, then fetches batches from it by row id.
"""
server = MockPermutationServer()
with mock_remote_table(server) as table:
ds = StreamingDataset(table, num_splits=2, shuffle_seed=SHUFFLE_SEED)
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"
)
assert len(server.scans) == 1, "the permutation is built with one row-id scan"
assert server.takes, "rows must be fetched with row-id takes"
assert_server_safe_row_id_requests(server)