From 357a6dd4679e99aefc31a2e36e78e9a0e5a64d2b Mon Sep 17 00:00:00 2001
From: Alexander Schepanovski <suor.web@gmail.com>
Date: Mon, 9 Mar 2020 17:13:07 +0700
Subject: [PATCH] dvc: optimize building graph (#3443)

Use trie structure to efficiently check out, dep and stage paths being
prefixes of each other.
---
 dvc/repo/__init__.py | 73 +++++++++++++++++++++++++-------------------
 setup.py             |  1 +
 2 files changed, 42 insertions(+), 32 deletions(-)

diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
index 5912d5dc5..3cbfb8590 100644
--- a/dvc/repo/__init__.py
+++ b/dvc/repo/__init__.py
@@ -1,12 +1,11 @@
 import os
 from contextlib import contextmanager
 from functools import wraps
-from itertools import chain
 
 from dvc.ignore import CleanTree
 from dvc.compat import fspath_py35
 
-from funcy import cached_property, cat
+from funcy import cached_property, cat, first
 
 from dvc.config import Config
 from dvc.exceptions import (
@@ -305,6 +304,7 @@ class Repo(object):
             CyclicGraphError: resulting graph has cycles
         """
         import networkx as nx
+        from pygtrie import Trie
         from dvc.exceptions import (
             OutputDuplicationError,
             StagePathAsOutputError,
@@ -314,48 +314,57 @@ class Repo(object):
         G = nx.DiGraph()
         stages = stages or self.stages
         stages = [stage for stage in stages if stage]
-        outs = {}
+        outs = Trie()  # Use trie to efficiently find overlapping outs and deps
 
         for stage in stages:
             for out in stage.outs:
-                if out.path_info in outs:
-                    dup_stages = [stage, outs[out.path_info].stage]
+                out_key = out.path_info.parts
+
+                # Check for dup outs
+                if out_key in outs:
+                    dup_stages = [stage, outs[out_key].stage]
                     raise OutputDuplicationError(str(out), dup_stages)
-                outs[out.path_info] = out
+
+                # Check for overlapping outs
+                if outs.has_subtrie(out_key):
+                    parent = out
+                    overlapping = first(outs.values(prefix=out_key))
+                else:
+                    parent = outs.shortest_prefix(out_key).value
+                    overlapping = out
+                if parent and overlapping:
+                    msg = (
+                        "Paths for outs:\n'{}'('{}')\n'{}'('{}')\n"
+                        "overlap. To avoid unpredictable behaviour, "
+                        "rerun command with non overlapping outs paths."
+                    ).format(
+                        str(parent),
+                        parent.stage.relpath,
+                        str(overlapping),
+                        overlapping.stage.relpath,
+                    )
+                    raise OverlappingOutputPathsError(parent, overlapping, msg)
+
+                outs[out_key] = out
 
         for stage in stages:
-            for out in stage.outs:
-                for p in out.path_info.parents:
-                    if p in outs:
-                        msg = (
-                            "Paths for outs:\n'{}'('{}')\n'{}'('{}')\n"
-                            "overlap. To avoid unpredictable behaviour, "
-                            "rerun command with non overlapping outs paths."
-                        ).format(
-                            str(outs[p]),
-                            outs[p].stage.relpath,
-                            str(out),
-                            out.stage.relpath,
-                        )
-                        raise OverlappingOutputPathsError(outs[p], out, msg)
+            out = outs.shortest_prefix(PathInfo(stage.path).parts).value
+            if out:
+                raise StagePathAsOutputError(stage, str(out))
 
+        # Building graph
+        G.add_nodes_from(stages)
         for stage in stages:
-            stage_path_info = PathInfo(stage.path)
-            for p in chain([stage_path_info], stage_path_info.parents):
-                if p in outs:
-                    raise StagePathAsOutputError(stage, str(outs[p]))
-
-        for stage in stages:
-            G.add_node(stage)
-
             for dep in stage.deps:
                 if dep.path_info is None:
                     continue
 
-                for out_path_info, out in outs.items():
-                    if out_path_info.overlaps(dep.path_info):
-                        G.add_node(out.stage)
-                        G.add_edge(stage, out.stage)
+                dep_key = dep.path_info.parts
+                overlapping = list(n.value for n in outs.prefixes(dep_key))
+                if outs.has_subtrie(dep_key):
+                    overlapping.extend(outs.values(prefix=dep_key))
+
+                G.add_edges_from((stage, out.stage) for out in overlapping)
 
         check_acyclic(G)
 
diff --git a/setup.py b/setup.py
index 8dea1df6d..1f7f67e84 100644
--- a/setup.py
+++ b/setup.py
@@ -77,6 +77,7 @@ install_requires = [
     "speedcopy>=2.0.1; sys_platform == 'win32'",
     "flatten_json>=0.1.6",
     "texttable>=0.5.2",
+    "pygtrie==2.3.2",
 ]