Skip to content

Commit

Permalink
Improve performance of infer primary key
Browse files Browse the repository at this point in the history
  • Loading branch information
gshank committed Sep 25, 2024
1 parent 359a2c0 commit 6314d8a
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 12 deletions.
6 changes: 6 additions & 0 deletions .changes/unreleased/Fixes-20240925-160543.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
kind: Fixes
body: Improve performance of infer primary key
time: 2024-09-25T16:05:43.59536-04:00
custom:
Author: gshank
Issue: "10781"
30 changes: 18 additions & 12 deletions core/dbt/parser/manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -1150,10 +1150,15 @@ def process_saved_queries(self, config: RuntimeConfig):

def process_model_inferred_primary_keys(self):
"""Processes Model nodes to populate their `primary_key`."""
model_to_generic_test_map: Dict[str, List[GenericTestNode]] = {}

Check warning on line 1153 in core/dbt/parser/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/manifest.py#L1153

Added line #L1153 was not covered by tests
for node in self.manifest.nodes.values():
if not isinstance(node, ModelNode):
continue
generic_tests = self._get_generic_tests_for_model(node)
if node.created_at < self.started_at:
continue
if not model_to_generic_test_map:
model_to_generic_test_map = self.build_model_to_generic_tests_map()
generic_tests = model_to_generic_test_map[node.unique_id]

Check warning on line 1161 in core/dbt/parser/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/manifest.py#L1157-L1161

Added lines #L1157 - L1161 were not covered by tests
primary_key = node.infer_primary_key(generic_tests)
node.primary_key = sorted(primary_key)

Expand Down Expand Up @@ -1425,23 +1430,24 @@ def write_perf_info(self, target_path: str):
write_file(path, json.dumps(self._perf_info, cls=dbt.utils.JSONEncoder, indent=4))
fire_event(ParsePerfInfoPath(path=path))

def _get_generic_tests_for_model(
def build_model_to_generic_tests_map(
self,
model: ModelNode,
) -> List[GenericTestNode]:
) -> Dict[str, List[GenericTestNode]]:
"""Return a list of generic tests that are attached to the given model, including disabled tests"""
tests = []
model_to_generic_tests_map: Dict[str, List[GenericTestNode]] = {}

Check warning on line 1438 in core/dbt/parser/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/manifest.py#L1438

Added line #L1438 was not covered by tests
for _, node in self.manifest.nodes.items():
if isinstance(node, GenericTestNode) and node.attached_node == model.unique_id:
tests.append(node)
if isinstance(node, GenericTestNode) and node.attached_node:
if node.attached_node not in model_to_generic_tests_map:
model_to_generic_tests_map[node.attached_node] = []
model_to_generic_tests_map[node.attached_node].append(node)

Check warning on line 1443 in core/dbt/parser/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/manifest.py#L1440-L1443

Added lines #L1440 - L1443 were not covered by tests
for _, nodes in self.manifest.disabled.items():
for disabled_node in nodes:
if (
isinstance(disabled_node, GenericTestNode)
and disabled_node.attached_node == model.unique_id
):
tests.append(disabled_node)
return tests
if isinstance(disabled_node, GenericTestNode) and disabled_node.attached_node:
if disabled_node.attached_node not in model_to_generic_tests_map:
model_to_generic_tests_map[disabled_node.attached_node] = []
model_to_generic_tests_map[disabled_node.attached_node].append(disabled_node)
return model_to_generic_tests_map

Check warning on line 1450 in core/dbt/parser/manifest.py

View check run for this annotation

Codecov / codecov/patch

core/dbt/parser/manifest.py#L1446-L1450

Added lines #L1446 - L1450 were not covered by tests


def invalid_target_fail_unless_test(
Expand Down

0 comments on commit 6314d8a

Please sign in to comment.