如何高效读取分区Parquet数据集的指定子集？

核心问题分析 #

你的问题本质是PyArrow默认扫描逻辑未充分利用Hive分区元数据做精准路径过滤，导致扫描了大量无关分区文件；而手动定位目标文件夹直接命中数据，所以速度差极大。

推荐的高效实现方式 #

1. 优化PyArrow Dataset扫描逻辑 #

你当前使用的ds.dataset本身支持精准分区过滤，但需要明确配置分区规则，确保过滤条件直接命中分区列：

import pyarrow.dataset as ds
import pyarrow as pa

# 显式定义分区Schema，让Dataset精准识别分区列及类型
partition_schema = pa.schema([
    ("partition_col1", pa.string()),  # 替换为你的实际分区列和类型
    ("partition_col2", pa.int64())
])

# 创建Dataset时指定分区规则细节
dataset = ds.dataset(
    dataset_storage_root,
    format="parquet",
    partitioning=ds.partitioning(
        schema=partition_schema,
        flavor="hive"
    )
)

# 构建基于分区列的过滤条件（必须用ds.field指定）
filters = ds.field("partition_col1") == "target_value"

# 启用多线程+分区剪枝，只读取目标数据
result = dataset.scanner(
    columns=columns,
    filter=filters,
    use_threads=True,
    batch_size=1024*1024
).to_table().to_pandas()

2. 用Dask实现并行分区扫描 #

如果过滤涉及多分区或需要更高并行度，Dask可自动识别Hive分区并做并行读取，效率接近手动定位：

import dask.dataframe as dd

# Dask自动识别Hive分区，直接指定过滤条件和分区列
ddf = dd.read_parquet(
    dataset_storage_root,
    engine="pyarrow",
    filters=filters,
    columns=columns,
    partition_on=["partition_col1", "partition_col2"]
)

# 转换为Pandas DataFrame
result = ddf.compute()

3. 验证分区剪枝是否生效 #

可以通过以下代码检查实际扫描的文件路径，确认是否仅读取目标分区：

files = list(dataset.get_fragments(filter=filters))
for fragment in files:
    print(fragment.path)

若输出均为目标分区的文件路径，说明分区剪枝已生效。

为什么之前的方法慢？ #

• pd.read_parquet和pq.read_table默认会扫描所有文件的元数据，若分区列识别不精准，会遍历300GB数据集的全部元数据，导致耗时剧增。
• 你之前的ds.dataset调用未显式指定分区Schema，可能导致PyArrow无法正确识别分区列，无法跳过无关分区，只能全量扫描后再过滤。

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