PySpark 中的 Python 字典查找

Posted 2023-04-15

【中文标题】PySpark 中的 Python 字典查找

【英文标题】：Python dictionary lookup in PySpark

【发布时间】：2021-05-03 15:12:45

【问题描述】：

在 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 的新手，我真的很难让它发挥作用。你知道吗？

【参考方案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|
+------+------+------+--------------+