fix(python): make Permutation picklable for PyTorch multiprocessing#3335
Merged
Conversation
A Permutation could not be passed to a PyTorch DataLoader with num_workers > 0 because the underlying Rust PermutationReader and the LanceTable references it held could not be pickled. Worker processes would fail with "cannot pickle 'builtins.PermutationReader' object". Track the base table, permutation table, split, offset, and limit on the Permutation directly and rebuild the reader lazily, so __getstate__ can capture the base table by URI/name and the in-memory permutation table as Arrow IPC bytes. __setstate__ reopens the base table and rematerializes the permutation table from those bytes. Use a separate async _reader_async() inside coroutines to avoid re-entering the background event loop and deadlocking. Also drops the persist() method from the Python permutation builder bindings — breaking change. The permutation table is now always in-memory; the underlying Rust persist() API is untouched and can be re-exposed later. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
westonpace
commented
Apr 29, 2026
Comment on lines
-54
to
-62
| def persist( | ||
| self, database: Union[DBConnection, AsyncConnection], table_name: str | ||
| ) -> "PermutationBuilder": | ||
| """ | ||
| Persist the permutation to the given database. | ||
| """ | ||
| self._async.persist(database, table_name) | ||
| return self | ||
|
|
Contributor
Author
There was a problem hiding this comment.
This is a breaking change. We're just removing this capability for now. We could add it back in the future but we'd have to think about it some more to figure out how we can pickle this.
…ling The URI-introspection __getstate__ path doesn't work for non-local connections: RemoteDBConnection consumes api_key/region/host_override at construction and never stores them, and LanceDBConnection only retains storage_options (not read_consistency_interval, session, etc). Add an optional connection_factory: Callable[[str], LanceTable] that the user attaches via with_connection_factory(). When set, __getstate__ pickles the factory in place of the URI; __setstate__ calls factory(name) to recover the base table in the worker. The factory must be picklable (top-level function or functools.partial — not lambdas or closures). The URI-introspection fallback stays for the common OSS-file case where it works; remote-style connections that lack a usable uri/storage_options now raise a clear ValueError pointing at with_connection_factory. Docstring includes three examples: native + functools.partial, native with namespace_path, and LanceDB Cloud reading auth from environment variables at worker startup. New test exercises the explicit factory path with a top-level _open_native_table helper. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
Previously, pickling a Permutation whose base table lived in an in-memory database raised ValueError because the worker had no way to reopen it. Mirror what we already do for the in-memory permutation table: serialize the base table to Arrow IPC stream bytes and rebuild it as an in-memory LanceTable in __setstate__. This adds dataset-sized cost to pickles of memory-backed permutations, so users with large in-memory base tables should still prefer persisting them or supplying a connection_factory. New test confirms round-trip behavior for a memory-backed identity permutation. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
The lazy reader build introduced a sync `reader` property and a parallel
async `_reader_async()` to dodge a re-entrant `LOOP.run` deadlock when
the property was touched from inside a coroutine. The lazy path bought
nothing in practice — every `_clone()` rebuilt anyway because we never
shared `_reader` across clones — so the cost (two accessors, deadlock
trap, `_ = new_perm.reader` band-aids in with_skip/with_take, validation
deferred to first access) wasn't justified.
Switch to eager building:
- `__init__` accepts an optional `_reader=` kwarg; if absent it builds
via `LOOP.run(_build_reader())`. `self.reader` is a plain attribute.
- `from_tables` passes the reader it already built via `_reader=` so
there is no double build.
- `__setstate__` rebuilds eagerly after restoring tables.
- `_clone()` checks whether the overrides touch reader-affecting fields
({base_table, permutation_table, split, offset, limit}) — if not, it
passes the existing `self.reader` through, so with_batch_size,
_with_selection, with_transform, and with_connection_factory share
the parent's reader instead of rebuilding.
- `_reader_async()` is gone; schema/iter/__getitems__ just
`await self.reader.foo(...)` directly.
- with_skip/with_take drop their explicit `_ = new_perm.reader` since
__init__ already builds eagerly.
Net -13 lines and the deadlock surface is removed. All 65 tests green.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
_clone was just sugar around the constructor. Replace it with copy.copy(self) at each call site: - Reader-preserving updates (with_batch_size, _with_selection, with_transform, with_connection_factory) just shallow-copy and mutate the one field. The reader is preserved automatically. - Reader-rebuilding updates (with_skip, with_take) shallow-copy, mutate offset/limit, then explicitly reassign self.reader from a fresh _build_reader(). Removes the _READER_AFFECTING_FIELDS set and the field-list duplication between __init__ and _clone. -14 net lines. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
The second example was using namespace_path on open_table(), which is
not what "namespace connection" means in lancedb. Replace it with a
proper lancedb.connect_namespace() usage that takes an implementation
name and a properties dict (e.g. "dir" with {"root": ...} or "rest"
with a URL), partial-bound so the worker rebuilds the same namespace
connection on unpickle.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
westonpace
marked this pull request as ready for review
May 1, 2026 12:28
wjones127
approved these changes
May 4, 2026
wjones127
left a comment
Contributor
There was a problem hiding this comment.
Looks good. I think the approach works. I think we can later make Connection and Table pickle-able. I think instead of a connection_factory, for those we'll just need an optional credential_factory. Otherwise we can just reconstruct stuff from the pickled data.
Comment on lines
+357
to
+368
| def _serialize_arrow_table(table: pa.Table) -> bytes: | ||
| """Serialize a pyarrow Table to Arrow IPC stream bytes.""" | ||
| sink = pa.BufferOutputStream() | ||
| with pa.ipc.new_stream(sink, table.schema) as writer: | ||
| writer.write_table(table) | ||
| return sink.getvalue().to_pybytes() | ||
|
|
||
|
|
||
| def _deserialize_arrow_table(data: bytes) -> pa.Table: | ||
| """Read an Arrow IPC stream back into a pyarrow Table.""" | ||
| with pa.ipc.open_stream(pa.BufferReader(data)) as reader: | ||
| return reader.read_all() |
Contributor
There was a problem hiding this comment.
note: I had thought pa.Table was already pickle-able.
Contributor
Author
There was a problem hiding this comment.
Good catch. Removed these helpers.
pyarrow.Table is picklable natively (its __reduce__ uses Arrow IPC under the hood), so wrapping it in custom serialize/deserialize helpers in __getstate__/__setstate__ added indirection without benefit. Store the pa.Table objects in the state dict directly and let Python's pickle machinery handle them. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
When pytorch is used with multiprocessing and the mp mode is spawn then the Permutation needs to be pickled. It could not be pickled because
TableandConnectionare not serializable. This PR adds pickle support to Permutation without adding general pickle support toTableorConnection. To add general support we probably need to start by adding serialization in the namespace client.In the meantime this PR enable pickling by adding special cases for:
If a user is not using one of the above cases (e.g. using a remote connection) then they will need to provide a connection factory that can be pickled.
Breaking change
PermutationBuilder.persist(...)is removed from the Python bindings; the permutation table is now always in-memory. The underlying RustPermutationBuilder::persistAPI is untouched and can be re-exposed later if needed. It probably won't make sense to do that until we have a way to serializeTableandConnection.🤖 Generated with Claude Code