如何在PySpark中通过多OR条件过滤并修改ZIPCODE字段？

原需求逻辑：当COUNTRY_TABLE表中STATE字段值为TN、DEL、UK、UP、HP、JK、MP，且ZIPCODE字段长度小于5时，将ZIPCODE设为"0"。

原SQL代码 #

UPDATE COUNTRY_TABLE SET COUNTRY_TABLE.ZIPCODE = "0" 
WHERE (((COUNTRY_TABLE.STATE)="TN" Or
(COUNTRY_TABLE.STATE)="DEL" Or
(COUNTRY_TABLE.STATE)="UK" Or 
(COUNTRY_TABLE.STATE)="UP" Or 
(COUNTRY_TABLE.STATE)="HP" Or 
(COUNTRY_TABLE.STATE)="JK" Or 
(COUNTRY_TABLE.STATE)="MP") AND ((Len([ZIPCODE]))<"5"));

尝试的PySpark代码 #

df=df.withColumn('ZIPCODE', substring('ZIPCODE', 1,5)) -- take only first five character
df=df.withColumn('length_ZIP',length(df.ZIPCODE))

df=df.withColumn('ZIPCODE', F.when( (col('COUNTRY_TABLE.STATE') == 'TN') | (col('COUNTRY_TABLE.STATE') == 'DEL') \
| (col('COUNTRY_TABLE.STATE') == 'UK') | (col('COUNTRY_TABLE.STATE') == 'UP') | (col('COUNTRY_TABLE.STATE') == 'HP')  \
| (col('COUNTRY_TABLE.STATE') == 'JK') | (col('COUNTRY_TABLE.STATE') == 'MP') & (col('length_ZIP') < '5'), '0')
.otherwise(df.ZIPCODE))

问题分析与修正代码 #

错误点说明 #

1. 字段引用错误：PySpark中无需加表名前缀COUNTRY_TABLE.，直接使用字段名STATE即可
2. 运算符优先级问题：&的优先级高于|，所有STATE的条件需要用括号整体包裹后，再和长度条件做与运算
3. 类型匹配错误：长度是数值类型，应和数字5比较，而非字符串'5'
4. 冗余操作：原SQL是判断原ZIPCODE的长度，无需先截取前5位，该操作会破坏原本不需要更新的ZIPCODE值

修正后的PySpark代码 #

from pyspark.sql import functions as F
from pyspark.sql.functions import col, length

# 定义目标STATE列表，简化条件判断
target_states = ["TN", "DEL", "UK", "UP", "HP", "JK", "MP"]

# 直接更新ZIPCODE字段
df = df.withColumn(
    "ZIPCODE",
    F.when(
        (col("STATE").isin(target_states)) & (length(col("ZIPCODE")) < 5),
        "0"
    ).otherwise(col("ZIPCODE"))
)

代码说明 #

• 用isin()替代多个==和|，代码更简洁易维护
• 修正逻辑运算符优先级，确保条件判断符合原SQL逻辑
• 直接基于原ZIPCODE字段计算长度，保留无需更新的原始值
• 去掉不必要的中间字段，减少冗余操作

内容的提问来源于stack exchange，提问作者BigData Lover