Repository: iterative/dvc

Issue:
Add `key` to matrix
The addition of `matrix` functionality simplified quite a lot for me. However, contrary to `foreach` stages, there is no `key` property available. I often find myself composing output paths using the parametrization fields.

Let me give an example – first, using foreach:

```yaml
train_model:
  foreach:
    backbone_A_dim_100:
      backbone: A
      dim: 100
   backbone_B_dim_100:
     backbone: B
     dim: 100
   backbone_A_dim_200:
     ....
  do:
    cmd: python train_model.py --backbone ${item.backbone} --dim ${item.dim}
    outs:
    - data/models/${key}/
```

This can be simplified into a `matrix` stage:

```yaml
train_model:
  matrix:
    backbone: ["A", "B"]
   dim: [100, 200]
  cmd: python train_model.py --backbone ${item.backbone} --dim ${item.dim}
  outs:
  - data/models/${item.backbone}_${item.dim}/
```

The output path is composed of both the `item.backbone` and `item.dim`, which in this simplified case might not be that problematic, but it becomes more demanding and error prone with more `matrix` parameters. An easier way would be to include the `key` of the `matrix` stage into the templating context, i.e.:

```yaml
train_model:
  matrix:
    backbone: ["A", "B"]
   dim: [100, 200]
  cmd: python train_model.py --backbone ${item.backbone} --dim ${item.dim}
  outs:
  - data/models/${key}/
```

I already implemented the requested functionality ~~and will link it in a moment.~~ – see [PR#10044](https://github.com/iterative/dvc/pull/10044)


Discussion:

Resolving patch:
diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py
--- a/dvc/parsing/__init__.py
+++ b/dvc/parsing/__init__.py
@@ -499,8 +499,26 @@ def _resolve_matrix_data(self) -> Dict[str, List]:
             iterable = self.context.resolve(self.matrix_data, unwrap=False)
         except (ContextError, ParseError) as exc:
             format_and_raise(exc, f"'{self.where}.{self.name}.matrix'", self.relpath)
+
+        # Matrix entries will have `key` and `item` added to the context.
+        # Warn users if these are already in the context from the global vars.
+        self._warn_if_overwriting([self.pair.key, self.pair.value])
         return iterable
 
+    def _warn_if_overwriting(self, keys: List[str]):
+        warn_for = [k for k in keys if k in self.context]
+        if warn_for:
+            linking_verb = "is" if len(warn_for) == 1 else "are"
+            logger.warning(
+                (
+                    "%s %s already specified, "
+                    "will be overwritten for stages generated from '%s'"
+                ),
+                " and ".join(warn_for),
+                linking_verb,
+                self.name,
+            )
+
     @cached_property
     def normalized_iterable(self) -> Dict[str, "DictStrAny"]:
         """Convert sequence to Mapping with keys normalized."""
@@ -540,7 +558,7 @@ def _each_iter(self, key: str) -> "DictStrAny":
         with reraise(KeyError, EntryNotFound(err_message)):
             value = self.normalized_iterable[key]
 
-        temp_dict = {self.pair.value: value}
+        temp_dict = {self.pair.key: key, self.pair.value: value}
         with self.context.set_temporarily(temp_dict, reserve=True):
             # optimization: item and key can be removed on __exit__() as they
             # are top-level values, and are not merged recursively.


Test patch:
diff --git a/tests/func/parsing/test_matrix.py b/tests/func/parsing/test_matrix.py
--- a/tests/func/parsing/test_matrix.py
+++ b/tests/func/parsing/test_matrix.py
@@ -50,3 +50,44 @@ def test_matrix_interpolated(tmp_dir, dvc, matrix):
         "build@linux-3.8-dict1-list0": {"cmd": "echo linux 3.8 --arg2 2 -- out1 out11"},
         "build@linux-3.8-dict1-list1": {"cmd": "echo linux 3.8 --arg2 2 -- out2 out22"},
     }
+
+
+@pytest.mark.parametrize(
+    "matrix",
+    [
+        MATRIX_DATA,
+        {
+            "os": "${os}",
+            "pyv": "${pyv}",
+            "dict": "${dict}",
+            "list": "${list}",
+        },
+    ],
+)
+def test_matrix_key_present(tmp_dir, dvc, matrix):
+    (tmp_dir / "params.yaml").dump(MATRIX_DATA)
+    resolver = DataResolver(dvc, tmp_dir.fs_path, {})
+    data = {
+        "matrix": matrix,
+        "cmd": "echo ${key}",
+    }
+    definition = MatrixDefinition(resolver, resolver.context, "build", data)
+
+    assert definition.resolve_all() == {
+        "build@win-3.7-dict0-list0": {"cmd": "echo win-3.7-dict0-list0"},
+        "build@win-3.7-dict0-list1": {"cmd": "echo win-3.7-dict0-list1"},
+        "build@win-3.7-dict1-list0": {"cmd": "echo win-3.7-dict1-list0"},
+        "build@win-3.7-dict1-list1": {"cmd": "echo win-3.7-dict1-list1"},
+        "build@win-3.8-dict0-list0": {"cmd": "echo win-3.8-dict0-list0"},
+        "build@win-3.8-dict0-list1": {"cmd": "echo win-3.8-dict0-list1"},
+        "build@win-3.8-dict1-list0": {"cmd": "echo win-3.8-dict1-list0"},
+        "build@win-3.8-dict1-list1": {"cmd": "echo win-3.8-dict1-list1"},
+        "build@linux-3.7-dict0-list0": {"cmd": "echo linux-3.7-dict0-list0"},
+        "build@linux-3.7-dict0-list1": {"cmd": "echo linux-3.7-dict0-list1"},
+        "build@linux-3.7-dict1-list0": {"cmd": "echo linux-3.7-dict1-list0"},
+        "build@linux-3.7-dict1-list1": {"cmd": "echo linux-3.7-dict1-list1"},
+        "build@linux-3.8-dict0-list0": {"cmd": "echo linux-3.8-dict0-list0"},
+        "build@linux-3.8-dict0-list1": {"cmd": "echo linux-3.8-dict0-list1"},
+        "build@linux-3.8-dict1-list0": {"cmd": "echo linux-3.8-dict1-list0"},
+        "build@linux-3.8-dict1-list1": {"cmd": "echo linux-3.8-dict1-list1"},
+    }