Pyspark: как кодировать сложные вычисления Dataframe - PullRequest
Новая
       157

Pyspark: как кодировать сложные вычисления Dataframe

1 голос
/ 06 августа 2020

Фрейм данных уже отсортирован по дате,

col1 == 1 значение уникально,

и col1 == 1 переданы, приращение увеличится на 1 (например, 1 , 2,3,4,5,6,7 ...) и только -1 являются дубликатами.

У меня фрейм данных выглядит так, назовите его df 

TEST_schema = StructType([StructField("date", StringType(), True),\
                          StructField("col1", IntegerType(), True),\
                          StructField("col2", IntegerType(), True)])
TEST_data = [('2020-08-01',-1,-1),('2020-08-02',-1,-1),('2020-08-03',-1,3),('2020-08-04',-1,2),('2020-08-05',1,4),\
             ('2020-08-06',2,1),('2020-08-07',3,2),('2020-08-08',4,3),('2020-08-09',5,-1)]
rdd3 = sc.parallelize(TEST_data)
TEST_df = sqlContext.createDataFrame(TEST_data, TEST_schema)
TEST_df.show()



+--------+----+----+
    date |col1|col2|
+--------+----+----+
2020-08-01| -1|  -1|
2020-08-02| -1|  -1|
2020-08-03| -1|   3|
2020-08-04| -1|   2|
2020-08-05| 1 |   4|
2020-08-06| 2 |   1|
2020-08-07| 3 |   2|
2020-08-08| 4 |   3|
2020-08-09| 5 |  -1|
+--------+----+----+

условие - когда col1 == 1, тогда мы начинаем добавлять в обратном направлении от col2 == 4 (например, 4,5,6,7,8, ...), а после col2 == 4 полностью возвращаем 0 (например, . 4,0,0,0,0 ...)

Итак, мой результирующий df будет выглядеть примерно так. 

   +--------+----+----+----+
        date |col1|col2|want
    +--------+----+----+----+
    2020-08-01| -1|  -1|  8 |
    2020-08-02| -1|  -1|  7 |
    2020-08-03| -1|   3|  6 |
    2020-08-04| -1|   2|  5 |
    2020-08-05| 1 |   4|  4 |
    2020-08-06| 2 |   1|  0 |
    2020-08-07| 3 |   2|  0 |
    2020-08-08| 4 |   3|  0 |
    2020-08-09| 5 |  -1|  0 |
   +---------+----+----+----+

Улучшение : Я хочу чтобы добавить дополнительное условие, где col2 == -1, когда col1 == 1 (в 2020-08-05), а col2 == -1 идет последовательно .. тогда я хочу подсчитать последовательный -1, а затем добавить, где последовательные разрывы col2 ==? стоимость. так вот пример для очистки. 

    +--------+----+----+----+
        date |col1|col2|want
    +--------+----+----+----+
    2020-08-01| -1|  -1|  11|
    2020-08-02| -1|  -1|  10|
    2020-08-03| -1|   3|  9 |
    2020-08-04| -1|   2|  8 |
    2020-08-05| 1 |  -1|  7*|
    2020-08-06| 2 |  -1|  0 |
    2020-08-07| 3 |  -1|  0 |
    2020-08-08| 4 |  4*|  0 |
    2020-08-09| 5 |  -1|  0 |
   +---------+----+----+----+

Итак, мы видим 3 последовательных -1 (начиная с 2020-08-05, мы заботимся только о первых последовательных -1), а после последовательного у нас есть 4 (в 2020-08-08 обозначено *), тогда у нас будет 4+ 3 = 7 в строке col1 == 1. возможно ли это?

** МОЯ ПЕРВАЯ ПОПЫТКА **

TEST_df = TEST_df.withColumn('cumsum', sum(when( col('col1') < 1, col('col1') ) \
                 .otherwise( when( col('col1') == 1, 1).otherwise(0))).over(Window.partitionBy('col1').orderBy().rowsBetween(-sys.maxsize, 0)))
TEST_df.show()

+----------+----+----+------+
|      date|col1|col2|cumsum|
+----------+----+----+------+
|2020-08-01|  -1|  -1|    -1|
|2020-08-02|  -1|  -1|    -2|
|2020-08-03|  -1|   3|    -3|
|2020-08-04|  -1|   2|    -4|
|2020-08-05|   1|   4|     1|
|2020-08-07|   3|   2|     0|
|2020-08-09|   5|  -1|     0|
|2020-08-08|   4|   3|     0|
|2020-08-06|   2|   1|     0|
+----------+----+----+------+

w1 = Window.orderBy(desc('date'))
w2 =Window.partitionBy('case').orderBy(desc('cumsum'))

TEST_df.withColumn('case', sum(when( (col('cumsum') == 1) & (col('col2') != -1) , col('col2')) \
       .otherwise(0)).over(w1)) \
  .withColumn('rank', when(col('case') != 0, rank().over(w2)-1).otherwise(0)) \
  .withColumn('want', col('case') + col('rank')) \
  .orderBy('date') \
