带有case语句的for循环

Question

我正在尝试寻找一种方法来运行 for 循环，以更好地优化我的案例语句的脚本。

下面显示的脚本没有错误，但是我觉得这太冗长了，可能会在下次维护时造成混乱。

df = df.withColumn('Product', when(df.where('input_file_name LIKE "%CAD%"'), 'Cash and DUE').
                   when(df.where('input_file_name LIKE "%TP%"'), 'Trade Product').
                   when(df.where('input_file_name LIKE "%LNS%"'), 'Corp Loans').
                   when(df.where('input_file_name LIKE "%DBT%"'), 'Debt').
                   when(df.where('input_file_name LIKE "%CRD%"'), 'Retail Cards').
                   when(df.where('input_file_name LIKE "%MTG%"'), 'Mortage').
                   when(df.where('input_file_name LIKE "%OD%"'), 'Overdraft').
                   when(df.where('input_file_name LIKE "%PLN%"'), 'Retail Personal Loan').
                   when(df.where('input_file_name LIKE "%CLN%"'), 'CLN').
                   when(df.where('input_file_name LIKE "%CAT%"'), 'Custody and Trust').
                   when(df.where('input_file_name LIKE "%DEP%"'), 'Deposits').
                   when(df.where('input_file_name LIKE "%STZ%"'), 'Securitization').
                   when(df.where('input_file_name LIKE "%SECZ%"'), 'Security Securitization').
                   when(df.where('input_file_name LIKE "%SEC%"'), 'Securities').
                   when(df.where('input_file_name LIKE "%MTSZ%"'), 'Retail Mortage Securitization').
                   when(df.where('input_file_name LIKE "%PLSZ%"'), 'Retail Personal Loan Securitization').
                   when(df.where('input_file_name LIKE "%CCSZ%"'), 'Retail Cards Securitization').
                   when(df.where('input_file_name LIKE "%CMN%"'), 'Cash Management').
                   when(df.where('input_file_name LIKE "%OTC%"'), 'Over-the-counter').
                   when(df.where('input_file_name LIKE "%SFT%"'), 'Securities Financing Transactions').
                   when(df.where('input_file_name LIKE "%ETD%"'), 'Excahnge Traded Deriative').
                   when(df.where('input_file_name LIKE "%DEF%"'), 'Default Products').
                   when(df.where('input_file_name LIKE "%FFS%"'), 'Not Required').
                   when(df.where('input_file_name LIKE "%hdfs%"'), 'Not Required').
                   otherwise('feed_name'));

我想过运行一个循环，下面是一个例子（脚本不正确，这是为了演示）

product_code = ['%CAD%','%TP%','%LNS%','%DBT%','%CRD%','%MTG%','%OD%','%PLN%','%CLN%','%CAT%','%DEP%','%STZ%','%SECZ%','%SEC%','%MTSZ%','%PLSZ%','%CCSZ%','%CMN%','%OTC%','%SFT%','%ETD%','%DEF%','%FFS%','%hdfs%']
product_name = ['Cash and Due','Trade Product','Corp Loans','Debt','Retail Cards','Mortage','Overdraft','Retail Personal Loan','CLN','Custody and Trust','Deposits','Securitization','Securities Securitization','Securities','Retail Mortage Securitization','Retail Personal Loan Securitization','Retail Cards Securitization','Cash Management','Over-the-counter','Securities Finanacing Transactions','Exchange Traded Derivative','Default Products','Not Required','Not Required']
   
##Both product_code & product name have the same number of index

lastIndex = len(product_code)    
    for x in product_code:
       # Logic i thought df.withColumn('Product', when(df.where('input_file_name LIKE "%'product_code[x]'%"'), product_name[x])
       if(product_code[lastIndex]): 
      #otherwise('feed_name')

如果在 spark 中可以为 when(df.where()).otherwise 的 case 语句运行循环，或者我可以查看另一种方法或用例，则需要一些建议

更新

我已经按照建议使用该方法实现了，查询在条件集上返回正确，但我想知道为什么它没有在下面的脚本中返回正确的值，否则称为 lit()，而是删除了不正确的行满足条件

Sample DF:
product_code = ['%CMN%','%TP%','%LNS%']
product_name = ['Cash and Due','Trade Product']
feed_name = ['farid','arshad','jimmy']   

df = spark.createDataFrame(
     list(zip(inp_file,feed_name)),
     ['input_file_name','feed_name']
)

+---------------+---------+
|input_file_name|feed_name|
+---------------+---------+
|sdasdasdasd    |bob      |
|_CMN_BD        |arshad   |
|_CMN_BD_WS     |jimmy    |
+---------------+---------+

