python - 在 pyspark 中保留至少一个元素满足条件的组

Question

我一直在尝试在 pyspark 中重现在 Pandas 中相当容易做到的事情，但我已经挣扎了一段时间了。 假设我有以下数据框:

df = pd.DataFrame({'a':[1,2,2,1,1,2], 'b':[12,5,1,19,2,7]})
print(df)
   a   b
0  1  12
1  2   5
2  2   1
3  1  19
4  1   2
5  2   7

还有列表

l = [5,1]

我想做的是按a分组，如果b中的任何元素在列表中，则返回True 对于组中的所有值。然后我们可以使用结果来索引数据帧。 Pandas 的等价物是:

df[df.b.isin(l).groupby(df.a).transform('any')]

   a  b
1  2  5
2  2  1
5  2  7

pyspark 中可重现的数据帧:

from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()

df = pd.DataFrame({'a':[1,2,2,1,1,2], 'b':[12,5,1,19,2,7]})
sparkdf = spark.createDataFrame(df)

我目前正在朝着按 a 分组并应用 pandasUDF 的方向，尽管肯定有更好的方法仅使用 Spark 来做到这一点。

Answer 1

我已经找到了一个足够简单的解决方案。第一步是使用 isin 和 filter 过滤出列表中 b 中的值所在的行，然后保留唯一的分组键(a) 在列表中。

然后通过与 a 上的数据帧合并，我们将组保留在列表中:

unique_a = (sparkdf.filter(f.col('b').isin(l))
                   .select('a').distinct())
sparkdf.join(unique_a, 'a').show()

+---+---+
|  a|  b|
+---+---+
|  2|  5|
|  2|  1|
|  2|  7|
+---+---+