Multi-partition inputs & IO managers

[sryza]

Sometimes the input of an asset computation is fed by multiple partitions of an upstream asset. For example:

• A non-partitioned asset depends on a partitioned asset
• A daily-partitioned asset depends on an hourly-partitioned asset
• A single-dimensional-partitioned asset depends on a multi-partitioned asset
• A daily-partitioned asset has a rolling window dependency on a daily partitioned asset

This discussion is about what Python types should be passed between the IO manager and the asset input in this case. I.e. it concerns what should go in the question-marks below:

class MyIOManager(IOManager):
    def load_input(self, context) -> ???:
        partition_keys = context.asset_partition_keys
        return ???

@asset(...)
def my_asset(partitioned_upstream_asset: ???):
    ...

Status quo

• The Pandas and PySpark type handlers of the DB IO managers (Snowflake, DuckDB, BigQuery) always return a single DataFrame, which can includes values from all the partitions.
• When loading an input that corresponds to multiple partitions, the UPathIOManager returns a dictionary that maps each input partition key to the input value for that partition key.

The reason these are different is that the DB IO managers load all the partitions using a single SQL query, but the UPathIOManager loads each partition independently from a separate file.

While it's a little weird, I believe it's fine that they handle these in different ways. Different IO managers generally operate on different types.

Problem with the status quo that led me to create this Discussion

What input type should your asset expect if its inputs are loaded using the UPathIOManager, and it sometimes receives multiple partitions and other times receives a single partition?

Currently, you would need to check if the input type is a dict, but even that isn't foolproof, because what if the value of an individual partition is a dict?

I have encountered this in two situations:

• Writing my own dynamic asset partitioning example.
• This issue (not exactly the same, but related): Dynamic type checking fails when downstream asset refers to a group of upstream partitions #12426.

Proposal

Introduce a PartitionedInput type that extends dict. Then here's how the IO manager and asset would get implemented:

class MyIOManager(IOManager):
    def load_input(self, context):
        inputs_by_partition_key = read_files(context.asset_partition_keys)

        if context.dagster_type.typing_type == PartitionedInput:
            return inputs_by_partition_key
        else:
            if len(inputs_by_partition_key) == 1:
                return next(inputs_by_partition_key.values())
            else:
                # Or should we raise an error??
                return inputs_by_partition_key


@asset(...)
def my_asset(partitioned_upstream_asset: PartitionedInput[int]):
    for partition, value in partitioned_upstream_asset.items():
        ...

[harrylojames]

This looks great to me. Preserves existing behaviour while allowing a way to more elegantly handle partitioned inputs. I hadn't considered "Currently, you would need to check if the input type is a dict, but even that isn't foolproof, because what if the value of an individual partition is a dict?"!

I would be interested to see how this is documented to help users avoid this pitfall as it's subtle that adding the type changes the iomanager behaviour.

[smackesey]

Some questions:

(1) Is this a mistake? Why does load_input return a partition key (not a loaded value) here?

            if len(inputs_by_partition_key) == 1:
                return next(context.asset_partition_keys)

(2) You write:

Currently, you would need to check if the input type is a dict, but even that isn't foolproof, because what if the value of an individual partition is a dict?

Can't you currently call context.asset_partition_key_range_for_input for a foolproof check?

(3) Is the idea that PartitionedInput should only be used for UPathIOManager (and descendants)?

If that is the case, IMO that's confusing. I would expect an abstraction with this generic a name to be usable with any IO manager.

Given that we already have the available information in the framework to determine whether a loaded value represents multiple partitions or a single partition, why not have the framework do the wrapping in PartitionedInput? The less requirements we can impose on user code (IO manager load_input/handle_output), the better. For backcompat, we could only wrap if the dagster_type on the input is annotated as a PartitionedInput, similarly to what is shown in your IO manager code.

In this case PartitionedInput would be a dead simple wrapper:

