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
+3 -1
View File
@@ -268,7 +268,9 @@ impl PyPermutationReader {
.await
.infer_error()?
} else {
PermutationReader::identity(base_table).await
PermutationReader::identity(base_table)
.await
.infer_error()?
};
Ok(Self::from_reader(reader))
})