如何将DataFrame的Key列子串按x/a拆分至指定列？

我有如下DataFrame：

+--------+
|     Key|
+--------+
|  x10x60|
|x1x19x33|
|   x12x6|
|   a14x4|
|x1x1x1x6|
|x2a23x30|
+--------+

我希望得到如下输出：将Key列按每个x元素拆分，依次放入xa/xb/xc/xd列；若存在a元素，则将其依次放入ta/tb/tc/td列。

+--------+-----+------+-----+-----+-----+----+----+-----+
|     Key|   xa|    xb|   xc|   xd|   ta|  tb|  tc|   td|
+--------+-----+------+-----+-----+-----+----+----+-----+
|  x10x60|  x10|   x60|     |     |     |    |    |     |
|x1x19x33|   x1|   x19|  x33|     |     |    |    |     |
|   x12x6|  x12|    x6|     |     |     |    |    |     |
|   a14x4|     |    x4|     |     |  a14|    |    |     |
|x1x1x1x6|   x1|    x1|   x1|   x6|     |    |    |     |
|x2a23x30|   x2|      |  x30|     |     | a23|    |     |
+--------+-----+------+-----+-----+-----+----+----+-----+

我尝试使用substr()或substring()函数，但无法得到预期输出，拆分过程出现问题。

可以通过正则表达式提取目标片段，拆分分类后展开为对应列，具体步骤如下：

1. 提取所有x/a开头的数字片段 #

用regexp_extract_all函数从Key列中提取所有符合[xa]\d+格式的片段：

from pyspark.sql import functions as F

# 假设原DataFrame名为df
df = df.withColumn("all_parts", F.regexp_extract_all("Key", r"[xa]\d+"))

2. 拆分x和a类型的片段 #

分别过滤出x开头和a开头的片段，生成两个独立的列表列：

df = df.withColumn(
    "x_parts",
    F.expr("filter(all_parts, part -> startsWith(part, 'x'))")
).withColumn(
    "a_parts",
    F.expr("filter(all_parts, part -> startsWith(part, 'a'))")
)

3. 将列表展开为指定列 #

用element_at函数按位置取出列表元素，对应到xa/xb/xc/xd和ta/tb/tc/td列（位置从1开始，超出列表长度则返回空）：

# 处理x系列列
for idx, col_name in enumerate(["xa", "xb", "xc", "xd"], start=1):
    df = df.withColumn(col_name, F.element_at("x_parts", idx))

# 处理a系列列
for idx, col_name in enumerate(["ta", "tb", "tc", "td"], start=1):
    df = df.withColumn(col_name, F.element_at("a_parts", idx))

# 清理中间临时列
df = df.drop("all_parts", "x_parts", "a_parts")

执行以上代码后，即可得到符合需求的输出格式。

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