Добавить столбец с ближайшими значениями в PySpark Dataframe

Question

У меня есть фрейм данных PySpark (скажем, df), который имеет два столбца (Name и Score). Ниже приведен пример фрейма данных:

+------+-----+
|  Name|Score|
+------+-----+
| name1|11.23|
| name2|14.57|
| name3| 2.21|
| name4| 8.76|
| name5|18.71|
+------+-----+

У меня есть массив Numpy (скажем, bin_array), который имеет значения, близкие к числовым значениям в столбце с именем Score фрейма данных PySpark,

Ниже приведен вышеупомянутый массив numpy:

bin_array = np.array([0, 5, 10, 15, 20])

Я хочу сравнить значение из каждой строки столбца Score со значениями в bin_array и сохранить ближайшее значение(получено из bin_array) в отдельном столбце в фрейме данных PySpark.

Ниже приведен порядок отображения моего нового фрейма данных (скажем, df_new).

+------+-----+------------+
|  Name|Score| Closest_bin|
+------+-----+------------+
| name1|11.23|      10.0  |
| name2|14.57|      15.0  |
| name3| 2.21|       0.0  |
| name4| 8.76|      10.0  |
| name5|18.71|      20.0  |
+------+-----+------------+

У меня есть нижеупомянутая функция, которая дает мне самые близкие значения от bin_array. Функция отлично работает, когда я тестирую ее с отдельными числами.

def find_nearest(array, value):
    array = np.asarray(array)
    idx = (np.abs(array - value)).argmin()
    return float(array[idx])

В моей работе у меня будут миллионы строк в datafrmae. Какой самый эффективный способ создания df_new?

Ниже приведены шаги, которые я пытался использовать для создания пользовательской функции (udf) и нового фрейма данных (* 1034). *).

closest_bin_udf = F.udf( lambda x: find_nearest(array, x) )
df_new = df.withColumn( 'Closest_bin' , closest_bin_udf(df.Score)  )

Но при попытке df_new.show() я получил ошибки. Часть ошибки показана ниже.

---------------------------------------------------------------------------
Py4JJavaError                             Traceback (most recent call last)
<ipython-input-11-685c9b7e25d9> in <module>()
----> 1 df_new.show()

/usr/lib/spark/python/pyspark/sql/dataframe.py in show(self, n, truncate, vertical)
    376         """
    377         if isinstance(truncate, bool) and truncate:
--> 378             print(self._jdf.showString(n, 20, vertical))
    379         else:
    380             print(self._jdf.showString(n, int(truncate), vertical))

Вы можете использовать указанные ниже шаги для создания вышеупомянутого кадра данных:

from pyspark.sql import *
import pyspark.sql.functions as F
import numpy as np

Stats = Row("Name", "Score")

stat1 = Stats('name1', 11.23)
stat2 = Stats('name2', 14.57)
stat3 = Stats('name3', 2.21)
stat4 = Stats('name4', 8.76)
stat5 = Stats('name5', 18.71)

stat_lst = [stat1 , stat2, stat3, stat4, stat5]
df = spark.createDataFrame(stat_lst)

df.show()

Answer 1

Вы также можете pandas_udf, хотя я бы посоветовал вам проверить скорость и потребление памяти при увеличении

from pyspark.sql.functions import pandas_udf, PandasUDFType
import numpy as np
import pandas as pd
df = spark.createDataFrame(zip(["name_"+str(i) for i in range(1,6)], [11.23, 14.57, 2.21, 8.76, 18.71]), ["Name", "Score"])

bin_array = np.array([0, 5, 10, 15, 20])

@pandas_udf('double', PandasUDFType.SCALAR)
def find_nearest(value):
    res = bin_array[np.newaxis, :] - value.values[:, np.newaxis]
    ret_vals = [bin_array[np.argmin(np.abs(i))] for i in res]
    return pd.Series(ret_vals)

df.withColumn('v2', find_nearest(df.Score)).show()

Вывод

+------+-----+----+
|  Name|Score|  v2|
+------+-----+----+
|name_1|11.23|10.0|
|name_2|14.57|15.0|
|name_3| 2.21| 0.0|
|name_4| 8.76|10.0|
|name_5|18.71|20.0|
+------+-----+----+

Answer 2

Вы можете использовать bucketizer из pyspark.mllib

from pyspark.sql import *
import pyspark.sql.functions as F
import numpy as np

Stats = Row("Name", "Score")

stat_lst = [Stats('name1', 11.23) , Stats('name2', 14.57), Stats('name3', 2.21), Stats('name4', 8.76), Stats('name5', 18.71)]
df = spark.createDataFrame(stat_lst)

from pyspark.ml.feature import Bucketizer

"""
Bucketizer creates bins like 0-5:0, 5-10:1, 10-15:2, 15-20:3
As I see, your expected output wants the closest numbered bin, so you might 
have to change your buckets or the variable `t` below accordingly.
"""
bucket_list = [0, 5, 10, 15, 20]

bucketizer = Bucketizer(splits=bucket_list, inputCol="Score", outputCol="buckets")
df_buck = bucketizer.setHandleInvalid("keep").transform(df)

df_buck.show()

Я все еще работаю над получением ближайшей корзины, я обновлю свой ответ.

Если вам нужны значения вашего массива для каждого сегмента, вы можете использовать udf для создания нового столбца с именами блоков

from pyspark.sql.functions import udf
from pyspark.sql.types import *

t = dict(zip(range(len(bucket_list)), bucket_list))
udf_foo = udf(lambda x: t[x], IntegerType())
df_buck = df_buck.withColumn("score_bucket", udf_foo("buckets"))

Вывод

>>> df_buck.show()

+-----+-----+-------+------------+
| Name|Score|buckets|score_bucket|
+-----+-----+-------+------------+
|name1|11.23|    2.0|          10|
|name2|14.57|    2.0|          10|
|name3| 2.21|    0.0|           0|
|name4| 8.76|    1.0|           5|
|name5|18.71|    3.0|          15|
+-----+-----+-------+------------+

РЕДАКТИРОВАТЬ: Исправление сегментов оценки:

# Not dynamic, but please try to figure out this business logic according to your use-case
df_buck = df_buck.withColumn("correct_buckets", F.when(df_buck.Score-df_buck.score_bucket > 5/2, F.col("score_bucket") + 5).otherwise(F.col("score_bucket"))).drop("buckets", "score_bucket")

Теперь вывод соответствует ожидаемому:

+-----+-----+---------------+
| Name|Score|correct_buckets|
+-----+-----+---------------+
|name1|11.23|             10|
|name2|14.57|             15|
|name3| 2.21|              0|
|name4| 8.76|             10|
|name5|18.71|             20|
+-----+-----+---------------+