Skip to content

fix(python): make Permutation picklable for PyTorch multiprocessing#3335

Merged
westonpace merged 7 commits into
lancedb:mainfrom
westonpace:fix-pickle-permutation
May 5, 2026
Merged

fix(python): make Permutation picklable for PyTorch multiprocessing#3335
westonpace merged 7 commits into
lancedb:mainfrom
westonpace:fix-pickle-permutation

Conversation

@westonpace

@westonpace westonpace commented Apr 29, 2026

Copy link
Copy Markdown
Contributor

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 Table and Connection are not serializable. This PR adds pickle support to Permutation without adding general pickle support to Table or Connection. 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:

  • In-memory tables (just serialize as Arrow IPC)
  • Native tables (serialize the URI)

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 Rust PermutationBuilder::persist API 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 serialize Table and Connection.

🤖 Generated with Claude Code

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]>
@github-actions github-actions Bot added bug Something isn't working Python Python SDK labels 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

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

westonpace and others added 5 commits April 30, 2026 21:41
…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
westonpace marked this pull request as ready for review May 1, 2026 12:28
@sarwarbhuiyan
sarwarbhuiyan requested a review from wjones127 May 1, 2026 13:32

@wjones127 wjones127 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 thread python/python/lancedb/permutation.py Outdated
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()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

note: I had thought pa.Table was already pickle-able.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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]>
@westonpace
westonpace merged commit 1fc23e5 into lancedb:main May 5, 2026
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working Python Python SDK

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants