Question

我有一个看起来像这样的数据框：

Filename        Type
file1.A.txt 
file2.A.txt 
file3.B.txt 
file4.A.txt
file5.B.txt
...

我想添加另一列Type，这取决于文件名。如果文件名中有A，则添加A，如果有B，则添加B。

我在Add column to Data Frame conditionally in Pyspark中看到了与此相似的东西，但是我看不到如何在我的情况下应用。

我可以通过df = df.withColumn('NewCol', lit('a'))向Spark添加常量，但是在某些情况下如何使用正则表达式更改此表达式以添加特定字符串，而在其他情况下添加另一个字符串？

这与链接的问题Spark Equivalent of IF Then ELSE相似，但是Michael West的答案更容易输入，并且更具体地针对该问题。但是，我认为它仍然可以解决问题（尽管更难阅读）。

Answer 1

类似的事情应该起作用

from pyspark.sql.functions import regexp_extract

df = spark.createDataFrame([
    ("file1.A.txt",),
    ("file2.A.txt",),
    ("file3.B.txt",),
    ("file4.A.txt",),
    ("file5.B.txt",)
  ],
  ["filenames"]
)

df.withColumn('A_or_B', regexp_extract('filenames', '^\w+\.(.*)\.txt$', 1)).show()

+-----------+------+
|  filenames|A_or_B|
+-----------+------+
|file1.A.txt|     A|
|file2.A.txt|     A|
|file3.B.txt|     B|
|file4.A.txt|     A|
|file5.B.txt|     B|
+-----------+------+

Answer 2

替代答案，不太笼统。将文件名拆分为基于的数组。分隔符

 m <-cbind(A,B,C)
#      A B  C
#[1,]  1 2  3
#[2,]  2 3  4
#[3,]  3 4  5
#[4,]  4 5 NA
#[5,]  5 6 NA
#[6,] NA 7 NA
#[7,] NA 8 NA

如何在PySpark数据框中添加变量/条件列

2 个答案: