pyspark: withcolumn Ошибка анализа после преобразования столбца в нижний регистр

Question

У меня есть фрейм данных, который выглядит следующим образом:

+------------+------+
|        food|pounds|
+------------+------+
|       bacon|   4.0|
|STRAWBERRIES|   3.5|
|       Bacon|   7.0|
|STRAWBERRIES|   3.0|
|       BACON|   6.0|
|strawberries|   9.0|
|Strawberries|   1.0|
|      pecans|   3.0|
+------------+------+

И ожидаемый результат равен

+------------+------+---------+
|        food|pounds|food_type|
+------------+------+---------+
|       bacon|   4.0|     meat|
|STRAWBERRIES|   3.5|    fruit|
|       Bacon|   7.0|     meat|
|STRAWBERRIES|   3.0|    fruit|
|       BACON|   6.0|     meat|
|strawberries|   9.0|    fruit|
|Strawberries|   1.0|    fruit|
|      pecans|   3.0|    other|
+------------+------+---------+

Так что я по сути определил new_column на основе моей логики и применил его к .withcolumn

new_column = when((col('food') == 'bacon') | (col('food') == 'BACON') | (col('food') == 'Bacon'), 'meat'
                   ).when((col('food') == 'STRAWBERRIES') | (col('food') == 'strawberries') | (col('food') == 'Strawberries'), 'fruit'
                   ).otherwise('other')

А потом

df.withColumn("food_type", new_column).show()

Что отлично работает.Но я хотел обновить оператор new_column с меньшим количеством кода, поэтому переписать, как показано ниже

new_column = when(lower(col('food') == 'bacon') , 'meat'
                   ).when(lower(col('food') == 'strawberries'), 'fruit'
                   ).otherwise('other')

Теперь, когда я df.withColumn("food_type", new_column).show()

я получаю ошибку

AnalysisException: "cannot resolve 'CASE WHEN lower(CAST((`food` = 'bacon') AS STRING)) THEN 'meat' WHEN lower(CAST((`food` = 'strawberries') AS STRING)) THEN 'fruit' ELSE 'other' END' due to data type mismatch: WHEN expressions in CaseWhen should all be boolean type, but the 1th when expression's type is lower(cast((food#165 = bacon) as string));;\n'Project [food#165, pounds#166, CASE WHEN lower(cast((food#165 = bacon) as string)) THEN meat WHEN lower(cast((food#165 = strawberries) as string)) THEN fruit ELSE other END AS food_type#197]\n+- Relation[food#165,pounds#166] csv\n"

Чего мне не хватает?

Answer 1

упрощено:

new_column = когда (нижний (col («еда»)) == «бекон», «мясо»). Когда (нижний (col («еда»)) == 'клубника',' fruit '). иначе (' other ')

df.withColumn ("food_type", new_column) .show ()

Answer 2

Я хотел бы поделиться другим подходом, который более похож на SQL-запросы и может быть более подходящим для более сложных и вложенных условий.

from pyspark.sql.functions import *
cond = """case when lower(food) in ('bacon') then 'meat'
            else case when lower(food) in ('strawberries') then 'fruit'
                 else 'other'
                end
            end"""

newdf = df.withColumn("food_type", expr(cond))

Надеюсь, это поможет.

С уважением,

Neeraj

Answer 3

Ваши скобки не совпадают.

new_column = when(lower(col('food')) == 'bacon' , 'meat').when(lower(col('food')) == 'strawberries', 'fruit').otherwise('other')