如何在Palantir Foundry的PySpark测试中使用现有数据集？

直接在测试函数上使用@transform装饰器会将其转换为Foundry Transform对象，破坏PyTest对测试用例的识别逻辑（依赖test_前缀+无特殊装饰器的函数）。以下是两种实用的替代方案：

方案1：测试函数内部手动加载数据集 #

不用@transform，直接通过Foundry的Dataset API在测试函数内读取目标数据集，同时保留PyTest的测试用例识别能力。

import pytest
from pyspark.sql import SparkSession
from foundry.dataset import Dataset

# 定义SparkSession fixture，供测试复用
@pytest.fixture(scope="session")
def spark():
    return SparkSession.builder.getOrCreate()

def test_compute(spark):
    # 直接加载Foundry现有数据集
    weightings_df = Dataset("/path/to/your/target-weightings-dataset").read_spark()
    
    # 结合自定义测试数据编写逻辑
    custom_test_data = spark.createDataFrame(
        [(1, 0.6), (2, 0.4)], 
        ["user_id", "weight_score"]
    )
    combined_data = weightings_df.union(custom_test_data)
    
    # 执行断言验证
    assert combined_data.count() == weightings_df.count() + 2

方案2：将数据集加载逻辑封装为Pytest Fixture #

把数据集加载逻辑抽成可复用的fixture，提高代码整洁度，同时支持多个测试函数共享数据集。

import pytest
from foundry.dataset import Dataset

@pytest.fixture(scope="module")
def weightings_dataset(spark):
    # 加载并可选预处理数据集（如采样、过滤）
    raw_df = Dataset("/path/to/your/target-weightings-dataset").read_spark()
    return raw_df.filter(raw_df["weight_score"] > 0.1)  # 示例预处理

def test_compute_with_weightings(weightings_dataset, spark):
    # 使用fixture传入的数据集+自定义测试数据
    test_input = spark.createDataFrame([(3, 0.3)], ["user_id", "weight_score"])
    processed_result = your_transform_function(weightings_dataset, test_input)
    
    # 验证结果
    assert processed_result.select("weight_score").sum().first()[0] > 0

补充说明 #

• 确保测试环境具备目标数据集的读取权限
• 如果数据集规模较大，建议在加载时做采样或过滤，避免测试运行过慢
• 若需基于其他Transform的输出做测试，可先定义一个负责预处理数据的Transform，再在测试函数中读取该Transform的输出数据集

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