Pyspark从文件名获取子字符串并存储为新列

Question

我正在使用pyspark处理来自S3的CSV文件，但是我希望将文件名作为新列并入我正在使用的以下代码中：

spark.udf.register("filenamefunc", lambda x: x.rsplit('/', 1)[-2])
df=spark.read.csv("s3a://exportcsv-battery/S5/243/101*",sep=',',header=True,inferSchema=True)
df=df.withColumn("filename", 'filenamefunc(input_file_name())')

但是我想要一个子字符串，而不是文件名，例如，如果这是input_file_name：-

s3a：// exportcsv-battery / S5 / 243 /101_002932_243_AAA_A_T01_AAA_AAA_0_0_0_0_0_2_10Hz.csv

我只希望将243提取并存储在新列中，为此我将UDF定义为：

spark.udf.register("filenamefunc", lambda x: x.rsplit('/', 1)[-2])

但是它似乎不起作用。有什么我可以解决的方法或其他方法吗？谢谢！

Answer 1

您可以使用split()功能

import pyspark.sql.functions as f

[...]

df = df.withColumn('filename', f.split(f.input_file_name(), '/')[4])