Я пишу udf, который возьмет два столбца dataframe вместе с дополнительным параметром (постоянное значение) и должен добавить новый столбец в dataframe. Моя функция выглядит так:
def udf_test(column1, column2, constant_var):
if column1 == column2:
return column1
else:
return constant_var
Кроме того, я делаю следующее, чтобы пройти в нескольких столбцах:
apply_test = udf(udf_test, StringType())
df = df.withColumn('new_column', apply_test('column1', 'column2'))
Это не работает прямо сейчас, если я не удаляю constant_var
как третий аргумент функций, но мне это действительно нужно. Поэтому я попытался сделать что-то вроде следующего:
constant_var = 'TEST'
apply_test = udf(lambda x: udf_test(x, constant_var), StringType())
df = df.withColumn('new_column', apply_test(constant_var)(col('column1', 'column2')))
а также
apply_test = udf(lambda x,y: udf_test(x, y, constant_var), StringType())
Ни один из вышеперечисленных не работал для меня. Я получил эти идеи на основе этого и этих сообщений stackoverflow, и я думаю, что очевидно, что мой вопрос отличается от обоих. Любая помощь приветствуется.
ПРИМЕЧАНИЕ. Я упростил функцию здесь только ради обсуждения, и фактическая функция более сложна. Я знаю, что эта операция может быть сделана с помощью, when
и в otherwise
заявления.
Вам не нужно использовать определенную пользователем функцию. Вы можете использовать функции, когда() и иначе():
from pyspark.sql import functions as f
df = df.withColumn('new_column',
f.when(f.col('col1') == f.col('col2'), f.col('col1'))
.otherwise('other_value'))
Другой способ сделать это - создать определенную пользователем функцию. Однако использование udf
оказывает негативное влияние на производительность, поскольку данные должны быть (де) сериализованы в и из python. Чтобы создать пользовательскую функцию, вам понадобится функция, которая возвращает (определяемую пользователем) функцию. Например:
def generate_udf(constant_var):
def test(col1, col2):
if col1 == col2:
return col1
else:
return constant_var
return f.udf(test, StringType())
df = df.withColumn('new_column',
generate_udf('default_value')(f.col('col1'), f.col('col2')))
when()
и otherwise()
не являются операторами, они являются функциями .
.when()
и.otherwise()
, верно?