product_code = ['%CAD_%','%TP%','%LNS%','%DBT%','%CRD%','%MTG%','%_OD_%','%PLN%','%CLN%','%CAT%','%DEP%','%STZ%','%SECZ%','%SEC%','%MTSZ%','%PLSZ%','%CCSZ%','%CMN%','%OTC%','%SFT%','%ETD%','%DEF%','%FFS%','%hdfs%']
product_name = ['Cash and Due','Trade Product','Corp Loans','Debt','Retail Cards','Mortage','Overdraft','Retail Personal Loan','CLN','Custody and Trust','Deposits','Securitization','Securities Securitization','Securities','Retail Mortage Securitization','Retail Personal Loan Securitization','Retail Cards Securitization','Cash Management','Over-the-counter','Securities Finanacing Transactions','Exchange Traded Derivative','Default Products','Not Required','Not Required']
   
## -- Create spark dataframe and with list tuple     
## -- Lit is used to add new column

product_ref_df = spark.createDataFrame(
     list(zip(product_code, product_name)),
     ["product_code", "product_name"]
)
    

def tempDF(df,targetField,columnTitle,condition,targetResult,show=False):
    product_ref_df = spark.createDataFrame(
         list(zip(condition,targetResult)),
         ["condition", "target_result"]
    )
    
    df.join(broadcast(product_ref_df), expr(""+targetField+" like condition")) \
    .withColumn(columnTitle, coalesce(col("target_result"), lit("feed_name"))) \
    .drop('condition','target_result') \
    .show()
    
    return df

product_ref_df = tempDF(df,'input_file_name','Product',product_code,product_name)

脚本触发时，没有错误，返回结果如图，

+---------------+---------+------------+
|input_file_name|feed_name|     Product|
+---------------+---------+------------+
|        _CMN_BD|   arshad|Cash and Due|
|     _CMN_BD_WS|    jimmy|Cash and Due|
+---------------+---------+------------+

由于我们没有删除任何行，因此结果是否应该包含第一行，

+---------------+---------+------------+
|input_file_name|feed_name|     Product|
+---------------+---------+------------+
|    sdasdasdasd|      bob|bob         |
|        _CMN_BD|    jimmy|Cash and Due|
|     _CMN_BD_WS|    jimmy|Cash and Due|
+---------------+---------+------------+
+---------------+---------+------------+

Answer 1

您可以使用coalesce 将所有when 语句合并为一个。 coalesce 将选择第一个非空列，when 仅在条件匹配时才给出非空值；否则它给出 null （没有otherwise 条件）。

product_code = ['%CAD%','%TP%','%LNS%','%DBT%','%CRD%','%MTG%','%OD%','%PLN%','%CLN%','%CAT%','%DEP%','%STZ%','%SECZ%','%SEC%','%MTSZ%','%PLSZ%','%CCSZ%','%CMN%','%OTC%','%SFT%','%ETD%','%DEF%','%FFS%','%hdfs%']
product_name = ['Cash and Due','Trade Product','Corp Loans','Debt','Retail Cards','Mortage','Overdraft','Retail Personal Loan','CLN','Custody and Trust','Deposits','Securitization','Securities Securitization','Securities','Retail Mortage Securitization','Retail Personal Loan Securitization','Retail Cards Securitization','Cash Management','Over-the-counter','Securities Finanacing Transactions','Exchange Traded Derivative','Default Products','Not Required','Not Required']

import pyspark.sql.functions as F

df2 = df.withColumn(
    'Product',
    F.coalesce(
        *[F.when(F.col('input_file_name').like(code), F.lit(name))
        for (code, name) in zip(product_code, product_name)]
    )
)

Answer 2

您可以从这些产品名称引用创建一个新的 DataFrame 并与原始 df 连接以获取产品名称：

from pyspark.sql.functions import expr, col, broadcast, coalesce

product_ref_df = spark.createDataFrame(
     list(zip(product_code, product_name)),
     ["product_code", "product_name"]
)

df.join(broadcast(product_ref_df), expr("input_file_name like product_code"), "left") \
  .withColumn("Product", coalesce(col("product_name"), col("Feed_name"))) \
  .drop("product_code", "product_name") \
  .show()

或使用functools.reduce 链接case/when 条件，如下所示：

import functools

from pyspark.sql.functions import lit, col, when

case_conditions = list(zip(product_code, product_name))

product_col = functools.reduce(
    lambda acc, x: acc.when(col(f"input_file_name").like(x[1]), lit(x[1])),
    case_conditions[1:],
    when(col("input_file_name").like(case_conditions[0][0]), lit(case_conditions[0][1]))
).otherwise(col("Feed_name"))

df.withColumn("Product", product_col).show()