在 Databricks 中使用 PySpark 获取 S3 中的文件列表答案

【问题标题】：Get a list of files in S3 using PySpark in Databricks在 Databricks 中使用 PySpark 获取 S3 中的文件列表
【发布时间】：2021-12-23 11:55:13
【问题描述】：

我正在尝试生成存储桶/文件夹中所有 S3 文件的列表。文件夹中通常有数百万个文件。我现在使用 boto，它每分钟能够检索大约 33k 个文件，即使是一百万个文件，也需要半个小时。我还将这些文件加载到数据框中，但生成并使用此列表来跟踪正在处理的文件。

我注意到，当我要求 Spark 读取文件夹中的所有文件时，它会列出自己的列表，并且能够比 boto 调用更快地列出它们，然后处理这些文件。我在 PySpark 中查找了一种方法，但没有找到好的示例。我得到的最接近的是一些 Java 和 Scala 代码来列出使用 HDFS 库的文件。

有没有一种方法可以在 Python 和 Spark 中做到这一点？作为参考，我正在尝试复制以下代码 sn-p：

def get_s3_files(source_directory, file_type="json"):
    s3_resource = boto3.resource("s3")

    file_prepend_path = f"/{'/'.join(source_directory.parts[1:4])}"
    bucket_name = str(source_directory.parts[3])
    prefix = "/".join(source_directory.parts[4:])

    bucket = s3_resource.Bucket(bucket_name)

    s3_source_files = []

    for object in bucket.objects.filter(Prefix=prefix):
        if object.key.endswith(f".{file_type}"):
            s3_source_files.append(
                (
                    f"{file_prepend_path}/{object.key}",
                    object.size,
                    str(source_directory),
                    str(datetime.now()),
                )
            )

    return s3_source_files

【问题讨论】：

看看这个stackoverflow.com/a/40258750/1386551
@blackbishop - 我做了，但它似乎特定于 HDFS，因此不确定我们是否也可以将它应用到 S3。

标签： python apache-spark pyspark databricks aws-databricks

【解决方案1】：

这可以通过dbutils 非常简单地实现。

def get_dir_content(ls_path):
  dir_paths = dbutils.fs.ls(ls_path)
  subdir_paths = [get_dir_content(p.path) for p in dir_paths if p.isDir() and p.path != ls_path]
  flat_subdir_paths = [p for subdir in subdir_paths for p in subdir]
  return list(map(lambda p: p.path, dir_paths)) + flat_subdir_paths
    

paths = get_dir_content('s3 location')
[print(p) for p in paths]

【讨论】：

我们以前用dbutils，但是太慢了。我们发现 boto 更快，但正在尝试看看我们是否可以让它更快。
表示，您正在寻找选项，列出来自给定 s3 位置路径的文件，而不是读取的数据。我对吗？请确认。
是的，只要我可以复制我使用 PySpark 发布的代码（并让它运行得更快），那就完美了。也就是说，我使用 AWS CLI 来列出对象，而且速度明显更快。奇怪的是 CLI 会比 boto 花费更少的时间。

【解决方案2】：

由于某种原因，使用 AWS CLI 命令比使用 boto 快大约 15 倍（！）。不知道为什么会这样，但这是我目前正在使用的代码，以防有人发现它很方便。基本上，使用s3api 列出对象，然后使用jq 操作输出并将其变成我喜欢的形式。

def get_s3_files(source_directory, schema, file_type="json"):

    file_prepend_path = f"/{'/'.join(source_directory.parts[1:4])}"
    bucket = str(source_directory.parts[3])
    prefix = "/".join(source_directory.parts[4:])

    s3_list_cmd = f"aws s3api list-objects-v2 --bucket {bucket} --prefix {prefix} | jq -r '.Contents[] | select(.Key | endswith(\".{file_type}\")) | [\"{file_prepend_path}/\"+.Key, .Size, \"{source_directory}\", (now | strftime(\"%Y-%m-%d %H:%M:%S.%s\"))] | @csv'"

    s3_list = subprocess.check_output(s3_list_cmd, shell=True, universal_newlines=True)

    with open(f"s3_file_paths.csv", "w") as f:
        f.truncate()
        f.write(s3_list)

    s3_source_files_df = spark.read.option("header", False).schema(schema).csv(f"s3_file_paths.csv")

    return s3_source_files_df

【讨论】：