如何从 PySpark DataFrame 的列中分离特定字符并使用它们形成一个新列？

Question

我对 PySpark 非常陌生，我真的很难理解如何使用它。我有一个案例，我必须对仅包含 column1 的 PySpark DataFrame 应用操作，如下所示，并返回一个包含 column1、column2 和 column3 的新数据帧：

column1              column2      column3
'A123'                '123'         'A'
'321B'   --------->   '321'         'B'
'C875'                '875'         'C'

要分隔的字符是 ['A', 'B', 'C'] 并且它们可以在 'column1' 中的字符串的前面或末尾。 column1 中的每个字符串都有一个，而且只有一个。

Answer 1

您可以使用 UDF 来执行此操作。或者你可以使用内置的regex_extract函数来提取

from pyspark.sql import Row
from pyspark.sql.types import *
from pyspark import SparkContext, SQLContext
from pyspark.sql.functions import *

sc = SparkContext('local')
sqlContext = SQLContext(sc)

data1 = [
    ('A123',),
    ('321B',),
    ('C875',) ,
      ]


df1Columns = ["column1"]
df1 = sqlContext.createDataFrame(data=data1, schema = df1Columns)
df1.show(20, truncate=False)


import re

def findChar(input_string):
    result = "".join(re.findall("[A-C]", input_string))
    return result

findChar_udf = udf(findChar, StringType())

def findNumber(input_string):
    result = "".join(re.findall("[0-9]+", input_string))
    return result

findNumber_udf = udf(findNumber, StringType())

print("Using udf")
df2 = df1.withColumn("column2", findChar_udf("column1")).withColumn("column3", findNumber_udf("column1"))
df2.show(20, truncate=False)

print("Using regex_extract")
df3 = df1.withColumn("column2", regexp_extract("column1", "[A-C]", 0)).withColumn("column3", regexp_extract("column1", "[0-9]+", 0))
df3.show(20, truncate=False)

以下是输出：

+-------+
|column1|
+-------+
|A123   |
|321B   |
|C875   |
+-------+

Using udf
+-------+-------+-------+
|column1|column2|column3|
+-------+-------+-------+
|A123   |A      |123    |
|321B   |B      |321    |
|C875   |C      |875    |
+-------+-------+-------+

Using regex_extract
+-------+-------+-------+
|column1|column2|column3|
+-------+-------+-------+
|A123   |A      |123    |
|321B   |B      |321    |
|C875   |C      |875    |
+-------+-------+-------+

Answer 2

您可以使用正则表达式来分隔两个条件并获取子字符串。

import pyspark.sql.functions as F

df.select('column1', 

F.when(
    F.col('column1').rlike('^[ABC]'), 
    F.substring(F.col('column1'), 2, 3)
).when(
    F.col('column1').rlike('[ABC]$'),
    F.substring(F.col('column1'), 1, 3)
).alias('column2'),

F.when(
    F.col('column1').rlike('^[ABC]'), 
    F.substring(F.col('column1'), 1, 1)
).when(
    F.col('column1').rlike('[ABC]$'),
    F.substring(F.col('column1'), 4, 1)
).alias('column3')

)