如何检查Pyspark中count的值？

Question

我正在尝试创建一个查询，该查询将允许我获取一个表，其中ID至少出现3个计数，并且这些ID在k列中的值分别为0、3、4。

+---+---+
|  i|  k|
+---+---+
| 1 |  0|
| 1 |  3|
| 1 |  4|
| 2 |  0|
| 2 |  3|
| 2 |  3|
+---+---+

我想要的输出是：

+---+---+
|  i|  k|
+---+---+
| 1 |  0|
+---+---+
| 1 |  3|
+---+---+
| 1 |  4|
+---+---+  

这是我目前的代码。但是它只显示一个ID出现3次的表，我不确定如何检查这些计数中至少1个是0、3还是4

    sample= sample.join(
        sample.groupBy('i').count().where(('count == 3')).drop('count'), on=['i']
    )

Answer 1

在分组计数之前，您可以在k = 0, 3, 4上进行不同的选择。

df = df.join(
    df.select("i", "k").where("k in (0, 3, 4)").distinct()
      .groupby("i").count().where("count == 3").drop("count"), on=["i"]
)

Answer 2

我又添加了一个元素(1,5)进行测试。

from pyspark.sql.functions import count,collect_list, min, col
from pyspark.sql.window import Window
#values = [(1,0),(1,3),(1,4),(2,0),(2,3),(2,3)]
values = [(1,0),(1,3),(1,4),(1,5),(2,0),(2,3),(2,3)]
df = sqlContext.createDataFrame(values,['i','k'])
df.show()
+---+---+
|  i|  k|
+---+---+
|  1|  0|
|  1|  3|
|  1|  4|
|  1|  5|
|  2|  0|
|  2|  3|
|  2|  3|
+---+---+

#Creating a list of all elements grouped by 'i'
w = Window().partitionBy('i')
df = df.withColumn('count',count(col('i')).over(w)).where(col('count') >= 3).drop('count')\
       .withColumn('list',collect_list(col('k')).over(w))
df.show()
+---+---+------------+
|  i|  k|        list|
+---+---+------------+
|  1|  0|[0, 3, 4, 5]|
|  1|  3|[0, 3, 4, 5]|
|  1|  4|[0, 3, 4, 5]|
|  1|  5|[0, 3, 4, 5]|
|  2|  0|   [0, 3, 3]|
|  2|  3|   [0, 3, 3]|
|  2|  3|   [0, 3, 3]|
+---+---+------------+

# Creating a UDF to check if 'k' contains ATLEAST one of these three values (0,3,4), irrespective 
# of whether there are other values like 7,8 and so on.
check_for_atleast_0_3_4 = udf(lambda row: set([0,3,4]).issubset(row))
df1 = df.withColumn('bool',check_for_atleast_0_3_4(col('list'))).where(col('bool')==True).drop('bool','list')
df1.show()
+---+---+
|  i|  k|
+---+---+
|  1|  0|
|  1|  3|
|  1|  4|
|  1|  5|
+---+---+

# Creating a UDF to check if 'k' contains ATLEAST and ONLY one of these three values (0,3,4). 
# If 'k' has value 7, then all rows corresponding to that 'i' will be removed.
check_for_atleast_and_only_0_3_4 = udf(lambda row: set(list(set(row))).issubset([0,3,4]) and set([0,3,4]).issubset(row))
df1 = df.withColumn('bool',check_for_atleast_and_only_0_3_4(col('list'))).where(col('bool')==True).drop('bool','list')
df1.show()
+---+---+
|  i|  k|
+---+---+
+---+---+