Pyspark：根据条件在列的所有值中添加零作为前缀

Question

我在 pyspark 数据框中有一列 A，其值如下所示

Column A      Column B
1234           AB
45685          CD
2345           KN
789899         MN
12335577       LN

如果值小于 8 位，我需要在 A 列的开头添加零以使值变为 8 位。 所以结果应该如下所示 -

Column A      Column B
00001234           AB
00045685           CD
00002345           KN
00789899           MN
12335577           LN

我怎样才能实现它？

Answer 1

使用内置函数lpad

df.withColumn("Column A", F.lpad(F.col('Column A'), 8, '0') ).show()

Answer 2

您可以采用两种方法：首先，编写一个 UDF。 其次，使用when函数。

UDF

如果您使用的是 UDF，您可以使用常规 Python 逻辑来添加预期的0数。

from pyspark.sql import functions as F
from pyspark.sql import types as T

@F.udf(returnType=T.StringType())
def get_string_to_length(text):
    nulls_to_fill = 8 - len(text) 
    return '0' * nulls_to_fill+text    

df = df.withColumn('Column A', get_string_to_length(F.col('Column A')))

当函数

使用 when 函数，您可以动态地将多个when语句相互链接。

from pyspark.sql import functions as F

when_function = F.when(
    F.length('Column A') == 8,
    F.col('Column A')
)

for i in range(1, 8):
    when_function = when_function.when(
        F.length('Column A') == F.lit(8 - i),
        F.concat(F.lit('0'*i), F.col('Column A'))
    )

df = df.withColumn('Column A', when_function)