Dask重分区保存Parquet后无法恢复索引与分区信息

背景 #

从Spark迁移哈希分区的Parquet DataFrame到Dask后，连接操作性能极差。计划将数据按source_id列重新分区并建立排序索引，划分为N个按source_id范围划分的分区，保存为Parquet以实现后续快速查询和同索引数据集的高效连接，但保存后读取时无法恢复索引和分区信息。

重分区并保存代码 #

# Config
num_new_partitions = 20
in_dataset = '/mnt/foo/part.1.parquet' # 实际有数千个文件，当前用单个测试
out_dataset = '/mnt/bar/'

# Read in
ddf = dd.read_parquet(in_dataset)

# Sort by source_id
ddf = ddf.sort_values('source_id')

# Define a sorted index on source_id
ddf = ddf.set_index('source_id', sorted=True)

# Repartition
ddf = ddf.repartition(npartitions=num_new_partitions)

# Write to disk
ddf.to_parquet(out_dataset, write_index=True)

保存后的检查结果 #

执行代码生成20个文件后，检查结果如下：

# Check the _meta, known_divisions etc
print(ddf._meta)
Empty DataFrame
Columns: [List of columns, redacted for brevity but source_id is not one of them]
Index: []

[0 rows x 25 columns]

print(ddf.known_divisions)
True

print(ddf.divisions)
(List of divisions, redacted for brevity)

print(ddf.index.name)
source_id

疑问1 #

ddf.index.name显示为“source_id”，但ddf._meta的Index为空，这是为什么？

通过pyarrow检查Parquet元数据，发现“source_id”被列为列，且pandas元数据显示其为索引列。

多次读取尝试（均失败） #

尝试1：常规读取 #

dataset = '/mnt/bar/*.parquet'

ddf = dd.read_parquet(
    dataset,
)

结果：

print(ddf._meta)
Empty DataFrame
Columns: [Redacted, source_id not present]
Index: []

print(ddf.index.name)
source_id

print(ddf.known_divisions)
False

print(ddf.divisions)
(None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None)

尝试2：读取时指定索引 #

ddf = dd.read_parquet(
    dataset,
    index='source_id'
)

结果：

print(ddf._meta)
Empty DataFrame
Columns: [Redacted, source_id not present]
Index: []

print(ddf.index.name)
source_id

print(ddf.known_divisions)
False

print(ddf.divisions)
(None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None)

尝试3：读取后重新设置索引 #

ddf = dd.read_parquet(
    dataset,
    index='source_id'
)

ddf = ddf.set_index('source_id', sorted=True)

结果：

print(ddf._meta)
Empty DataFrame
Columns: [Redacted, source_id not present]
Index: []

print(ddf.index.name)
source_id

print(ddf.known_divisions)
False

print(ddf.divisions)
(None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None)

尝试4：持久化计算 #

ddf = dd.read_parquet(
    dataset,
    index='source_id'
)

ddf = ddf.set_index('source_id', sorted=True).persist()

# Wait until above has completed

结果：

print(ddf._meta)
Empty DataFrame
Columns: [Redacted, source_id not present]
Index: []

print(ddf.index.name)
source_id

print(ddf.known_divisions)
False

print(ddf.divisions)
(None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None)

目标 #

将数据集划分为N个文件，以source_id为索引并排序，实现与同索引排序数据集的快速连接，且无需每次加载都重新排序分区。

解决方案 #

1. 修正保存流程的顺序问题 #

原代码中repartition在set_index之后执行，会打乱已有的有序分区。正确做法是通过set_index一步完成排序和分区：

# Config
num_new_partitions = 20
in_dataset = '/mnt/foo/part.1.parquet'
out_dataset = '/mnt/bar/'

# Read in
ddf = dd.read_parquet(in_dataset)

# 一步完成排序、设索引、分区
ddf = ddf.set_index('source_id', sorted=True, npartitions=num_new_partitions)

# 用pyarrow引擎写入，确保元数据完整
ddf.to_parquet(out_dataset, write_index=True, engine='pyarrow', compression='snappy')

解释：set_index指定sorted=True和npartitions时，会直接按source_id排序并划分成指定数量的有序分区，避免后续操作破坏分区顺序。

2. 正确读取已分区数据集 #

读取时需指定pyarrow引擎，且读取整个数据集目录而非单个文件，让Dask识别分区元数据：

ddf = dd.read_parquet(
    out_dataset,
    engine='pyarrow',
    index='source_id'
)

# 验证分区信息
print(ddf.known_divisions)  # 应返回True
print(ddf.divisions)        # 显示实际分区边界

3. 关于_meta的疑问解答 #

ddf._meta是描述数据结构的空DataFrame，Index为空是正常的——它仅用来定义列和索引的结构，不存储实际数据。只要ddf.index.name正确且known_divisions为True，就说明索引和分区信息有效。

4. 实现快速连接的关键 #

• 确保两个待连接的数据集按相同索引列排序且分区边界完全对齐
• 读取后确认known_divisions为True，Dask会自动使用高效的有序连接算法，避免全量数据shuffle

内容的提问来源于stack exchange，提问作者es-code-bar