fix(dpmodel): reject incompatible LMDB merge type maps

deepmd/dpmodel/utils/lmdb_data.py

@@ -2685,20 +2685,61 @@ def get_test(self) -> dict[str, Any]:
         return self._inner.get_test(nloc=self._nloc)
 
 
+def _validate_merge_type_maps(
+    source_metadata: list[tuple[str, dict[str, Any]]],
+) -> list[str] | None:
+    """Return the shared type map required for byte-for-byte frame merging.
+
+    ``merge_lmdb`` does not decode and rewrite atom-type arrays, so every
+    source must use exactly the same index-to-species mapping. All-missing
+    legacy metadata remains supported, but mixing explicit and missing maps is
+    rejected because compatibility cannot be established.
+    """
+    source_type_maps = [(path, meta.get("type_map")) for path, meta in source_metadata]
+    explicit_type_maps = [
+        (path, list(type_map))
+        for path, type_map in source_type_maps
+        if type_map is not None
+    ]
+    if not explicit_type_maps:
+        return None
+
+    formatted_maps = ", ".join(
+        f"{path}: {list(type_map)!r}" if type_map is not None else f"{path}: missing"
+        for path, type_map in source_type_maps
+    )
+    if len(explicit_type_maps) != len(source_type_maps):
+        raise ValueError(
+            "Cannot merge LMDB datasets with mixed type_map metadata because "
+            f"raw atom-type indices cannot be validated ({formatted_maps})"
+        )
+
+    canonical_type_map = explicit_type_maps[0][1]
+    if any(type_map != canonical_type_map for _, type_map in explicit_type_maps[1:]):
+        raise ValueError(
+            "Cannot merge LMDB datasets with incompatible type_map values "
+            f"because frames are copied without remapping ({formatted_maps})"
+        )
+    return canonical_type_map
+
+
 def _copy_lmdb_source(
     src_path: str,
+    metadata: dict[str, Any],
     dst_env: lmdb.Environment,
     dst_format: str,
     frame_idx: int,
     frame_nlocs: list[int],
     frame_system_ids: list[int],
     system_id_offset: int,
-) -> tuple[int, dict, list[str] | None, int]:
-    """Copy one source under a ref-counted environment lease."""
+) -> tuple[int, dict, int]:
+    """Copy one validated source under a ref-counted environment lease.
+
+    The caller supplies metadata collected during the validation preflight so
+    this copy pass does not read and decode it a second time.
+    """
     src_env = _open_lmdb(src_path)
     try:
-        with src_env.begin() as transaction:
-            metadata = _read_metadata(transaction)
         nframes, src_format, natoms_per_type = _parse_metadata(metadata)
         fallback_natoms = sum(natoms_per_type)
         source_nlocs = metadata.get("frame_nlocs")
