在 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|
+-----+----+--------+