How to add variable/conditional column in PySpark data frame

user_3973175 0
Viewed 1582

I have a data frame which looks like this:

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

I want to add another column, Type, which will depend on the filename. If there is an A in the filename, add A, if there is a B, add B.

I've seen something vaguely similar to this in Add column to Data Frame conditionally in Pyspark but I can't see how I can apply this in my case.

I can add constants to Spark by df = df.withColumn('NewCol', lit('a')) but how can I alter this expression, using regex, to add a certain string in some cases, and another string in other cases?

This is similar to the linked question Spark Equivalent of IF Then ELSE but Michael West's answer is easier to type out and more specific to the problem. However, I think it could still solve the problem (though would be more difficult to read).

pyspark apache-spark python
2 Answers

Something like this should work

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|
+-----------+------+
user_21540510

Alternate answer, less general. Split the filename into an array based on . delimiter

from pyspark.sql.functions import split, col

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

df.withColumn('a_or_b', split(col("filenames"), "\.")[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|
+-----------+------+
user_21540510
Related