diff --git a/python/python/lancedb/permutation.py b/python/python/lancedb/permutation.py index ac2ebfa4d..bd0dbdd75 100644 --- a/python/python/lancedb/permutation.py +++ b/python/python/lancedb/permutation.py @@ -11,7 +11,7 @@ import pyarrow as pa from ._lancedb import async_permutation_builder, PermutationReader from .table import LanceTable, Table from .background_loop import LOOP -from .util import batch_to_tensor, batch_to_tensor_rows +from .util import batch_to_tensor, batch_to_tensor_dict, batch_to_tensor_rows from typing import Any, Callable, Iterator, Literal, Optional, TYPE_CHECKING, Union if TYPE_CHECKING: @@ -946,6 +946,7 @@ class Permutation: "pandas", "arrow", "torch", + "torch_row", "torch_col", "polars", ], @@ -961,15 +962,19 @@ class Permutation: - "python_col" - the batch will be a dict of lists (one entry per column) - "pandas" - the batch will be a pandas DataFrame - "arrow" - the batch will be a pyarrow RecordBatch - - "torch" - the batch will be a list of tensors, one per row + - "torch" - the batch will be a list of per-row dicts mapping column + name to a 0-D torch tensor. Works with the default + ``torch.utils.data.DataLoader`` collate, which stacks the per-row + dicts back into a dict of batched tensors. + - "torch_row" - the batch will be a list of tensors, one per row - "torch_col" - the batch will be a 2D torch tensor (first dim indexes columns) - "polars" - the batch will be a polars DataFrame Conversion may or may not involve a data copy. Lance uses Arrow internally and so it is able to zero-copy to the arrow and polars formats. - Conversion to torch_col will be zero-copy but will only support a subset of data - types (numeric types). + Conversion to torch and torch_col will be zero-copy but will only support a + subset of data types (numeric types). Conversion to numpy and/or pandas will typically be zero-copy for numeric types. Conversion of strings, lists, and structs will require creating python @@ -990,6 +995,8 @@ class Permutation: elif format == "arrow": return self.with_transform(Transforms.arrow2arrow) elif format == "torch": + return self.with_transform(batch_to_tensor_dict) + elif format == "torch_row": return self.with_transform(batch_to_tensor_rows) elif format == "torch_col": return self.with_transform(batch_to_tensor) diff --git a/python/python/lancedb/util.py b/python/python/lancedb/util.py index cd170c221..932a3e982 100644 --- a/python/python/lancedb/util.py +++ b/python/python/lancedb/util.py @@ -515,3 +515,38 @@ def batch_to_tensor_rows(batch: pa.RecordBatch): stacked = torch.tensor(numpy.column_stack(columns)) rows = list(stacked.unbind(dim=0)) return rows + + +def batch_to_tensor_dict(batch: pa.RecordBatch): + """ + Convert a PyArrow RecordBatch to a list of per-row dicts of PyTorch Tensors. + + Each column is first converted to a 1-D tensor (zero-copy via DLPack), then + sliced per row. The result is a list whose length is ``batch.num_rows`` and + whose items are dicts keyed by column name. This shape composes directly + with PyTorch's default ``DataLoader`` collate, which stacks the per-row + dicts back into a dict of batched tensors. + + Fails if torch is not installed. + Fails if a column's data type is not supported by PyTorch. + + Parameters + ---------- + batch : pa.RecordBatch + The record batch to convert. + + Returns + ------- + list[dict[str, torch.Tensor]] + One per-row dict per row in the batch. Each dict maps column name to a + 0-D tensor view into the column. + """ + torch = attempt_import_or_raise("torch", "torch") + tensors = { + name: torch.from_dlpack(col) + for name, col in zip(batch.schema.names, batch.columns) + } + return [ + {name: tensor[i] for name, tensor in tensors.items()} + for i in range(batch.num_rows) + ] diff --git a/python/python/tests/test_permutation.py b/python/python/tests/test_permutation.py index e58a49624..0da12c7dc 100644 --- a/python/python/tests/test_permutation.py +++ b/python/python/tests/test_permutation.py @@ -935,14 +935,41 @@ def test_transform_fn(mem_db): try: import torch - torch_result = list( - permutation.with_format("torch").iter(10, skip_last_batch=False) + # "torch" returns a list of per-row dicts. Default DataLoader collate + # stacks the per-row dicts back into a dict of batched tensors. + torch_perm = permutation.with_format("torch") + torch_batch = list(torch_perm.iter(10, skip_last_batch=False))[0] + assert isinstance(torch_batch, list) + assert len(torch_batch) == 10 + assert isinstance(torch_batch[0], dict) + assert set(torch_batch[0].keys()) == {"id", "value"} + assert isinstance(torch_batch[0]["id"], torch.Tensor) + assert torch_batch[0]["id"].dtype == torch.int64 + + rows = torch_perm.__getitems__([0, 1, 2]) + assert isinstance(rows, list) + assert len(rows) == 3 + assert isinstance(rows[0], dict) + assert set(rows[0].keys()) == {"id", "value"} + assert isinstance(rows[0]["id"], torch.Tensor) + + # "torch_row" returns a list of tensors, one per row. + torch_rows = list( + permutation.with_format("torch_row").iter(10, skip_last_batch=False) )[0] - assert isinstance(torch_result, list) - assert len(torch_result) == 10 - assert isinstance(torch_result[0], torch.Tensor) - assert torch_result[0].shape == (2,) - assert torch_result[0].dtype == torch.int64 + assert isinstance(torch_rows, list) + assert len(torch_rows) == 10 + assert isinstance(torch_rows[0], torch.Tensor) + assert torch_rows[0].shape == (2,) + assert torch_rows[0].dtype == torch.int64 + + # "torch_col" stacks columns into a single 2D tensor. + torch_col = list( + permutation.with_format("torch_col").iter(10, skip_last_batch=False) + )[0] + assert isinstance(torch_col, torch.Tensor) + assert torch_col.shape == (2, 10) + assert torch_col.dtype == torch.int64 except ImportError: # Skip check if torch is not installed pass diff --git a/python/python/tests/test_torch.py b/python/python/tests/test_torch.py index 8337568a3..38a367330 100644 --- a/python/python/tests/test_torch.py +++ b/python/python/tests/test_torch.py @@ -148,15 +148,47 @@ def test_permutation_dataloader(mem_db): for batch in dataloader: assert batch["a"].size(0) == 10 - permutation = permutation.with_format("torch") - dataloader = torch.utils.data.DataLoader(permutation, batch_size=10, shuffle=True) + # "torch" produces a list of per-row dicts per batch. The default + # DataLoader collate stacks the per-row dicts back into a batched dict. + torch_perm = permutation.with_format("torch") + batch = next(torch_perm.iter(10, skip_last_batch=False)) + assert isinstance(batch, list) + assert len(batch) == 10 + assert isinstance(batch[0], dict) + assert isinstance(batch[0]["a"], torch.Tensor) + rows = torch_perm.__getitems__([0, 1, 2]) + assert isinstance(rows, list) + assert len(rows) == 3 + assert isinstance(rows[0], dict) + assert isinstance(rows[0]["a"], torch.Tensor) + dataloader = torch.utils.data.DataLoader(torch_perm, batch_size=10, shuffle=True) + for batch in dataloader: + assert isinstance(batch, dict) + assert batch["a"].shape == (10,) + # Spawn-based workers exercise the pickle round-trip path: the new + # transform-as-list shape must survive pickling so workers produce the + # same per-row dicts the parent does. + spawn_loader = torch.utils.data.DataLoader( + torch_perm, + batch_size=10, + num_workers=2, + multiprocessing_context="spawn", + ) + for batch in spawn_loader: + assert isinstance(batch, dict) + assert batch["a"].shape == (10,) + + # "torch_row" returns a list of row tensors. Works with the default + # DataLoader collate (stacks rows into 2D). + row_perm = permutation.with_format("torch_row") + dataloader = torch.utils.data.DataLoader(row_perm, batch_size=10, shuffle=True) for batch in dataloader: assert batch.size(0) == 10 assert batch.size(1) == 1 - permutation = permutation.with_format("torch_col") + col_perm = permutation.with_format("torch_col") dataloader = torch.utils.data.DataLoader( - permutation, collate_fn=lambda x: x, batch_size=10, shuffle=True + col_perm, collate_fn=lambda x: x, batch_size=10, shuffle=True ) for batch in dataloader: assert batch.size(0) == 1