class PartitionedInput(Generic[T]):

    def __init__(self, value: T, partition_keys: Sequence[str]):
         self.value = value
         self.partition_keys = partition_keys

    @property
    def is_single_partition():
        return len(self.partition_keys) == 1

Advantages:

1. This allows users to opt into PartitionedInput on any IO manager without no breaking changes or even changes to UPathIOManager.
2. This mirrors how handle_output already works. Even though ops may return an Output or DynamicOutput, handle_output on any IO manager never receives these objects-- instead the framework "unwraps" the value and passes it to handle_output. That provides a nice abstraction barrier between framework code and IO managers, which can work with whatever domain objects the user is working with rather than Dagster abstractions (i.e. Output).

The other half of this is PartitionedOutput, as discussed here. Like PartitionedInput in the above scheme, PartitionedOutput belongs to the op layer and framework internals. It is optionally returned from a multi-partition-returning op, but unwrapped (like its cousins Output and DynamicOutput) before passing the contents to the IO manager.

[sryza]

Is this a mistake? Why does load_input return a partition key (not a loaded value) here?

Oops, yes. Edited to fix.

Can't you currently call context.asset_partition_key_range_for_input for a foolproof check?

That's a good point that I hadn't considered. However, I think this is too arcane. I.e. the user ends up needing to know lore about our protocol rather than interacting with a set of self-explanatory types.

Is the idea that PartitionedInput should only be used for UPathIOManager (and descendants)?

This is how I had originally envisioned it, but open to making it more deeply embedded into the framework.

If that is the case, IMO that's confusing. I would expect an abstraction with this generic a name to be usable with any IO manager.

That's a good point about the name. If we made it specific to UPathIOManager, we'd probably want to give it a different name, maybe ValuesByPartition.

[smackesey]

This is how I had originally envisioned it, but open to making it more deeply embedded into the framework.

I strongly favor the "framework-does-the-wrapping" approach I outlined above-- any thoughts on that?

[sryza]

Suppose I want to write code like this inside my asset:

@asset
def asset1(asset0: <something>):
    for partition_key, partition_value in <something>:
        # do something with that partition value

and I want this code to work both if there's a single partition loaded or multiple partitions loaded. What would go inside the <something>s in the "framework-does-the-wrapping" approach?

[smackesey]

@asset
def asset1(asset0: PartitionedInput[Mapping[str, <type of value>]`):
    for partition_key, partition_value in asset0.value.items():
        ...

IIUC the difference between our proposals is that:

• In your proposal, the IO manager needs to deal directly with PartitionedInput (both returning the value and checking the dagster_type).
• In my proposal, the IO manager doesn't concern itself with the PartitionedInput concept and deals directly with "raw" values. But the framework knows (a) whether the value corresponds to multiple partitions or not; and (b) whether the dagster type of the target input is a PartitionedInput, then it is able to wrap the value returned by the IO manager before passing as input.

Alternatively we could use dict as the default assumed type but allow users to subclass:

@asset
def asset1(asset0: PartitionedInput[int]`):  # default assumes the value is a dict
    for partition_key, partition_value in asset0.items():
        ...

class PartitionedPandasInput(PartitionedInput):

    value: pd.DataFrame

    def items():
        # ... yield item pairs from dataframe ..

@asset
def asset3(asset2: PartitionedPandasInput[<sometype>]):
        for partition_key, partition_value in asset0.items():
            ...

[sryza]

In your proposal, what type would PickledObjectFilesystemIOManager.load_input return in each of these cases:

• The input corresponds to a single partition
• The input corresponds to multiple partitions

[smackesey]

The input corresponds to a single partition

A bare value. This will be packaged in PartitionedInput if the user is using PartitionedInput in their annotation.

The input corresponds to multiple partitions

A dict.

[sryza]

A bare value. This will be packaged in PartitionedInput if the user is using PartitionedInput in their annotation.

Would this line from your example above execute successfully then? asset0.value would be the bare value, not a dictionary?

    for partition_key, partition_value in asset0.value.items():