PySpark-熊猫UDF的顺序计数

Question

I have a column dataset that has increasing values any columns for the same month, and then resets for the next month.

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

The problem is that their date is always the same, so I want to do some sort of counting up, and then resetting the count when we approach a different day.

I've written a Pandas UDF Function:

@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)

However, the output does not seem to be continuous:

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

Which is to be expected, because I guess the dataframe is split off into different machines (a few on DataBrick's community edition), and each have their own array to maintain.

Is there a way to perform a sequential counting up?

Answer 1

Combination of Window and row_number functions should solve it for you. I have used all of the columns for ordering as you've said

dataset that has increasing values any columns for the same month...

but you can use only one column or many which has the increasing values.

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)

which should give you

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