How to concurrently do a multi stage data merge with Python

The problem is to find a more elegant solution to concurrently process a multi stage data merge/pipeline (function process_concurrent() below).

Error handling was omitted for brevity, but assume that rows for which there is an exception at any stage would be skipped from the output. The order of producing the output "files" is not important.

Beside being complicated the solution in process_concurrent() has several shortcomings:

• Can potentially accumulate large amount of data in memory (for real world problems). Some management is done with del statements, but player_age dict for example will keep growing, if we want to allow for a player to be in more than one team (i.e. player one is in teams A and C).

• Although phase 4 (simulates output) starts as soon as first team's data is ready, the result from that phase will not be known until all teams data are ready. This makes it impossible to abort early (or retry) in case I/O fails.

• The above statement can be generalized for every phase in that although its tasks begin as soon as submitted, to check results of completed tasks we wait for completion of the entire phase before it.

I am not "married" to concurrent.futures idioms. This problem can be probably expressed as a multi-stage map/reduce.

File conc.py:

"""Produces a CSV table of team players with age. 

Simulation for testing concurrent solution to a multi stage 
data merge problem.

>>> import conc
>>> conc.process_sequential(teams=['A', 'C', 'B'])

Team_A.csv:
team,player,age
A,1,10
A,2,15
A,3,20

Team_C.csv:
team,player,age
C,8,20
C,1,10

Team_B.csv:
team,player,age
B,4,25
B,5,10
B,7,15
>>>
>>> conc.process_concurrent(teams=['A', 'C', 'B'])

Team_B.csv:
team,player,age
B,4,25
B,5,10
B,7,15

Team_C.csv:
team,player,age
C,8,20
C,1,10

Team_A.csv:
team,player,age
A,1,10
A,2,15
A,3,20
>>>
"""

from concurrent import futures
from time import sleep
from random import random


team_players = {
    'A': ['1', '2', '3'],
    'B': ['4', '5', '7'],
    'C': ['8', '1']
}

player_age = {
    '1': 10,
    '2': 15,
    '3': 20,
    '4': 25,
    '5': 10,
    '6': 20,
    '7': 15,
    '8': 20
}


def phase1_lookup_team_players(team):
    """lookup list of players for given team"""
    sleep(random())
    return team_players[team]


def phase2_lookup_player_age(player):
    """lookup player's age"""
    sleep(random())
    return player_age[player]


def phase3_merge(team, players, ages):
    """pulls together three data sources into a list of dicts"""
    sleep(random())
    recs = [{'team': team, 'player': p, 'age': a} for p, a in zip(players, ages)]
    return recs


def phase4_make_csv_table(team, team_records):
    """I/O phase, simulates output to file or database"""
    sleep(random())
    rs = f'\nTeam_{team}.csv:\nteam,player,age'
    for r in team_records:
        rs += f'\n{r["team"]},{r["player"]},{r["age"]}'
    return rs


def process_sequential(teams):
    """simple sequential solution"""
    for team in teams:
        players = phase1_lookup_team_players(team)
        ages = []
        for player in players:
            age = phase2_lookup_player_age(player)
            ages.append(age)
        team_records = phase3_merge(team, players, ages)
        csvstr = phase4_make_csv_table(team, team_records)
        print(csvstr)


def process_concurrent(teams):
    """concurrent solution with overy complicated state management"""
    team_players = {}
    player_age = {}
    player_count = {}
    tp_jobs = {}
    pa_jobs = {}
    m_jobs = {}
    render_jobs = {}

    with futures.ProcessPoolExecutor(4) as tpe:

        # submit team player lookup jobs
        for t in teams:
            tp_jobs[tpe.submit(phase1_lookup_team_players, t)] = t

        # print(f'tp_jobs: {tp_jobs.values()}')
        for tj in futures.as_completed(tp_jobs):
            team = tp_jobs[tj]
            players = tj.result()
            player_count[team] = len(players)
            team_players[team] = players
            # print(f'tp_job completed: {team}: {players}')
            del tp_jobs[tj]

            for player in players:
                # submit player age lookup jobs
                pa_jobs[tpe.submit(phase2_lookup_player_age, player)] = team, player
            del players
        assert len(tp_jobs) == 0

        # accumulate ages list for each player in team
        # print(f'pa_jobs: {pa_jobs.values()}')
        for aj in futures.as_completed(pa_jobs):
            team, player = pa_jobs[aj]
            age = aj.result()
            # print(f'pa_job completed {team}, {player}, {age}')
            player_age[player] = age
            player_count[team] -= 1  # reduce the number of players remaining to process for the team
            del pa_jobs[aj]

            # once all players' age is collected for a team, submit to the merge function
            if player_count[team] == 0:
                ages = [player_age[p] for p in team_players[team]]
                players = team_players[team]
                del team_players[team]
                del player_count[team]
                m_jobs[tpe.submit(phase3_merge, team=team, players=players, ages=ages)] = team
        assert len(pa_jobs) == 0

        # retrieve merged records and submit for output
        for mj in futures.as_completed(m_jobs):
            team = m_jobs[mj]
            team_records = mj.result()
            del m_jobs[mj]
            render_jobs[tpe.submit(phase4_make_csv_table, team=team, team_records=team_records)] = team
        assert len(m_jobs) == 0

        for rj in futures.as_completed(render_jobs):
            team = render_jobs[rj]
            csvstr = rj.result()
            # print(f'render_job completed {team}')
            print(csvstr)
            del render_jobs[rj]