Как проверить значение счета в Pyspark?

Question

Я пытаюсь создать запрос, который позволил бы мне получить таблицу, в которой идентификатор имел бы не менее 3-х отсчетов, и эти значения имели бы значение 0, 3, 4 в столбце k.

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

Мой желаемый результат:

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

В настоящее время это мой код. Но он показывает только таблицу, в которой идентификатор встречается 3 раза, и я не совсем уверен, как проверить, является ли хотя бы 1 из этих счетчиков 0, 3 или 4

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

Answer 1

Я добавил еще один элемент (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|
+---+---+
+---+---+

Answer 2

Вы можете сделать отдельное выделение на 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"]
)