Skip to content

Commit

Permalink
Fix two other O(n^2) input callsites (#26493)
Browse files Browse the repository at this point in the history
Summary:
Spiritual siblings to #26482
just in a slightly different callstack

BK, speedscope on large asset graph

## Summary & Motivation

## How I Tested These Changes

## Changelog

> Insert changelog entry or delete this section.
  • Loading branch information
gibsondan authored Dec 16, 2024
1 parent b9a5069 commit 7cf69b9
Showing 1 changed file with 10 additions and 9 deletions.
19 changes: 10 additions & 9 deletions python_modules/dagster/dagster/_core/snap/dep_snapshot.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from collections import defaultdict
from functools import cached_property
from typing import DefaultDict, Dict, List, Mapping, NamedTuple, Sequence

import dagster._check as check
Expand Down Expand Up @@ -140,11 +141,7 @@ def get_upstream_outputs(self, node_name: str, input_name: str) -> Sequence["Out
check.str_param(node_name, "node_name")
check.str_param(input_name, "input_name")

for input_dep_snap in self.get_invocation(node_name).input_dep_snaps:
if input_dep_snap.input_name == input_name:
return input_dep_snap.upstream_output_snaps

check.failed(f"Input {input_name} not found for node {node_name}")
return self.get_invocation(node_name).input_dep_snap(input_name).upstream_output_snaps

def get_upstream_output(self, node_name: str, input_name: str) -> "OutputHandleSnap":
check.str_param(node_name, "node_name")
Expand Down Expand Up @@ -235,9 +232,13 @@ def __new__(
is_dynamic_mapped=check.bool_param(is_dynamic_mapped, "is_dynamic_mapped"),
)

@cached_property
def input_dep_map(self) -> Mapping[str, InputDependencySnap]:
return {inp_snap.input_name: inp_snap for inp_snap in self.input_dep_snaps}

def input_dep_snap(self, input_name: str) -> InputDependencySnap:
for inp_snap in self.input_dep_snaps:
if inp_snap.input_name == input_name:
return inp_snap
inp_snap = self.input_dep_map.get(input_name)
if inp_snap:
return inp_snap

check.failed(f"No input found named {input_name}")
check.failed(f"Input {input_name} not found for node {self.node_name}")

0 comments on commit 7cf69b9

Please sign in to comment.