pass down asset partitions to ops in a graph

[dagsir[bot]]

Hi! How can I pass down asset partitions to ops in a graph?

from dagster import DailyPartitionsDefinition, OpExecutionContext, graph_asset, op
import pandas as pd

@op(
    required_resource_keys={"snowflake_airbyte_manager"},
)
def create_table(context: OpExecutionContext) -> bool:
    context.log.info(context.asset_partitions_time_window_for_output())
    return True

@op(
    required_resource_keys={"snowflake_airbyte_manager"},
)
def replace_table(context: OpExecutionContext, tbl) -> None:
    context.log.info(context.asset_partitions_time_window_for_output())
    return True

@graph_asset(
    group_name="kunal_test",
    partitions_def=DailyPartitionsDefinition(start_date="2023-01-01"),
)
def kunal_test_graph_asset() -> pd.DataFrame:
    return replace_table(create_table())

In this simple example, I'm able to get the asset partitions in the Op replace_table but the next Op in the dependency chain create_table doesn't get access to the partitions at all:

dagster._core.errors.DagsterExecutionStepExecutionError: Error occurred while executing op "kunal_test_graph_asset.create_table"::

ValueError: Tried to get asset partitions for an output that does not correspond to a partitioned asset.

How can I propagate the partitions to create_table in this case?

---

The question was originally asked in Dagster Slack.

[clairelin135]

Hi there -- I believe the context.asset... methods may only work when invoked within asset execution, and not within op execution.

Within the op you can call context.partition_time_window instead.

[juicemanjames]

thank you so much @clairelin135 . I tried to replace those context partitions but got this

dagster._core.errors.DagsterInvariantViolationError: Cannot access partition_key for a partitioned run with a range of partitions. Call asset_partition_key_range instead.

what should we do here?

[clairelin135]

That's odd -- locally I tried running the following on a partition range and was able to execute successfully:

from dagster import DailyPartitionsDefinition, OpExecutionContext, graph_asset, op, BackfillPolicy
import pandas as pd

@op
def create_table(context: OpExecutionContext) -> None:
    context.log.info(context.partition_time_window)

@op
def replace_table(context: OpExecutionContext, tbl) -> None:
    context.log.info(context.partition_time_window)

@graph_asset(
    partitions_def=DailyPartitionsDefinition(start_date="2023-01-01"),
    backfill_policy=BackfillPolicy.single_run()
)
def kunal_test_graph_asset() -> pd.DataFrame:
    return replace_table(create_table())

What version are you on? Do you have a code snippet so I can repro?

[juicemanjames]

thanks @clairelin135 . It might be a version thing, as I got this error when running your code:

ImportError: cannot import name 'BackfillPolicy' from 'dagster' (/Users/juicemanjames/Library/Caches/pypoetry/virtualenvs/spindl-dagster-N9Wxfy4A-py3.8/lib/python3.8/site-packages/dagster/__init__.py)

My code snippet for the error in the above thread was literally replacing context.asset... with context.partition_time_window

{"pypi": {"package": "dagster==1.3.13"}},

[clairelin135]

@wangry-james ah yeah I think the context.partition_time_window fix landed in version 1.4.17 so you would have to upgrade to that version or above