Hey everyone - how do i perform a conditional split on a single input

Hi! What do you mean by conditional “split” — as in branching task trees based on the value of an input? Does https://docs.prefect.io/api/latest/tasks/control_flow.html#functions do what you want?

I have a pandas df with Group A and Group B. I need to pass group A to task 1 and group B to task 2

Is it possible to implement your situation such that the dataframe is returned from the upstream task already split up into two, like

{'a': df1, 'b': df2}

, and then the downstream tasks index that result like

task1(result['a']

and

task2(result['b']

? Basically doing the splitting itself in the upstream task

Can you provide a full code sample

This is more psuedocode but this is the idea; where you split it yourself in an upstream task based on whatever rules you want to split it by:

@task
def dataframe_stuff() -> dict:
    df = pd.DataFrame([['a', 1], ['b', 2]], columns=['group','data'])
    # add/do other stuff to your dataframe
    groups = df.groupby('group')
    return {k:v for (k,v) in groups}

@task 
task1(a: 'pandas.DataFrame' -> 'pandas.DataFrame'):
	# do stuff special to a
    return a

@task 
task2(b: 'pandas.DataFrame') -> 'pandas.DataFrame':
	# do stuff special to b
    return b

with Flow("my flow") as flow:
    split_dataframes = dataframe_stuff()
    result_on_a = task1(split_dataframes['a'])
    result_on_b = task2(split_dataframe['b'])

Alternatively, if

task1

and

task2

do the same thing so calling either is the same, you can probably use mapping instead and return a list of dataframes from

dataframe_stuff()

instead of a dictionary. Either way you’ll still need to return some structure from the upstream task that is already split up however it is that you want the dataframe to be split up by for its downstream tasks.