Столбец Pyspark UDF на Датафрейме

1

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

df = pd.DataFrame([[0,1,0],[1,0,0],[1,1,1]],columns = ['Foo','Bar','Baz'])

spark_df = spark.createDataFrame(df)

def get_profile():
    if 'Foo'==1:
        return 'Foo'
    elif 'Bar' == 1:
        return 'Bar'
    elif 'Baz' ==1 :
        return 'Baz'

spark_df = spark_df.withColumn('get_profile', lit(get_profile()))
spark_df.show()

   Foo  Bar  Baz get_profile
    0    1    0        None
    1    0    0        None
    1    1    1        None

Я бы ожидал, что столбец get_profile будет заполнен для всех строк.

Я также пробовал:

spark_udf = udf(get_profile,StringType())

spark_df = spark_df.withColumn('get_profile', spark_udf())
print(spark_df.toPandas())

к тому же эффекту.

  • 1
    Вы сравниваете строки с числами. 'Foo' != 1 , то же самое для других условий. Вот почему вы ничего не получите. UDF ожидает столбцы в качестве аргументов, в то время как get_profile имеет нулевые аргументы.
  • 1
    идти с когда / в противном случае встроенная функция вместо функции UDF
Показать ещё 2 комментария
Теги:
apache-spark
pyspark

1 ответ

1
Лучший ответ

udf не знает, что такое имена столбцов. Поэтому он проверяет каждое из ваших условий в блоке if/elif и все они оценивают значение False. Таким образом, функция вернет None.

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

from pyspark.sql.functions import udf

def get_profile(foo, bar, baz):
    if foo == 1:
        return 'Foo'
    elif bar == 1:
        return 'Bar'
    elif baz == 1 :
        return 'Baz'

spark_udf = udf(get_profile, StringType())
spark_df = spark_df.withColumn('get_profile',spark_udf('Foo', 'Bar', 'Baz'))
spark_df.show()
#+---+---+---+-----------+
#|Foo|Bar|Baz|get_profile|
#+---+---+---+-----------+
#|  0|  1|  0|        Bar|
#|  1|  0|  0|        Foo|
#|  1|  1|  1|        Foo|
#+---+---+---+-----------+

Если у вас много столбцов и вы хотите передать их (по порядку):

spark_df = spark_df.withColumn('get_profile', spark_udf(*spark_df.columns))

В более общем плане вы можете распаковать любой упорядоченный список столбцов:

cols_to_pass_to_udf = ['Foo', 'Bar', 'Baz']
spark_df = spark_df.withColumn('get_profile', spark_udf(*cols_to_pass_to_udf ))

Но для этой конкретной операции не требуется udf. Я бы сделал это так:

from pyspark.sql.functions import coalesce, when, col, lit

spark_df.withColumn(
    "get_profile",
    coalesce(*[when(col(c)==1, lit(c)) for c in spark_df.columns])
).show()
#+---+---+---+-----------+
#|Foo|Bar|Baz|get_profile|
#+---+---+---+-----------+
#|  0|  1|  0|        Bar|
#|  1|  0|  0|        Foo|
#|  1|  1|  1|        Foo|
#+---+---+---+-----------+

Это работает, потому что pyspark.sql.functions.when() будет возвращать значение null по умолчанию, если условие принимает значение False и не otherwise. Тогда в представлении списка pyspark.sql.functions.coalesce будет возвращен первый ненулевой столбец.

Обратите внимание, что это эквивалентно udf ТОЛЬКО, если порядок столбцов совпадает с get_profile который оценивается в функции get_profile. Чтобы быть более явным, вам следует:

spark_df.withColumn(
    "get_profile",
    coalesce(*[when(col(c)==1, lit(c)) for c in ['Foo', 'Bar', 'Baz'])
).show()
  • 0
    Понял - высоко ценится. Так что это упрощенная версия фактической функции. В действительности, существует намного больше столбцов / условий для присваивания значений, и вложенные, когда структура не будет практичной. Есть ли способ передать все столбцы в UDF в качестве аргумента?
  • 0
    @flyingmeatball вы также можете сделать spark_udf(*spark_df.columns) , но вы должны убедиться, что порядок столбцов совпадает с порядком аргументов вашего udf.
Показать ещё 2 комментария

Ещё вопросы

Сообщество Overcoder
Наверх
Меню