Использование литерала MapType для создания нового столбца

Question

У меня есть следующее pyspark.DataFrame

+---+--------+--------+--------------+
|SEX|_AGEG5YR|_IMPRACE|       _LLCPWT|
+---+--------+--------+--------------+
|  2|    11.0|     1.0| 79.4259469451|
|  2|    10.0|     1.0| 82.1648291655|
|  2|    11.0|     2.0| 55.7851100058|
|  2|    13.0|     1.0|115.9818718258|
|  2|    12.0|     1.0|194.7566575195|
+---+--------+--------+--------------+

Я хочу создать новый столбец на основе столбца SEX В соответствии с предложением из этого предыдущего ответа , я определили литерал MapType следующим образом

brfss_mapping = {
    "SEX": {
        1: "Male",
        2: "Female",
        9: "Refused"
    }
}
brfss_sex_mapping = create_map(
    [lit(x) for x in chain(*brfss_mapping["SEX"].items())]
)

Теперь, когда я использую withColumn и brfss_sex_mapping.getItem(...) с постоянным значением, например ниже

brfss_dmy = brfss_dmy.withColumn(
    "SEX_2",
    brfss_sex_mapping.getItem(1)
)

, я получаю ожидаемый результат

+---+--------+--------+--------------+-----+                                    
|SEX|_AGEG5YR|_IMPRACE|       _LLCPWT|SEX_2|
+---+--------+--------+--------------+-----+
|  1|    13.0|     1.0|381.8001043164| Male|
|  2|    10.0|     1.0| 82.1648291655| Male|
|  1|    11.0|     1.0|279.1864457296| Male|
|  1|    10.0|     1.0| 439.024136158| Male|
|  2|     8.0|     1.0| 372.921644978| Male|
+---+--------+--------+--------------+-----+

Однако, когда я пытаюсь пропустить соответствующий столбец следующим образом (снова, как это предлагается в предыдущем ответе)

brfss_dmy = brfss_dmy.withColumn(
    "SEX_2",
    brfss_sex_mapping.getItem(col("SEX"))
)

, я получаю следующее java.lang.RuntimeException: Unsupported literal type class org.apache.spark.sql.Column SEX

Answer 1

Похоже, что в Spark 3.0 мы больше не можем передавать столбец в функцию getItem, но я не смог найти никаких ссылок в коде или документации.

Вы можете использовать element_at вместо:

df.withColumn("SEX_2", element_at(brfss_sex_mapping, col("SEX")).show()

Или получить доступ к значению в виде массива:

df.withColumn("SEX_2", brfss_sex_mapping[col("SEX")]).show()

В Scala:

df.withColumn("SEX_2", element_at(brfss_sex_mapping, $"SEX")).show()
df.withColumn("SEX_2", brfss_sex_mapping($"SEX")).show()