Python多线程实现大规模数据集高效转移

背景

在处理大规模数据集时，通常需要在不同存储设备、不同服务器或文件夹之间高效地传输数据。如果采用单线程传输方式，当数据量非常大时，整个过程会非常耗时。因此，通过多线程并行处理可以大幅提升数据传输效率。

本文将分享一个基于Python多线程实现的高效数据传输工具，通过遍历源文件夹中的所有文件，将它们移动到目标文件夹。

工具和库

这个数据集转移工具主要依赖于以下Python标准库：

• os：用于文件系统操作，如遍历目录、判断文件和文件夹是否存在等。
• shutil：用于执行文件和文件夹的移动、删除操作。
• concurrent.futures.ThreadPoolExecutor：用于实现多线程并发操作，提升文件传输的效率。

代码分享

下面是整个程序的代码实现：

import os
import shutil
from concurrent.futures import ThreadPoolExecutor

def move_file(src, dst):
    """移动单个文件"""
    try:
        if os.path.exists(dst):
            if os.path.isdir(dst):
                shutil.rmtree(dst)  # 如果目标是文件夹，递归删除
            else:
                os.remove(dst)  # 如果目标是文件，先删除
        shutil.move(src, dst)  # 移动文件
        print(f"Moved: {src} -> {dst}")
    except Exception as e:
        print(f"Error moving {src} to {dst}: {e}")

def gather_files_in_directory(directory, src_dir, dst_dir):
    """遍历目录并收集文件的源路径和目标路径"""
    file_pairs = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            src_file = os.path.join(root, file)  # 源文件路径
            relative_path = os.path.relpath(src_file, src_dir)  # 相对路径
            dst_file = os.path.join(dst_dir, relative_path)  # 目标文件路径
            file_pairs.append((src_file, dst_file))
    return file_pairs

def gather_all_files(src_dir, dst_dir):
    """多线程遍历文件夹并收集所有文件的路径和目标路径"""
    all_file_pairs = []
    sub_dirs = [os.path.join(src_dir, d) for d in os.listdir(src_dir) if os.path.isdir(os.path.join(src_dir, d))]

    with ThreadPoolExecutor(max_workers=20) as executor:
        futures = []
        for sub_dir in sub_dirs:
            futures.append(executor.submit(gather_files_in_directory, sub_dir, src_dir, dst_dir))

        for future in futures:
            all_file_pairs.extend(future.result())

    return all_file_pairs

def move_all_files(file_pairs):
    """多线程传输所有文件"""
    with ThreadPoolExecutor(max_workers=20) as executor:
        for src, dst in file_pairs:
            dst_dir = os.path.dirname(dst)
            if not os.path.exists(dst_dir):
                os.makedirs(dst_dir)  # 如果目标文件夹不存在，则创建
            executor.submit(move_file, src, dst)

def main():
    # 提示用户输入源文件夹和目标文件夹
    src_base_dir = input("请输入源文件夹路径：").strip()
    dst_base_dir = input("请输入目标文件夹路径：").strip()

    # 检查源文件夹是否存在
    if not os.path.exists(src_base_dir):
        print(f"源文件夹不存在：{src_base_dir}")
        return

    # 检查目标文件夹是否存在
    if not os.path.exists(dst_base_dir):
        print(f"目标文件夹不存在，正在创建：{dst_base_dir}")
        os.makedirs(dst_base_dir)

    # 收集文件路径
    print(f"正在遍历源文件夹 '{src_base_dir}'，并准备移动文件至 '{dst_base_dir}'...")
    file_pairs = gather_all_files(src_base_dir, dst_base_dir)

    # 如果没有找到文件，结束程序
    if not file_pairs:
        print(f"在源文件夹 '{src_base_dir}' 中未找到任何文件。")
        return

    print(f"已找到 {len(file_pairs)} 个文件准备移动。")

    # 开始移动文件
    print("开始移动文件...")
    move_all_files(file_pairs)

    print(f"所有文件已成功移动至 '{dst_base_dir}'。")
    print("操作完成。")

if __name__ == "__main__":
    main()

代码讲解

1. move_file 函数

该函数负责将单个文件从源路径移动到目标路径。它首先检查目标位置是否已经存在文件或文件夹，如果存在，则先删除旧的文件或文件夹，然后再执行移动操作。shutil.move() 是核心的文件移动函数，负责将文件从源路径移动到目标路径。

2. gather_files_in_directory 函数

这是一个递归遍历函数，用于遍历指定目录，并将每个文件的源路径和目标路径作为一对元组收集起来。os.walk() 用于递归遍历目录树，os.path.relpath() 用于生成相对路径，以确保源路径和目标路径的文件结构一致。

3. gather_all_files 函数

该函数使用多线程并发遍历子目录，通过 ThreadPoolExecutor 实现并行处理多个目录的文件收集工作。每个子目录都提交给线程池中的一个线程，并且通过 futures 的结果合并所有收集到的文件路径对。

4. move_all_files 函数

同样，这里使用了多线程来并行移动文件。每个文件对（源路径和目标路径）都会提交给线程池，并行处理。如果目标路径的目录不存在，程序会自动创建对应的目录。

性能优化

1. 多线程并发：由于文件移动操作是 I/O 密集型任务，使用 ThreadPoolExecutor 来并发执行任务，能够充分利用 CPU 资源并显著减少文件传输的时间。

2. 避免重复创建目录：在文件移动之前，通过检查目标路径的文件夹是否存在，避免重复的 os.makedirs() 调用，从而提高效率。

3. 减少文件移动冲突：通过先删除目标文件或文件夹，避免因同名文件或文件夹导致的移动冲突。

4. 合理设置线程数：根据硬件资源（CPU 核心数量、硬盘 I/O 性能）设置合理的线程数。在程序中，我们使用了20个线程，这个数量可以根据系统性能做适当调整。

总结

通过多线程并发执行，我们能够显著加速大规模文件的传输过程。本文的 Python 程序展示了如何高效地遍历目录并移动文件，适用于需要在本地存储、远程服务器、不同硬盘之间转移大量数据的场景。

如果你有任何问题或者建议，欢迎在评论区留言讨论。希望这篇文章能够对你有所帮助，感谢阅读！