How to depend on only one specific key from input with MultiPartitionsDefinition?

[rprechelt]

I have a MultiPartitionsDefinition that is used by an upstream asset:

date_and_country = MultiPartitionsDefinition(
    {
        "date": DailyPartitionsDefinition(start_date="2022-03-12"),
        "country": StaticPartitionsDefinition(["USA", "Canada", "Australia"])
    })

@asset(partitions_def=date_and_country)
def upstream():
    pass

I have another MultiPartitionsDefinition that is used by a downstream asset:

date_and_color = MultiPartitionsDefinition(
    {
        "date": DailyPartitionsDefinition(start_date="2022-03-12"),
        "color": StaticPartitionsDefinition(["red", "green", "blue"])
    })

@asset(partitions_def=date_and_color)
def downstream(upstream):
    pass

I want downstream to receive the upstream asset for only a single "specific" key from the "country" dimension in upstream, i.e. downstream depends (and receives as input) only the "USA" slice of the upstream asset (while preserving the one-to-one mapping of the "date" dimension). I can't quite get this to work with the MultiPartitionMapping's that I've tried.

How can I achieve this?

[clairelin135]

@rprechelt In your case, you can provide a SpecificPartitionsPartitionMapping to MultiPartitionMapping that defines that each color partition depends on a specific partition in the upstream "country" dimension:

@asset(
    partitions_def=date_and_color,
    ins={
        "upstream": AssetIn(
            partition_mapping=MultiPartitionMapping(
                {
                    "date": DimensionPartitionMapping(
                        "date", TimeWindowPartitionMapping()
                    ),
                    "country": DimensionPartitionMapping(
                        "color", SpecificPartitionsPartitionMapping(["USA"])
                    ),
                }
            ),
        )
    },
)
def downstream(context, upstream):
    context.log.info(context.asset_partition_keys_for_input("upstream"))
    ...

[rprechelt]

Great! This worked exactly as described.