在 PySpark Dataframe 中新建一列，其中包含另一列的 Substring

Question

我有一個 dataframe，在該 dataframe 中有一列包含字符串值。 每當存在一組特定字符時，我需要從該列中提取 substring 並將其轉換為新列。 我希望能夠做到這一點，同時也不過濾，以免丟失所有其他行。 對於不包含該特定字符串值的任何行，我希望相應的新列讀作“null”。 例如，假設我有以下 dataframe：

+---------------------------------------+----------+---------+
|id                                     |compliance|workflow |
+---------------------------------------+----------+---------+
|account/product/rule-id/r-1879bajhdfd80|PASS      |      NEW|
|account/product/rule-id/r-198Hhfu89421s|PASS      |      NEW|
|account/product/test/run/date/YYYYMMDD |FAIL      |      NEW|
+---------------------------------------+----------+---------+

我希望能夠識別 substring 'rule-id' 並創建一個名為 'rule-id' 的新列，對於沒有 substring 的行，我希望該值為 null。例如 output 應該看起來像這樣：

+---------------------------------------+----------+---------+---------------+
|id                                     |compliance|workflow |rule-id        |
+---------------------------------------+----------+---------+---------------+
|account/product/rule-id/r-1879bajhdfd80|PASS      |      NEW|r-1879bajhdfd80|
|account/product/rule-id/r-198Hhfu89421s|PASS      |      NEW|r-198Hhfu89421s|
|account/product/test/run/date/YYYYMMDD |FAIL      |      NEW|null           |
+---------------------------------------+----------+---------+---------------+

我知道我可以使用 substring() function 來提取我想要的字符串部分，但這將對所有行執行此操作，給我一些奇怪的規則 ID 值。

df2 = df1.withColumn("rule-id", substring("id", 25, 15))

但是，我該如何編寫它，以便如果 'rule-id' substring 出現在 'id' 字符串值中，它會提取我正在尋找的 substring 但僅針對那些行，而 rest 將獲得“空”值新的“規則 ID”列？

Answer 1

如果它只需要處理描述的情況（所以 id 的長度不會改變並且模式將一直相似）你可以添加 when/otherwise 另一個 substring 檢查

import pyspark.sql.functions as F
from pyspark.sql import Window

inputData = [
    ("account/product/rule-id/r-1879bajhdfd80", "PASS", "NEW"),
    ("account/product/rule-id/r-198Hhfu89421s", "PASS", "NEW"),
    ("account/product/test/run/date/YYYYMMDD", "FAIL", "NEW"),
]
df1 = spark.createDataFrame(inputData, schema=["id", "compliance", "workflow"])
df1.show()

df2 = df1.withColumn(
    "rule-id",
    F.when(
        F.substring("id", 17, 7) == F.lit("rule-id"), F.substring("id", 25, 15)
    ).otherwise(None),
).show()

output

+--------------------+----------+--------+---------------+
|                  id|compliance|workflow|        rule-id|
+--------------------+----------+--------+---------------+
|account/product/r...|      PASS|     NEW|r-1879bajhdfd80|
|account/product/r...|      PASS|     NEW|r-198Hhfu89421s|
|account/product/t...|      FAIL|     NEW|           null|
+--------------------+----------+--------+---------------+

如果它應該更靈活先 substring 我們應該改變模式檢查