Repository: dbt-labs/dbt-core
Pull request: Improve tree traversal of select_children
Resolves #10434

### Problem

DBT core performance was reported as degraded between `1.5` and `1.8` for some users with selectors of the style `tag:mytag+`. For selectors like this and for some large DAGs, the current implementation of getting descendants makes unnecessarily many calls to find children get the edge type, which I believe is a new check in finding children between `1.5` and `1.8`. Many nodes may have the same descendents; this implementation takes care not to check the same edge twice.

### Solution

This searches for descendents by stepping at a depth at 1 at a time, de-duplicating descendents by nature of Python sets, and removes nodes at each step that have already had their descendents searched.

I am looking for user testing to see if this improves performance and to apply this solution to parents as well if it is helpful.

I am seeing improved performance and will post my results in the issue thread.

### Checklist

- [X] I have read [the contributing guide](https://github.com/dbt-labs/dbt-core/blob/main/CONTRIBUTING.md) and understand what's expected of me.
- [X] I have run this code in development, and it appears to resolve the stated issue.
- [x] This PR includes tests, or tests are not required or relevant for this PR.
- [x] This PR has no interface changes (e.g., macros, CLI, logs, JSON artifacts, config files, adapter interface, etc.) or this PR has already received feedback and approval from Product or DX.
- [X] This PR includes [type annotations](https://docs.python.org/3/library/typing.html) for new and modified functions.


Problem statement:
[Regression] 1.8.2 slower to build than 1.5.9 when tag+ includes many nodes
### Is this a regression in a recent version of dbt-core?

- [X] I believe this is a regression in dbt-core functionality
- [X] I have searched the existing issues, and I could not find an existing issue for this regression

### Current Behavior

`dbt build -s tag:my_tag+` takes about 20 minutes longer to start on dbt 1.8.2 than it does on 1.5.9 with the same tag. The tag used has a lot of downstream nodes in our project, about 11k. Generally we’re seeing better performance on 1.8 so we were surprised to see this big regression in performance. 

### Expected/Previous Behavior

Previously building everything downstream of a tag with lots of nodes would take a couple minutes of startup time and then begin running queries against our db. Now it takes 20+ minutes. 

### Steps To Reproduce

1. Set up dbt project with a tag that has about 11k downstream nodes
2. Install dbt 1.8.2
3. Run dbt build -s tag:my_tag+

### Relevant log output

_No response_

### Environment

```markdown
- OS: MacOS 14.5 and Ubuntu 22.04
- Python: 3.9.12
- dbt (working version): 1.5.9
- dbt (regression version): 1.8.2
```


### Which database adapter are you using with dbt?

bigquery

### Additional Context

The reason we need this is because we have a selector used in CI/CD that excludes everything downstream of a tag that is upstream of many nodes. The example given is a simpler version of the original issue we found with that selector. (We tested the simpler version and found to also have the same performance issue.) The selector was something like
```
- name: my_selector
  definition: 
    union:
      - state:modified+
      - exclude:
        - tag:my_tag+
```

Hints from issue discussion:
Thanks @cajubelt! Is this only for the `build` command (i.e. not `run`)?

I suspect this might be due to the large number of tests in your project (per our conversation), and the additional time that dbt spends "linking" the DAG (adding edges between test on upstream model -> downstream models, so that they skip on test failure). While it's not immediately clear to me what change we made between v1.5 -> v1.8 to that logic, if you only see this slowdown on `build`, it would be a strong indication that therein lies the problem.
@jtcohen6 yes I just confirmed that dbt run doesn't have any noticeable slowdown on 1.8 with the same selector. 
@cajubelt Okay! So I think the hypothesis is:
- A project with a lot of tests (~5k models and ~15k data tests) is meaningfully slower at startup for `dbt build` in v1.8 compared to v1.5
- My suspicion is that this slowdown is happening within (or closely adjacent to) the [`add_test_edges`](https://github.com/dbt-labs/dbt-core/blob/1aeff2c58ffd608bb35320b7d0ef35850a88da42/core/dbt/compilation.py#L193) method
If you're up for it, there are two ways to try confirming that hypothesis by profiling dbt's performance:
- Using snakeviz as documented here: https://docs.getdbt.com/reference/global-configs/record-timing-info (this way is a bit older, but more familiar to me)
- Using `py-spy` (snazzier) in the way @peterallenwebb recommends [here](https://github.com/dbt-labs/dbt-core/issues/10127#issuecomment-2105107655):
```
pip install py-spy
sudo py-spy record -s -f speedscope -- dbt build
```
We have the ability to exclude tests from the build resource types, but that doesn't affect whether or not we pass "add_test_edges" to compile. It would be possible to check the resource types and if tests are excluded, pass "add_test_edges" = False.

Of course if you do want tests to run, that doesn't help.
We do want to run tests if possible.

@jtcohen6 here's a screenshot of a search through py-spy's output for the longest segments of the flame graph. Looks like `generic_bfs_edges` is adding 12:41, not sure if that's particularly revealing however

<img width="735" alt="Screenshot 2024-07-14 at 8 43 37 PM" src="https://github.com/user-attachments/assets/4f5c8996-4535-446d-b4aa-bb8170c3b183">

@jtcohen6 Not super sure if related to this issue but we're seeing an extreme slowdown when `build`ing a large chunk of a fairly big graph with `--indirect-selection buildable` in CI/CD. The project has `2573 models, 13 snapshots, 7 analyses, 2254 data tests, 20 seeds, 1 operation, 359 sources, 44 exposures, 1408 macros, 1 group, 13 unit tests`

`cautious` and `eager` end up being fine, but with `buildable`, the process "hangs" for up to 20 minutes after it prints the project stats "Found 2573 models ..." etc. and before the "Concurrency: 4 threads (target='test')" message. No other output, even on `debug` log level. 

I tried to figure out what's going on using the `trace` module but it was a bit too much spam :). Will try add some trace output when I find a minute
here's me waiting for about 5minutes for the `generic_bfs_edges` traversal

<img width="1338" alt="image" src="https://github.com/user-attachments/assets/70c39911-c0a4-490c-9725-f6fb6739b453">

In a "good" run with `eager` selection this takes about two seconds with the exact same selector (the rest is dbt actually starting to execute, which I didn't wait for in the other run)

<img width="1306" alt="image" src="https://github.com/user-attachments/assets/e59c7fb7-bc43-403a-837b-1adf520d87c5">

What percentage of your graph (including tests) is selected by the tag, and what % is selected by the selector?

I had someone on my team try using the `@` decorator on a selector (which I don't normally do) and we both had DBT hang. I kind of pushed it out of my head, will be thinking on this now.

I don't really understand the `py-spy` output (worked with SnakeViz).... if I had to guess what might be happening:
* for some reason, `generic_bfs_edges` is either really slow or being called a lot
* `collect_specified_neighbors` in `selector.py` calls `self.graph.select_children` when passed certain selector patterns
* `self.graph.select_children` applies the `self.descendants` method for each node, which is perhaps a lot of times, or maybe it uses up a lot of memory.
* the slow part based on your output is `filter_edges_by_type`, which is calling `get_edge_data` in `networkx`,

If you can share the snakeviz would help out, also maybe check your memory usage? Last bug I worked on for this graph optimization stuff was having performance issues from running out of memory when the DAG got too big.

If my stab in the dark here is in the right direction, maybe there is a pruning that could be done to reduce calling `descendants` so much. Like, if you have a ton of nodes in your selector for which many are already descendents of each other, you don't need to re-tread and evaluate each node. could first get a list of all descendents for each node with a depth of `1`, then take the difference of that and the list of nodes in the selector to restrict to searching for descendents of unsearched nodes, find descendents of depth 1 from that list, and iterate.

For your graph has 20k nodes, if 1k have the tag, 10k are descendents, and each descendent is on average the descendent of 70% of the tagged nodes, this takes the number of times each of the descendents has to be evaluated by ~700. 700*10,000 is a lot of calls.

After writing this out this seems actually pretty likely XD. Think an update to `select_children` with this logic would be the place to implement this. This seems similar to https://github.com/dbt-labs/dbt-core/pull/7194

Also if the slow part is `filtered_graph = self.exclude_edge_type("parent_test")`, maybe could cache that or pass it as a parameter into `descendants`?

will try tests on my own project with ~500 models and ~5000 tests next week.
I added a PR! https://github.com/dbt-labs/dbt-core/pull/10526

I believe I was able to recreate the issue in my project, which has: `Found 442 models, 43 snapshots, 5112 data tests, 32 seeds, 2 operations, 134 sources, 23 exposures, 1075 macros`

Running `dbt build --select tag:silver+` on my DAG on 1.8.4, it takes 132 seconds to start building. `dbt build` just takes 11 seconds.

My PR applies the pruning strategy I described. I saw the runtime improve for `dbt build --select tag:silver+` to 66 seconds. Using SnakeViz, I identified that the number of calls to `get_edge_data` reduced from `14,332,569` to `383,916`.

I am unsure how much DBT 1.9/main branch contains performance improvements... the total runtime improvement was 66 seconds, but `select_children` only improved from 80 seconds to 45 seconds (35 seconds improvement). Going to try to run `main` branch for performance check.

Feedback welcome, especially if you can test/review the change!

Commit message: remove unused function

Patch to review:
diff --git a/.changes/unreleased/Under the Hood-20240809-130234.yaml b/.changes/unreleased/Under the Hood-20240809-130234.yaml
new file mode 100644
index 00000000000..964dd2fedf2
--- /dev/null
+++ b/.changes/unreleased/Under the Hood-20240809-130234.yaml	
@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Improve speed of tree traversal when finding children, increasing build speed for some selectors
+time: 2024-08-09T13:02:34.759905-07:00
+custom:
+    Author: ttusing
+    Issue: "10434"
diff --git a/core/dbt/graph/graph.py b/core/dbt/graph/graph.py
index 87ba64414fb..54984f51102 100644
--- a/core/dbt/graph/graph.py
+++ b/core/dbt/graph/graph.py
@@ -59,18 +59,40 @@ def select_childrens_parents(self, selected: Set[UniqueId]) -> Set[UniqueId]:
     def select_children(
         self, selected: Set[UniqueId], max_depth: Optional[int] = None
     ) -> Set[UniqueId]:
-        descendants: Set[UniqueId] = set()
-        for node in selected:
-            descendants.update(self.descendants(node, max_depth))
-        return descendants
+        """Returns all nodes which are descendants of the 'selected' set.
+        Nodes in the 'selected' set are counted as children only if
+        they are descendants of other nodes in the 'selected' set."""
+        children: Set[UniqueId] = set()
+        i = 0
+        while len(selected) > 0 and (max_depth is None or i < max_depth):
+            next_layer: Set[UniqueId] = set()
+            for node in selected:
+                next_layer.update(self.graph.successors(node))
+            next_layer = next_layer - children  # Avoid re-searching
+            children.update(next_layer)
+            selected = next_layer
+            i += 1
+
+        return children
 
     def select_parents(
         self, selected: Set[UniqueId], max_depth: Optional[int] = None
     ) -> Set[UniqueId]:
-        ancestors: Set[UniqueId] = set()
-        for node in selected:
-            ancestors.update(self.ancestors(node, max_depth))
-        return ancestors
+        """Returns all nodes which are ancestors of the 'selected' set.
+        Nodes in the 'selected' set are counted as parents only if
+        they are ancestors of other nodes in the 'selected' set."""
+        parents: Set[UniqueId] = set()
+        i = 0
+        while len(selected) > 0 and (max_depth is None or i < max_depth):
+            next_layer: Set[UniqueId] = set()
+            for node in selected:
+                next_layer.update(self.graph.predecessors(node))
+            next_layer = next_layer - parents  # Avoid re-searching
+            parents.update(next_layer)
+            selected = next_layer
+            i += 1
+
+        return parents
 
     def select_successors(self, selected: Set[UniqueId]) -> Set[UniqueId]:
         successors: Set[UniqueId] = set()


Code review comments:
diff_hunk: @@ -59,18 +59,40 @@ def select_childrens_parents(self, selected: Set[UniqueId]) -> Set[UniqueId]:
     def select_children(
         self, selected: Set[UniqueId], max_depth: Optional[int] = None
     ) -> Set[UniqueId]:
-        descendants: Set[UniqueId] = set()
-        for node in selected:
-            descendants.update(self.descendants(node, max_depth))
-        return descendants
+        """Returns all nodes which are descendants of the 'selected' set.
+        Nodes in the 'selected' set are counted as children only if
+        they are descendants of other nodes in the 'selected' set."""
+        children: Set[UniqueId] = set()
+        i = 0
+        while len(selected) > 0 and (max_depth is None or i < max_depth):
+            next_layer: Set[UniqueId] = set()
+            for node in selected:
+                next_layer.update(self.graph.successors(node))
+            next_layer = next_layer - children  # Avoid re-searching
+            children.update(next_layer)
+            selected = next_layer
+            i += 1
+
+        return children
 
     def select_parents(
         self, selected: Set[UniqueId], max_depth: Optional[int] = None
     ) -> Set[UniqueId]:
-        ancestors: Set[UniqueId] = set()
-        for node in selected:
-            ancestors.update(self.ancestors(node, max_depth))
-        return ancestors
+        """Returns all nodes which are ancestors of the 'selected' set.
+        Nodes in the 'selected' set are counted as parents only if
+        they are ancestors of other nodes in the 'selected' set."""
+        parents: Set[UniqueId] = set()
+        i = 0
+        while len(selected) > 0 and (max_depth is None or i < max_depth):
+            next_layer: Set[UniqueId] = set()
+            for node in selected:
+                next_layer.update(self.graph.predecessors(node))
path: core/dbt/graph/graph.py
text: @user1:
And this one should read: `next_layer.update(self.ancestors(node, 1))`

diff_hunk: @@ -59,18 +59,40 @@ def select_childrens_parents(self, selected: Set[UniqueId]) -> Set[UniqueId]:
     def select_children(
         self, selected: Set[UniqueId], max_depth: Optional[int] = None
     ) -> Set[UniqueId]:
-        descendants: Set[UniqueId] = set()
-        for node in selected:
-            descendants.update(self.descendants(node, max_depth))
-        return descendants
+        """Returns all nodes which are descendants of the 'selected' set.
+        Nodes in the 'selected' set are counted as children only if
+        they are descendants of other nodes in the 'selected' set."""
+        children: Set[UniqueId] = set()
+        i = 0
+        while len(selected) > 0 and (max_depth is None or i < max_depth):
+            next_layer: Set[UniqueId] = set()
+            for node in selected:
+                next_layer.update(self.graph.successors(node))
path: core/dbt/graph/graph.py
text: @user1:
It was my mistake. This line should read: `next_layer.update(self.descendants(node, 1))`

@author:
I made this same mistake on first pass :P. Fixing!