如何使用pyspark rdd处理多个输入文件并将输出保存在各个文件夹中，文件名为文件夹名

我有5个输入文件，比如A、B、C、D、E。我想将这些文件加载到pyspark rdd并进行一些处理。最后，我想将输出保存在一个文件夹中，并将相应的文件名作为文件夹名。这在spark cluster模式下可能吗？

由于rdd/dataframe本质上是分布在多个分区上的一组行，因此从多个源读入数据后，您将无法跟踪数据的来源。因此，我的简单解决方案是为跟踪其原点的行指定一个附加值。使用dataframe API：

from pyspark.sql.functions import lit, col
from pyspark.sql import DataFrame
from functools import reduce

fnames = ['file_A.csv','file_B.csv','file_C.csv']

dfs = []

# 1. Read in data from individual sources and assign the filename as a string to a column
for fname in fnames:
    dfs.append(spark.read.format('csv')
                    .option("header", "true")
                    .load(fname)
                    .withColumn('origin',lit(fname))
              )

# 2. Combine data 
df = reduce(DataFrame.unionAll,dfs)

# +---+---+---+----------+
# |  A|  B|  C|    origin|
# +---+---+---+----------+
# |  1|  1|  1|file_A.csv|
# |  1|  1|  1|file_A.csv|
# |  1|  1|  1|file_A.csv|
# |  0|  0|  0|file_B.csv|
# |  0|  0|  0|file_B.csv|
# |  0|  0|  0|file_B.csv|
# |  2|  2|  2|file_C.csv|
# |  2|  2|  2|file_C.csv|
# |  2|  2|  2|file_C.csv|
# +---+---+---+----------+

# 3. Do processing
# ...

# 4. Subset the combined data by origin and write out each subset to file
for fname in fnames:
    
    out_fname = '_new.'.join(fname.split('.'))
    
    df.filter(col('origin')==fname)\
      .write.format('csv')\
      .option('header',True)\
      .save(out_fname)

---

应该可以做到这一点-是否有一个特定的步骤被卡住了？在保存输出时，spark rdd如何知道文件的相应内容？如何使用文件名将处理后的数据保存在相应的文件夹中？