I have a pyspark dataframe with a column FullPath.
How can I use the function os.path.splitext(FullPath) to extract the extension of each entry in the FullPath column and put them in a new column?
Thanks.
You can use pyspark.sql.functions.regexp_extract() to extract the file extension:
import pyspark.sql.functions as f
data = [
('/tmp/filename.tar.gz',)
]
df = sqlCtx.createDataFrame(data, ["FullPath"])
df.withColumn("extension", f.regexp_extract("FullPath", "\.[0-9a-z]+$", 0)).show()
#+--------------------+---------+
#| FullPath|extension|
#+--------------------+---------+
#|/tmp/filename.tar.gz| .gz|
#+--------------------+---------+
However if you wanted to use os.path.splittext(), you would need to use a udf (which will be slower than the above alternative):
import os
splittext = f.udf(lambda FullPath: os.path.splitext(FullPath)[-1], StringType())
df.withColumn("extension", splittext("FullPath")).show()
#+--------------------+---------+
#| FullPath|extension|
#+--------------------+---------+
#|/tmp/filename.tar.gz| .gz|
#+--------------------+---------+
There is a split function the SQL functions module, so you can split your full path on the "." character and take the last element. Assuming that there is only one "." in each filepath string.
import pyspark.sql.functions as F
myDataFrame = myDataFrame.withColumn("pathArray", F.split(myDataFrame.FullPath, ".")
myDataFrame = myDataFrame.withColumn("FileExtension", myDataFrame.pathArray.getItem(1))