Spark加载数据并将文件名添加为数据框列 [英] Spark load data and add filename as dataframe column

问题描述

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

I am loading some data into Spark with a wrapper function:

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一次加载一堆文件:

specifically, I am using a glob to load a bunch of files at once:

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

文件为:

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

我希望主机名"列实际包含要加载的文件的真实名称，而不是glob(即host1，host2等，而不是*).

I would like the column 'hostname' to actually contain the real name of the file being loaded rather than the glob (ie host1, host2 etc, rather than *).

我该怎么做?

推荐答案

您可以使用input_file_name其中:

为当前Spark任务的文件名创建一个字符串列.

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())

Scala中的相同内容

Same thing in Scala:

import org.apache.spark.sql.functions.input_file_name

df.withColumn("filename", input_file_name)

这篇关于Spark加载数据并将文件名添加为数据框列的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！