Fix gather on sliced nullable string and list tensors (#1289)

> # Summary
>
> Fix wrong values when gathering or materializing nullable
`StringTensor`/`VarLenTensor` built from a sliced Arrow array or cuDF
series.
>
> ```python
> import pyarrow as pa, torch
> from sdm import StringTensor
>
> s = StringTensor.from_arrow(pa.array(["LONGPREFIX", "a", None,
"b"]).slice(1))
> s.tolist()                        # ['a', None, 'b']
> s[torch.tensor([0, 2])].tolist() # main: ['LONGPREFIX', '']; expected:
['a', 'b']
> ```
>
> # Description of the change
>
> Root cause: for nullable inputs, `StringTensor.from_arrow`,
`StringTensor.from_cudf` and `VarLenTensor.from_arrow` slice the
offsets, so `_offset` is a view with its own nonzero storage offset.
`_materialize`, `_layout_view` and `_valid` pass only the wrapper's
storage offset to `torch.as_strided`, which addresses the underlying
storage absolutely, so they read offsets (and validity) from before the
slice.
>
> Fix: add the inner tensor's storage offset in those `as_strided` calls
and subtract it again in `_from_layout_view`. This keeps inner buffers
zero-copy and adds no branches, so it stays compile-friendly.
Normalizing inner tensors in the constructor (cloning views) was the
alternative; it copies buffers and adds a data-dependent branch for
symbolic storage offsets.
>
> # Validation
>
> - New regression tests for sliced nullable Arrow input (`string`,
`large_string`, list) on CPU and sliced nullable cuDF input on CUDA; the
CPU tests fail on `main`.
> - CPU: `test/tensor` (166 passed), `test/processing test/relational
test/nn` (580 passed).
> - GPU (torch 2.14.1+cu130, cuDF): `test/tensor` (219 passed),
`test/processing test/relational test/nn` (826 passed), cuDF repro above
returns `['a', 'b']`.

Author: aw471

sdm/tensor/var_len.py

@@ -1009,11 +1009,13 @@ def _valid(inp: VarLenTensor) -> Tensor | None:
     if inp._valid is None:
         return None
 
+    # `as_strided` offsets are absolute, but `_valid` may be a view:
+    storage_offset = inp._valid.storage_offset() + int(inp.storage_offset())
     return torch.as_strided(
         inp._valid,
         size=inp.size(),
         stride=inp.stride(),
-        storage_offset=int(inp.storage_offset()),
+        storage_offset=storage_offset,
     )
 
 
@@ -1440,22 +1442,25 @@ def _span_len(size: Sequence[int], stride: Sequence[int]) -> int:
 
 
 def _layout_view(inp: VarLenTensor) -> Tensor:
+    # `as_strided` offsets are absolute, but `_offset` may be a view:
+    storage_offset = inp._offset.storage_offset() + int(inp.storage_offset())
     return torch.as_strided(
         inp._offset,
         size=inp.size(),
         stride=inp.stride(),
-        storage_offset=int(inp.storage_offset()),
+        storage_offset=storage_offset,
     )
 
 
 def _from_layout_view(inp: VarLenTensor, view: Tensor) -> VarLenTensor:
+    storage_offset = view.storage_offset() - inp._offset.storage_offset()
     return inp.__class__(
         data=inp._data,
         offset=inp._offset,
         valid=inp._valid,
         size=view.size(),
         stride=view.stride(),
-        storage_offset=0 if view.numel() == 0 else int(view.storage_offset()),
+        storage_offset=0 if view.numel() == 0 else int(storage_offset),
     )
 
 
@@ -1494,12 +1499,7 @@ def _materialize(
     non_blocking: bool = False,
 ) -> VarLenTensor:
     # Use PyTorch's own memory-format semantics to materialize data:
-    start = torch.as_strided(
-        inp._offset,
-        size=inp.size(),
-        stride=inp.stride(),
-        storage_offset=int(inp.storage_offset()),
-    )
+    start = _layout_view(inp)
     start = function(start)
     assert start.storage_offset() == 0
     assert _span_len(start.size(), start.stride()) == start.numel()
@@ -1508,7 +1508,9 @@ def _materialize(
         inp._offset,
         size=inp.size(),
         stride=inp.stride(),
-        storage_offset=int(inp.storage_offset()) + 1,
+        storage_offset=(
+            inp._offset.storage_offset() + int(inp.storage_offset()) + 1
+        ),
     )
     end = function(end)
     assert end.storage_offset() == 0
@@ -1529,14 +1531,8 @@ def _materialize(
         stride=(1,),
         storage_offset=end.storage_offset(),
     )
-    valid: Tensor | None = None
-    if inp._valid is not None:
-        valid = torch.as_strided(
-            inp._valid,
-            size=inp.size(),
-            stride=inp.stride(),
-            storage_offset=int(inp.storage_offset()),
-        )
+    valid = _valid(inp)
+    if valid is not None:
         valid = function(valid).contiguous().view(-1)
 
     offset, index = _compact(start, end)

test/tensor/test_string.py

@@ -113,6 +113,19 @@ def test_from_arrow_zero_string_chunks(arrow_type: pa.DataType) -> None:
     assert tensor.to_arrow().to_pylist() == []
 
 
+@pytest.mark.parametrize("arrow_type", [pa.string(), pa.large_string()])
+def test_from_arrow_sliced_nullable(arrow_type: pa.DataType) -> None:
+    tensor = StringTensor.from_arrow(
+        pa.array(["x", "hi", None, "yo", "z"], type=arrow_type)[1:],
+    )
+
+    assert tensor.tolist() == ["hi", None, "yo", "z"]
+    assert tensor[torch.tensor([0, 2])].tolist() == ["hi", "yo"]
+    transposed = tensor.view(2, 2).t()
+    assert not transposed.is_contiguous()
+    assert transposed.contiguous().tolist() == [["hi", "yo"], [None, "z"]]
+
+
 @onlyCUDA
 def test_from_cudf() -> None:
     cudf = pytest.importorskip("cudf")
@@ -148,6 +161,17 @@ def test_from_cudf_sliced_values() -> None:
     )
 
 
