feat(data): carry semantic temporal parquet side channels

Author: apstenku123

cppmega_mlx/data/parquet_dataset.py

@@ -68,6 +68,15 @@
     "hunk_id_per_token",
     "edit_op_per_token",
 )
+_FAMILY_TOKEN_SIDE_CHANNEL_COLUMNS: Mapping[str, tuple[str, ...]] = {
+    "semantic_graph": _TOKEN_SEMANTIC_METADATA_COLUMNS,
+    "temporal_diff": _TOKEN_TEMPORAL_METADATA_COLUMNS,
+}
+_FAMILY_TOKEN_SIDE_CHANNEL_COLUMNS_FLAT = tuple(
+    column
+    for columns in _FAMILY_TOKEN_SIDE_CHANNEL_COLUMNS.values()
+    for column in columns
+)
 _SOURCE_PLATFORM_IDS_COLUMN = "source_platform_ids"
 _VALID_TOKEN_COUNT_COLUMN = "valid_token_count"
 _ROW_METADATA_COLUMNS = (
@@ -179,6 +188,12 @@ class _SideChannelColumns:
     skipped: Sequence[Mapping[str, str | None]]
 
 
+@dataclass(frozen=True)
+class _FamilySideChannelColumns:
+    channels: Mapping[str, Mapping[str, np.ndarray]]
+    sources: Mapping[str, Mapping[str, Mapping[str, str | None]]]
+
+
 @dataclass(frozen=True)
 class _ModelMetadataColumns:
     channels: Mapping[str, np.ndarray]
@@ -288,6 +303,11 @@ def __init__(
         )
         side_channel_columns = _side_channel_windows(columns, token_rows, seq_len)
         side_channels = side_channel_columns.channels
+        family_side_channel_columns = _family_side_channel_windows(
+            columns,
+            token_rows,
+            seq_len,
+        )
         model_metadata_columns = _model_metadata_windows(columns, token_rows, seq_len)
         batch_metadata_columns = _resolve_batch_metadata_columns(
             columns,
@@ -325,6 +345,13 @@ def __init__(
             key: _to_side_channel_values(key, value)
             for key, value in side_channels.items()
         }
+        self._family_side_channels = {
+            family: {
+                column: value.astype(np.int32, copy=False)
+                for column, value in family_columns.items()
+            }
+            for family, family_columns in family_side_channel_columns.channels.items()
+        }
         self._model_metadata_channels = {
             key: value.astype(np.int32, copy=False)
             for key, value in model_metadata_columns.channels.items()
@@ -342,6 +369,7 @@ def __init__(
             document_id_source=document_id_source,
             side_channel_sources=side_channel_columns.sources,
             skipped_side_channels=side_channel_columns.skipped,
+            family_side_channel_sources=family_side_channel_columns.sources,
             model_metadata_sources=model_metadata_columns.sources,
             batch_metadata_columns=batch_metadata_columns,
         )
@@ -441,6 +469,13 @@ def _make_batch(self, sample_idx: np.ndarray) -> LMTokenBatch:
                 for key, value in self._model_metadata_channels.items()
             }
         )
+        family_side_channels = {
+            family: {
+                column: mx.array(value[sample_idx])
+                for column, value in family_columns.items()
+            }
+            for family, family_columns in self._family_side_channels.items()
+        }
         return LMTokenBatch(
             tokens=mx.array(self._tokens[sample_idx]),
             target_tokens=None
@@ -452,6 +487,7 @@ def _make_batch(self, sample_idx: np.ndarray) -> LMTokenBatch:
             document_ids=None
             if self._document_ids is None
             else mx.array(self._document_ids[sample_idx]),
+            side_channels=family_side_channels or None,
             metadata=self._make_batch_metadata(sample_idx),
             **kwargs,
         )
@@ -496,6 +532,7 @@ def _candidate_parquet_columns(
         candidates.append(text_key)
     candidates.extend(_LOSS_MASK_COLUMN_ALIASES)
     candidates.extend(_DOCUMENT_ID_COLUMN_ALIASES)
+    candidates.extend(_FAMILY_TOKEN_SIDE_CHANNEL_COLUMNS_FLAT)
     candidates.append(_SOURCE_PLATFORM_IDS_COLUMN)
     candidates.append(_VALID_TOKEN_COUNT_COLUMN)
     for aliases in _SIDE_CHANNEL_COLUMN_ALIASES.values():
@@ -721,6 +758,41 @@ def _side_channel_windows(
     return _SideChannelColumns(channels=channels, sources=sources, skipped=skipped)
 
 
+def _family_side_channel_windows(
+    columns: ParquetColumns,
+    token_rows: list[list[int]],
+    seq_len: int,
+) -> _FamilySideChannelColumns:
+    channels: dict[str, dict[str, np.ndarray]] = {}
+    sources: dict[str, dict[str, dict[str, str | None]]] = {}
+    for family, family_columns in _FAMILY_TOKEN_SIDE_CHANNEL_COLUMNS.items():
+        for column in family_columns:
+            if column not in columns.values:
+                continue
+            _reject_non_integer_parquet_type(
+                columns,
+                column,
+                f"{column} {family} side-channel IDs",
+            )
+            rows = [
+                _coerce_token_row(value, label=f"{column} {family} side-channel")
+                for value in columns.require(column)
+            ]
+            if not _rows_are_token_aligned(rows, token_rows):
+                raise ValueError(
+                    f"{column} rows must be token-aligned with token IDs"
+                )
+            channels.setdefault(family, {})[column] = _fixed_windows_from_rows(
+                rows,
+                seq_len,
+            )
+            sources.setdefault(family, {})[column] = {
+                "column": column,
+                "type": columns.type_label(column),
+            }
+    return _FamilySideChannelColumns(channels=channels, sources=sources)
+
+
 def _model_metadata_windows(
     columns: ParquetColumns,
     token_rows: list[list[int]],
@@ -1137,6 +1209,10 @@ def _parquet_receipt(
     document_id_source: str | None,
     side_channel_sources: Mapping[str, Mapping[str, str | None]],
     skipped_side_channels: Sequence[Mapping[str, str | None]],
+    family_side_channel_sources: Mapping[
+        str,
+        Mapping[str, Mapping[str, str | None]],
+    ],
     model_metadata_sources: Mapping[str, Mapping[str, str | None]],
     batch_metadata_columns: Sequence[str],
 ) -> dict[str, Any]:
@@ -1173,6 +1249,14 @@ def _parquet_receipt(
         receipt["model_metadata_sources"] = {
             key: dict(value) for key, value in sorted(model_metadata_sources.items())
         }
+    if family_side_channel_sources:
+        receipt["family_side_channel_sources"] = {
+            family: {
+                column: dict(source)
+                for column, source in sorted(columns.items())
+            }
+            for family, columns in sorted(family_side_channel_sources.items())
+        }
     training_sources = {
         "target_tokens": target_source,
         "loss_mask": loss_mask_source,

tests/test_parquet_dataset.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pytest
+import mlx.core as mx
 import mlx.nn as nn
 
 from cppmega_mlx.data.parquet_dataset import TokenParquetDataset
@@ -334,6 +335,114 @@ def test_all_metadata_excludes_token_content_and_slices_token_level_fields() ->
     assert metadata_windows[1]["repo"] == "llvm"
 
 
+def test_parquet_token_semantic_and_temporal_metadata_reach_side_channel_map(
+    tmp_path,
+) -> None:
+    pa = pytest.importorskip("pyarrow")
+    pq = pytest.importorskip("pyarrow.parquet")
+    path = tmp_path / "semantic_temporal.parquet"
+    table = pa.table(
+        {
+            "token_ids": pa.array(
+                [[0, 1, 2, 3, 4, 5, 6, 7]],
+                type=pa.large_list(pa.int32()),
+            ),
+            "token_symbol_ids": pa.array(
+                [[10, 11, 12, 13, 14, 15, 16, 17]],
+                type=pa.large_list(pa.int32()),
+            ),
+            "token_call_targets": pa.array(
+                [[20, 21, 22, 23, 24, 25, 26, 27]],
+                type=pa.large_list(pa.int32()),
+            ),
+            "token_change_mask_post": pa.array(
+                [[0, 0, 1, 1, 0, 1, 0, 1]],
+                type=pa.large_list(pa.int32()),
+            ),
+            "edit_op_per_token": pa.array(
+                [[0, 0, 2, 2, 0, 3, 0, 3]],
+                type=pa.large_list(pa.int32()),
+            ),
+        }
+    )
+    pq.write_table(table, path)
+
+    dataset = TokenParquetDataset(path, seq_len=4, batch_size=2, token_key="token_ids")
+    batch = next(dataset.iter_batches())
+
+    assert batch.side_channels is not None
+    assert tuple(batch.model_kwargs()) == ()
+    assert set(batch.side_channels) == {"semantic_graph", "temporal_diff"}
+    np.testing.assert_array_equal(
+        np.array(batch.side_channels["semantic_graph"]["token_symbol_ids"]),
+        [[10, 11, 12, 13], [14, 15, 16, 17]],
+    )
+    np.testing.assert_array_equal(
+        np.array(batch.side_channels["semantic_graph"]["token_call_targets"]),
+        [[20, 21, 22, 23], [24, 25, 26, 27]],
+    )
+    np.testing.assert_array_equal(
+        np.array(batch.side_channels["temporal_diff"]["token_change_mask_post"]),
+        [[0, 0, 1, 1], [0, 1, 0, 1]],
+    )
+    np.testing.assert_array_equal(
+        np.array(batch.side_channels["temporal_diff"]["edit_op_per_token"]),
+        [[0, 0, 2, 2], [0, 3, 0, 3]],
+    )
+
+    dropped = batch.with_side_channel_dropout(
+        {"semantic_graph": 1.0, "temporal_diff": 1.0},
+        seed=7,
+    )
+    assert dropped.side_channels is not None
+    assert float(
+        mx.sum(
+            mx.abs(dropped.side_channels["semantic_graph"]["token_symbol_ids"])
+        ).item()
+    ) == 0.0
+    assert float(
+        mx.sum(
+            mx.abs(dropped.side_channels["temporal_diff"]["edit_op_per_token"])
+        ).item()
+    ) == 0.0
+    assert (
+        dataset.parquet_receipt["family_side_channel_sources"]["semantic_graph"][
+            "token_symbol_ids"
+        ]["column"]
+        == "token_symbol_ids"
+    )
+    assert (
+        dataset.parquet_receipt["family_side_channel_sources"]["temporal_diff"][
+            "edit_op_per_token"
+        ]["column"]
+        == "edit_op_per_token"
+    )
+
+
+def test_parquet_token_semantic_side_channel_fails_closed_when_not_token_aligned(
+    tmp_path,
+) -> None:
+    pa = pytest.importorskip("pyarrow")
+    pq = pytest.importorskip("pyarrow.parquet")
+    path = tmp_path / "bad_semantic.parquet"
+    table = pa.table(
+        {
+            "token_ids": pa.array(
+                [[0, 1, 2, 3, 4, 5, 6, 7]],
+                type=pa.large_list(pa.int32()),
+            ),
+            "token_symbol_ids": pa.array(
+                [[10, 11]],
+                type=pa.large_list(pa.int32()),
+            ),
+        }
+    )
+    pq.write_table(table, path)
+
+    with pytest.raises(ValueError, match="token_symbol_ids.*token-aligned"):
+        TokenParquetDataset(path, seq_len=4, batch_size=1, token_key="token_ids")
+
+
 def test_platform_ids_parquet_column_threads_to_model_kwargs(tmp_path) -> None:
     pa = pytest.importorskip("pyarrow")
     pq = pytest.importorskip("pyarrow.parquet")