Preserve staged cache source entry metadata

Author: francois-drielsma

src/spine/io/read/stage_hdf5.py

@@ -401,6 +401,7 @@ def get(self, idx: int) -> dict[str, object]:
         data: dict[str, object] = {
             "file_index": file_idx,
             "file_entry_index": entry_idx,
+            "source_file_entry_index": entry_idx,
         }
         data.update(self._source_info.get(file_idx, {}))
         product_stage_map = self._resolved_products[file_idx]

test/test_io/test_read/test_hdf5.py

@@ -290,6 +290,34 @@ def test_stage_hdf5_reader_loads_one_stage(tmp_path):
     reader.close()
 
 
+def test_stage_hdf5_reader_fills_missing_source_entry_index(tmp_path):
+    """Older/minimal stage caches should expose source entry metadata."""
+    path = tmp_path / "stage_cache.h5"
+    writer = StageHDF5Writer(str(path), overwrite=True)
+    writer.write_stage(
+        "deghosting",
+        {
+            "index": np.asarray([0, 1]),
+            "source_file_name": np.asarray(["source.root", "source.root"]),
+            "source_file_size": np.asarray([10, 10]),
+            "source_file_mtime_ns": np.asarray([20, 20]),
+            "dummy_data": [
+                np.asarray([[1.0, 2.0]]),
+                np.asarray([[3.0, 4.0]]),
+            ],
+        },
+    )
+    writer.finalize_stage("deghosting")
+    writer.close()
+
+    reader = StageHDF5Reader("deghosting", str(path), build_classes=False)
+    entry = reader.get(1)
+    assert entry["file_entry_index"] == 1
+    assert entry["source_file_entry_index"] == 1
+    np.testing.assert_array_equal(entry["dummy_data"], np.asarray([[3.0, 4.0]]))
+    reader.close()
+
+
 def test_stage_hdf5_reader_rejects_incomplete_stages(tmp_path):
     """Incomplete stages should be rejected by default."""
     path = tmp_path / "stage_cache.h5"

test/test_io/test_write/test_hdf5.py

@@ -529,6 +529,36 @@ def test_stage_hdf5_writer_respects_explicit_keys(tmp_path):
         assert "extra_tensor" not in events.dtype.names
 
 
+def test_stage_hdf5_writer_preserves_source_entry_with_explicit_keys(tmp_path):
+    """Stage caches should preserve source entry provenance from file_entry_index."""
+    cache_path = tmp_path / "cache.h5"
+    writer = StageHDF5Writer(
+        str(cache_path),
+        overwrite=True,
+        keys=["dummy_data"],
+    )
+    writer.write_stage(
+        "deghosting",
+        {
+            "index": np.asarray([0]),
+            "file_entry_index": np.asarray([7]),
+            "source_file_name": np.asarray(["source.root"]),
+            "source_file_size": np.asarray([10]),
+            "source_file_mtime_ns": np.asarray([20]),
+            "dummy_data": [np.asarray([[1.0, 2.0]])],
+            "extra_tensor": [np.asarray([[3.0, 4.0]])],
+        },
+    )
+    writer.finalize_stage("deghosting")
+    writer.close()
+
+    with h5py.File(cache_path, "r") as out_file:
+        stage_group = out_file["stages"]["deghosting"]
+        assert "source_file_entry_index" in stage_group
+        assert "source_file_entry_index" in stage_group["events"].dtype.names
+        np.testing.assert_array_equal(stage_group["source_file_entry_index"][:], [7])
+
+
 def test_stage_hdf5_writer_rejects_mismatched_source(tmp_path):
     """Writing a later stage with different source provenance should fail."""
     source_a = tmp_path / "source_a.root"