@@ -2741,7 +2782,6 @@ def _copy_lmdb_source(
         return (
             frame_idx,
             metadata.get("system_info", {}),
-            metadata.get("type_map"),
             system_id_offset,
         )
     finally:
@@ -2772,30 +2812,47 @@ def merge_lmdb(
     -------
     str
         Path to the created LMDB.
+
+    Raises
+    ------
+    ValueError
+        If sources use different explicit type maps, or mix explicit type-map
+        metadata with legacy metadata where the mapping is missing.
     """
     import os
     import shutil
 
+    # Validate every source before replacing or creating the destination. A
+    # type-map validation failure must not destroy an existing dataset or
+    # leave a partial output.
+    source_metadata: list[tuple[str, dict[str, Any]]] = []
+    for src_path in src_paths:
+        src_env = _open_lmdb(src_path)
+        try:
+            with src_env.begin() as txn:
+                source_metadata.append((src_path, _read_metadata(txn)))
+        finally:
+            _close_lmdb(src_path)
+    merged_type_map = _validate_merge_type_maps(source_metadata)
+
     if os.path.exists(dst_path):
         shutil.rmtree(dst_path)
-
     dst_env = lmdb.open(dst_path, map_size=map_size)
     frame_idx = 0
     fmt = "012d"
     frame_nlocs: list[int] = []
     frame_system_ids: list[int] = []
     first_system_info: dict | None = None
-    first_type_map: list[str] | None = None
     sys_id_offset = 0
     try:
-        for src_path in src_paths:
+        for src_path, metadata in source_metadata:
             (
                 frame_idx,
                 source_system_info,
-                source_type_map,
                 sys_id_offset,
             ) = _copy_lmdb_source(
                 src_path,
+                metadata,
                 dst_env,
                 fmt,
                 frame_idx,
@@ -2805,8 +2862,6 @@ def merge_lmdb(
             )
             if first_system_info is None:
                 first_system_info = source_system_info
-            if first_type_map is None:
-                first_type_map = source_type_map
 
         merged_meta = {
             "nframes": frame_idx,
@@ -2815,8 +2870,8 @@ def merge_lmdb(
             "frame_nlocs": frame_nlocs,
             "frame_system_ids": frame_system_ids,
         }
-        if first_type_map is not None:
-            merged_meta["type_map"] = first_type_map
+        if merged_type_map is not None:
+            merged_meta["type_map"] = merged_type_map
         with dst_env.begin(write=True) as transaction:
             transaction.put(
                 b"__metadata__",

source/tests/pt/test_lmdb_dataloader.py

@@ -1068,6 +1068,57 @@ def test_merge_preserves_type_map(self, tmp_path):
         env.close()
         assert meta.get("type_map") == ["O", "H"]
 
+        reader = LmdbDataReader(dst, ["O", "H"])
+        expected_atype = np.array([0, 0, 0, 1, 1, 1])
+        np.testing.assert_array_equal(reader[0]["atype"], expected_atype)
+        np.testing.assert_array_equal(reader[5]["atype"], expected_atype)
+
+    @pytest.mark.parametrize(
+        "second_type_map",
+        [["H", "O"], ["O", "H", "N"]],
+        ids=["reordered", "prefix-compatible-superset"],
+    )
+    def test_merge_rejects_incompatible_type_maps_before_creating_output(
+        self, tmp_path, second_type_map
+    ):
+        """Raw frames cannot be shared under two different type index maps."""
+        src1, src2 = str(tmp_path / "tm1.lmdb"), str(tmp_path / "tm2.lmdb")
+        _create_lmdb_with_system_ids(
+            src1, system_frames=[1], natoms=6, type_map=["O", "H"]
+        )
+        # Even a prefix-compatible superset is rejected: merge_lmdb deliberately
+        # requires identical metadata instead of proving frame-by-frame safety.
+        _create_lmdb_with_system_ids(
+            src2, system_frames=[1], natoms=6, type_map=second_type_map
+        )
+        dst = tmp_path / "incompatible.lmdb"
+        dst.mkdir()
+        marker = dst / "existing-data"
+        marker.write_text("preserve me")
+
+        with pytest.raises(ValueError, match="incompatible type_map values") as exc:
+            merge_lmdb([src1, src2], str(dst))
+
+        assert src1 in str(exc.value)
+        assert src2 in str(exc.value)
+        assert marker.read_text() == "preserve me"
+
+    def test_merge_rejects_mixed_explicit_and_missing_type_maps(self, tmp_path):
+        """A legacy source without a map cannot be proven index-compatible."""
+        src_without_map = str(tmp_path / "legacy.lmdb")
+        src_with_map = str(tmp_path / "typed.lmdb")
+        _create_test_lmdb(src_without_map, nframes=1, natoms=6)
+        _create_lmdb_with_system_ids(
+            src_with_map, system_frames=[1], natoms=6, type_map=["O", "H"]
+        )
+        dst = tmp_path / "mixed_metadata.lmdb"
+
+        with pytest.raises(ValueError, match="mixed type_map metadata") as exc:
+            merge_lmdb([src_without_map, src_with_map], str(dst))
+
+        assert "missing" in str(exc.value)
+        assert not dst.exists()
+
 
 # ============================================================
 # Multitask LMDB training