Handling MultiPartitions with DuckDB IO Manager

[dagsir[bot]]

I'm setting up a Dagster pipeline in a local environment, and using the DuckDBPandasIOManager because it seems to be pretty good for local environments. I've got an upstream asset with a MultiPartitionsDefinition (a MonthlyPartitionsDefinition and StaticPartitionsDefinition) that will feed into a downstream asset with a single MonthlyPartitionsDefinition partition. The downstream asset is defined as follows:

@asset(
    group_name="statements",
    description="The statements",
    ins={
        "upstream": AssetIn(
            key=upstream.key,
            partition_mapping=MultiToSingleDimensionPartitionMapping(
                partition_dimension_name="date"
            )
        ),
    },
    partitions_def=monthly_partition_def,
)
def downstream(upstream: pd.DataFrame):

Whenever I try to materialize this asset I get the following error:
dagster._check.CheckError: Failure condition: Tried to access partition key for asset 'AssetKey(['upstream'])', but the number of input partitions != 1

#13290 This looks to be very similar to this error, which has its own issue. It seems the DuckDBPandasIOManager may also not support inputs with multiple partitions. For now I just want to build this pipeline without worrying about where the data will be written to. Is there an I/O manager that supports inputs with multiple partitions and that works well in a local environment? Am I better writing my own version of the DuckDB I/O manager that has this support, as was originally done in that issue?

---

The question was originally asked in Dagster Slack.

[jamiedemaria]

This looks like a bug. Making an issue for this so that we can track #17838

[archy-bold]

Is there a way this can be worked around until the bug is fixed?

[archy-bold]

If anyone finds this discussion and wants to know a potential workaround, you could create an intermediary asset between the upstream and downstream assets, just ensure this asset has the same partitions as the upstream, and uses an I/O manager that supports multi-partition mappings ie one of the Pickle I/O managers such as FilesystemIOManager.

Here's an example of the asset I created:

@asset(
    io_manager_key="fs_io_manager",
    partitions_def=upstream_partitions_def,
)
def upstream_file(upstream: pd.DataFrame):
    return upstream

Obviously this shouldn't be permanent, I'm just using it until the issue is fixed.

[jamiedemaria]

Unfortunately, I don't think so while still using DuckDB. I'm trying to take a look at fixing it soon though