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