Я пытаюсь создать новый столбец на фреймворке данных на основе значений некоторых столбцов. Он возвращает нуль во всех случаях. Кто-нибудь знает, что происходит с этим простым примером?
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())
к тому же эффекту.
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()
spark_udf(*spark_df.columns)
, но вы должны убедиться, что порядок столбцов совпадает с порядком аргументов вашего udf.
'Foo' != 1
, то же самое для других условий. Вот почему вы ничего не получите. UDF ожидает столбцы в качестве аргументов, в то время какget_profile
имеет нулевые аргументы.