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.

python/python/lancedb/permutation.py

@@ -41,21 +41,15 @@ class PermutationBuilder:
     The permutation is stored in memory and will be lost when the program exits.
     """
 
-    def __init__(self, table: LanceTable):
+    def __init__(self, table: Table):
         """
         Creates a new permutation builder for the given table.
 
         By default, the permutation builder will create a single split that contains all
         rows in the same order as the base table.
+
+        Tables with an LSM write spec are rejected: unflushed rows have no row id.
         """
-        if not hasattr(table, "_inner"):
-            raise TypeError(
-                f"PermutationBuilder requires a local LanceTable, "
-                f"got {type(table).__name__}. "
-                "The permutation API is not supported on remote tables. "
-                "Remote tables connect to LanceDB Cloud or Enterprise and do not have "
-                "direct access to the underlying Lance dataset needed for permutations."
-            )
         self._async = async_permutation_builder(table)
 
     def split_random(
@@ -231,7 +225,7 @@ class PermutationBuilder:
         return LOOP.run(do_execute())
 
 
-def permutation_builder(table: LanceTable) -> PermutationBuilder:
+def permutation_builder(table: Table) -> PermutationBuilder:
     return PermutationBuilder(table)
 
 
@@ -248,7 +242,7 @@ class Permutations:
 
     Attributes
     ----------
-    base_table: LanceTable
+    base_table: Table
         The base table that the permutations are based on.
     permutation_table: LanceTable
         The permutation table that defines the splits.
@@ -282,7 +276,7 @@ class Permutations:
     {'train': 0, 'test': 1}
     """
 
-    def __init__(self, base_table: LanceTable, permutation_table: LanceTable):
+    def __init__(self, base_table: Table, permutation_table: LanceTable):
         self.base_table = base_table
         self.permutation_table = permutation_table
 

python/python/tests/test_elastic_dataloader.py

@@ -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)

python/python/tests/test_permutation.py

@@ -8,6 +8,11 @@ import pytest
 from lancedb import DBConnection, Table, connect
 from lancedb.background_loop import LOOP
 from lancedb.permutation import Permutation, Permutations, permutation_builder
+from utils import (
+    MockPermutationServer,
+    assert_server_safe_row_id_requests,
+    mock_remote_table,
+)
 
 
 def test_split_random_ratios(mem_db):
@@ -1214,3 +1219,57 @@ def test_remove_rowid_after_select(some_permutation: Permutation):
     perm_without_rowid = perm_with_rowid.remove_columns(["_rowid"])
     assert "_rowid" not in perm_without_rowid.column_names
     assert perm_without_rowid.column_names == ["id"]
+
+
+def test_permutation_is_stable_when_remote_scan_order_varies():
+    """Splits are assigned by scan position, and every rank builds its own
+    permutation, so two ranks seeing different scan orders must still agree."""
+    server = MockPermutationServer(num_rows=16, vary_scan_order=True)
+
+    def split_of_each_row(permutation_tbl):
+        # Sequential splits are assigned by position, so a reversed scan would put
+        # the last rows in split 0.  Compare the mapping rather than the table order,
+        # which the split-id sort does not pin down.
+        rows = permutation_tbl.search(None).to_arrow().to_pydict()
+        return dict(zip(rows["row_id"], rows["split_id"]))
+
+    with mock_remote_table(server) as table:
+        first = split_of_each_row(
+            permutation_builder(table).split_sequential(fixed=2).execute()
+        )
+        second = split_of_each_row(
+            permutation_builder(table).split_sequential(fixed=2).execute()
+        )
+
+    assert server.scan_calls == 2, "both builds must have scanned"
+    assert first == second
+    assert first[0] == 0 and first[server.num_rows - 1] == 1, first
+
+
+def test_permutation_over_remote_table():
+    """The permutation API accepts a remote table, addressing rows by `_rowid` just
+    as `take_row_ids` does.  Also pins the request shapes sent to the server.
+    """
+    server = MockPermutationServer()
+
+    with mock_remote_table(server) as table:
+        permutation_tbl = permutation_builder(table).split_sequential(fixed=2).execute()
+        assert permutation_tbl.count_rows() == server.num_rows
+
+        permutation = Permutation.from_tables(table, permutation_tbl, 0)
+        assert permutation.num_rows == server.num_rows // 2
+
+        # Compare against the permutation's own order; the split-id sort is not stable.
+        rows = permutation_tbl.search(None).to_arrow().to_pydict()
+        split0 = [
+            row_id
+            for row_id, split in zip(rows["row_id"], rows["split_id"])
+            if not split
+        ]
+        # The mock table's `id` equals its `_rowid`.
+        assert permutation.take_offsets([2, 0]) == [
+            {"id": split0[2]},
+            {"id": split0[0]},
+        ]
+
+    assert_server_safe_row_id_requests(server)

