Spark2.4下如何高效对目录中每个文件单独执行WordCount？

核心思路 #

使用wholeTextFiles读取目标目录，它会返回键值对RDD：(文件完整路径, 文件内容)，正好满足按单个文件处理的需求。Spark会分布式分区处理这些文件，不会一次性加载所有文件到内存，可通过调整分区数平衡性能与内存占用。

Python实现代码 #

from pyspark.sql import SparkSession
from collections import Counter

# 初始化Spark会话
spark = SparkSession.builder.appName("PerFileWordCount").getOrCreate()
sc = spark.sparkContext

def process_file(file_tuple):
    full_path, content = file_tuple
    # 提取文件名（从完整路径中截取最后一段）
    filename = full_path.split("/")[-1]
    # 分割内容为单词（处理换行、多空格）
    words = content.strip().split()
    # 统计单词频次
    word_counts = Counter(words)
    return (filename, list(word_counts.items()))

# 读取目标目录，设置合理分区数（根据文件数量/集群资源调整，数万个文件可设为100-500）
input_dir = "path/to/demotxt"
file_rdd = sc.wholeTextFiles(input_dir, minPartitions=200)

# 处理每个文件
result_rdd = file_rdd.map(process_file)

# 打印结果示例
for filename, counts in result_rdd.collect():
    print(f"{filename}:  {', '.join([str(item) for item in counts])}")

# 批量保存到数据库（推荐用foreachPartition减少DB连接数）
def save_partition(partition):
    # 这里替换为你的DB连接逻辑，以PostgreSQL为例
    import psycopg2
    conn = psycopg2.connect("dbname=your_db user=your_user password=your_pwd")
    cur = conn.cursor()
    for filename, counts in partition:
        for word, count in counts:
            cur.execute(
                "INSERT INTO word_counts (filename, word, count) VALUES (%s, %s, %s)",
                (filename, word, count)
            )
    conn.commit()
    cur.close()
    conn.close()

result_rdd.foreachPartition(save_partition)

# 停止Spark会话
spark.stop()

Scala实现代码 #

import org.apache.spark.sql.SparkSession
import scala.collection.mutable

object PerFileWordCount {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder.appName("PerFileWordCount").getOrCreate()
    val sc = spark.sparkContext

    val inputDir = "path/to/demotxt"
    // 读取目录并设置分区数
    val fileRDD = sc.wholeTextFiles(inputDir, minPartitions = 200)

    val resultRDD = fileRDD.map { case (fullPath, content) =>
      // 提取文件名
      val filename = fullPath.split("/").last
      // 分割内容为单词
      val words = content.trim.split("\\s+")
      // 统计单词频次
      val wordCounts = mutable.HashMap[String, Int]()
      words.foreach(word => wordCounts(word) = wordCounts.getOrElse(word, 0) + 1)
      (filename, wordCounts.toList)
    }

    // 打印结果
    resultRDD.collect().foreach { case (filename, counts) =>
      println(s"$filename:  ${counts.mkString(", ")}")
    }

    // 批量保存到数据库
    resultRDD.foreachPartition { partition =>
      // 替换为你的JDBC连接信息
      val conn = java.sql.DriverManager.getConnection(
        "jdbc:postgresql://localhost/your_db",
        "your_user",
        "your_pwd"
      )
      val stmt = conn.createStatement()
      partition.foreach { case (filename, counts) =>
        counts.foreach { case (word, count) =>
          val sql = s"INSERT INTO word_counts (filename, word, count) VALUES ('$filename', '$word', $count)"
          stmt.executeUpdate(sql)
        }
      }
      conn.commit()
      stmt.close()
      conn.close()
    }

    spark.stop()
  }
}

关键注意事项 #

1. 分区数调整：针对数万个小文件，设置minPartitions为100-500（根据集群核心数调整，建议每个核心对应2-4个分区），避免单分区处理过慢，同时防止分区过多消耗资源。
2. 内存控制：wholeTextFiles按分区读取文件，不会一次性加载所有数据到内存，只要分区数合理，不会出现内存溢出问题。若单个文件过大，可调整spark.files.maxPartitionBytes参数拆分大文件。
3. 数据库优化：禁止在map操作中直接创建DB连接，必须使用foreachPartition，每个分区创建一次连接处理批量数据，大幅降低DB连接压力。

内容的提问来源于stack exchange，提问作者AppleCEO