Dep. between multi-partitioned asset and two single-dimension partitioned assets

[shiukaheng]

Dependency between multi-partitioned asset and two single-dimension partitioned assets

Context

I’m trying to model a very simple scenario in Dagster:

• A set of datasets
• A set of models
• An evaluation matrix that should run for every (dataset, model) pair

So the structure is:

datasets  (partitioned by dataset)
models    (partitioned by model)
evaluation_matrix  (partitioned by [dataset, model])

---

Minimal Example

import dagster as dg

dataset_partitions = dg.StaticPartitionsDefinition(["dataset_1", "dataset_2"])
model_partitions = dg.StaticPartitionsDefinition(["model_a", "model_b"])

@dg.asset(partitions_def=dataset_partitions)
def datasets(context: dg.AssetExecutionContext):
    print(f"Materialized dataset {context.partition_key}")

@dg.asset(partitions_def=model_partitions)
def models(context: dg.AssetExecutionContext):
    print(f"Materialized model {context.partition_key}")

dataset_model_matrix = dg.MultiPartitionsDefinition({
    "dataset": dataset_partitions,
    "model": model_partitions,
})

@dg.asset(
    partitions_def=dataset_model_matrix,
    deps=[datasets, models],
)
def evaluation_matrix(context: dg.AssetExecutionContext):
    keys = context.partition_key.keys_by_dimension
    print(f"Evaluating dataset={keys['dataset']} with model={keys['model']}")

---

What I Expect

• Each (dataset, model) partition of evaluation_matrix should depend on:
  • the corresponding dataset partition
  • the corresponding model partition

So (dataset_1, model_a) → datasets(dataset_1) + models(model_a)

---

What Actually Happens

• Dagster doesn’t seem to support “splitting” a multi-partition dependency across two single-partitioned upstreams.
• Theres MultiToSingleDimensionPartitionMapping described in discussion #15675 which I am trying to use but it doesn't seem to work as I would expect:

import dagster as dg

dataset_partitions = dg.StaticPartitionsDefinition(["dataset_1", "dataset_2"])
model_partitions = dg.StaticPartitionsDefinition(["model_a", "model_b"])
multi_partitions = dg.MultiPartitionsDefinition({
    "dataset": dataset_partitions,
    "model": model_partitions,
})

@dg.asset(partitions_def=dataset_partitions)
def datasets(context):
    print(f"Materialized dataset {context.partition_key}")

@dg.asset(partitions_def=model_partitions)
def models(context):
    print(f"Materialized model {context.partition_key}")

@dg.asset(
    partitions_def=multi_partitions,
    deps=[
        dg.AssetDep(
            "datasets",
            partition_mapping=dg.MultiToSingleDimensionPartitionMapping("dataset"),
        ),
        dg.AssetDep(
            "models",
            partition_mapping=dg.MultiToSingleDimensionPartitionMapping("model"),
        ),
    ]
)
def evaluation_matrix(context):
    keys = context.partition_key.keys_by_dimension
    print(f"Evaluating dataset={keys['dataset']} with model={keys['model']}")

Running dg launch --assets "+evaluation_matrix" --partition "dataset_1|model_a" gives:
dagster_shared.check.functions.CheckError: Failure condition: There is no PartitionsDefinition shared by all the provided assets

---

Question

Is there a way to express:

“This multi-partitioned asset depends on asset A along dimension X, and asset B along dimension Y”

Or is this currently unsupported?