PySpark - 使用Pandas UDF的顺序计数

时间:2018-04-17 21:40:51

标签: python pandas apache-spark pyspark

我有一个列数据集,它在同一个月的任何列中都会增加值,然后在下个月重置。

+----------+------+-----------+----+-----------+------------+
|      Date|column|column_2   |co_3|column_4   |column_5    |
+----------+------+-----------+----+-----------+------------+
|2016-12-14|     0|          0|   0|         14|           0|
|2016-12-14|     0|          0|   0|         14|           0|
|2016-12-14|     0|          0|   0|         18|           0|
|2016-12-14|     0|          0|   0|         19|           0|
|2016-12-14|     0|          0|   0|         20|           0|
|2016-12-14|     0|          0|   0|         26|           0|
|2016-12-14|     0|          0|   0|         60|           0|
|2016-12-14|     0|          0|   0|         63|           0|
|2016-12-14|     0|          0|   0|         78|           0|
|2016-12-14|     0|          0|   0|         90|           0|
+----------+------+-----------+----+-----------+------------+

问题是他们的日期总是一样的,所以我想做一些计数,然后当我们接近不同的一天时重置计数。

我写了一个Pandas UDF函数:

@pandas_udf('int', PandasUDFType.SCALAR)
def get_counts_up(v):
  prev = None
  series = []
  count = 0
  for i in v:
    if prev != i:
      count = 0
      prev = i
    series.append(count)
    count += 1
  return pd.Series(series)

但是,输出似乎不是连续的:

sdf.filter(sdf.Date == "2016-12-14").sort("Date_Count").show()

+------------+----------+------+-----------+----+-----------+------------+---------+----------+--------+----------+-----+----------+
|Date_Convert|      Date|column|column_____|col_|column_____|column______|Date_Year|Date_Month|Date_Day|Date_Epoch|count|Date_Count|
+------------+----------+------+-----------+----+-----------+------------+---------+----------+--------+----------+-----+----------+
|  2016-12-14|2016-12-14|     0|          0|   0|         14|           0|     2016|        12|      14|1481673600|14504|         0|
|  2016-12-14|2016-12-14|     0|          0|   0|         18|           0|     2016|        12|      14|1481673600|14504|         0|
|  2016-12-14|2016-12-14|     0|          0|   0|         14|           0|     2016|        12|      14|1481673600|14504|         1|
|  2016-12-14|2016-12-14|     0|          0|   0|         18|           0|     2016|        12|      14|1481673600|14504|         1|
|  2016-12-14|2016-12-14|     0|          0|   0|         18|           0|     2016|        12|      14|1481673600|14504|         2|
|  2016-12-14|2016-12-14|     0|          0|   0|         14|           0|     2016|        12|      14|1481673600|14504|         2|
|  2016-12-14|2016-12-14|     0|          0|   0|         14|           0|     2016|        12|      14|1481673600|14504|         3|
+------------+----------+------+-----------+----+-----------+------------+---------+----------+--------+----------+-----+----------+

这是可以预料到的,因为我猜数据框被拆分为不同的机器(DataBrick的社区版中的一些机器),每个机器都有自己的阵列来维护。

有没有办法进行顺序计数?

1 个答案:

答案 0 :(得分:1)

Windowrow_number函数的组合应该为您解决。我已经使用了所有列进行排序,如你所说

  
    

同一月份任何列的值增加的数据集...

  

但您只能使用一列或多列具有增加值的列。

from pyspark.sql import window as w
windowSpec = w.Window.partitionBy("Date").orderBy("column", "column_2", "co_3", "column_4", "column_5")

from pyspark.sql import functions as f
df.withColumn('inc_count', f.row_number().over(windowSpec)).show(truncate=False)

应该给你

+----------+------+--------+----+--------+--------+---------+
|Date      |column|column_2|co_3|column_4|column_5|inc_count|
+----------+------+--------+----+--------+--------+---------+
|2016-12-14|0     |0       |0   |14      |0       |1        |
|2016-12-14|0     |0       |0   |14      |0       |2        |
|2016-12-14|0     |0       |0   |18      |0       |3        |
|2016-12-14|0     |0       |0   |19      |0       |4        |
|2016-12-14|0     |0       |0   |20      |0       |5        |
|2016-12-14|0     |0       |0   |26      |0       |6        |
|2016-12-14|0     |0       |0   |60      |0       |7        |
|2016-12-14|0     |0       |0   |63      |0       |8        |
|2016-12-14|0     |0       |0   |78      |0       |9        |
|2016-12-14|0     |0       |0   |90      |0       |10       |
+----------+------+--------+----+--------+--------+---------+