Skip to content

Commit f404fb4

Browse files
committed
Make migration node frozen
1 parent a3a491f commit f404fb4

File tree

1 file changed

+14
-15
lines changed

1 file changed

+14
-15
lines changed

src/databricks/labs/ucx/assessment/sequencing.py

Lines changed: 14 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ class MigrationStep:
4242
MigrationNodeKey = tuple[str, str]
4343

4444

45-
@dataclass
45+
@dataclass(frozen=True)
4646
class MigrationNode:
4747
# TODO: @JCZuurmond the prefixes look a bit redundant
4848
node_id: int = field(compare=False)
@@ -236,20 +236,20 @@ def generate_steps(self) -> Iterable[MigrationStep]:
236236
# pre-compute incoming keys for best performance of self._required_step_ids
237237
incoming = self._invert_outgoing_to_incoming()
238238
incoming_counts = self._compute_incoming_counts(incoming)
239-
key_queue = self._create_key_queue(incoming_counts)
240-
key = key_queue.get()
239+
queue = self._create_node_queue(incoming_counts)
240+
node = queue.get()
241241
step_number = 1
242-
sorted_steps: list[MigrationStep] = []
243-
while key is not None:
244-
step = self._nodes[key].as_step(step_number, sorted(n.node_id for n in incoming[key]))
245-
sorted_steps.append(step)
242+
ordered_steps: list[MigrationStep] = []
243+
while node is not None:
244+
step = node.as_step(step_number, sorted(n.node_id for n in incoming[node.key]))
245+
ordered_steps.append(step)
246246
# Update queue priorities
247-
for dependency in self._outgoing[key]:
247+
for dependency in self._outgoing[node.key]:
248248
incoming_counts[dependency.key] -= 1
249-
key_queue.update(incoming_counts[dependency.key], dependency)
249+
queue.update(incoming_counts[dependency.key], dependency)
250250
step_number += 1
251-
key = key_queue.get()
252-
return sorted_steps
251+
node = queue.get()
252+
return ordered_steps
253253

254254
def _invert_outgoing_to_incoming(self) -> dict[MigrationNodeKey, set[MigrationNode]]:
255255
result: dict[MigrationNodeKey, set[MigrationNode]] = defaultdict(set)
@@ -266,14 +266,13 @@ def _compute_incoming_counts(
266266
result[node_key] = len(incoming[node_key])
267267
return result
268268

269-
@staticmethod
270-
def _create_key_queue(incoming_counts: dict[tuple[str, str], int]) -> PriorityQueue:
271-
"""Create a priority queue given the keys and their incoming counts.
269+
def _create_node_queue(self, incoming_counts: dict[MigrationNodeKey, int]) -> PriorityQueue:
270+
"""Create a priority queue for their nodes using the incoming count as priority.
272271
273272
A lower number means it is pulled from the queue first, i.e. the key with the lowest number of keys is retrieved
274273
first.
275274
"""
276275
priority_queue = PriorityQueue()
277276
for node_key, incoming_count in incoming_counts.items():
278-
priority_queue.put(incoming_count, node_key)
277+
priority_queue.put(incoming_count, self._nodes[node_key])
279278
return priority_queue

0 commit comments

Comments
 (0)