在 Palantir Foundry 代码库中定义 Pandas UDF 的正确方法是什么

问题描述

我想在 Palantir Foundry 代码库中定义以下 pandas_udf。

@pandas_udf("long",PandasUDFType.GROUPED_AGG)
def percentile_95_udf(v):
    return v.quantile(0.95)

但是当我尝试在全局范围内定义此 udf 时，出现错误：

AttributeError: 'nonetype' object has no attribute '_jvm'

但是，如果我在由我的转换调用的函数中定义相同的函数，则代码运行良好，如下所示：

from pyspark.sql import functions as F
from pyspark.sql.functions import pandas_udf,PandasUDFType
from transforms.api import transform,Input,Output

@transform(
    data_out=Output("output path"),data_in=Input("input path")
)
def percentile_95_transform(data_in,data_out):
    data_out.write_dataframe(percentile_95(data_in.dataframe()))

def percentile_95(df):
    @pandas_udf("long",PandasUDFType.GROUPED_AGG)
    def percentile_95_udf(v):
        return v.quantile(0.95)

    # group rows for each interface into 1 day periods
    grp_by = df.groupBy(df.objectId,F.window("TimeCaptured","1 day"))

    stats = [
        percentile_95_udf(df.ReceivedWidgets),percentile_95_udf(df.TransmittedWidgets),]
    result = grp_by.agg(*stats)

    cleaned = result.withColumn("Day",F.col("window").start).drop("window")
    return cleaned

为什么我的 pandas_udf 在全局范围内不起作用，但在另一个函数中定义时却起作用？另外，是否有更好的方法来定义 pandas_udf？将其定义为嵌套函数会阻止我重用我的 udf。

作为参考，我在 Palantir Foundry 中的代码库具有以下结构：

transforms-python
    conda_recipe
        Meta.yaml
    src
        myproject
            datasets
                __init__.py
                percentile_95.py
            __init__.py
            pipeline.py
        setup.cfg
        setup.py

解决方法

原因的根本原因类似于此问题：PySpark error: AttributeError: 'NoneType' object has no attribute '_jvm'

当您在全局级别进行调用时，您正在尝试在设置 spark 之前执行 spark 命令（在您的情况下通过 pandas）。当您在转换中调用时，spark 可用，因此它可以工作。

这里的主要问题是注释本身在顶层被调用，而 spark 仅在转换运行时设置。当您从 def percentile_95(df): 中调用它时，您实际上是从此处的转换中调用注释：

@transform(
    data_out=Output("output path"),data_in=Input("input path")
)
def percentile_95_transform(data_in,data_out):
    data_out.write_dataframe(
                             percentile_95(  # <-- here we're inside a transform
                                           data_in.dataframe())) 

如果您想在多个地方重用这些 UDF，也许您可​​以将它们包装在一个函数或一个类中，并在您要使用的每个转换中初始化。