apache-spark - 使用包含至少一个特定值的窗口过滤数据框

Question

假设您有一个包含两列的 pyspark 数据框，id和 val .
您要过滤行，其中单个 id -window 至少有一个 val具有一定值的条目。
例如。如果我们想获取行，其中 id -window 至少有一个 5 val 中的值柱子:

# Input:

df = spark.createDataFrame(
    [(1, '0'), 
     (1, '5'),
     (2, '2'),
     (2, '5'),    
     (2, '5'),
     (3, '1'),
     (3, '0'),],
     ['id', 'val']
)

# Desired output:

+---+----+
| id| val|
+---+----+
|  1|   0|
|  1|   5|
|  2|   2|
|  2|   5|
|  2|   5|
+---+----+

我想过以某种方式使用窗口函数吗？

Answer 1

使用 array_contains() 函数来自 Spark-2.4然后加入df只得到身份证 有 5 val 在他们。
Example:

from pyspark.sql.functions import *
from pyspark.sql.types import *

df1=df.groupBy("id").\
agg(array_contains(collect_set(col("val")).cast("array<int>"),5).alias("has_5")).\
filter(col("has_5")).\
drop('has_5')

df.join(df1,['id'],'inner').show()
#+---+---+
#| id|val|
#+---+---+
#|  1|  0|
#|  1|  5|
#|  2|  2|
#|  2|  5|
#|  2|  5|
#+---+---+

Another way Using Window function:

import sys
from pyspark.sql import *

w=Window.partitionBy("id").orderBy("val").rowsBetween(-sys.maxsize,sys.maxsize)

df.withColumn("has_5",array_contains(collect_set(col("val")).over(w).cast("array<int>"),5)).\
filter(col("has_5")).\
drop("has_5").\
show()
#+---+---+
#| id|val|
#+---+---+
#|  1|  0|
#|  1|  5|
#|  2|  2|
#|  2|  5|
#|  2|  5|
#+---+---+