PySpark拆分含C/D标识的Transaction列为金额与借贷标识列

现有如下测试数据，需要将Transaction列拆分为数值型交易金额Amount、末尾C/D标识的借贷标记CreditOrDebit两列：

+-----+--------------------+-----------+
|Sr No|             User Id|Transaction|
+-----+--------------------+-----------+
|1    |paytm 111002203@p...|       100D|
|2    |paytm 111002203@p...|        50C|
|3    |paytm 111002203@p...|        20C|
|4    |paytm 111002203@p...|        10C|
|5    |                null|         1C|
+-----+--------------------+-----------+

初始固定位置截取金额的逻辑因为金额长度不固定，返回全空值；尝试用instr动态定位截取位置时触发TypeError: Column is not iterable报错。

• 代码里写的'C' or 'D'是Python原生布尔逻辑判断，返回值固定为'C'，不属于Spark SQL支持的列运算表达式，无法被Spark解析执行
• 即使instr部分不报错，substring函数的第三个入参是截取长度，不是结束位置，直接传入instr返回的字符位置会把C/D标识也截进金额字段，遇到D结尾的行还会因为instr找不到C返回0，得到空值。

方案1：按字符串长度动态截取（适配当前固定末尾为标识的场景，性能最好） #

你之前取最后1位作为借贷标识的逻辑完全正确，金额部分只需要截取从第一位开始、长度为总字符串长度减1的子串即可，不需要额外做字符定位：

from pyspark.sql.functions import substring, length, col

df_result = df_sample.withColumn(
    "CreditOrDebit", substring("Transaction", -1, 1)
).withColumn(
    "Amount", substring("Transaction", 1, length("Transaction") - 1).cast("int")
)

df_result.show()

执行返回结果：

+-----+--------------------+-----------+-------------+------+
|Sr No|             User Id|Transaction|CreditOrDebit|Amount|
+-----+--------------------+-----------+-------------+------+
|    1|paytm 111002203@p...|       100D|            D|   100|
|    2|paytm 111002203@p...|        50C|            C|    50|
|    3|paytm 111002203@p...|        20C|            C|    20|
|    4|paytm 111002203@p...|        10C|            C|    10|
|    5|                null|         1C|            C|     1|
+-----+--------------------+-----------+-------------+------+

方案2：正则提取（适配格式不规则的场景，鲁棒性最强） #

如果后续数据格式变动，借贷标识不一定在字符串末尾，可以用正则表达式分别提取数字部分和C/D标识，不需要依赖固定位置：

from pyspark.sql.functions import regexp_extract

df_result = df_sample.withColumn(
    "CreditOrDebit", regexp_extract("Transaction", r"([CD])", 1)
).withColumn(
    "Amount", regexp_extract("Transaction", r"(\d+)", 1).cast("int")
)

方案3：修正instr逻辑写法 #

如果需要保留instr定位的实现思路，可以分别计算C、D两个字符的位置，取非零的位置值减1作为截取长度：

from pyspark.sql.functions import substring, instr, greatest

df_result = df_sample.withColumn(
    "CreditOrDebit", substring("Transaction", -1, 1)
).withColumn(
    "Amount",
    substring(
        "Transaction",
        1,
        greatest(instr("Transaction", "C"), instr("Transaction", "D")) - 1
    ).cast("int")
)

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