python/python/tests/utils.py

@@ -1,7 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright The LanceDB Authors
+import contextlib
+import http.server
+import json
+import re
+import threading
+
+import lancedb
+import pyarrow as pa
 import pytest
 
+ARROW_FILE_CONTENT_TYPE = "application/vnd.apache.arrow.file"
+
 
 def exception_output(e_info: pytest.ExceptionInfo):
     import traceback
@@ -9,3 +19,199 @@ def exception_output(e_info: pytest.ExceptionInfo):
     # skip traceback part, since it's not worth checking in tests
     lines = traceback.format_exception_only(e_info.type, e_info.value)
     return "".join(lines).strip()
+
+
+def parse_in_list(filter_sql: str) -> list[int]:
+    """Pull the integers out of a `<col> IN (a, b, c)` predicate.
+
+    Scoped to the parenthesised list so a cast in the SQL adds no phantom values.
+    """
+    match = re.search(r"\bIN\s*\(([^)]*)\)", filter_sql, re.IGNORECASE)
+    assert match is not None, f"expected an IN list, got: {filter_sql}"
+    return [int(m) for m in re.findall(r"-?\d+", match.group(1))]
+
+
+def is_row_id_take(body) -> bool:
+    """True when a query body fetches specific rows by row id."""
+    return "_rowid" 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()
+    with pa.ipc.new_file(sink, table.schema) as writer:
+        writer.write_table(table)
+    return sink.getvalue().to_pybytes()
+
+
+class MockPermutationServer:
+    """A stand-in LanceDB server hosting one table whose ``id`` equals its ``_rowid``.
+
+    Records every ``/query/`` body so tests can assert on the request shapes sent to
+    the server, which is the part that has to stay compatible.
+    """
+
+    def __init__(self, name="remote_data", num_rows=8, vary_scan_order=False):
+        self.name = name
+        self.num_rows = num_rows
+        self.query_bodies = []
+        # Stand in for a distributed scan that answers in no fixed order.
+        self.vary_scan_order = vary_scan_order
+        self.scan_calls = 0
+
+    def __call__(self, request):
+        path = request.path
+        if path == f"/v1/table/{self.name}/describe/":
+            return self._json(
+                request,
+                {
+                    "version": 1,
+                    "schema": {
+                        "fields": [
+                            {"name": "id", "type": {"type": "int64"}, "nullable": False}
+                        ]
+                    },
+                },
+            )
+        if path == f"/v1/table/{self.name}/get_lsm_write_spec/":
+            self._read_body(request)
+            # Null spec: this table has no LSM write path.
+            return self._json(request, {"lsm_write_spec": None})
+        if path == f"/v1/table/{self.name}/count_rows/":
+            self._read_body(request)
+            return self._json(request, self.num_rows)
+        if path == f"/v1/table/{self.name}/query/":
+            return self._query(request, self._read_body(request))
+
+        # Drain first, so an unexpected route cannot desync a keep-alive connection.
+        self._read_body(request)
+        request.send_response(404)
+        request.end_headers()
+
+    @property
+    def scans(self):
+        """Bodies of the permutation build scan: the row id column, nothing else."""
+        return [b for b in self.query_bodies if b.get("columns") == ["_rowid"]]
+
+    @property
+    def takes(self):
+        """Bodies of the row-id takes the loader fetches batches with.
+
+        Keyed on `_rowid`, not "has a filter": the schema probe also has a predicate.
+        """
+        return [b for b in self.query_bodies if is_row_id_take(b)]
+
+    @staticmethod
+    def _read_body(request):
+        content_len = int(request.headers.get("Content-Length") or 0)
+        return json.loads(request.rfile.read(content_len)) if content_len else {}
+
+    @staticmethod
+    def _json(request, payload):
+        request.send_response(200)
+        request.send_header("Content-Type", "application/json")
+        request.end_headers()
+        request.wfile.write(json.dumps(payload).encode())
+
+    @staticmethod
+    def _arrow(request, table):
+        body = arrow_file_bytes(table)
+        request.send_response(200)
+        request.send_header("Content-Type", ARROW_FILE_CONTENT_TYPE)
+        request.send_header("Content-Length", str(len(body)))
+        request.end_headers()
+        request.wfile.write(body)
+
+    def _query(self, request, body):
+        self.query_bodies.append(body)
+
+        if is_row_id_take(body):
+            # A row-id take. Answer ascending, so tests prove the client reorders.
+            row_ids = sorted(parse_in_list(body["filter"]))
+            return self._arrow(
+                request,
+                pa.table(
+                    {
+                        "id": pa.array(row_ids, pa.int64()),
+                        "_rowid": pa.array(row_ids, pa.uint64()),
+                    }
+                ),
+            )
+
+        if body.get("columns") == ["_rowid"]:
+            # The permutation build scan: row ids and nothing else.
+            row_ids = list(range(self.num_rows))
+            if self.vary_scan_order and self.scan_calls % 2:
+                row_ids.reverse()
+            self.scan_calls += 1
+            return self._arrow(
+                request,
+                pa.table({"_rowid": pa.array(row_ids, pa.uint64())}),
+            )
+
+        # The schema probe: filtered to nothing, so it carries schema and no rows.
+        return self._arrow(request, pa.table({"id": pa.array([], pa.int64())}))
+
+
+def _make_handler(serve):
+    class MockLanceDBHandler(http.server.BaseHTTPRequestHandler):
+        def do_GET(self):
+            serve(self)
+
+        def do_POST(self):
+            serve(self)
+
+        def log_message(self, *args):
+            pass  # keep pytest output readable
+
+    return MockLanceDBHandler
+
+
+@contextlib.contextmanager
+def mock_remote_table(server):
+    """Run ``server`` on a local port and yield an open remote table against it.
+
+    Threading: the loader fans out fetch threads a single-threaded server would
+    serialize, hiding the prefetch overlap under test.
+    """
+    with http.server.ThreadingHTTPServer(
+        ("localhost", 0), _make_handler(server)
+    ) as srv:
+        thread = threading.Thread(target=srv.serve_forever)
+        thread.start()
+        try:
+            db = lancedb.connect(
+                "db://dev",
+                api_key="fake",
+                host_override=f"http://localhost:{srv.server_address[1]}",
+                client_config={"timeout_config": {"connect_timeout": 5}},
+            )
+            yield db.open_table(server.name)
+        finally:
+            srv.shutdown()
+            thread.join()
+
+
+def assert_server_safe_row_id_requests(server):
+    """Assert the loader fetched rows by row id and bounded everything else.
+
+    `.get`, not `[...]`, so a dropped field reads as the assertion, not a KeyError.
+    """
+    for body in server.takes:
+        # The fetch needs the row id back to restore the requested order.
+        assert body.get("with_row_id") is True, body
+        assert "_rowid" in body["filter"], body
+
+    # Only the one-off permutation scan may scan the whole table; the schema probe is
+    # built once per split per epoch. `k == 0` counts as unbounded: lance reads a zero
+    # limit as "no limit".
+    def is_unbounded(body):
+        if is_row_id_take(body):
+            return False
+        k = body.get("k")
+        return k is None or k == 0 or k > server.num_rows
+
+    unbounded = [b for b in server.query_bodies if is_unbounded(b)]
+    assert unbounded == server.scans, (
+        f"only the permutation scan may be unbounded, got {unbounded}"
+    )

python/src/permutation.rs

@@ -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))
         })