Pyspark: как применить функцию, чтобы получить уклон в кадре данных?

Question

У меня есть такой фрейм данных искры (столбцы x и y, каждый с 6 точками данных). Я хочу иметь возможность извлечь наклон путем подгонки простой линии регрессии (в основном, чтобы увидеть тренд у, увеличивая или уменьшая).

+------------------+----------------------------------------------------------------------------------------------------------------------------------------+
|x                 |y                                                                                                                                       |
+------------------+----------------------------------------------------------------------------------------------------------------------------------------+
|[1, 2, 3, 4, 5, 6]|[3.42102562348192E-6, 4.2159323917750995E-6, 3.924587540944015E-6, 4.167182871752131E-6, 4.109192066532302E-6, 4.297804458327455E-6]    |
|[1, 2, 3, 4, 5, 6]|[1.384402399630826E-5, 9.913141993957704E-6, 1.1145077060247102E-5, 1.1005472165326649E-5, 1.1004462921073546E-5, 1.1004462921073546E-5]|
+------------------+----------------------------------------------------------------------------------------------------------------------------------------+

Я хочу использовать эту функцию, но получаю ошибки, когда применяя к данным. Как правильно применить эту функцию, которая принимает 2 массива?

import numpy as np
def trendline(x,y, order=1):
    coeffs = np.polyfit(x, y, order)
    slope = coeffs[-2]
    return float(slope)

#example to run
x_example=[1,2,3,4,5,6]
y_example=[3.42102562348192,4.2159323917750995,3.924587540944015,4.167182871752131,4.109192066532302,4.297804458327455 ]
slope=trendline(x_example,y_example)
print(slope)  

эта функция работает как в примере выше, но у меня возникают проблемы при применении к фрейму данных, я хочу создать новый столбец в фрейм данных с возвращенным уклоном.

спасибо !!

Я пробовал это, и это не сработало

def get_slope_func(x,y, order=1):
    coeffs = np.polyfit(x, y, order)
    slope = coeffs[-2]
    return float(slope)
get_slope = pandas_udf(get_slope_func, returnType=LongType())

df.select(get_slope(col("x"), col("y"))).show()

Answer 1

Итак, вы возвращаете float в get_slope_func и при регистрации udf вы установили тип возврата как LongType(), который в основном является BigInt в SQL. Установите returnType на DoubleType().

from pyspark.sql import functions as F
from pyspark.sql.type import *
import numpy as np
def get_slope_func(x,y,order=1):
    coeffs = np.polyfit(x, y, order)
    slope = coeffs[-2]
    return float(slope)
get_slope = F.udf(get_slope_func, returnType=DoubleType())

df.select(get_slope(F.col("x"), F.col("y")).alias("slope")).show(truncate=False)
#+----------------------+
#|slope                 |
#+----------------------+
#|1.2303624369449727E-7 |
#|-3.1609849970704353E-7|
#+----------------------+