in built-in object IO managers, handle loading objects for multiple partitions

If a daily asset depends on an hourly asset, then each partition of the daily asset will correspond to 24 partitions of the hourly asset.

When materializing the daily asset and load_input is called to load the contents of the hourly asset, we should return the contents of those 24 partitions.

For the Pandas type handler of the Snowflake IO manager, we handle this by returning a single dataframe that contains the concatenated contents of all the hourly partitions.

For the built-in object store IO managers, we could just return a list of the pickled objects. Or potentially a dictionary keyed by partition?

What we’ve heard:

• 2022-10-20: https://dagster.slack.com/archives/C01U954MEER/p1666175125167899

Here’s a start on the implementation:

from dagster import io_manager, IOManager, AssetKey


@io_manager
def my_io_manager():
    return MyIOManager()


class MyIOManager(IOManager):
    def handle_output(self, context, obj):
        ...

    def _load(self, asset_key: AssetKey):
        ...

    def _load_partition(self, asset_key: AssetKey, partition_key: str):
        ...

    def load_input(self, context):
        if context.has_asset_partitions:
            partition_key_range = context.asset_partition_key_range
            if partition_key_range.start == partition_key_range.end:
                return self._load_partition(context.asset_key, partition_key_range.start)
            else:
                return [
                    self._load_partition(context.asset_key, partition_key)
                    for partition_key in context.asset_partition_keys
                ]
        else:
            return self._load(context.asset_key)

[sryza]

@sryza would a generator make more sense?

Interesting - this isn’t a pattern that we’ve used before with IO managers, but I think it makes sense.

Testing it using materialize_to_memory is proving tricky since all assets need to share the same partitions_def.

You could do something like this:

materialize([*hourly_asset.to_source_assets(), daily_asset], partition_key="2022-01-01")

That will run load_input for daily_asset without trying to materialize hourly_asset in the same run.

[sryza]

I also think a syntax to be more specific for higher frequency asset dependencies would be useful, but out of scope for this ticket. In some cases not all overlapping partitions are required, it would be nice to specify e.g. hourly_asset[‘-3H’] or something similar

You might already be aware of this, but you could define a custom PartitionMapping to express this