PySpark 中的 Python 字典查找

Question

在 PySpark 中遇到以下问题。我有一本 Python 字典，如下所示：

COUNTRY_MAP = {
    "AND": "AD", "ARE": "AE", "AFG": "AF", "ATG": "AG", "AIA": "AI", ... };

我现在想要建立一个由 3 列组成的值，例如 value1、value2 和 value3。问题是 value3 需要使用上面的查找将 3 字母代码转换为 2 字母代码，如果不存在，则应使用“NONE”，即

from pyspark.sql import functions as sf

combined = sf.trim(sf.concat(sf.col("value1"), sf.lit(":"), sf.col("value2"), sf.lit(":"),
                                 sf.coalesce(sf.col("value3"), "NONE")))

tmp = (df.withColumn('COMBINED_FIELD', combined)
       ...<other stuff>
       )

这给了我像“abc:4545:AND”、“def:7789:ARE”和“ghi:1122:NONE”这样的值。我现在需要：“abc:4545:AD”、“def:7789:AE”和“ghi:1122:NONE”。

作为 PySpark 的新手，我真的很难让它发挥作用。你知道吗？

Answer 1

您可以将字典转换为映射类型列并使用value3作为键获取值：

import pyspark.sql.functions as F

COUNTRY_MAP = {"AND": "AD", "ARE": "AE", "AFG": "AF", "ATG": "AG", "AIA": "AI"}

result = df.withColumn(
    'combined_field', 
    F.trim(
        F.concat_ws(':', 
            'value1', 'value2', 
            F.coalesce(
                F.create_map(*sum([[F.lit(k), F.lit(v)] for (k,v) in COUNTRY_MAP.items()], []))[F.col('value3')], 
                F.lit('NONE')
            )
        )
    )
)

result.show()
+------+------+------+--------------+
|value1|value2|value3|combined_field|
+------+------+------+--------------+
|   abc|  4545|   AND|   abc:4545:AD|
|   def|  7789|   ARE|   def:7789:AE|
|   ghi|  1122|  NONE| ghi:1122:NONE|
+------+------+------+--------------+