@@ -48,14 +48,14 @@ def split(
48
48
process_graph = process ["process_graph" ]
49
49
50
50
# Extract necessary back-ends from `load_collection` usage
51
- backend_usage = collections . Counter (
52
- self .backend_for_collection (node [ "arguments" ][ "id" ] )
53
- for node in process_graph . values ()
54
- if node ["process_id" ] == "load_collection"
55
- )
56
- _log . info (
57
- f"Extracted backend usage from `load_collection` nodes: { backend_usage } "
58
- )
51
+ backend_per_collection : Dict [ str , str ] = {
52
+ cid : self .backend_for_collection (cid )
53
+ for cid in (
54
+ node [ "arguments" ][ "id" ] for node in process_graph . values () if node ["process_id" ] == "load_collection"
55
+ )
56
+ }
57
+ backend_usage = collections . Counter ( backend_per_collection . values ())
58
+ _log . info ( f"Extracted backend usage from `load_collection` nodes: { backend_usage = } { backend_per_collection = } " )
59
59
60
60
primary_backend = backend_usage .most_common (1 )[0 ][0 ] if backend_usage else None
61
61
secondary_backends = {b for b in backend_usage if b != primary_backend }
@@ -70,7 +70,7 @@ def split(
70
70
71
71
for node_id , node in process_graph .items ():
72
72
if node ["process_id" ] == "load_collection" :
73
- bid = self . backend_for_collection ( node ["arguments" ]["id" ])
73
+ bid = backend_per_collection [ node ["arguments" ]["id" ]]
74
74
if bid == primary_backend and not (
75
75
self ._always_split and primary_has_load_collection
76
76
):
0 commit comments