从字典中的数组列中查找值Pyspark

7qhs6swi  于 5个月前  发布在  Spark
关注(0)|答案(1)|浏览(45)

我在Pyspark中有这样一个框架:

data = [("definitely somewhere",), ("Las Vegas",), ("其他",), (None,), ("",), ("Pucela Madrid Langreo, España",), ("Trenches, With Egbon Adugbo",)]
df = spark.createDataFrame(data, ["address"])
city_country = {
    'las vegas': 'US',
    'lagos': 'NG',
    'España': 'ES'
}
cities_name_to_code = spark.sparkContext.broadcast(city_country )
df_with_codes = df.withColumn('cities_array', F.lower(F.col('address'))) \
                          .withColumn('cities_array', F.split(F.col('cities_array'), ', '))

字符串
我想在cities_array中找到cities_name_to_code中每个元素的所有键(得到一个值数组)。问题是我不想使用UDF。

kuhbmx9i

kuhbmx9i1#

对于这个用例,你可以使用transform高阶函数,并将case作为它内部的函数传递。
这里有一个例子

# create case when builder function
case_whens = lambda c: reduce(lambda x, y: x.when(c == y[0].lower(), y[1]), city_country.items(), func)

# test case when builder
# case_whens(func.lit('bork'))
# Column<'CASE WHEN (bork = las vegas) THEN US WHEN (bork = lagos) THEN NG WHEN (bork = españa) THEN ES END'>

# use case when inside the `transform`
df_with_codes_sdf = data_sdf. \
    withColumn('cities_array', func.lower(func.col('address'))). \
    withColumn('cities_array', func.split(func.col('cities_array'), ', ')). \
    withColumn('city_codes_array', func.transform('cities_array', lambda a: case_whens(a))). \
    show(truncate=False)

# +-----------------------------+-------------------------------+----------------+
# |address                      |cities_array                   |city_codes_array|
# +-----------------------------+-------------------------------+----------------+
# |definitely somewhere         |[definitely somewhere]         |[null]          |
# |Las Vegas                    |[las vegas]                    |[US]            |
# |其他                         |[其他]                         |[null]          |
# |null                         |null                           |null            |
# |                             |[]                             |[null]          |
# |Pucela Madrid Langreo, España|[pucela madrid langreo, españa]|[null, ES]      |
# |Trenches, With Egbon Adugbo  |[trenches, with egbon adugbo]  |[null, null]    |
# +-----------------------------+-------------------------------+----------------+

字符串

相关问题