根据第一个字符的出现，将字符串列分开

Question

我希望将 Spark DataFrame 的列根据第一次出现的字符（在本例中是下划线“_”）分成两个不同的列。

我准备了一个可以100％复制的示例：

模拟的 Spark DataFrame 如下：

df = spark.createDataFrame(
    [
     (1, 1.8, 'newyork_3434_north'), 
     (4, 2.6, 'la_432432432_south'), 
     (6, 3.3, 'boston_234324_east'), 
     (8, 4.1, 'detroit_6757_west'), 
     (2, 5.7, 'miami_133123_north'), 
     (3, 6.2, 'atlanta_093394_west'), 
     (1, 6.1, 'houston_87342_east')
    ],
    ('ranking', "coordenate", "city")
)

上述代码创建了一个表格，格式如下：

ranking  coordenate  city
1        1.8         newyork_3434_north
4        2.6         la_432432432_south
6        3.3         boston_234324_east
8        4.1         detroit_6757_west
2        5.7         miami_133123_north
3        6.2         atlanta_093394_west 
1        6.1         houston_87342_east

我希望做的是根据从左到右第一个下划线的位置，将城市列分为两个不同的列。

最终期望的表格应该像这样：

ranking  coordenate  city       code
1        1.8         newyork    3434_north
4        2.6         la         432432432_south
6        3.3         boston     234324_east
8        4.1         detroit    6757_west
2        5.7         miami      133123_north
3        6.2         atlanta    093394_west
1        6.1         houston    87342_east

我看到了关于这个主题的几个帖子，但它们并没有谈论字符第一次出现的情况(链接1，链接2等)，而是根据字符串上的所有特定字符进行分割；或者通过字符在字符串中的特定位置来分割。

我还尝试了Python Pandas方法，但是如预期的那样，在PySpark中不能通过扩展或类比应用该方法 (链接3)。

非常感谢您的帮助。

- NuValue

---

你的数据框有多大？你真的需要用Spark吗？ - acushner

嗨@acushner，感谢您的回复。我同意您的观点，这似乎非常容易，在Pandas中解决方案很简单（我已经知道如何在本地执行），但数据集大小为123 GB。我已经为此目的调整了一个非常简单易懂的SO数据集，但它是一个巨大的数据集，即使只选择我想要的列（“city”）。提前致谢！ - NuValue

使用 split 和 concat_ws 函数。 - Ramesh Maharjan

是否总是恰好有两个下划线？ - pault

嗨@pault，答案是否定的。这只是我举例说明时的巧合；我在我的SparkDF中看到了有超过2个下划线的注册表。谢谢你的问题。 - NuValue

Answer 1

我认为在这里最好的选择是使用 pyspark.sql.functions.regexp_extract() 和 pyspark.sql.functions.regexp_replace()：

import pyspark.sql.functions as f

df.select(
    "ranking",
    "coordenate",
    f.regexp_extract("city", pattern="^[A-Za-z]+(?=_)", idx=0).alias('city'),
    f.regexp_replace("city", "^[A-Za-z]+_", "").alias("code")
).show()
#+-------+----------+----------+---------------+
#|ranking|coordenate|      city|           code|
#+-------+----------+----------+---------------+
#|      1|       1.8|   newyork|     3434_north|
#|      4|       2.6|        la|432432432_south|
#|      6|       3.3|    boston|    234324_east|
#|      8|       4.1|   detroit|      6757_west|
#|      2|       5.7|     miami|   133123_north|
#|      3|       6.2|   atlanta|    093394_west|
#|      1|       6.1|   houston|     87342_east|
#+-------+----------+----------+---------------+

在这两种情况下，模式本质上是相同的：

- `^[A-Za-z]+`：匹配字符串开头的任意数量字母 - `(?=_)`：正向先行断言匹配下划线

对于“城市”，我们找到这个模式并提取第一个匹配项。对于“代码”，我们将“look-ahead”更改为匹配，并用空字符串替换该模式。

---

如果难以找到合适的正则表达式模式，则以下方法是 Spark 2.1 及以上版本中可行的替代方法：

- 获取“城市”很简单 - 可以使用 `pyspark.sql.functions.split()` 在下划线上拆分字符串，然后使用 `getItem(0)` 获取分割列表的第一个元素。 - 对于“代码”部分，将“城市”拆分为下划线，并使用 `pyspark.sql.functions.posexplode()` 将结果数组展开。然后过滤出 `pos > 0`，按原始列分组，并使用 `pyspark.sql.functions.concat_ws` 连接收集到的标记。

df.select(
        "*",
        f.posexplode(f.split("city", "_")).alias("pos", "token")
    )\
    .where("pos > 0")\
    .groupBy("ranking", "coordenate", "city")\
    .agg(f.concat_ws("_" ,f.collect_list("token")).alias("code"))\
    .select(
        "ranking",
        "coordenate",
        f.split("city", "_").getItem(0).alias("city"),
        "code"
    )\
    .show()

- pault

---

感谢 @pault 的反馈。非常棒的答案。我会将其标记为问题的答案。 - NuValue

Answer 2

@pault已经使用regex、split和concat_ws内置函数给出了一个很棒的答案。

这里有一个简单的替代方法，可以使用udf函数来实现：

from pyspark.sql.functions import col, udf
from pyspark.sql.types import ArrayType, StringType

@udf(ArrayType(StringType()))
def splitUdf(x):
    splitted = x.split('_')
    return [splitted[0], '_'.join(splitted[1:])]

df.withColumn('city', splitUdf(col('city')))\
    .select(col('ranking'), col('coordenate'), col('city')[0].alias('city'), col('city')[1].alias('code'))\
    .show()

这应该给你

+-------+----------+-------+---------------+
|ranking|coordenate|   city|           code|
+-------+----------+-------+---------------+
|      1|       1.8|newyork|     3434_north|
|      4|       2.6|     la|432432432_south|
|      6|       3.3| boston|    234324_east|
|      8|       4.1|detroit|      6757_west|
|      2|       5.7|  miami|   133123_north|
|      3|       6.2|atlanta|    093394_west|
|      1|       6.1|houston|     87342_east|
+-------+----------+-------+---------------+

我希望这个答案对你有所帮助

- Ramesh Maharjan