如何 F.when 基於可變數量的條件提供 pyspark

Question

我正在嘗試基於可變數量的條件構建一系列F.when 。 如何使用循環構建下面的邏輯，在該循環中提供要測試的項目列表（即[1,2,3]遵循下面的示例）？

我問的原因是因為我希望能夠在列表中使用可變數量的測試項來構建這些條件。循環邏輯應該構建如下所示的內容，但是通過傳遞一個數字列表來測試， [1,2,3] 。

F.when(F.col("test") == 1, "out_" + str(1) ).when(F.col("test") == 2, "out_" + str(2)).when(F.col("test") == 3, "out_" + str(3)).otherwise(-1)

我曾嘗試使用reduce來做到這一點，但之前沒有弄清楚這一點。 有人有建議嗎？

reduce(lambda x, i: x.when(F.col("test") == i , "out_" + str(i)),  
              output_df, 
              F).otherwise(-1)

我預期的 output 應該提供與以下相同的邏輯：

Column<b'CASE WHEN (test = 1) THEN out_1 WHEN (test = 2) THEN out_2 WHEN (test = 3) THEN out_3 ELSE -1 END'>

Answer 1

你幾乎明白了，你需要將測試用例列表作為第二個參數傳遞給reduce function：

from functools import reduce
import pyspark.sql.functions as F


tests = [1, 2, 3]

new_col = reduce(
    lambda x, i: x.when(F.col("test") == i, "out_" + str(i)),
    tests,
    F
).otherwise(-1)

print(new_col)

#Column<'CASE WHEN (test = 1) THEN out_1 WHEN (test = 2) THEN out_2 WHEN (test = 3) THEN out_3 ELSE -1 END'>

Answer 2

由於您的支票與您的 output 具有相同的價值，但附加out_ 。 您可以檢查該值是否在預定義列表中，以及是否只是添加out_ 。

例子：

from pyspark.sql import SparkSession
import pyspark.sql.functions as F


data = [
    {"test": 1},
    {"test": 2},
    {"test": 3},
    {"test": 4},
    {"test": 5},
]

test_ints = [1, 2, 3]

spark = SparkSession.builder.getOrCreate()
df = spark.createDataFrame(data)
df = df.withColumn(
    "result",
    F.when(
        F.col("test").isin(test_ints),
        F.concat(F.lit("out_"), F.col("test")),
    ).otherwise(-1),
)

結果：

+----+------+                                                                   
|test|result|
+----+------+
|1   |out_1 |
|2   |out_2 |
|3   |out_3 |
|4   |-1    |
|5   |-1    |
+----+------+