Add bounded native variable-length payload selection

sdm/tensor/var_len.py

@@ -1577,3 +1577,49 @@ def _materialize(
     )
     out._max_length = inp._max_length
     return out
+
+
+def _select_entries(inp: VarLenTensor, index: Tensor) -> VarLenTensor:
+    """Compact selected entries into bounded backing storage.
+
+    The input and index are one-dimensional. Padding after the final offset
+    keeps the output capacity static without truncating selected values.
+    """
+    index = index.where(index >= 0, index + inp.numel())
+    positions = index * inp.stride(0) + inp._logical_storage_offset
+    start = inp._offset[positions]
+    end = inp._offset[positions + 1]
+    max_length = inp._max_length
+    capacity = index.numel() * max_length
+    if capacity > torch.iinfo(inp._offset.dtype).max:
+        return cast(VarLenTensor, inp[index])
+    offset = torch.cat((start.new_zeros(1), (end - start).cumsum(0)))
+    offset = offset.to(inp._offset.dtype)
+    if index.numel() == 0 or inp._data.numel() == 0:
+        data = inp._data.new_empty(0)
+    elif inp._data.device.type == "cpu":
+        byte = torch.arange(max_length, device=inp.device, dtype=offset.dtype)
+        used = byte < (end - start)[:, None]
+        source = (start[:, None] + byte).clamp_max(inp._data.numel() - 1)
+        target = torch.where(used, offset[:-1, None] + byte, capacity)
+        values = inp._data[source].where(used, 0)
+        # Padding writes only to the extra byte, outside the logical payload.
+        data = inp._data.new_zeros(capacity + 1).scatter_(
+            0, target.flatten().long(), values.flatten()
+        )[:-1]
+    else:
+        byte = torch.arange(capacity, device=inp.device, dtype=offset.dtype)
+        owner = torch.searchsorted(offset[1:], byte, right=True)
+        owner = owner.clamp_max(index.numel() - 1)
+        source = start[owner] + byte - offset[owner]
+        used = byte < offset[-1]
+        source = source.where(used, 0)
+        data = inp._data[source].where(used, 0)
+    out = inp.__class__(
+        data=data,
+        offset=offset,
+        valid=inp._valid[positions] if inp._valid is not None else None,
+        size=index.shape,
+    )
+    out._max_length = max_length
+    return out

test/tensor/test_var_len.py

@@ -9,7 +9,7 @@
 from torch import Tensor
 
 from sdm import StringTensor, VarLenTensor
-from sdm.tensor.var_len import _compact
+from sdm.tensor.var_len import _compact, _select_entries
 from sdm.testing import onlyCUDA, withCUDA
 
 
@@ -1034,3 +1034,88 @@ def test_reshape_inference_mode() -> None:
 
     assert not torch.is_inference(view)
     assert torch.is_inference(copied)
+
+
+@withCUDA
+@pytest.mark.parametrize("offset_dtype", [torch.int32, torch.int64])
+@pytest.mark.parametrize("indices", [[], [2, 0, 2, 1], [0, 1, 2], [-1, 0]])
+@pytest.mark.parametrize("layout", ["slice", "stride", "broadcast"])
+def test_native_string_selection(
+    device: torch.device,
+    offset_dtype: torch.dtype,
+    indices: list[int],
+    layout: str,
+) -> None:
+    values = StringTensor.from_arrow(
+        pa.array(["unused", "é🙂", None, "", "long尾"])
+    )
+    assert values._valid is not None
+    values = StringTensor(
+        data=values._data.to(device),
+        offset=values._offset.to(device=device, dtype=offset_dtype),
+        valid=values._valid.to(device),
+        size=values.shape,
+    )
+    if layout == "slice":
+        values = cast(StringTensor, values[1:4])
+    elif layout == "stride":
+        values = cast(StringTensor, values[::2])
+    else:
+        values = cast(StringTensor, values[1:2].expand(3))
+    index = torch.tensor(indices, device=device, dtype=torch.int64)
+    expected = cast(StringTensor, values[index])
+    actual = _select_entries(values, index)
+    assert actual.tolist() == expected.tolist()
+    assert actual.device == expected.device
+    assert actual.data_offset[1].dtype == expected.data_offset[1].dtype
+    assert actual.data_offset[1].diff().equal(expected.data_offset[1].diff())
+
+
+def test_native_string_selection_fullgraph(compilation: None) -> None:
+    values = StringTensor.from_arrow(pa.array(["é🙂", None, "", "long尾"]))
+    index = torch.tensor([3, 0, 3, 1, 2])
+    compiled = torch.compile(_select_entries, fullgraph=True)
+    actual = compiled(values, index)
+    assert actual.tolist() == values[index].tolist()
+
+
+def test_native_string_selection_meta() -> None:
+    values = StringTensor(
+        data=torch.empty(12, dtype=torch.uint8, device="meta"),
+        offset=torch.empty(4, dtype=torch.int32, device="meta"),
+        valid=torch.empty(3, dtype=torch.bool, device="meta"),
+        size=(3,),
+    )
+    index = torch.empty(4, dtype=torch.int64, device="meta")
+    actual = _select_entries(values, index)
+    assert actual.shape == index.shape
+    assert actual.dtype == torch.uint8
+    assert actual.device == index.device
+    assert actual._offset.shape == (index.numel() + 1,)
+    assert actual._offset.dtype == torch.int32
+    assert actual._data.numel() == index.numel() * 12
+    assert actual.valid is not None
+    assert actual.valid.shape == index.shape
+
+
+@pytest.mark.parametrize("empty", [False, True])
+def test_native_string_selection_empty_payload(empty: bool) -> None:
+    values = StringTensor.from_arrow(
+        pa.array([], type=pa.string()) if empty else pa.array(["", None, ""])
+    )
+    index = torch.tensor([] if empty else [2, 0, 1, 2], dtype=torch.int64)
+    actual = _select_entries(values, index)
+    assert actual.tolist() == values[index].tolist()
+
+
+def test_native_string_selection_null_payload() -> None:
+    values = StringTensor(
+        data=torch.tensor([97, 255, 98], dtype=torch.uint8),
+        offset=torch.tensor([0, 1, 2, 3], dtype=torch.int32),
+        valid=torch.tensor([True, False, True]),
+        size=(3,),
+    )
+    index = torch.tensor([1, 2, 0, 1])
+    actual = _select_entries(values, index)
+    assert actual.tolist() == values[index].tolist()
+    assert actual.data_offset[1].diff().tolist() == [1, 1, 1, 1]