python - 如何根据 pyspark 提供的条件的可变数量进行 F.when

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)

我的预期输出应提供与以下相同的逻辑:

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 函数:

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'>