@@ -48,14 +48,14 @@ def split(
4848 process_graph = process ["process_graph" ]
4949
5050 # 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 = } " )
5959
6060 primary_backend = backend_usage .most_common (1 )[0 ][0 ] if backend_usage else None
6161 secondary_backends = {b for b in backend_usage if b != primary_backend }
@@ -70,7 +70,7 @@ def split(
7070
7171 for node_id , node in process_graph .items ():
7272 if node ["process_id" ] == "load_collection" :
73- bid = self . backend_for_collection ( node ["arguments" ]["id" ])
73+ bid = backend_per_collection [ node ["arguments" ]["id" ]]
7474 if bid == primary_backend and not (
7575 self ._always_split and primary_has_load_collection
7676 ):
0 commit comments