+----------+----+----+------+----+----+----+
|date      |col1|col2|cumsum|case|rank|want|
+----------+----+----+------+----+----+----+
|2020-08-01|-1  |-1  |-1    |4   |1   |5   |
|2020-08-02|-1  |-1  |-2    |4   |2   |6   |
|2020-08-03|-1  |3   |-3    |4   |3   |7   |
|2020-08-04|-1  |2   |-4    |4   |4   |8   |
|2020-08-05|1   |4   |1     |4   |0   |4   |
|2020-08-06|2   |1   |0     |0   |0   |0   |
|2020-08-07|3   |2   |0     |0   |0   |0   |
|2020-08-08|4   |3   |0     |0   |0   |0   |
|2020-08-09|5   |-1  |0     |0   |0   |0   |
+----------+----+----+------+----+----+----+

Вы видите, что ранг 1,2,3,4, если я смогу сделать это 4,3,2,1 это будет выглядеть как полученный фрейм данных .... как это изменить? Я пробовал как orderby как c, так и des c ... и, конечно, это до улучшения

1 Ответ

1 голос
/ 07 августа 2020

IIU C, вы можете попробовать следующее:

  1. groupby и создать collect_list всех связанных строк (vals в коде ниже), отсортируйте список по дате в в порядке убывания ( Примечание: измените groupby(lit(1)) на любые столбцы, которые вы можете использовать для разделения данных на независимое подмножество.

  2. найдите индекс массива idx, который имеет col1 == 1

  3. если col2==-1 в idx, тогда найдите смещение от idx до начала списка с первой строкой, имеющей col2 != -1 ( Примечание: в текущем коде смещение может быть NULL, если все столбцы col2 до idx равны -1, вам нужно будет решить, что вы хотите. Например, используйте coalesce(IF(...),0))

  4. после того, как у нас есть смещение и idx, столбец want можно рассчитать следующим образом: 

    IF(i<idx, 0, vals[idx-offset].col2 + offset + i - idx)

  5. используйте Spark SQL function inline , чтобы взорвать массив структур.

Примечание: Тот же logi c может быть применен с использованием функции Window в случае, если существует слишком много столбцов t в вашем производственном фрейме данных.

Код ниже: 

from pyspark.sql.functions import sort_array, collect_list, struct, expr, lit

TEST_df = spark.createDataFrame([
  ('2020-08-01', -1, -1), ('2020-08-02', -1, -1), ('2020-08-03', -1, 3),
  ('2020-08-04', -1, 2), ('2020-08-05', 1, -1), ('2020-08-06', 2, -1),
  ('2020-08-07', 3, -1), ('2020-08-08', 4, 4), ('2020-08-09', 5, -1)
], ['date', 'col1', 'col2'])

# list of column used in calculation
cols = ["date", "col1", "col2"]

df_new = TEST_df \
    .groupby(lit(1)) \
    .agg(sort_array(collect_list(struct(*cols)),False).alias('vals')) \
    .withColumn('idx', expr("filter(sequence(0,size(vals)-1), i -> vals[i].col1=1)[0]")) \
    .withColumn('offset', expr("""
        coalesce(IF(vals[idx].col2=-1, filter(sequence(1,idx), i -> vals[idx-i].col2 != -1)[0],0),0)
     """)).selectExpr("""
       inline(
         transform(vals, (x,i) -> named_struct(
             'date', x.date,
             'col1', x.col1,
             'col2', x.col2,
             'want', IF(i<idx, 0, vals[idx-offset].col2 + offset + i - idx)
           )
         )
    )""")

Вывод: 

df_new.orderBy('date').show()
+----------+----+----+----+
|      date|col1|col2|want|
+----------+----+----+----+
|2020-08-01|  -1|  -1|  11|
|2020-08-02|  -1|  -1|  10|
|2020-08-03|  -1|   3|   9|
|2020-08-04|  -1|   2|   8|
|2020-08-05|   1|  -1|   7|
|2020-08-06|   2|  -1|   0|
|2020-08-07|   3|  -1|   0|
|2020-08-08|   4|   4|   0|
|2020-08-09|   5|  -1|   0|
+----------+----+----+----+

Изменить: Per комментарии, добавлена ​​альтернатива для использования агрегатной функции Window вместо groupby: 

from pyspark.sql import Window

# WindowSpec to cover all related Rows in the same partition
w1 = Window.partitionBy().orderBy('date').rowsBetween(Window.unboundedPreceding,Window.unboundedFollowing)

cols = ["date", "col1", "col2"]

# below `cur_idx` is the index for the current Row in array `vals`
df_new = TEST_df.withColumn('vals', sort_array(collect_list(struct(*cols)).over(w1),False)) \
    .withColumn('idx', expr("filter(sequence(0,size(vals)-1), i -> vals[i].col1=1)[0]")) \
    .withColumn('offset', expr("IF(vals[idx].col2=-1, filter(sequence(1,idx), i -> vals[idx-i].col2 != -1)[0],0)")) \
    .withColumn("cur_idx", expr("array_position(vals, struct(date,col1,col2))-1")) \
    .selectExpr(*TEST_df.columns, "IF(cur_idx<idx, 0, vals[idx-offset].col2 + offset + cur_idx - idx) as want")
...