我需要创建一个 pyspark UDF，它从带有比较的查询中输出一个表

如何解决我需要创建一个 pyspark UDF，它从带有比较的查询中输出一个表

我正在使用 kaggle 上的 IBM 损耗数据集。我想要做的是将分类变量的出现次数计算为 Attrition == 'Yes' 和 Attrition == 'No'，并采用简单的比率来查看分类变量的哪个级别更有可能发生损耗。现在我可以在 Pandas 中做到这一点，就像这样：

def cal_ratio(x):
    n_1 = sum(x['Attrition'].values == 'Yes')
    n_0 = sum(x['Attrition'].values == 'No')
    return n_1/n_0

或者我可以很容易地编写一个执行它的 spark.sql 查询，并为我想要比较的每个分类变量重新编写它。像 Pandas 这样的函数会让我的生活更轻松，但我找不到任何关于如何创建此类 UDF 或如何注册的真正指导。

编辑：如果我还问这在带有 UDF 的 pyspark 中如何工作，可能会有所帮助？

b = data.groupby('BusinessTravel').apply(cal_ratio)

解决方法

不确定这是最好的解决方案，但你可以试试这个：

# My sample dataframe
df.show()
+---------+                                                                     
|Attrition|
+---------+
|      Yes|
|      Yes|
|      Yes|
|      Yes|
|      Yes|
|       No|
|       No|
+---------+

from pyspark.sql import functions as F

result = (
    df.agg(
        F.sum(F.when(F.col("Attrition") == "Yes",1)).alias("Yes"),F.sum(F.when(F.col("Attrition") == "No",1)).alias("No"),)
    .select((F.col("Yes") / F.col("No")).alias("ratio"))
    .first()
)

print(result.ratio)
> 2.5

当然，您可以通过用变量替换硬编码值来将 result 事物转换为函数。

def cal_ratio(df):
    result = (
        df.agg(
            F.sum(F.when(F.col("Attrition") == "Yes",)
        .select((F.col("Yes") / F.col("No")).alias("ratio"))
        .first()
    )
    return result.ratio

---

编辑：如果您需要按列分组，则需要将 first 替换为 collect：

def cal_ratio(df):
    result = (
        df.groupBy("BusinessTravel")
        .agg(
            F.sum(F.when(F.col("Attrition") == "Yes",)
        .select((F.col("Yes") / F.col("No")).alias("ratio"))
        .collect()
    )
    return result