How might I implement a map-reduce style graph_asset

[cmpadden]

I made a graph asset that has a downstream dependency, and according to the graph, the downstream dependency also got chunked.

The @graph_asset is generate_flowcharts. It has three ins and it has two ops: one that yields many DynamicOutput and one that is the map Callable. All in all, the @graph_asset does a chunk->map->collect operation and return result.

The downstream dependency as you can see in the image is generate_single_csv. It is a single @asset with deps=[AssetKey("generate_flowcharts")] . This should not be chunked, but it looks like it is.

From community Slack member: https://dagster.slack.com/archives/C01U954MEER/p1719926132336219

[cmpadden]

Here is example code to accomplish this -- note the @op that is used to coalesce our chunked results.

from itertools import chain

from dagster import Definitions, DynamicOut, DynamicOutput, asset, graph_asset, op


@asset
def upstream_asset() -> list:
    return list(range(1, 100))


@op(out=DynamicOut(list))
def split_into_chunks(upstream: list):
    chunk_size = 10
    for i in range(0, len(upstream), chunk_size):
        chunk: list[int] = upstream[i : i + chunk_size]
        yield DynamicOutput(chunk, mapping_key=str(i // chunk_size))


@op
def process_chunk(chunk: list) -> list:
    return [i**2 for i in chunk]


@op
def combine_chunks(chunks: list) -> list:
    return list(chain.from_iterable(chunks))


@graph_asset
def map_reduce_asset(upstream_asset: list) -> list:
    chunks = split_into_chunks(upstream_asset)
    processed_chunks = chunks.map(process_chunk)
    return combine_chunks(processed_chunks.collect())


@asset
def downstream_asset(map_reduce_asset):
    print(map_reduce_asset)


defs = Definitions(assets=[upstream_asset, map_reduce_asset, downstream_asset])