Support partitioned assets for fetching latest AssetMaterializations

[jamiedemaria]

Summary & Motivation

In #18971 we piggy-back on the versioning code path to provide AssetMaterializations for the direct upstream dependencies for a currently materializing asset on the context. However, due to scaling constraints, the versioning code path always fetches the latest AssetMaterialization without considering the partitions that are relevant to the current execution.

This means that for partitioned assets, the information provided via latest_materialization_for_upstream_asset could be incorrect. Consider the following scenario:

partitions_def = DailyPartitionsDefinition(start_date="2024-01-01")

@asset(partitions_def=partitions_def)
def upstream(context):
    return MaterializeResult(metadata={"partition_key": context.partition_key})

@asset(partitions_def=partitions_def, deps=[upstream])
def downstream(context: AssetExecutionContext):
    mat = context.latest_materialization_for_upstream_asset("upstream")
    assert mat is not None
    assert mat.metadata["partition_key"].value == context.partition_key

If we materialized partition 2024-01-01 of upstream, then materialized partition 2024-01-02 of upstream. At this point the latest AssetMaterialization for upstream is the one for the 2024-01-02 partition. If we then materialized partition 2024-01-01 of downstream, the final assertion would fail. Reasonable behavior to expect would be that the AssetMaterialization returned by latest_materailization_for_upstream_asset is the one for the most recent materialization of the partition of upstream that the currently materializing partition of downstream depends on. But the latest AssetMaterialization for upstream is for a different partition, so we get that AssetMaterialization instead.

This PR proposes adding additional logic to latest_materialization_for_upstream_asset so that for partitioned assets, we get the latest AssetMaterialization for the correct partition. This unfortunately requires a call to the DB. Renaming the function fetch_latest_materialization_for_upstream_asset should help indicate this to the user. For non-partitioned assets, we still use the AssetMaterializations fetched during the versioning code path, so those calls should remain efficient.

How I Tested These Changes

[jamiedemaria]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

Current dependencies on/for this PR:

• master
  • PR Only have one kind of context for direct invocation #17554
    • PR duplicate op context methods on asset context  #18573
      • PR deprecate methods on AssetExecutionContext that can be found via the DagsterRun object #17339
        • PR Direct invocation asset execution context #18549
          • PR deprecate asset_partition_*_for_output on AssetExecutionContext #19436
            • PR [docs] remove asset_partition_*_for_output from docs #19437
              • PR deprecate op related methods from AssetExecutionContext #19441
          • PR Make upstream asset materialization events available on the context #18971
            • PR [wip] store the tags fetched by data versioning so we can use them for upath io manager #19324
              • PR [UPath I/O managers] special case handling of None outputs #18820
            • PR Make input partitions methods based on asset key not input name #19027
              • PR Support partitioned assets for fetching latest AssetMaterializations #19286 👈
              • PR [do not merge] base example to showcase new apis #19199
                • PR [prototype] sub-context for each upstream dependency  #19283
                • PR [prototype] AssetExecutionContext partition methods on sub context #19236
                • PR [prototype] AssetExecutionContext partition methods given as named tuple #19037
                • PR [prototype] AssetExecutionContext partition methods directly on the context #19032
            • PR [wip] two execution contexts with minimal changes #18543
              • PR [wip] see wht needs to change when asset context is no longer an op context #18553
        • PR support Asset and Op ExecutionContexts in standard execution path #17972
          • PR [wip] add AssetExecutionContext to direct invocation #18044

This stack of pull requests is managed by Graphite.

[jamiedemaria]

starting a thread to discuss adding this DB access path

[jamiedemaria]

my major concern with not adding this access is that we could be providing inaccurate information to users who are working with partitioned assets

[jamiedemaria]

and a concern with adding it is that a user could iterate through a big list of assets and fetch the latest materialization for each and break something

[schrockn]

What if we prefetched this information for partitioned execution? Doesn't seem too crazy.

Also what about the behavior for parition-ranged backfills?

[jamiedemaria]

@schrockn and @smackesey just want to bump this convo since it probably got lost during the offsite

[jamiedemaria]

converting back to draft to get this out of review queues until context work is re-prioritized