feat(python)!: align Permutation.with_format("torch") with HuggingFace set_format("torch") (#3369)

Closes #3245.

> **BREAKING CHANGE:** `with_format("torch")` no longer returns a list
of stacked row tensors. It now returns per-row dicts so PyTorch's
default `DataLoader` collate stacks them into `{col: tensor(B,)}`.
Switch to `with_format("torch_row")` to keep the old shape.

### What changed

`"torch"` now returns a list of per-row dicts (`[{col: tensor}, ...]`)
at every indexed access path. The default `DataLoader` collate stacks
them into a column-keyed batched dict, no custom `collate_fn` needed.
The old shape is preserved under a new `"torch_row"` literal.
`"torch_col"` is unchanged.

The unbatching lives inside the transform (`batch_to_tensor_dict`), not
`__getitems__`, so the shape survives pickling and works under
`DataLoader(num_workers>0, multiprocessing_context="spawn")`.

### Format comparison

| Format | `iter(batch_size=N)` | `__getitems__([0,1,2])` | `DataLoader`
default collate |
|---|---|---|---|
| `"torch"` (new) | `list[{col: tensor}]` length N | `list[{col:
tensor}]` length 3 | `{col: tensor(B,)}` |
| `"torch_row"` (old `"torch"` behavior) | `list[tensor(n_cols,)]`
length N | `list[tensor(n_cols,)]` length 3 | `tensor(B, n_cols)` |
| `"torch_col"` (unchanged) | `tensor(n_cols, N)` | `tensor(n_cols, 3)`
| needs `collate_fn=lambda x: x` |

Output matches HuggingFace `Dataset.set_format("torch")` on container
shape, keys, and values at every access path. The only divergence:
HuggingFace downcasts `float64` to `torch.float32` by default, LanceDB
preserves dtype. Verified by `scripts/verify_torch_format.py`.

### Migration

```python
# Old default — column names lost, shape was tensor(B, n_cols)
DataLoader(Permutation.identity(table).with_format("torch"))

# New default — column names preserved
DataLoader(Permutation.identity(table).with_format("torch"))     # {col: tensor(B,)}

# Keep old behavior
DataLoader(Permutation.identity(table).with_format("torch_row")) # tensor(B, n_cols)
```

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)

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)
+    ]

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

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