apache-spark – spark加载数据并将文件名添加为dataframe列

我正在使用包装函数将一些数据加载到spark中:

def load_data( filename ):
    df = sqlContext.read.format("com.databricks.spark.csv")\
        .option("delimiter", "\t")\
        .option("header", "false")\
        .option("mode", "DROPMALFORMED")\
        .load(filename)
    # add the filename base as hostname
    ( hostname, _ ) = os.path.splitext( os.path.basename(filename) )
    ( hostname, _ ) = os.path.splitext( hostname )
    df = df.withColumn('hostname', lit(hostname))
    return df

具体来说,我使用glob来一次加载一堆文件:

df = load_data( '/scratch/*.txt.gz' )

文件是:

/scratch/host1.txt.gz
/scratch/host2.txt.gz
...

我希望列’hostname’实际上包含正在加载的文件的真实名称而不是glob(即host1,host2等,而不是*).我怎样才能做到这一点?

最佳答案
您可以使用input_file_name:

Creates a string column for the file name of the current Spark task.

from  pyspark.sql.functions import input_file_name

df.withColumn("filename", input_file_name())

转载注明原文:apache-spark – spark加载数据并将文件名添加为dataframe列 - 代码日志