#3226 [Core feature] map_task should be able to handle a partitioned StructuredDataset Issue created by cosmicBboy Motivation: Why do you think this is important? As a data practitioner, I should be able to apply a

map_task

to a partitioned StructuredDataset automatically so that I can process the partitions in an embarrassingly parallel fashion without too much extra code. Goal: What should the final outcome look like, ideally? Suppose we have a task that produces a

StructuredDataset

@task
def make_df() -> StructuredDataset:
    df = pd.DataFrame.from_records([
        {
            "id": i,
            "partition": (i % 10) + 1,
            "name": "".join(
                random.choices(string.ascii_uppercase + string.digits, k=10)
            )
        }
        for i in range(1000)
    ])
    return StructuredDataset(dataframe=df, partition_col=["partition"])

Ideally, I should be able to do something like this:

@task
def process_df(dataset: StructuredDataset) -> StructuredDataset:
    df = structured_dataset.open(pd.DataFrame).read_partition()  # read the partition
    ... # do stuff

@task
def use_processed_df(dataset: List[StructuredDataset]) -> ...:
    ...

@workflow
def wf() -> StructuredDataset:
    structured_dataset = make_df()
    # where structured_dataset.partitions is a list of unpartitioned StructuredDatasets
    results: List[StructuredDataset] = map_task(process_df)(dataset=structured_dataset.partitions)
    return use_processed_df(dataset=results)

Note that in this example code a few magical things are happening: 1. we pass in

structured_dataset.partitions

into the map task, which indicates that we want to apply

process_df

to each of the partitions defined in

make_df

2. The fact that

map_task(process_df)

returns a

StructuredDataset

implies that using map tasks with structured datasets does an implicit reduction, i.e. the outputs of

map_task(process_df)

are written to the same blob store prefix. Ideally the solution enables processing of

StructuredDataset

without having to manually handle reading in of partitions in the map task, and automatically reduces the results into a

StructuredDataset

without having to explicitly write a coalense/reduction task. Describe alternatives you've considered Users would have to roll their own way of processing partitions of a structured dataset using dynamic tasks. Propose: Link/Inline OR Additional context Slack context: https://flyte-org.slack.com/archives/CP2HDHKE1/p1673380243923279 Related to #3219 Are you sure this issue hasn't been raised already? ☑︎ Yes Have you read the Code of Conduct? ☑︎ Yes flyteorg/flyte