在 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|
+---+---+