Dagster资产间数据流转与用户自定义参数的正确姿势

# 错误示例片段
import pandas as pd
from dagster import asset, Config

# ... (generate_dataset 函数定义省略，与正确代码相同)

class fruit_config(Config):
    fruit_select: str

@asset(deps=[generate_dataset]) # 这里的deps虽然声明了，但内部调用方式是错误的
def filter_data(config: fruit_config):
    # 错误：直接调用generate_dataset()会重新执行，且无法获取其物化结果
    df = generate_dataset()
    df2 = df[df['fruit'] == config.fruit_select]
    return df2

@asset(deps=[filter_data]) # 这里的deps虽然声明了，但内部调用方式是错误的
def filter_again():
    # 错误：直接调用filter_data()同样会导致问题
    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

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

    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("Generated Dataset Head:\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:
    """
    根据用户配置的水果类型筛选数据集。
    参数：
        generate_dataset (pd.DataFrame): 来自上游 generate_dataset 资产的原始数据。
        config (fruit_config): 包含用户选择水果的配置对象。
    """
    # generate_dataset 参数会自动接收上游资产 generate_dataset 的输出
    filtered_df = generate_dataset[generate_dataset['fruit'] == config.fruit_select]
    print(f"Filtered Data (fruit_select='{config.fruit_select}') Head:\n", filtered_df.head())
    return filtered_df

# 4. 定义再次筛选的资产，接收上游筛选后的数据
@asset
def filter_again(filter_data: pd.DataFrame) -> pd.DataFrame:
    """
    在已筛选的数据集上，进一步筛选单位数量大于5的记录。
    参数：
        filter_data (pd.DataFrame): 来自上游 filter_data 资产的已筛选数据。
    """
    # filter_data 参数会自动接收上游资产 filter_data 的输出
    final_df = filter_data[filter_data['units'] > 5]
    print("Final Filtered Data (units > 5) Head:\n", final_df.head())
    return final_df

# 如果需要在本地直接运行测试，可以使用 materialize
if __name__ == "__main__":
    # 示例：如何通过代码提供配置并物化资产
    # 注意：在Dagster UI中运行，配置将通过UI界面输入
    result = materialize(
        [generate_dataset, filter_data, filter_again],
        run_config={
            "ops": { # 对于资产，配置是在 "ops" 键下，尽管现在推荐使用 "assets" 键，但旧版本或某些情况下仍兼容
                "filter_data": {
                    "config": {
                        "fruit_select": "Banana"
                    }
                }
            }
        }
    )
    assert result.success
    print("\nPipeline execution successful!")