如何 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)

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

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

因为您的支票与您的输出具有相同的值，但附加了 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    |
+----+------+

Answer 2

你差不多明白了，你需要将测试用例列表作为第二个参数传递给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'>

如何 F.when 基于条件的可变数量来提供 pyspark

How to F.when based on variable number of a conditions to supply with pyspark

python

apache-spark

pyspark

apache-spark-sql