How to best add metadata to partitions in a backfill situation

[derHeinzer]

Hi everybody!

I've encountered a challenge in Dagster and would like to share my thoughts with the community.

I'm working on adding metadata to a partitioned asset, and it works smoothly when materializing partition by partition. However, I also want it to work seamlessly in backfill situations, such as when initially materializing an asset that represents a table in a database. One specific metadata information I want the asset to hold is the record count by partition.

While I can yield additional AssetObservations with this metadata by partition, it doesn't show up in the web server when visualizing metadata plots by partition (although it does show up in plots by execution time, which is not my desired outcome).

Alternatively, I've managed to make the plots work by using additional AssetMaterialization events. However, these must be yielded after the original single output object. Otherwise, the last materialization triggered by the single output event would win, and the plots still won't work. The drawback is that the data provenance will break since the latest materialization events will have no tags with data provenance.

To work around this issue, I've implemented a somewhat hacky function that I call in the asset body:

    def attach_partition_information_to_output(
        self, context, orig_output, dataset, table, asset_key,
    ):
        bq_partition_id_start = datetime.strptime(
            context.partition_keys[0], "%Y-%m-%d"
        ).strftime("%Y%m%d")
        bq_partition_id_end = datetime.strptime(
            context.partition_keys[-1], "%Y-%m-%d"
        ).strftime("%Y%m%d")

        df_partitions = (
            self.get_partition_info(
                context,
                dataset,
                table,
                left_bound=bq_partition_id_start,
                right_bound=bq_partition_id_end,
            )
            .to_dataframe(dtypes={"partition_key": str})
            .set_index("partition_key")
        )
        if len(context.partition_keys) == 1:
            orig_output.metadata.update(
                df_partitions.loc[context.partition_keys[0]].to_dict()
            )
            yield orig_output
        else:
            from dagster._core.execution.plan.execute_step import _get_output_asset_materializations

            # TODO!!!!! hacky! ask community for better way!
            for am in _get_output_asset_materializations(asset_key, context.partition_keys, orig_output, None, {}, context._step_execution_context):
                am.metadata.update(
                    df_partitions.loc[am.partition].to_dict()
                )

                yield am

In a single partition scenario, I attach the partition information to the output metadata. However, in a multi-partition scenario, this approach doesn't work since the Output object can't hold metadata by partition. Hence, I manually call _get_output_asset_materializations and attach the right metadata to each yielded event.

I'm wondering if this is an acceptable way to handle it. I would prefer a cleaner solution but haven't found one yet. Perhaps it would be best to add a feature where the Output object can hold partition-specific metadata?

Thank you very much for your thoughts on this one!

[evzen-sirek-emq]

Hi, I am also interested in a solution for this!