diff --git a/docs/superpowers/plans/2026-07-27-ragged-string-getitem.md b/docs/superpowers/plans/2026-07-27-ragged-string-getitem.md new file mode 100644 index 0000000..011ae0b --- /dev/null +++ b/docs/superpowers/plans/2026-07-27-ragged-string-getitem.md @@ -0,0 +1,546 @@ +# String-under-axis integer indexing (issue #71) Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make integer indexing on a string-under-axis `Ragged` return one element per string instead of silently concatenating the whole group into one blob. + +**Architecture:** Integer indexing peels one real ragged level. For a string-under-axis leaf (`offsets` non-empty **and** `str_offsets` set) the peel lands on the standalone opaque-string layout (`offsets == []`, `str_offsets` set, `shape == (k,)`) — which Spec C already defines as the zero-real-level special case of the same layout. Two call sites construct this: the plain `__getitem__` integer branch and the record-row integer branch. Both mirror the narrowing that `_slice_contig_string` already performs for slices. + +**Tech Stack:** Python 3.9+, NumPy, pytest. Pure Python layer — no Rust (`src/`, `crates/`) changes. + +## Global Constraints + +- Spec: `docs/superpowers/specs/2026-07-27-ragged-string-getitem-design.md`. +- Data must stay **zero-copy**: the result's `data` is a view of the parent's buffer. Only the small `(k+1,)` offsets slice is copied (to rebase to zero, which `is_contiguous` requires for opaque strings — `python/seqpro/rag/_core.py:319-321`). +- The **standalone/flat** opaque-string case (`offsets == []`) is unchanged: `flat[0]` still returns `bytes`. The `self._layout.offsets` guard distinguishes it. +- No Python loops over elements — this is a per-batch-item accessor (repo rule: "No naive NumPy in hot paths"). +- Breaking change. `major_version_zero = true` in `pyproject.toml:64`, so a `!` conventional commit on 0.x produces a **minor** bump (0.21.2 → 0.22.0), which is what the spec calls for. Do **not** hand-edit `CHANGELOG.md` — commitizen generates it on bump. +- Test command (run from the worktree root; the worktree has no `.pixi` of its own, so borrow the main checkout's dev env and point `PYTHONPATH` at the worktree source): + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/test_ragged_core.py -q +``` + +## File Structure + +| File | Responsibility | Change | +|---|---|---| +| `python/seqpro/rag/_core.py:722-731` | `Ragged.__getitem__` integer branch — plain string-under-axis | Modify (Task 1) | +| `python/seqpro/rag/_core.py:1216-1226` | `_getitem_record_rows` integer branch — opaque-string record fields | Modify (Task 2) | +| `tests/test_ragged_core.py` | Existing home of the string-under-axis test section (`test_string_under_axis_integer_index` at line 733) | Modify + add (Tasks 1, 2) | +| `skills/seqpro/SKILL.md` | Public API skill doc; CLAUDE.md requires an update for any breaking change | Modify (Task 3) | + +`_getitem_record_rows_r2` (`_core.py:1243`) delegates to `Ragged(fl)[where]` and inherits Task 1's fix — no separate change, but Task 2 covers it with a test. + +--- + +### Task 1: Plain string-under-axis integer indexing + +**Files:** +- Modify: `python/seqpro/rag/_core.py:722-731` +- Test: `tests/test_ragged_core.py` (replace `test_string_under_axis_integer_index` at line 733; add new tests after it) + +**Interfaces:** +- Consumes: `RaggedLayout(data=..., offsets=..., shape=..., str_offsets=...)` from `python/seqpro/rag/_layout.py` (already imported at `_core.py:10`). +- Produces: `Ragged.__getitem__(int)` on a string-under-axis `Ragged` returns `Ragged` with `offsets == []`, `str_offsets` set, `shape == (k,)`, `is_string is True`. Task 2 relies on this same construction shape. + +- [ ] **Step 1: Replace the test that pins the old behavior** + +`tests/test_ragged_core.py:732-741` currently reads: + +```python +def test_string_under_axis_integer_index(): + rag = Ragged.from_offsets( + np.frombuffer(b"TTGG", "S1"), + (2, None), + np.array([0, 1, 2]), + str_offsets=np.array([0, 2, 4]), + ) + assert rag[0] == b"TT" + assert rag[1] == b"GG" +``` + +It uses exactly one string per group, so it can never distinguish concatenation from per-string indexing — that is why the bug survived. Replace it with: + +```python +def test_string_under_axis_integer_index(): + """Peeling one group yields the standalone opaque-string layout (Spec C).""" + rag = Ragged.from_offsets( + np.frombuffer(b"TTGG", "S1"), + (2, None), + np.array([0, 1, 2]), + str_offsets=np.array([0, 2, 4]), + ) + row = rag[0] + assert isinstance(row, Ragged) + assert row.is_string and row.shape == (1,) + assert len(row) == 1 + assert row[0] == b"TT" + assert rag[1][0] == b"GG" +``` + +- [ ] **Step 2: Add the failing tests from the issue** + +Append immediately after the test above: + +```python +def _issue71_pair(): + """String-under-axis and numeric Ragged sharing one offsets object. + + Groups: 0 -> ('A', 'GG'), 1 -> ('TC',). + """ + data = np.frombuffer(b"AGGTC", dtype="S1") + outer = np.array([0, 2, 3], dtype=OFFSET_TYPE) # group -> string index + inner = np.array([0, 1, 3, 5], dtype=OFFSET_TYPE) # string -> byte index + s = Ragged.from_offsets(data, (2, None), outer, str_offsets=inner) + n = Ragged.from_offsets(np.array([10, 20, 30], dtype=np.int32), (2, None), outer) + return s, n + + +def test_string_under_axis_index_preserves_boundaries(): + """Issue #71: interior str_offsets boundaries must survive an integer index.""" + s, _ = _issue71_pair() + row = s[0] + assert len(row) == 2 + assert row[0] == b"A" + assert row[1] == b"GG" + assert list(s[1]) == [b"TC"] + + +def test_string_under_axis_index_matches_lengths(): + """len(s[i]) must equal s.lengths[i] and the numeric row length.""" + s, n = _issue71_pair() + for i in range(len(s)): + assert len(s[i]) == int(s.lengths[i]) + assert len(s[i]) == len(n[i]) + + +def test_string_under_axis_index_is_zero_copy(): + s, _ = _issue71_pair() + assert np.shares_memory(s[0].data, s.data) + + +def test_string_under_axis_index_empty_group(): + """A group holding zero strings peels to a length-0 result.""" + rag = Ragged.from_offsets( + np.frombuffer(b"AC", "S1"), + (2, None), + np.array([0, 0, 2], dtype=OFFSET_TYPE), # group 0 empty + str_offsets=np.array([0, 1, 2], dtype=OFFSET_TYPE), + ) + assert len(rag[0]) == 0 + assert len(rag[1]) == 2 + + +def test_string_under_axis_index_agrees_with_to_chars(): + s, _ = _issue71_pair() + chars = s.to_chars() + for i in range(len(s)): + for j in range(len(s[i])): + assert s[i][j] == chars[i][j].tobytes() + + +def test_string_under_axis_index_negative_and_oob(): + s, _ = _issue71_pair() + assert list(s[-1]) == [b"TC"] + with pytest.raises(IndexError): + s[5] + + +def test_string_under_axis_index_multidim(): + """(batch, ploidy, ~variants): reaches the flat branch via _getitem_multidim.""" + data = np.frombuffer(b"AGGTCNNAC", dtype="S1") + o0 = np.array([0, 2, 3, 4, 6], dtype=OFFSET_TYPE) # 4 segments -> string idx + i0 = np.array([0, 1, 3, 5, 7, 9], dtype=OFFSET_TYPE) # 6 boundaries -> bytes + rag = Ragged.from_offsets(data, (2, 2, None), o0, str_offsets=i0) + row = rag[0] # -> (2, None) string-under-axis + assert list(row[0]) == [b"A", b"GG"] + assert list(row[1]) == [b"TC"] + + +def test_standalone_string_index_still_returns_bytes(): + """Regression guard: the zero-real-level case is the terminal peel.""" + flat = Ragged.from_offsets( + np.frombuffer(b"cathithere", "S1"), (3,), np.array([0, 3, 5, 10]) + ) + assert flat[0] == b"cat" + assert flat[-1] == b"there" + assert list(flat) == [b"cat", b"hi", b"there"] +``` + +- [ ] **Step 3: Run the tests to verify they fail** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/test_ragged_core.py -q -k "string_under_axis or standalone_string" +``` + +Expected: `test_string_under_axis_index_preserves_boundaries` fails (`len(row)` raises `TypeError`/returns 3 rather than 2 — the current return is `bytes`), `test_string_under_axis_integer_index` fails on `isinstance(row, Ragged)`, and the other new per-string tests fail. `test_standalone_string_index_still_returns_bytes` must PASS already. + +- [ ] **Step 4: Implement the fix** + +In `python/seqpro/rag/_core.py`, replace the integer branch at lines 722-731: + +```python + if isinstance(where, (int, np.integer)): + lo, hi = int(starts[where]), int(stops[where]) + if self._rl.str_offsets is not None and self._layout.offsets: + # string-under-axis: outer offsets index variants -> map to bytes via str_offsets + so = self._rl.str_offsets + return self._rl.data[int(so[lo]) : int(so[hi])].tobytes() + row = self._rl.data[lo:hi] + if self._rl.is_string: + return row.tobytes() + return row +``` + +with: + +```python + if isinstance(where, (int, np.integer)): + lo, hi = int(starts[where]), int(stops[where]) + if self._rl.str_offsets is not None and self._layout.offsets: + # string-under-axis: peel the real level -> standalone opaque + # string (k,), preserving the per-string boundaries that live in + # str_offsets. Concatenating here would drop them (issue #71). + return Ragged(_peel_string_row(self._rl, lo, hi)) + row = self._rl.data[lo:hi] + if self._rl.is_string: + return row.tobytes() + return row +``` + +Then add this module-level helper next to the other layout helpers — place it immediately above `class Ragged` in `python/seqpro/rag/_core.py` (Task 2 reuses it, which is why it is a free function rather than a method: `_getitem_record_rows` operates on per-field `RaggedLayout`s, not on `self`): + +```python +def _peel_string_row( + rl: "RaggedLayout[Any]", lo: int, hi: int +) -> "RaggedLayout[Any]": + """Peel strings ``[lo, hi)`` off a string-under-axis leaf. + + Returns the standalone opaque-string layout (``offsets == []``, + ``str_offsets`` set, ``shape == (hi - lo,)``) — the zero-real-level special + case of string-under-axis (Spec C Section 2). The data buffer is a view; + only the ``(k + 1,)`` offsets slice is copied, rebased to zero as + ``is_contiguous`` requires. + """ + so = rl.str_offsets + assert so is not None # caller guarantees a string leaf + b0 = int(so[lo]) + return RaggedLayout( + data=rl.data[b0 : int(so[hi])], + offsets=[], + shape=(hi - lo,), + str_offsets=so[lo : hi + 1] - b0, + ) +``` + +Note `lo`/`hi` are indices in **string** space, so `so[lo : hi + 1]` is the right slice whether `offsets[0]` is 1-D canonical or a lazy `(2, M)` gather layout — `_starts_stops()` normalizes both. + +- [ ] **Step 5: Run the tests to verify they pass** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/test_ragged_core.py -q -k "string_under_axis or standalone_string" +``` + +Expected: all PASS. + +- [ ] **Step 6: Run the full ragged suite for regressions** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/ -q +``` + +Expected: no new failures. If another test asserts the old concatenating behavior, judge it the same way as `test_string_under_axis_integer_index` — update it to index one level deeper, and note it in the commit body. + +- [ ] **Step 7: Commit** + +```bash +git add python/seqpro/rag/_core.py tests/test_ragged_core.py +git commit -m "fix(rag)!: preserve string boundaries when indexing string-under-axis + +Integer indexing on a string-under-axis Ragged concatenated the whole +group into one bytes, dropping the per-string boundaries already held in +str_offsets. It now peels to the standalone opaque-string layout (k,), +so len(s[i]) == s.lengths[i] and s[i][j] is one string. + +BREAKING CHANGE: Ragged.__getitem__ with an integer on a string-under-axis +array returns a Ragged of strings, not one concatenated bytes. Use +b\"\".join(s[i]) for the old value. + +Refs #71" +``` + +--- + +### Task 2: Record-layout string fields + +**Files:** +- Modify: `python/seqpro/rag/_core.py:1216-1226` +- Test: `tests/test_ragged_core_records.py` + +**Interfaces:** +- Consumes: `_peel_string_row(rl, lo, hi) -> RaggedLayout` from Task 1. +- Produces: `_getitem_record_rows` integer branch returns `dict[str, NDArray | Ragged]` where opaque-string fields are `Ragged` (standalone opaque-string layout) and numeric/char fields stay `ndarray`; every entry has length `hi - lo`. + +- [ ] **Step 1: Write the failing tests** + +Append to `tests/test_ragged_core_records.py`: + +```python +def _issue71_record(): + """Record with an opaque-string field and a numeric field sharing offsets. + + Groups: 0 -> ('A', 'GG') / starts (1, 2), 1 -> ('TC',) / starts (3,). + Group 0's second allele is multi-byte, so a 1-byte coincidence cannot + mask a regression. + """ + outer = np.array([0, 2, 3], dtype=OFFSET_TYPE) + alt = Ragged.from_offsets( + np.frombuffer(b"AGGTC", dtype="S1"), + (2, None), + outer, + str_offsets=np.array([0, 1, 3, 5], dtype=OFFSET_TYPE), + ) + start = Ragged.from_offsets(np.array([1, 2, 3], dtype=np.int32), (2, None), outer) + return Ragged.from_fields({"alt": alt, "start": start}) + + +def test_record_row_string_field_preserves_boundaries(): + """Issue #71: a string field peeled from a record row keeps its boundaries.""" + row = _issue71_record()[0] + assert isinstance(row, dict) + assert list(row["alt"]) == [b"A", b"GG"] + np.testing.assert_array_equal(row["start"], np.array([1, 2], dtype=np.int32)) + + +def test_record_row_fields_have_matching_lengths(): + """Every field of a peeled row must have the same length, so zip aligns.""" + rec = _issue71_record() + for i in range(len(rec)): + row = rec[i] + assert len(row["alt"]) == len(row["start"]) + row0 = rec[0] + assert list(zip(row0["start"], row0["alt"])) == [(1, b"A"), (2, b"GG")] + + +def test_record_row_string_field_is_zero_copy(): + rec = _issue71_record() + assert np.shares_memory(rec[0]["alt"].data, rec["alt"].data) + + +def test_record_multidim_row_string_field_preserves_boundaries(): + """(batch, ploidy, ~variants) record: rec[0][h] routes via _getitem_record_rows_r2.""" + outer = np.array([0, 2, 3, 4, 6], dtype=OFFSET_TYPE) + alt = Ragged.from_offsets( + np.frombuffer(b"AGGTCNNAC", dtype="S1"), + (2, 2, None), + outer, + str_offsets=np.array([0, 1, 3, 5, 7, 9], dtype=OFFSET_TYPE), + ) + start = Ragged.from_offsets( + np.arange(6, dtype=np.int32), (2, 2, None), outer + ) + rec = Ragged.from_fields({"alt": alt, "start": start}) + row = rec[0][0] + assert list(row["alt"]) == [b"A", b"GG"] + assert len(row["alt"]) == len(row["start"]) +``` + +- [ ] **Step 2: Run the tests to verify they fail** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/test_ragged_core_records.py -q -k "issue71 or record_row or record_multidim" +``` + +Expected: FAIL. Today `rec[0]["alt"]` is `array([b'A', b'G', b'G'], dtype='|S1')` — a 3-element `S1` char array next to a 2-element numeric array. + +- [ ] **Step 3: Implement the fix** + +In `python/seqpro/rag/_core.py`, replace the integer branch at lines 1216-1226: + +```python + if isinstance(where, (int, np.integer)): + lo, hi = int(starts[where]), int(stops[where]) + out: dict[str, Any] = {} + for name, fl in rec.fields.items(): + if fl.str_offsets is not None: + so = fl.str_offsets + row = fl.data[int(so[lo]) : int(so[hi])] + else: + row = fl.data[lo:hi] + out[name] = row + return out +``` + +with: + +```python + if isinstance(where, (int, np.integer)): + lo, hi = int(starts[where]), int(stops[where]) + out: dict[str, Any] = {} + for name, fl in rec.fields.items(): + if fl.str_offsets is not None: + # Each field carries its own str_offsets (Spec C Section 5); + # peel it against the shared lo/hi so every field of the row + # has the same length (issue #71). + out[name] = Ragged(_peel_string_row(fl, lo, hi)) + else: + out[name] = fl.data[lo:hi] + return out +``` + +- [ ] **Step 4: Run the tests to verify they pass** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/test_ragged_core_records.py -q -k "issue71 or record_row or record_multidim" +``` + +Expected: all PASS. + +- [ ] **Step 5: Run the full suite** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/ -q +``` + +Expected: no new failures. + +- [ ] **Step 6: Commit** + +```bash +git add python/seqpro/rag/_core.py tests/test_ragged_core_records.py +git commit -m "fix(rag)!: preserve string boundaries in peeled record rows + +_getitem_record_rows returned a string field as the raw concatenated S1 +buffer, so a peeled row mixed a 3-char array with a 2-element numeric +array. String fields now peel to the standalone opaque-string layout, +giving every field of the row the same length. + +BREAKING CHANGE: peeling a record row returns opaque-string fields as a +Ragged of strings, not a concatenated S1 array. + +Refs #71" +``` + +--- + +### Task 3: Lint, typecheck, and skill docs + +**Files:** +- Modify: `skills/seqpro/SKILL.md` + +**Interfaces:** +- Consumes: the public behavior established in Tasks 1 and 2. +- Produces: nothing other tasks depend on. + +- [ ] **Step 1: Document the behavior in the "do this, not that" table** + +In `skills/seqpro/SKILL.md`, add this row to the `### Working with `Ragged` — do this, not that` table (the table starting at line ~92), after the `Count top-level rows` row: + +```markdown +| Index one group of an opaque-string `Ragged` | `rag[i]` → `Ragged` of `bytes`, one per string (`len(rag[i]) == rag.lengths[i]`); `rag[i][j]` is one `bytes` | `b"".join(rag[i])`-style concatenation — that was the pre-0.22 behavior and it dropped the per-string boundaries | +``` + +- [ ] **Step 2: Document the layout rule in the record section** + +In `skills/seqpro/SKILL.md`, append to the bullet list at the end of the `### Record-layout `Ragged` (multi-field)` section (after the `view` and `apply` bullet): + +```markdown +- Peeling a row (`rag[i]` where `i` is an integer) returns a **dict** whose entries all have the same length: numeric/char fields as `ndarray`, opaque-string fields as a `Ragged` of `bytes`. This is what makes `zip(row["start"], row["alt"])` correct. +``` + +- [ ] **Step 3: Run lint and typecheck** + +```bash +cd /carter/users/dlaub/projects/ML4GLand/SeqPro/.claude/worktrees/issue-71-string-under-axis-getitem +PY=/carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin +$PY/ruff check python/ tests/ && $PY/ruff format --check python/ tests/ +$PY/pyrefly check python +``` + +Expected: clean. Fix anything reported before committing. + +- [ ] **Step 4: Run the full suite one last time** + +```bash +PYTHONPATH=python /carter/users/dlaub/projects/ML4GLand/SeqPro/.pixi/envs/dev/bin/python -m pytest tests/ -q +``` + +Expected: all pass. + +- [ ] **Step 5: Commit** + +```bash +git add skills/seqpro/SKILL.md +git commit -m "docs(skill): document string-under-axis integer indexing + +Refs #71" +``` + +- [ ] **Step 6: Push and open a draft PR** + +```bash +git push -u origin worktree-issue-71-string-under-axis-getitem +gh pr create --draft --title "fix(rag)!: preserve string boundaries when indexing a string-under-axis Ragged" --body "$(cat <<'EOF' +Closes #71. + +Integer indexing on a string-under-axis `Ragged` concatenated a whole group +into one blob, dropping the per-string boundaries already sitting in +`str_offsets`. `s.lengths[0] == 2` but `len(s[0]) == 3`. + +Two sites had the defect: + +- `Ragged.__getitem__` integer branch — the one reported. +- `_getitem_record_rows` integer branch — worse, and the one GenVarLoader + hits: a peeled row mixed a concatenated `S1` char array (not even `bytes`) + with a per-variant numeric array in the same dict. + +Both now peel to the standalone opaque-string layout (`offsets == []`, +`str_offsets` set, `shape == (k,)`), which Spec C already defines as the +zero-real-level special case of string-under-axis. Zero-copy on data; only +the small offsets slice is copied. + +This makes `zip(rv.start[0][h], rv.alt[0][h])` correct in GenVarLoader +(mcvickerlab/GenVarLoader#330). + +**Breaking:** the integer index now returns a `Ragged` of strings rather than +one concatenated `bytes`. `b"".join(s[i])` recovers the old value. +`major_version_zero` is set, so this bumps 0.21.2 → 0.22.0. + +The existing `test_string_under_axis_integer_index` pinned the old behavior +but used one string per group, so it could never distinguish the two +interpretations — that is why the bug survived. It has been rewritten. + +Spec: `docs/superpowers/specs/2026-07-27-ragged-string-getitem-design.md` +Plan: `docs/superpowers/plans/2026-07-27-ragged-string-getitem.md` +EOF +)" +``` + +--- + +## Self-Review + +**Spec coverage:** + +| Spec item | Task | +|---|---| +| Site 1 — `__getitem__` integer branch | 1 | +| Site 2 — `_getitem_record_rows` | 2 | +| `_getitem_record_rows_r2` inherits site 1 | 2 (Step 1, `test_record_multidim_row_string_field_preserves_boundaries`) | +| Standalone/flat case unchanged | 1 (`test_standalone_string_index_still_returns_bytes`) | +| Test 1 — issue repro | 1 | +| Test 2 — numeric/string parity | 1 | +| Test 3 — multi-dim | 1 | +| Test 4 — record row with indel | 2 | +| Test 5 — zero-copy | 1 and 2 | +| Test 6 — empty group | 1 | +| Test 7 — `to_chars()` agreement | 1 | +| Test 8 — negative index / OOB | 1 | +| Test 9 — standalone regression | 1 | +| Minor bump via `!` commit | 1, 2 (commit messages) | +| `skills/seqpro/SKILL.md` update | 3 | + +**Type consistency:** `_peel_string_row(rl: RaggedLayout, lo: int, hi: int) -> RaggedLayout` is defined in Task 1 Step 4 and used with that exact signature in Task 2 Step 3. Both call sites wrap it in `Ragged(...)`. + +**Imports:** the new tests use `Ragged`, `OFFSET_TYPE`, `np`, and `pytest`. `tests/test_ragged_core.py` already imports all four. `tests/test_ragged_core_records.py` does **not** import `OFFSET_TYPE` — Task 2 Step 1 must extend its existing `from seqpro.rag._utils import lengths_to_offsets` (line 6) to `from seqpro.rag._utils import OFFSET_TYPE, lengths_to_offsets`. diff --git a/docs/superpowers/specs/2026-07-27-ragged-string-getitem-design.md b/docs/superpowers/specs/2026-07-27-ragged-string-getitem-design.md new file mode 100644 index 0000000..9bbbdb4 --- /dev/null +++ b/docs/superpowers/specs/2026-07-27-ragged-string-getitem-design.md @@ -0,0 +1,203 @@ +# Design: integer indexing on a string-under-axis `Ragged` (issue #71) + +## Problem + +For a **string-under-axis** `Ragged` — non-empty `offsets` plus `str_offsets` set, +per Spec C Section 2 — integer indexing collapses a whole group into one +concatenated blob, silently discarding the per-string boundaries that +`str_offsets` is already carrying. + +```python +data = np.frombuffer(b"AGGTC", dtype="S1") +outer = np.array([0, 2, 3]) # group -> string index +inner = np.array([0, 1, 3, 5]) # string -> byte index + +s = Ragged.from_offsets(data, (2, None), outer, str_offsets=inner) +n = Ragged.from_offsets(np.array([10, 20, 30], np.int32), (2, None), outer) +``` + +``` +s.lengths -> [2 1] +n[0] -> array([10, 20], dtype=int32) # 2 elements, matches lengths[0] +s[0] -> b'AGG' # ONE bytes; is that 'A'+'GG' or 'AG'+'G'? +``` + +`s.lengths[0] == 2` but `len(s[0]) == 3`. The boundary at `str_offsets[1] == 1` is +dropped. + +### It fails silently, data-dependently, and in two places + +Consumers zip a numeric field against a string field on the shared axis. While +every string is one byte the lengths coincide and the misalignment is invisible; +the first multi-byte entry (an indel) breaks it. + +The issue reports one site. There are **two**: + +| # | Site | Layout | +|---|---|---| +| 1 | `Ragged.__getitem__` integer branch, `python/seqpro/rag/_core.py:722-731` | plain string-under-axis | +| 2 | `_getitem_record_rows` integer branch, `python/seqpro/rag/_core.py:1216-1226` | opaque-string **field** of a record | + +Site 2 is the one GenVarLoader actually hits. `RaggedVariants` is a record +`Ragged` with shape `(batch, ploidy, ~variants)` whose `alt`/`ref` are +opaque-string fields sharing one offsets object with numeric `start`/`ilen` +(`genvarloader/_dataset/_rag_variants.py:193`). Peeling a row today gives: + +``` +rec[0] -> {'alt': array([b'A', b'G', b'G'], dtype='|S1'), # 3 chars + 'start': array([1, 2], dtype=int32)} # 2 variants +``` + +The string field is not even `bytes` — it is the raw concatenated `S1` char +buffer, sitting next to a per-variant numeric array in the same dict. This is +strictly worse than the reported symptom. + +`_getitem_record_rows_r2` (`_core.py:1243`) delegates to `Ragged(fl)[where]`, so +it inherits site 1's behavior and needs no separate change. + +### Why fix rather than document + +`to_numpy()` and `to_packed()` both *refuse* on this layout rather than return +something lossy (`_core.py:1777-1780`, `:1801-1805`, `:1282-1285`). Integer +indexing is the one accessor that quietly returns a wrong-shaped answer. No +information is lost — `str_offsets` holds the boundaries — so this is purely a +question of what the accessor hands back. + +## Decision + +**Integer indexing that peels the last real ragged level off a string-under-axis +`Ragged` returns the standalone opaque-string layout: `offsets == []`, +`str_offsets` set, `shape == (k,)`, `is_string == True`.** + +This is not a new return type. Spec C's own layout table +(`docs/superpowers/specs/2026-06-20-rust-ragged-nested-design.md:129-135`) defines +the standalone opaque string of Spec A/B as *"the zero-real-level special case"* +of string-under-axis: + +| | standalone string (Spec B) | string-under-axis (Spec C) | chars | +|---|---|---|---| +| `offsets` | `[]` | `[O0]` | `[O0, O1]` | +| `str_offsets` | set | set | `None` | +| `.shape` | `(N,)` | `(*leading, None)` | `(*leading, None, None)` | + +Peeling the one real level off `(N, ~var)` therefore lands exactly on `(k,)`. +"Return the layout with one fewer real level" is what integer indexing already +means everywhere else in the class; this makes the string case obey it too. + +The returned layout already supports everything a consumer needs (verified on +0.21.2): + +``` +flat.shape (3,) is_string True len 3 +flat[0] b'cat' flat[-1] b'there' +list(flat) -> [b'cat', b'hi', b'there'] +zip(...) -> [(1, b'cat'), (2, b'hi'), (3, b'there')] +flat[5] -> IndexError flat.lengths -> [3 2 5] +``` + +So `zip(rv.start[0][h], rv.alt[0][h])` — the GenVarLoader pattern from +mcvickerlab/GenVarLoader#330 — becomes correct, and `len(s[i]) == s.lengths[i]` +holds. + +### Alternatives rejected + +- **Object array of `bytes`** (issue option 1). Same ergonomics as the chosen + design but allocates `k` Python `bytes` objects on every row access, against + the repo's "no naive NumPy / no Python loops in hot paths" rule — and this is a + per-batch-item accessor. It also introduces a return type the layout algebra + does not otherwise use. Strictly dominated. +- **Char `Ragged`, i.e. `to_chars()[i]`** (issue option 2). Shape `(k, None)` of + `S1`; forces every consumer to re-assemble strings. Loses the "these are + strings" information the layout is carrying. +- **Raise `NotImplementedError`** (issue option 3). Follows the + `to_numpy()`/`to_packed()` precedent, but indexing is far more central than + those, and the information needed to answer correctly is present. Would leave + GenVarLoader with no ergonomic path at all. + +## Design + +### Site 1 — `_core.py:722-731` + +Mirror the existing `_slice_contig_string` (`_core.py:511-541`), which already +performs exactly this narrowing for the slice case: + +```python +if isinstance(where, (int, np.integer)): + lo, hi = int(starts[where]), int(stops[where]) + if self._rl.str_offsets is not None and self._layout.offsets: + # string-under-axis: peel the real level -> standalone opaque string (k,) + so = self._rl.str_offsets + b0 = int(so[lo]) + return Ragged(RaggedLayout( + data=self._rl.data[b0 : int(so[hi])], + offsets=[], + shape=(hi - lo,), + str_offsets=so[lo : hi + 1] - b0, + )) + ... +``` + +- Data is a view; only the small `(k+1,)` offsets slice is copied, to rebase to + zero as `is_contiguous` requires for opaque strings (`_core.py:319-321`). +- Correct whether `offsets[0]` is 1-D canonical or a lazy `(2, M)` gather layout: + `starts`/`stops` come from `_starts_stops()`, and a group is a contiguous run in + string-index space under either encoding. +- `lo == hi` (empty group) yields a well-formed length-0 result. + +### Site 2 — `_getitem_record_rows`, `_core.py:1216-1226` + +Apply the same construction per opaque-string field; numeric and char fields are +untouched. The returned dict then has one entry per field, all of length +`hi - lo`. + +Each field keeps its **own** `str_offsets` (Spec C Section 5), so fields are +narrowed independently against the shared `lo`/`hi`. + +### Explicitly unchanged + +- **The standalone/flat case** (`offsets == []`). `flat[0]` still returns `bytes` + — that is the terminal peel and what Spec A/B consumers depend on. The guard + `self._layout.offsets` already distinguishes the two. +- **Slicing, masking, fancy indexing.** Already correct; they preserve + `str_offsets`. +- **`_getitem_record_rows_r2`**, which delegates to site 1. +- **Rust (`src/`, `crates/`).** Python-layer accessor only. + +## Compatibility + +Breaking for anyone relying on the concatenation. Warrants a minor bump and a +CHANGELOG entry. The old value remains reachable as +`b"".join(s[i])` for a caller who genuinely wanted the concatenated group. + +`tests/test_ragged_core.py:733` (`test_string_under_axis_integer_index`) pins the +current behavior and must be updated. Note it uses exactly one string per group, +so it could never have distinguished the two interpretations — that is why the +bug survived. + +## Testing + +1. The issue's exact repro: `len(s[0]) == s.lengths[0]`, `s[0][0] == b'A'`, + `s[0][1] == b'GG'`, `list(s[1]) == [b'TC']`. +2. Numeric/string parity: `len(s[i]) == len(n[i])` for every `i`, for a + string-under-axis and a numeric `Ragged` sharing one offsets object. +3. Multi-dim `(batch, ploidy, ~variants)`: `s[0][h]` yields per-variant strings + (the shape that reaches site 1 via `_getitem_multidim`). +4. Record row: `rec[0]` — every field the same length, `zip` aligned, with at + least one multi-byte (indel) entry so 1-byte coincidence cannot mask a + regression. +5. Zero-copy: the result's data buffer shares memory with the parent's. +6. Empty group (`lo == hi`) → length-0 result, `len(...) == 0`. +7. Agreement with `to_chars()`: `s[i][j] == s.to_chars()[i][j].tobytes()`. +8. Negative index and out-of-bounds `IndexError` preserved. +9. Standalone string regression: `flat[0]` still returns `bytes`. + +## Out of scope + +- Changing `to_numpy()` / `to_packed()` behavior on this layout. +- The GenVarLoader-side fix (mcvickerlab/GenVarLoader#330) — tracked separately; + this change is what unblocks it. + +## Docs + +`skills/seqpro/SKILL.md` must be updated in the same PR — CLAUDE.md requires it +for any breaking change or public-behavior change. diff --git a/python/seqpro/rag/_core.py b/python/seqpro/rag/_core.py index 00fab44..37baf8c 100644 --- a/python/seqpro/rag/_core.py +++ b/python/seqpro/rag/_core.py @@ -59,6 +59,29 @@ def _build_layout( return RaggedLayout(data=data, offsets=off_list, shape=shape) +def _peel_string_row(rl: "RaggedLayout[Any]", lo: int, hi: int) -> "RaggedLayout[Any]": + """Peel strings ``[lo, hi)`` off a string-under-axis leaf. + + Returns the standalone opaque-string layout (``offsets == []``, + ``str_offsets`` set, ``shape == (hi - lo,)``) — the zero-real-level special + case of string-under-axis (Spec C Section 2). The data buffer is a view; + only the ``(k + 1,)`` offsets slice is copied, rebased to zero as + ``is_contiguous`` requires. + + ``lo``/``hi`` are indices in *string* space, so this is correct whether the + parent's ``offsets[0]`` is 1-D canonical or a lazy ``(2, M)`` gather layout. + """ + so = rl.str_offsets + assert so is not None # caller guarantees a string leaf + b0 = int(so[lo]) + return RaggedLayout( + data=rl.data[b0 : int(so[hi])], + offsets=[], + shape=(hi - lo,), + str_offsets=so[lo : hi + 1] - b0, + ) + + class Ragged(NDArrayOperatorsMixin, Generic[RDTYPE_co]): """A non-branching ragged array with a single ragged axis (Spec A).""" @@ -722,9 +745,10 @@ def _getitem(self, where: Any) -> Any: if isinstance(where, (int, np.integer)): lo, hi = int(starts[where]), int(stops[where]) if self._rl.str_offsets is not None and self._layout.offsets: - # string-under-axis: outer offsets index variants -> map to bytes via str_offsets - so = self._rl.str_offsets - return self._rl.data[int(so[lo]) : int(so[hi])].tobytes() + # string-under-axis: peel the real level -> standalone opaque + # string (k,), preserving the per-string boundaries that live in + # str_offsets. Concatenating here would drop them (issue #71). + return Ragged(_peel_string_row(self._rl, lo, hi)) row = self._rl.data[lo:hi] if self._rl.is_string: return row.tobytes() @@ -1218,11 +1242,12 @@ def _getitem_record_rows(self, where: Any) -> Any: out: dict[str, Any] = {} for name, fl in rec.fields.items(): if fl.str_offsets is not None: - so = fl.str_offsets - row = fl.data[int(so[lo]) : int(so[hi])] + # Each field carries its own str_offsets (Spec C Section 5); + # peel it against the shared lo/hi so every field of the row + # has the same length (issue #71). + out[name] = Ragged(_peel_string_row(fl, lo, hi)) else: - row = fl.data[lo:hi] - out[name] = row + out[name] = fl.data[lo:hi] return out sel_starts, sel_stops = self._row_gather(where) new_offsets = np.stack([sel_starts, sel_stops], 0) diff --git a/skills/seqpro/SKILL.md b/skills/seqpro/SKILL.md index 2bb8f8e..08490cc 100644 --- a/skills/seqpro/SKILL.md +++ b/skills/seqpro/SKILL.md @@ -89,6 +89,7 @@ rag = sp.rag.Ragged.empty((10, None, 4), dtype=np.uint8) # batch of 10 OHE seq | Bulk numeric op on the flat data | `rag.data[:] = ...` or `rag.data.view(...)` — zero-copy | Iterate `for seq in rag:` | | Apply a `np.ufunc` | Just call it: `np.exp(rag)` — dispatched via `__array_ufunc__` (NDArrayOperatorsMixin) to return a `Ragged` | Manually unpack and rebuild | | Count top-level rows | `len(rag)` — returns `shape[0]` (raises if `shape[0]` is the ragged axis) | `rag.shape[0]` with manual int-cast | +| Index one group of an opaque-string `Ragged` | `rag[i]` → a `Ragged` of `bytes`, one per string (`len(rag[i]) == rag.lengths[i]`); `rag[i][j]` is one `bytes` | Expect one concatenated `bytes` — that was the pre-0.22 behavior and it silently dropped the per-string boundaries | | Insert a leading size-1 axis | `rag[np.newaxis]` — returns `Ragged` with shape `(1, *old_shape)` | Manual `from_offsets` rebuild | | Reinterpret bytes/dtype | `rag.view(np.uint8)` | `np.asarray(rag).view(...)` (loses ragged structure) | | Reshape non-ragged axes | `rag.reshape(batch, None, k, 4)` | Touch `rag.data.shape` directly | @@ -127,6 +128,7 @@ The inputs **must share the same offsets object** (pass the same `shared_offsets - `rag["field"]` gives zero-copy single-field access and shares the parent's offsets object. Its `.data` is the flat NumPy buffer for that field. - `rag.to_numpy()` on a record layout returns a **dict `{field: dense ndarray}`** (raises if any field is still jagged — lengths must be uniform for a dense conversion). - `view` and `apply` are **not defined** on record layouts — operate per-field. +- Peeling a row (`rag[i]` with an integer `i`) returns a **dict** whose entries all have the same length: numeric/char fields as `ndarray`, opaque-string fields as a `Ragged` of `bytes`. That's what makes `zip(row["start"], row["alt"])` correct. ### Hashing strings diff --git a/tests/test_ragged_core.py b/tests/test_ragged_core.py index 8c97791..6666763 100644 --- a/tests/test_ragged_core.py +++ b/tests/test_ragged_core.py @@ -731,14 +731,103 @@ def test_r2_reshape_leading(): def test_string_under_axis_integer_index(): + """Peeling one group yields the standalone opaque-string layout (Spec C).""" rag = Ragged.from_offsets( np.frombuffer(b"TTGG", "S1"), (2, None), np.array([0, 1, 2]), str_offsets=np.array([0, 2, 4]), ) - assert rag[0] == b"TT" - assert rag[1] == b"GG" + row = rag[0] + assert isinstance(row, Ragged) + assert row.is_string and row.shape == (1,) + assert len(row) == 1 + assert row[0] == b"TT" + assert rag[1][0] == b"GG" + + +def _issue71_pair(): + """String-under-axis and numeric Ragged sharing one offsets object. + + Groups: 0 -> ('A', 'GG'), 1 -> ('TC',). + """ + data = np.frombuffer(b"AGGTC", dtype="S1") + outer = np.array([0, 2, 3], dtype=OFFSET_TYPE) # group -> string index + inner = np.array([0, 1, 3, 5], dtype=OFFSET_TYPE) # string -> byte index + s = Ragged.from_offsets(data, (2, None), outer, str_offsets=inner) + n = Ragged.from_offsets(np.array([10, 20, 30], dtype=np.int32), (2, None), outer) + return s, n + + +def test_string_under_axis_index_preserves_boundaries(): + """Issue #71: interior str_offsets boundaries must survive an integer index.""" + s, _ = _issue71_pair() + row = s[0] + assert len(row) == 2 + assert row[0] == b"A" + assert row[1] == b"GG" + assert list(s[1]) == [b"TC"] + + +def test_string_under_axis_index_matches_lengths(): + """len(s[i]) must equal s.lengths[i] and the numeric row length.""" + s, n = _issue71_pair() + for i in range(len(s)): + assert len(s[i]) == int(s.lengths[i]) + assert len(s[i]) == len(n[i]) + + +def test_string_under_axis_index_is_zero_copy(): + s, _ = _issue71_pair() + assert np.shares_memory(s[0].data, s.data) + + +def test_string_under_axis_index_empty_group(): + """A group holding zero strings peels to a length-0 result.""" + rag = Ragged.from_offsets( + np.frombuffer(b"AC", "S1"), + (2, None), + np.array([0, 0, 2], dtype=OFFSET_TYPE), # group 0 empty + str_offsets=np.array([0, 1, 2], dtype=OFFSET_TYPE), + ) + assert len(rag[0]) == 0 + assert len(rag[1]) == 2 + + +def test_string_under_axis_index_agrees_with_to_chars(): + s, _ = _issue71_pair() + chars = s.to_chars() + for i in range(len(s)): + for j in range(len(s[i])): + assert s[i][j] == chars[i][j].tobytes() + + +def test_string_under_axis_index_negative_and_oob(): + s, _ = _issue71_pair() + assert list(s[-1]) == [b"TC"] + with pytest.raises(IndexError): + s[5] + + +def test_string_under_axis_index_multidim(): + """(batch, ploidy, ~variants): reaches the flat branch via _getitem_multidim.""" + data = np.frombuffer(b"AGGTCNNAC", dtype="S1") + o0 = np.array([0, 2, 3, 4, 6], dtype=OFFSET_TYPE) # 4 segments -> string idx + i0 = np.array([0, 1, 3, 5, 7, 9], dtype=OFFSET_TYPE) # 6 boundaries -> bytes + rag = Ragged.from_offsets(data, (2, 2, None), o0, str_offsets=i0) + row = rag[0] # -> (2, None) string-under-axis + assert list(row[0]) == [b"A", b"GG"] + assert list(row[1]) == [b"TC"] + + +def test_standalone_string_index_still_returns_bytes(): + """Regression guard: the zero-real-level case is the terminal peel.""" + flat = Ragged.from_offsets( + np.frombuffer(b"cathithere", "S1"), (3,), np.array([0, 3, 5, 10]) + ) + assert flat[0] == b"cat" + assert flat[-1] == b"there" + assert list(flat) == [b"cat", b"hi", b"there"] # --------------------------------------------------------------------------- diff --git a/tests/test_ragged_core_records.py b/tests/test_ragged_core_records.py index 0d3a05f..802ee89 100644 --- a/tests/test_ragged_core_records.py +++ b/tests/test_ragged_core_records.py @@ -3,7 +3,7 @@ import pytest from seqpro.rag._core import Ragged from seqpro.rag._layout import RaggedLayout, RecordLayout, validate_layout -from seqpro.rag._utils import lengths_to_offsets +from seqpro.rag._utils import OFFSET_TYPE, lengths_to_offsets def _two_field_record(): @@ -680,3 +680,65 @@ def test_record_string_under_axis_to_packed(): packed = rec.to_packed() assert packed["ref"].to_ak().to_list() == [[b"A"], [b"CG"]] assert packed["alt"].to_ak().to_list() == [[b"TT"], [b"GG"]] + + +# --------------------------------------------------------------------------- +# Issue #71: string fields of a peeled record row keep their boundaries +# --------------------------------------------------------------------------- + + +def _issue71_record(): + """Record with an opaque-string field and a numeric field sharing offsets. + + Groups: 0 -> ('A', 'GG') / starts (1, 2), 1 -> ('TC',) / starts (3,). + Group 0's second allele is multi-byte, so a 1-byte coincidence cannot + mask a regression. + """ + outer = np.array([0, 2, 3], dtype=OFFSET_TYPE) + alt = Ragged.from_offsets( + np.frombuffer(b"AGGTC", dtype="S1"), + (2, None), + outer, + str_offsets=np.array([0, 1, 3, 5], dtype=OFFSET_TYPE), + ) + start = Ragged.from_offsets(np.array([1, 2, 3], dtype=np.int32), (2, None), outer) + return Ragged.from_fields({"alt": alt, "start": start}) + + +def test_record_row_string_field_preserves_boundaries(): + """Issue #71: a string field peeled from a record row keeps its boundaries.""" + row = _issue71_record()[0] + assert isinstance(row, dict) + assert list(row["alt"]) == [b"A", b"GG"] + np.testing.assert_array_equal(row["start"], np.array([1, 2], dtype=np.int32)) + + +def test_record_row_fields_have_matching_lengths(): + """Every field of a peeled row must have the same length, so zip aligns.""" + rec = _issue71_record() + for i in range(len(rec)): + row = rec[i] + assert len(row["alt"]) == len(row["start"]) + row0 = rec[0] + assert list(zip(row0["start"], row0["alt"])) == [(1, b"A"), (2, b"GG")] + + +def test_record_row_string_field_is_zero_copy(): + rec = _issue71_record() + assert np.shares_memory(rec[0]["alt"].data, rec["alt"].data) + + +def test_record_multidim_row_string_field_preserves_boundaries(): + """(batch, ploidy, ~variants) record: rec[0][h] routes via _getitem_record_rows_r2.""" + outer = np.array([0, 2, 3, 4, 6], dtype=OFFSET_TYPE) + alt = Ragged.from_offsets( + np.frombuffer(b"AGGTCNNAC", dtype="S1"), + (2, 2, None), + outer, + str_offsets=np.array([0, 1, 3, 5, 7, 9], dtype=OFFSET_TYPE), + ) + start = Ragged.from_offsets(np.arange(6, dtype=np.int32), (2, 2, None), outer) + rec = Ragged.from_fields({"alt": alt, "start": start}) + row = rec[0][0] + assert list(row["alt"]) == [b"A", b"GG"] + assert len(row["alt"]) == len(row["start"])