Add option to collapse foreach/matrix stages into a single stage (#10945)

Co-authored-by: skshetry <18718008+skshetry@users.noreply.github.com>

Author: fhausmann

dvc/commands/dag.py

@@ -127,19 +127,80 @@ def _filter(graph, targets, full):
     return new_graph
 
 
-def _build(repo, target=None, full=False, outs=False):
+def _is_foreach_matrix_stage(node, join_string):
+    if node.endswith(".dvc"):
+        return False
+    return join_string in node
+
+
+def _collapse_foreach_matrix_get_nodes(graph):
+    from dvc.parsing import JOIN
+
+    new_nodes = set()
+    nodes_to_remove = set()
+    for _node in list(graph.nodes):
+        if not _is_foreach_matrix_stage(_node, JOIN):
+            continue
+        nodes_to_remove.add(_node)
+        new_nodes.add(_node.split(JOIN)[0])
+    return new_nodes, nodes_to_remove
+
+
+def _collapse_foreach_matrix_get_edges(graph):
+    from dvc.parsing import JOIN
+
+    new_edges = set()
+    edges_to_remove = set()
+    for _e1, _e2 in list(graph.edges):
+        _replace = False
+        _new_e1 = _e1
+        _new_e2 = _e2
+        if _is_foreach_matrix_stage(_e1, JOIN):
+            _new_e1 = _e1.split(JOIN)[0]
+            _replace = True
+        if _is_foreach_matrix_stage(_e2, JOIN):
+            _new_e2 = _e2.split(JOIN)[0]
+            _replace = True
+        if _replace:
+            edges_to_remove.add((_e1, _e2))
+            new_edges.add((_new_e1, _new_e2))
+    return new_edges, edges_to_remove
+
+
+def _collapse_foreach_matrix(graph):
+    new_nodes, nodes_to_remove = _collapse_foreach_matrix_get_nodes(graph)
+    new_edges, edges_to_remove = _collapse_foreach_matrix_get_edges(graph)
+    new_graph = graph.copy()
+    new_graph.remove_edges_from(edges_to_remove)
+    new_graph.add_nodes_from(new_nodes)
+    new_graph.add_edges_from(new_edges)
+    new_graph.remove_nodes_from(nodes_to_remove)
+    return new_graph
+
+
+def _build(repo, target=None, full=False, outs=False, collapse_foreach_matrix=False):
     targets = _collect_targets(repo, target, outs)
     graph = _transform(repo.index, outs)
-    return _filter(graph, targets, full)
+    filtered_graph = _filter(graph, targets, full)
+    if collapse_foreach_matrix:
+        filtered_graph = _collapse_foreach_matrix(filtered_graph)
+    return filtered_graph
 
 
 class CmdDAG(CmdBase):
     def run(self):
+        from dvc.exceptions import InvalidArgumentError
+
+        if self.args.outs and self.args.collapse_foreach_matrix:
+            raise InvalidArgumentError(
+                "`--outs` and `--collapse-foreach-matrix` are mutually exclusive"
+            )
         graph = _build(
             self.repo,
             target=self.args.target,
             full=self.args.full,
             outs=self.args.outs,
+            collapse_foreach_matrix=self.args.collapse_foreach_matrix,
         )
 
         if self.args.dot:
@@ -181,6 +242,12 @@ def add_parser(subparsers, parent_parser):
         dest="markdown",
         help="Print DAG with mermaid format wrapped in Markdown block.",
     )
+    dag_parser.add_argument(
+        "--collapse-foreach-matrix",
+        action="store_true",
+        default=False,
+        help="Collapse foreach/matrix stages to single stage.",
+    )
     dag_parser.add_argument(
         "--full",
         action="store_true",

tests/unit/command/test_dag.py

@@ -2,7 +2,15 @@
 import pytest
 
 from dvc.cli import main, parse_args
-from dvc.commands.dag import CmdDAG, _build, _show_ascii, _show_dot, _show_mermaid
+from dvc.commands.dag import (
+    CmdDAG,
+    _build,
+    _collapse_foreach_matrix,
+    _show_ascii,
+    _show_dot,
+    _show_mermaid,
+)
+from dvc.parsing import JOIN
 
 
 @pytest.mark.parametrize(
@@ -49,10 +57,53 @@ def repo(tmp_dir, dvc):
     return dvc
 
 
+def test_collapse_foreach_matrix(repo):
+    graph = nx.DiGraph(
+        [
+            ("2", "1"),
+            ("3@a", "2"),
+            ("3@a", "1"),
+            ("4", "1"),
+            ("3@b", "4"),
+            ("3@b", "1"),
+            ("5", "3@a"),
+            ("6", "3@b"),
+            ("7", "5"),
+            ("7", "6"),
+        ]
+    )
+    expected_graph = nx.DiGraph(
+        [
+            ("2", "1"),
+            ("3", "2"),
+            ("4", "1"),
+            ("3", "4"),
+            ("3", "1"),
+            ("5", "3"),
+            ("6", "3"),
+            ("7", "5"),
+            ("7", "6"),
+        ]
+    )
+    collapsed_graph = _collapse_foreach_matrix(graph)
+    for node in collapsed_graph.nodes:
+        assert JOIN not in node
+    for n1, n2 in collapsed_graph.edges:
+        assert JOIN not in n1
+        assert JOIN not in n2
+    assert nx.is_isomorphic(collapsed_graph, expected_graph)
+
+
 def test_build(repo):
     assert nx.is_isomorphic(_build(repo), repo.index.graph)
 
 
+def test_build_collapse(repo):
+    assert nx.is_isomorphic(
+        _build(repo, collapse_foreach_matrix=True), repo.index.graph
+    )
+
+
 def test_build_target(repo):
     graph = _build(repo, target="3")
     assert set(graph.nodes()) == {"3", "b.dvc", "a.dvc"}