Pyspark 遍历行并使用结​​果列上的逻辑计算计数器

Question

我在 pyspark 数据框中有数据（这是一个非常大的表，有 900M 行）

这是我拥有的数据

+-------+---------+----------+
|    key|     time|      cond|
+-------+---------+----------+
|      6|     3704|      null|
|      6|    74967|      1062|
|      6|151565068|      null|
|      6|154999554|      null|
|      6|160595800|      null|
|      6|166192324|      null|
|      6|166549533|      null|
|      6|171318946|      null|
|      6|754759092|      null|
|      6|754999359|  18882624|
|      6|755171746|  11381128|
|      6|761097038|      null|
|      6|774496554|      null|
|      6|930609982|      null|
|      6|930809622|      null|
|      1|   192427|      null|
|      1|   192427|      2779|
|      1|   717931|      null|
|      1|  1110573|      null|
|      1|  1155854|      null|
|      1| 70049289|      null|
|      1| 70687548|      null|
|      1| 71222733|      null|
|      1| 85006084|      null|
|      1| 85029676|      null|
|      1| 85032605|   1424537|
|      1| 85240114|      null|
|      1| 85573757|      null|
|      1| 85710915|      null|
|      1| 85870370|      null|
+-------+---------+----------+

这是我需要对数据框执行的操作（中间步骤）：

+-------+---------+----------+--------+
|    key|     time|      cond|  result|
+-------+---------+----------+--------+
|      6|     3704|      null|       0|
|      6|    74967|      1062|       1|
|      6|151565068|      null|       0|
|      6|154999554|      null|       1|
|      6|160595800|      null|       2|
|      6|166192324|      null|       3|
|      6|166549533|      null|       4|
|      6|171318946|      null|       5|
|      6|754759092|      null|       6|
|      6|754999359|  18882624|       7|
|      6|755171746|  11381128|       0|
|      6|761097038|      null|       0|
|      6|774496554|      null|       1|
|      6|930609982|      null|       2|
|      6|930809622|      null|       3|
|      1|   192427|      null|       0|
|      1|   192427|      2779|       1|
|      1|   717931|      null|       0|
|      1|  1110573|      null|       1|
|      1|  1155854|      null|       2|
|      1| 70049289|      null|       3|
|      1| 70687548|      null|       4|
|      1| 71222733|      null|       5|
|      1| 85006084|      null|       6|
|      1| 85029676|      null|       7|
|      1| 85032605|   1424537|       8|
|      1| 85240114|      null|       0|
|      1| 85573757|      null|       1|
|      1| 85710915|      null|       2|
|      1| 85870370|      null|       3|
+-------+---------+----------+--------+

'result' 列的逻辑如下：每个键都有一个运行计数器，如果 'cond' 列不为空，则将计数器归零。

我们可以假设table是orderBy("key",asc("time"))

我的最终结果实际上是条件不为空的行上的结果（每个键）的平均值。 上面的数据应该是这样的（最终结果）：

+--------+--------------+
|    key |  avg_per_key |
+--------+--------------+
|       6|    2.66666665|  ==> (1+7+0)/3
|       1|           4.5|  ==> (1+8)/2
+--------+--------------+

我打算这样做：

df_results = df3[df3.cond.isNotNull()].groupby(['key']).agg(
    F.expr("avg(result)").alias("avg_per_key")
)

我认为它应该可以工作，但也许有更好的方法可以在没有中间步骤的情况下做到这一点。

如何在 pyspark 中有效地做到这一点？ （记住数据集很大）

Answer 1

这是我的解决方案，我并不是说它是最佳的，但它适用于我的情况，当其他尝试使集群崩溃时。 我是 spark 的初学者，所以我知道这种方法可能会导致问题，因为它将数据集放入内存。 如果我有更多时间玩它，我会尝试使用 sortWithinPartitions

def handleRow(row):
    temp = list(row[1])
    temp = np.array([temp[x:x+2] for x in range(0, len(temp),2)])
    temp[:,0] = temp[:,0].astype(float)
    temp = temp[temp[:,0].argsort()]
    avg_per_key= []
    counter=0
    for time,cond in temp:
        if cond!=None:
            avg_per_key.append(counter) 
            counter=0
        else:
            counter=counter+1

    return [(row[0],-1 if len(avg_per_key)==0 else np.mean(avg_per_key))]


count = df3.rdd.map(lambda x: (x.key, (x.time, x.cond)))\
    .reduceByKey(lambda a, b: a + b)\
    .flatMap(handleRow)\
    .collect()

Answer 2

试试这个。结果是通过在 conditions 上使用 incremental sum 来计算的，然后在另一个窗口中使用这些 groupings as partitionBy 作为 @ 987654324@ 以获得所需的result。 Filter 之前的 groupBy 应该由 reducing shuffle 执行。

from pyspark.sql import functions as F
from pyspark.sql.window import Window


w=Window().partitionBy("key").orderBy("time")
w1=Window().partitionBy("key","result").orderBy("time")

conditions=F.when((F.col("cond").isNotNull())&(F.col("lag").isNotNull()&\
                                                           F.col("lead").isNull()),F.lit(1))\
                        .when((F.col("cond").isNull())&(F.col("lag").isNotNull()),F.lit(1))\
                         .otherwise(F.lit(0))


df.withColumn("lag", F.lag("cond").over(w))\
  .withColumn("lead", F.lead("cond").over(w))\
  .withColumn("result",F.sum(conditions).over(w))\
  .withColumn("result", F.row_number().over(w1)-1).filter("cond is not null")\
  .groupBy("key").agg(F.mean(F.col("result")).alias("avg_per_key")).show()

#+---+------------------+
#|key|       avg_per_key|
#+---+------------------+
#|  6|2.6666666666666665|
#|  1|               4.5|
#+---+------------------+