Как добавить новую колонку относительно максимального элемента?

Question

У меня есть искровой фрейм данных из CSV-элемента.

csv_df = sqlContext.read.format("csv").option("header", "true").option("inferSchema", "false").load(csv_path)

Есть столбец со скоростью

Скорость: [3, 5, 4, 8, 7, 6, 1]

Я хочу добавить новый столбец, чтобы определить, является ли это ускорением или desc.

Чтобы достичь этого, мне нужно найти максимум предметов [8], пометить все предметы как «A» до этого и отметить начало «D»:

Accel: ["A", "A", "A", "A", "D", "D", "D"]

Я могу сделать это с помощью панд и библиотеки numpy

#find max location:
maxloc=np.argmax(csv_pandas_df.Speed)
csv_pandas_df.loc[0:maxloc,'Accel']='A'
csv_pandas_df.loc[maxloc:csv_pandas_df.shape[0],'Accel']='D'

Есть ли способ добиться этого с помощью встроенного Spark DataFrame?

Answer 1

Не уверен насчет собственной искры, но этого можно добиться с помощью пользовательской функции (UDF).

from pyspark.sql.functions import udf
import operator

def acc_dec(speeds):
  if len(speeds) is 0:
    return []
  index, value = max(enumerate(speeds), key=operator.itemgetter(1))
  return ["A"]*(index+1) + ["D"]*(len(speeds)-index-1)
acc_dec_udf = udf(acc_dec)

df_accel = csv_df.withColumn("Accel", acc_dec_udf("Speed"))