使用 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()

使用 MapType 文字创建新列

Using MapType literal to create new column

python

apache-spark

apache-spark-sql

pyspark

pyspark-sql