+@onlyCUDA
+def test_from_cudf_sliced_nullable() -> None:
+    cudf = pytest.importorskip("cudf")
+
+    tensor = StringTensor.from_cudf(cudf.Series(["x", "hi", None, "yo"])[1:])
+
+    assert tensor.tolist() == ["hi", None, "yo"]
+    index = torch.tensor([0, 2], device=tensor.device)
+    assert tensor[index].tolist() == ["hi", "yo"]
+
+
 @onlyCUDA
 def test_from_cudf_all_empty_values() -> None:
     cudf = pytest.importorskip("cudf")

test/tensor/test_var_len.py

@@ -216,6 +216,31 @@ def test_arrow() -> None:
     assert tensor.tolist() == [[1, 2], None, [3]]
 
 
+def test_arrow_sliced_nullable() -> None:
+    tensor = VarLenTensor.from_arrow(
+        pa.array([[0, 0, 0], [1], None, [2, 3]])[1:],
+    )
+
+    assert tensor.tolist() == [[1], None, [2, 3]]
+    assert tensor[torch.tensor([0, 2])].tolist() == [[1], [2, 3]]
+
+
+def test_sliced_offset_and_valid_views() -> None:
+    # Rows [[0, 0, 0], [1], None, [2, 3]]; inner tensors are views from row 1.
+    offset = torch.tensor([0, 3, 4, 4, 6])
+    valid = torch.tensor([True, True, False, True])
+    tensor = VarLenTensor(
+        data=torch.tensor([0, 0, 0, 1, 2, 3]),
+        offset=offset[1:],
+        valid=valid[1:],
+        size=(2,),
+        storage_offset=1,
+    )
+
+    assert tensor.tolist() == [None, [2, 3]]
+    assert tensor[torch.tensor([1, 0])].tolist() == [[2, 3], None]
+
+
 def test_arrow_bool() -> None:
     tensor = VarLenTensor.from_arrow(
         pa.array(