如何实现支持资产级参数配置的Dagster IO Manager？

方案1：利用资产Metadata传递时间列并区分分区粒度 #

这是最轻量化的实现方式，无需额外资源定义，直接通过资产的metadata字段绑定时间列和分区粒度：

步骤1：定义带Metadata的资产 #

from dagster import asset, HourlyPartitionsDefinition, DailyPartitionsDefinition

hourly_partitions = HourlyPartitionsDefinition(start_date="2024-01-01")
daily_partitions = DailyPartitionsDefinition(start_date="2024-01-01")

@asset(
    partitions_def=hourly_partitions,
    metadata={
        "time_column": "event_time",
        "partition_granularity": "hourly"
    }
)
def hourly_user_events(context):
    # 业务逻辑：生成/读取小时级分区数据
    ...

@asset(
    partitions_def=daily_partitions,
    metadata={
        "time_column": "event_time",
        "partition_granularity": "daily"
    }
)
def daily_user_events(context):
    # 业务逻辑：生成/读取日级分区数据
    ...

步骤2：修改IO Manager获取Metadata并生成唯一路径 #

在Parquet IO Manager中，从上下文提取资产Metadata中的时间列和分区粒度，同时生成区分粒度的路径避免文件名冲突：

from dagster import ConfigurableIOManager, InputContext, OutputContext
import pandas as pd
import os

class ParquetIOManager(ConfigurableIOManager):
    root_path: str = "/data/parquet"

    def _get_path(self, context: InputContext | OutputContext) -> str:
        asset_key_str = "_".join(context.asset_key.path)
        path_components = [self.root_path, asset_key_str]

        if context.has_asset_partitions:
            # 从Metadata获取分区粒度
            granularity = context.asset_key.metadata.get("partition_granularity", "daily")
            path_components.append(granularity)
            
            # 根据粒度生成时间格式的路径
            window_start = context.asset_partitions_time_window.start
            if granularity == "hourly":
                path_components.append(window_start.strftime("%Y%m%d/%H"))
            else:
                path_components.append(window_start.strftime("%Y%m%d"))
        
        return "/".join(path_components)

    def load_input(self, context: InputContext) -> pd.DataFrame:
        # 从Metadata获取时间列
        time_column = context.asset_key.metadata.get("time_column")
        if not time_column:
            raise ValueError(f"资产{context.asset_key}未指定time_column元数据")
        
        # 构造过滤条件
        window = context.asset_partitions_time_window
        filters = [
            (time_column, ">=", window.start),
            (time_column, "<", window.end)
        ]
        
        return pd.read_parquet(self._get_path(context), filters=filters)

    def handle_output(self, context: OutputContext, obj: pd.DataFrame):
        output_path = self._get_path(context)
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        obj.to_parquet(output_path)

方案2：通过Per-Asset IO Manager配置传递时间列 #

如果希望用更正式的配置方式而非Metadata，可以利用Dagster的io_manager_config为每个资产单独配置时间列：

步骤1：定义支持配置的IO Manager #

class ParquetIOManager(ConfigurableIOManager):
    root_path: str = "/data/parquet"
    # 定义可被资产覆盖的时间列默认配置
    time_column: str = "default_time"

    def _get_path(self, context: InputContext | OutputContext) -> str:
        asset_key_str = "_".join(context.asset_key.path)
        path_components = [self.root_path, asset_key_str]

        if context.has_asset_partitions:
            # 从分区定义名称获取粒度（需保证分区定义命名规范）
            granularity = context.partitions_def.name
            path_components.append(granularity)
            
            window_start = context.asset_partitions_time_window.start
            path_components.append(window_start.strftime("%Y%m%d/%H") if granularity == "hourly" else window_start.strftime("%Y%m%d"))
        
        return "/".join(path_components)

    def load_input(self, context: InputContext) -> pd.DataFrame:
        window = context.asset_partitions_time_window
        filters = [
            (self.time_column, ">=", window.start),
            (self.time_column, "<", window.end)
        ]
        return pd.read_parquet(self._get_path(context), filters=filters)

    def handle_output(self, context: OutputContext, obj: pd.DataFrame):
        output_path = self._get_path(context)
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        obj.to_parquet(output_path)

步骤2：为资产指定IO Manager配置 #

@asset(
    partitions_def=hourly_partitions,
    io_manager_config={"time_column": "event_time"}
)
def hourly_user_events(context):
    ...

@asset(
    partitions_def=daily_partitions,
    io_manager_config={"time_column": "event_time"}
)
def daily_user_events(context):
    ...

方案3：集中式时间列映射资源（适合大量资产） #

如果有大量资产需要配置时间列，可以创建一个集中管理的资源，避免重复配置：

步骤1：定义时间列映射资源 #

from dagster import ConfigurableResource
from typing import Dict

class TimeColumnMapping(ConfigurableResource):
    # 键为资产键的字符串形式，值为对应时间列名
    asset_time_columns: Dict[str, str] = {
        "hourly_user_events": "event_time",
        "daily_user_events": "event_time"
    }

步骤2：让IO Manager依赖该资源 #

class ParquetIOManager(ConfigurableIOManager):
    root_path: str = "/data/parquet"
    # 依赖时间列映射资源
    time_column_mapping: TimeColumnMapping

    def load_input(self, context: InputContext) -> pd.DataFrame:
        asset_key_str = "_".join(context.asset_key.path)
        time_column = self.time_column_mapping.asset_time_columns.get(asset_key_str)
        if not time_column:
            raise ValueError(f"未找到资产{asset_key_str}的时间列配置")
        
        window = context.asset_partitions_time_window
        filters = [
            (time_column, ">=", window.start),
            (time_column, "<", window.end)
        ]
        return pd.read_parquet(self._get_path(context), filters=filters)

    def _get_path(self, context: InputContext | OutputContext) -> str:
        # 同方案1的路径生成逻辑
        ...

    def handle_output(self, context: OutputContext, obj: pd.DataFrame):
        # 同方案1的输出逻辑
        ...

步骤3：注册资源和资产 #

from dagster import Definitions

defs = Definitions(
    assets=[hourly_user_events, daily_user_events],
    resources={
        "io_manager": ParquetIOManager(
            root_path="/data/parquet",
            time_column_mapping=TimeColumnMapping(
                asset_time_columns={
                    "hourly_user_events": "event_time",
                    "daily_user_events": "event_time"
                }
            )
        )
    }
)

内容的提问来源于stack exchange，提问作者guillaume latour