Another approach is to persist the collected codes in a Delta table and then read from this table in downstream tasks.

Make sure to add ample logging and counts.

Checkpointing also would help if you suspect the counts in set are not the same as what you see in the 

key = "codes".