首页 > 解决方案 > 循环一次 pyspark 数据框以查找包含值列表的行

问题描述

我在pyspark中有一个类似于下面的数据框

项目名称 项目编号
待定加拿大 10000000000029
待定中国 10000000000033
待定美国 10000000000974
美国潜水员 10000000000234
花2.0 10000000023947

我需要从包含 TBD 的 ProjectName 中提取国家/地区。我的维度表中还有一个国家/地区名称列表 country_list = ['Canada','United States',....., 'China'] 我在此列表中有大约 75 个国家/地区

我最初尝试使用 for 循环来解决它,我写了以下内容

country_list = ['Canada','United States',....., 'China']
# Set initial ExtractCountry column
df = df.withColumn("ExtractColumn", lit(None).cast("string"))
# Loop through data frame
for country in country_list:
    df = df.withColumn("ExtractColumn", when(df.ProjectName.contains("TBD") 
                                             & df.ProjectName.contains(country), lit(country))
                                       .otherwise(F.col("ExtractColumn")))

它可以工作,但运行需要很长时间,因为我的 df 太大了,这意味着它运行的数据帧长度为 75 个国家/地区。我知道 pyspark 没有像 pandas 数据帧那样的索引(我必须使用 pyspark 不能使用 pandas),但无论如何我是否可以只循环一次数据帧或以其他方式减少运行时间?

标签: pyspark

解决方案


你可以删除TBDfrom projectnamewithregexp_replace

from pyspark.sql import functions as F

(df
    .withColumn('ExtractColumn', F
        .when(F.col('projectname').startswith('TBD'), F.regexp_replace('projectname', 'TBD ', ''))
    )
    .show(10, False)
)

# Output
# +------------------+--------------+--------------+
# |projectname       |projectnumber |ExtractColumn |
# +------------------+--------------+--------------+
# |TBD Canada        |10000000000029|Canada        |
# |TBD China         |10000000000033|China         |
# |TBD United Kingdom|10000000000033|United Kingdom|
# |US Diver          |10000000000234|null          |
# |Flower 2.0        |10000000023947|null          |
# +------------------+--------------+--------------+

推荐阅读