如何在Palantir Foundry代码仓库用PySpark读取RID并提取数据集列

我是Palantir Foundry代码仓库的新用户，第一次发帖，要是表述有问题请多包涵。我现在需要完成这样一个任务：从存储了RID的数据集里读取RID（比如其中的backing_dataset_rid列存的是Foundry应用内的RID），加载这些RID对应的数据集，提取它们的列信息，然后把结果保存回最初提取RID的那个数据集。我之前看过一个相关思路，但想要用PySpark来实现。目前我的代码返回的是类似Transform(myproject.datasets.examples:extract_cols)</Foundry/OUTPUTFOLDER>的对象列表（敏感信息已隐藏），代码如下：

from pyspark.sql import functions as F
from transforms.api import transform, transform_df, Input, Output
 
def RID_extract(RID):
    @transform_df(
        Output('/folder_path/OutputDataset'),
        data=Input(RID)
    )
    def extract_cols(data):
        column_names = data.dataframe().columns
        return column_names
    return extract_cols

@transform_df(
    Output("/folder_path/OutputDataset"),
    source_df=Input("/folder_path/InputDataset")
)
def compute(source_df):
    df = source_df
    output_path = "/folder_path/OutputDataset"
    rows=df.collect()
    df2=[]
    print(source_df.columns)
    for row in rows:
        if row['backing_dataset_rid'] == (None):
            continue
        RID = row['backing_dataset_rid']
        print(RID)
        RID_transform = RID_extract(RID)
        df2.append(RID_transform)
        #df.rdd.map(RID_transform).collect()
        print(RID_transform)
    return df2

核心问题分析 #

你的代码有几个关键问题导致返回Transform对象而非预期数据：

• Transform嵌套定义：在compute函数里调用RID_extract动态创建新的transform_df，这不符合Foundry的Transform运行逻辑，会返回Transform对象而非实际数据集
• 全量数据拉取：用collect()把整个数据集拉到Driver节点处理，数据量大时会引发内存溢出，也违背了Spark分布式计算的原则
• 返回值类型错误：transform_df要求返回PySpark DataFrame，但你的代码返回的是Transform对象列表

正确PySpark实现 #

我们可以直接通过Spark读取Foundry格式的数据集，结合UDF或去重关联的方式，高效完成任务：

基础实现（单条RID处理）

from pyspark.sql import functions as F
from transforms.api import transform_df, Input, Output

@transform_df(
    Output("/folder_path/OutputDataset"),
    source_df=Input("/folder_path/InputDataset")
)
def compute(source_df):
    def get_dataset_columns(rid):
        if not rid:
            return None
        # 通过Foundry格式读取RID对应的数据集
        target_df = spark.read.format("foundry").load(rid)
        return target_df.columns
    
    # 注册返回数组类型的UDF
    get_cols_udf = F.udf(get_dataset_columns, "array<string>")
    
    # 给原数据集添加对应列名的数组列
    result_df = source_df.withColumn(
        "backing_dataset_columns",
        get_cols_udf(F.col("backing_dataset_rid"))
    )
    
    # 可选：将数组列展开为多行（如果需要每个列名单独一行）
    # result_df = result_df.select("*", F.explode("backing_dataset_columns").alias("column_name"))
    
    return result_df

优化实现（去重避免重复读取）

如果backing_dataset_rid存在大量重复值，先对RID去重再读取，减少重复IO：

from pyspark.sql import functions as F
from transforms.api import transform_df, Input, Output

@transform_df(
    Output("/folder_path/OutputDataset"),
    source_df=Input("/folder_path/InputDataset")
)
def compute(source_df):
    def get_dataset_columns(rid):
        if not rid:
            return None
        target_df = spark.read.format("foundry").load(rid)
        return target_df.columns
    
    # 提取唯一RID并读取对应列信息
    unique_rid_cols = source_df.filter(F.col("backing_dataset_rid").isNotNull()) \
                              .select("backing_dataset_rid").distinct() \
                              .rdd.map(lambda row: (row.backing_dataset_rid, get_dataset_columns(row.backing_dataset_rid))) \
                              .toDF(["backing_dataset_rid", "backing_dataset_columns"])
    
    # 和原数据集关联，保留所有原始行
    result_df = source_df.join(unique_rid_cols, on="backing_dataset_rid", how="left")
    
    return result_df

关键注意事项 #

1. 权限检查：确保Transform的服务账号有读取所有RID对应数据集的权限，否则会触发权限错误
2. 性能考量：如果涉及的数据集数量极多或数据量巨大，建议分批处理，或考虑使用Foundry Pipeline Builder管理依赖关系
3. 空值处理：代码中已经处理了空RID的情况，避免读取操作报错

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