在 Pyspark 中创建一个新列，该列在另一个可用列上计算

Question

我有一个如下的数据框：

>>> df
   Name    id
0   Tom   103
1  Jack   109
2  nick  9518
3  juli  1890

我想创建一个作为super_id的 ne 列，即 i) 如果 id 是 3 位数字，则 super_id 为零加上第一个整数 ii) 如果 id 是 4 位数字，则 super id 是前两个整数。

>>> df
   Name    id super_id
0   Tom   103       01
1  Jack   109       01
2  nick  9518       95
3  juli  1890       18

我有以下相同的python代码，但不确定如何将其转换为pyspark代码。

import pandas as pd

# initialise data of lists.
data = {'Name':['Tom', 'Jack', 'nick', 'juli'],
        'id':[103, 109, 9518, 1890]}
 
# Creates pandas DataFrame.
df = pd.DataFrame(data)

#Create super id
df['super_id'] = df.id.astype('int').astype('str').str.zfill(4).str[0:2]

在 pyspark 中尝试出错

df= df.withColumn('super_id', df['id'].astype('int').astype('str').str.zfill(4).str[0:2])

Answer 1

您需要使用spark 函数来做到这一点：

from pyspark.sql import functions as F

df.withColumn("super_id", F.substring(F.lpad("id", 4, "0"), 0, 2)).show()

+-----+----+--------+
| name|  id|super_id|
+-----+----+--------+
|  Tom| 103|      01|
| jack| 109|      01|
| nick|9518|      95|
|julie|1890|      18|
+-----+----+--------+