Dagster资产间数据传递与用户配置的最佳实践

from dagster import Config

class fruit_config(Config):
    fruit_select: str

# 错误示例：不正确的上游资产数据获取方式
@asset(deps=[generate_dataset]) 
def filter_data(config: fruit_config):
    df = generate_dataset() # ❌ 错误！直接调用上游资产函数无法获取其输出
    df2 = df[df['fruit'] == config.fruit_select]
    return df2

@asset(deps=[filter_data]) 
def filter_again():
    df2 = filter_data() # ❌ 错误！同样无法获取上游资产输出
    df3 = df2[df2['units'] > 5]
    return df3

import pandas as pd
import random
from datetime import datetime, timedelta
from dagster import asset, Config, materialize # 移除了 MaterializeResult, MetadataValue，因为在本例中未使用

# 1. 数据生成资产
@asset 
def generate_dataset() -> pd.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("生成的数据集:\n", df.head())
    return df

# 2. 用户配置类
class fruit_config(Config):
    """
    定义用户选择水果的配置。
    """
    fruit_select: str 

# 3. 数据筛选资产：接收上游资产输出和用户配置
@asset 
def filter_data(generate_dataset: pd.DataFrame, config: fruit_config) -> pd.DataFrame: # 关键改变：
    # 1. 移除了 deps=[generate_dataset]，因为数据依赖通过参数显式声明。
    # 2. 将 generate_dataset: pd.DataFrame 作为参数，Dagster会将 generate_dataset 资产的输出注入到此参数。
    # 3. config: fruit_config 接收用户配置。
    """
    根据用户配置筛选水果数据。
    """
    filtered_df = generate_dataset[generate_dataset['fruit'] == config.fruit_select]
    print(f"根据 '{config.fruit_select}' 筛选后的数据:\n", filtered_df.head())
    return filtered_df

# 4. 再次筛选资产：接收上游资产输出
@asset
def filter_again(filter_data: pd.DataFrame) -> pd.DataFrame: # 关键改变：
    # 1. 移除了 deps=[filter_data]。
    # 2. 将 filter_data: pd.DataFrame 作为参数，Dagster会将 filter_data 资产的输出注入到此参数。
    """
    对筛选后的数据进行二次筛选，保留单位大于5的记录。
    """
    final_df = filter_data[filter_data['units'] > 5]
    print("再次筛选（单位 > 5）后的数据:\n", final_df.head())
    return final_df

# 为了在本地测试，可以调用 materialize 函数
if __name__ == "__main__":
    # 示例运行，需要提供配置
    # 注意：在Dagster UI中运行时，UI会自动提示配置输入
    result = materialize(
        assets=[generate_dataset, filter_data, filter_again],
        run_config={
            "ops": { # 注意这里是 "ops" 即使在资产上下文，因为配置是针对底层操作
                "filter_data": {
                    "config": {
                        "fruit_select": "Banana"
                    }
                }
            }
        }
    )
    assert result.success
    print("\nDagster 管道执行成功！")