python - PySpark 中有等效的 .any() 吗？

Question

我想知道是否有办法在 Pyspark 中使用 .any() ？

我在Python中有以下代码，它本质上是搜索子集数据帧中感兴趣的特定列，如果这些列中的任何一个包含“AD”，我们不想处理它们.

这是Python代码:

index_list = [
    df.query("id == @id").index 
    for trial in unique_trial_id_list 
    if ~(df.query("id == @trial")["unit"].str.upper().str.contains("AD").any()]

这是 Pandas 中的示例数据框。

ID=1 具有与之关联的字符串 'AD'，因此我们希望将其从处理中排除。但是，ID=2 没有与之关联的字符串 'AD'，因此我们希望将其包含在进一步处理中。

data = [
    [1, "AD"],
    [1, "BC"],
    [1, "DE"],
    [1, "FG"],
    [2, "XY"],
    [2, "BC"],
    [2, "DE"],
    [2, "FG"],
]
df = pd.DataFrame(data, columns=["ID", "Code"])
df

问题是我不知道如何在 PySpark 中执行此等效功能。我已经能够对子集进行列表理解，并且能够使用 contains('AD') 进行子集化，但在涉及到 any 部分时陷入困境事情。

我想出的 PySpark 代码:

id = id_list[0] 
test = sdf.select(["ID", "Codes"]).filter(spark_fns.col("ID") == id).filter(~spark_fns.col("Codes").str.contains("AD"))

Answer 1

您可以使用窗口函数(如果至少有一个真值，则 bool 值的最大值为真):

from pyspark.sql import functions as F, Window

df1 = df.withColumn(
    "to_exclude",
    ~F.max(F.when(F.col("Code") == "AD", True).otherwise(False)).over(Window.partitionBy("ID"))
).filter(
    F.col("to_exclude")
).drop("to_exclude")

df1.show()
# +---+----+
# | ID|Code|
# +---+----+
# |  2|  XY|
# |  2|  BC|
# |  2|  DE|
# |  2|  FG|
# +---+----+

---

或者 groupby id 并使用 max 函数和 when 表达式来过滤包含 AD 的 id Code 列，然后与原始 df 联接:

from pyspark.sql import functions as F

filter_df = df.groupBy("id").agg(
    F.max(F.when(F.col("Code") == "AD", True).otherwise(False)).alias("to_exclude")
).filter(F.col("to_exclude"))

df1 = df.join(filter_df, ["id"], "left_anti")

---

在Spark 3+中，还有一个函数any:

from pyspark.sql import functions as F

filter_df = df.groupBy("id").agg(
    F.expr("any(Code = 'AD')").alias("to_exclude")
).filter(F.col("to_exclude"))

df1 = df.join(filter_df, ["id"], "left_anti")