CSV load into Dataframe with filename as additional column in pyspark

user_16130268 0
Viewed 391

I'm trying to create a dataframe from a directory full of csv files but I want to keep the filename of each file on the dataframe as an additional column, is that possible on pyspark without using pandas?, also I want to remove the path from the filename.

from pyspark.sql.functions import input_file_name

df = spark.read.option("delimiter", "\t").csv(mount_point_input)
df_.withColumn("filename", input_file_name())

I tried using input_file_name() but all rows on the dataframe got the same filename.

Input:

False    2021-06-05T14:45:09     Server       True
True     2021-06-02T21:32:42     Server       True

Output:

+-----+-----------------------+-------+-------+--------------------------------+
False  2021-06-05T14:45:09     Server   True   /2021-06-02-general/c32d3f47.csv
+-----+-----------------------+-------+-------+--------------------------------+
False  2021-06-02T21:32:42     Server   True   /2021-06-02-general/c32d3f47.csv
+-----+-----------------------+-------+-------+--------------------------------+

Expected output:

+-----+-----------------------+-------+-------+--------------------------------+
False  2021-06-05T14:45:09     Server   True   c32d3f47.csv
+-----+-----------------------+-------+-------+--------------------------------+
False  2021-06-02T21:32:42     Server   True   c32d3f48.csv
+-----+-----------------------+-------+-------+--------------------------------+
pyspark apache-spark-sql apache-spark python
1 Answers

You can use os.path.basename within a UDF -

>>> from pyspark.sql.functions import input_file_name,udf
>>> from pyspark.sql.types import StringType
>>> from os.path import basename
>>> 
>>> data = [("/home/user/test/File1.txt",10), 
...         ("/home/user/test/File2.txt",20), 
...         ("/home/user/test/File3.txt",30), 
...         ("/home/user/test/File4.txt",40),
...         ("/2021-06-02-general/c32d3f47.csv",50),
...         ("/2021-06-02-general/c32d3f47.csv",50)
...         ]
>>> 
>>> 
>>> cols = ["file_path","dummy_value"]
>>> testDF = spark.createDataFrame(data=data, schema = cols)
>>> 
>>> testDF.show(truncate=False)
+--------------------------------+-----------+
|file_path                       |dummy_value|
+--------------------------------+-----------+
|/home/user/test/File1.txt       |10         |
|/home/user/test/File2.txt       |20         |
|/home/user/test/File3.txt       |30         |
|/home/user/test/File4.txt       |40         |
|/2021-06-02-general/c32d3f47.csv|50         |
|/2021-06-02-general/c32d3f47.csv|50         |
+--------------------------------+-----------+

>>> 
>>> 
>>> @udf(StringType())
... def return_filename(inp):
...     if inp:
...       return basename(inp)
...     else:
...       return None
... 
>>> testDF = testDF.withColumn("file_name", return_filename('file_path'))
>>> testDF.show(truncate=False)
+--------------------------------+-----------+------------+
|file_path                       |dummy_value|file_name   |
+--------------------------------+-----------+------------+
|/home/user/test/File1.txt       |10         |File1.txt   |
|/home/user/test/File2.txt       |20         |File2.txt   |
|/home/user/test/File3.txt       |30         |File3.txt   |
|/home/user/test/File4.txt       |40         |File4.txt   |
|/2021-06-02-general/c32d3f47.csv|50         |c32d3f47.csv|
|/2021-06-02-general/c32d3f47.csv|50         |c32d3f47.csv|
+--------------------------------+-----------+------------+
user_91089120
Related