Dagster资产间数据传递与配置使用指南

import pandas as pd
import random
from datetime import datetime, timedelta
from dagster import asset, Config, materialize

# 资产1: 生成数据集
@asset 
def generate_dataset() -> pd.DataFrame:
    """
    生成一个包含随机水果销售数据的DataFrame。
    """
    def random_dates(start_date, end_date, n=10):
        date_range = end_date - start_date
        random_dates = [start_date + timedelta(days=random.randint(0, date_range.days)) for _ in range(n)]
        return random_dates

    random.seed(42) # 设置随机种子以保证可复现性
    num_rows = 100

    fruits = ['Apple', 'Banana', 'Orange', 'Grapes', 'Kiwi']
    fruit_column = [random.choice(fruits) for _ in range(num_rows)]
    units_column = [random.randint(1, 10) for _ in range(num_rows)]
    start_date = datetime(2022, 1, 1)
    end_date = datetime(2022, 12, 31)
    date_column = random_dates(start_date, end_date, num_rows)

    df = pd.DataFrame({
        'fruit': fruit_column,
        'units': units_column,
        'date': date_column
    })

    print("生成的数据集:")
    print(df.head())
    return df

# 配置类: 定义用户选择的水果参数
class FruitConfig(Config):
    fruit_select: str 

# 资产2: 根据用户配置筛选数据
@asset 
def filter_data(generate_dataset: pd.DataFrame, config: FruitConfig) -> pd.DataFrame:
    """
    根据用户在配置中选择的水果类型筛选数据集。

    Args:
        generate_dataset (pd.DataFrame): 上游资产 generate_dataset 的输出。
        config (FruitConfig): 用户自定义的配置对象，包含 fruit_select 参数。

    Returns:
        pd.DataFrame: 筛选后的数据集。
    """
    # 直接使用传入的 generate_dataset 参数，而不是重新调用函数
    df_filtered = generate_dataset[generate_dataset['fruit'] == config.fruit_select]
    print(f"\n根据 '{config.fruit_select}' 筛选后的数据:")
    print(df_filtered.head())
    return df_filtered

# 资产3: 进一步筛选数据
@asset
def filter_again(filter_data: pd.DataFrame) -> pd.DataFrame:
    """
    在已筛选的数据基础上，进一步筛选出销量大于5的记录。

    Args:
        filter_data (pd.DataFrame): 上游资产 filter_data 的输出。

    Returns:
        pd.DataFrame: 再次筛选后的数据集。
    """
    # 直接使用传入的 filter_data 参数
    df_final = filter_data[filter_data['units'] > 5]
    print("\n进一步筛选 (units > 5) 后的数据:")
    print(df_final.head())
    return df_final

# 如果需要在一个Job中运行这些资产
# from dagster import define_asset_job
# my_pipeline = define_asset_job("my_fruit_pipeline", selection="*")

# 示例：如何在本地物化（测试）
if __name__ == "__main__":
    # 运行所有资产，并提供配置
    # 注意：在Dagster UI中，配置会在运行时由UI提供
    # 在本地测试时，需要手动构建配置字典
    result = materialize(
        [generate_dataset, filter_data, filter_again],
        run_config={
            "ops": {
                "filter_data": { # 注意这里是资产名，不是函数名
                    "config": {
                        "fruit_select": "Banana"
                    }
                }
            }
        }
    )
    assert result.success
    print("\n所有资产成功物化！")