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作为组。然后，我们可以使用结果为数据帧建立索引。相当于熊猫的是：

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

我已经找到了足够简单的解决方案。第一步是使用b和isin过滤掉列表中filter中的值所在的行，然后将唯一的分组键（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|
+---+---+

将群组中至少一个满足条件的群组保留在pyspark

1 个答案: