Dask读取Parquet数据集时忽略分区划分信息的问题排查

问题描述 #

我有一个存储在目录dataset_path下的Parquet数据集，索引列为date，元数据由Dask生成。使用以下代码读取数据集：

import dask.dataframe as dd
ddf = dd.read_parquet("dataset_path", engine="pyarrow", calculate_divisions=True)

执行print(ddf.divisions)可确认Dask已获取正确的分区边界，但执行基于索引的操作（如ddf.loc[ddf.index == pd.Timestamp("2020-01-01")].compute()或索引合并）时，日志显示Dask仍会打开所有数据文件，而非仅访问匹配的分区。切换为fastparquet引擎后问题依旧。

以下是复现该问题的示例代码：

from datetime import date
from dateutil.relativedelta import relativedelta
import pandas as pd

def write_parquet_timeseries(path: str, start: date, end: date):
    for part, pdf in enumerate(generate_timeseries(start, end)):
        dd.to_parquet(
            dd.from_pandas(pdf, npartitions=1),
            path,
            engine="pyarrow",
            overwrite=part == 0,
            append=part > 0,
            write_metadata_file=True,
        )

def generate_timeseries(start: date, end: date):
    start0, end0 = None, start
    while end0 < end:
        start0, end0 = end0, end0 + relativedelta(months=1)
        yield timeseries(start0, end0)

def timeseries(start: date, end: date, num_rows: int = 2**16, num_cols: int = 2**4):
    index = pd.Index(pd.date_range(start, end, inclusive="left"), name="date").repeat(num_rows)
    return pd.DataFrame({f"x{i}": range(i, len(index) + i) for i in range(num_cols)}, index=index)

write_parquet_timeseries("dataset_path", date(2020, 1, 1), date(2021, 1, 1))

问题原因 #

核心问题出在写入数据集的方式：当前循环逐个写入单分区Dask DataFrame的操作，没有为Parquet文件添加date索引列的min/max统计信息，也没有用Dask的分区逻辑组织数据。Dask实现分区剪枝（仅访问匹配分区）需要两个关键条件：

1. 数据集的分区列与索引列一致；
2. 每个Parquet文件的元数据中包含索引列的范围统计值。
   缺少统计信息导致Dask无法判断哪些文件包含目标索引值，只能遍历所有文件。

解决方案 #

1. 重新按规范写入分区化数据集 #

不要手动循环追加文件，先生成完整的Dask DataFrame，再按date索引分区后写入，确保每个文件生成统计信息：

from datetime import date
from dateutil.relativedelta import relativedelta
import pandas as pd
import dask.dataframe as dd

def generate_timeseries(start: date, end: date):
    start0, end0 = None, start
    while end0 < end:
        start0, end0 = end0, end0 + relativedelta(months=1)
        yield timeseries(start0, end0)

def timeseries(start: date, end: date, num_rows: int = 2**16, num_cols: int = 2**4):
    index = pd.Index(pd.date_range(start, end, inclusive="left"), name="date").repeat(num_rows)
    return pd.DataFrame({f"x{i}": range(i, len(index) + i) for i in range(num_cols)}, index=index)

# 生成所有数据并转为Dask DataFrame
pdf_list = list(generate_timeseries(date(2020, 1, 1), date(2021, 1, 1)))
ddf = dd.from_pandas(pd.concat(pdf_list), npartitions=12)

# 设置date为索引并按索引分区
ddf = ddf.set_index("date").repartition(partition_size="100MB")

# 写入时启用统计信息生成
dd.to_parquet(
    ddf,
    "dataset_path",
    engine="pyarrow",
    write_metadata_file=True,
    overwrite=True,
    write_statistics=True  # fastparquet需显式设置，pyarrow默认开启
)

2. 正常读取并验证分区剪枝 #

读取时无需额外指定calculate_divisions=True，Dask会自动从Parquet元数据中获取分区边界并执行剪枝：

import dask.dataframe as dd
import pandas as pd

ddf = dd.read_parquet("dataset_path", engine="pyarrow")
# 此时执行索引查询只会访问匹配的分区
result = ddf.loc[pd.Timestamp("2020-01-01")].compute()

关键注意事项 #

• 必须通过set_index将date设为索引并完成分区，不能仅手动追加文件；
• 写入时确保开启统计信息（write_statistics=True），这是Dask判断文件是否匹配的核心依据；
• 读取时无需手动计算分区，元数据已包含完整的分区边界信息。

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