使用带有分区的 dask.dataframe.to_parquet() 时丢失索引信息

Question

当我将 dask=1.2.2 与 pyarrow 0.11.1 一起使用时，我没有观察到这种行为。更新后（dask=2.10.1 和 pyarrow=0.15.1），当我使用带有给定 partition_on 和 write_index 参数的 to_parquet 方法时，我无法保存索引。在这里，我创建了一个显示问题的最小示例：

from datetime import timedelta
from pathlib import Path

import dask.dataframe as dd
import pandas as pd

REPORT_DATE_TEST = pd.to_datetime('2019-01-01').date()
path = Path('/home/ludwik/Documents/YieldPlanet/research/trials/')

observations_nr = 3
dtas = range(0, observations_nr)
rds = [REPORT_DATE_TEST - timedelta(days=days) for days in dtas]
data_to_export = pd.DataFrame({
    'report_date': rds,
    'dta': dtas,
    'stay_date': [REPORT_DATE_TEST] * observations_nr,
    }) \
    .set_index('dta')

data_to_export_dask = dd.from_pandas(data_to_export, npartitions=1)

file_name = 'trial.parquet'
data_to_export_dask.to_parquet(path / file_name,
                               engine='pyarrow',
                               compression='snappy',
                               partition_on=['report_date'],
                               write_index=True
                              )

data_read = dd.read_parquet(path / file_name, engine='pyarrow')
print(data_read)

这给出了：

| | stay_date  |dta| report_date|
|0| 2019-01-01 | 2 | 2018-12-30 |
|0| 2019-01-01 | 1 | 2018-12-31 |
|0| 2019-01-01 | 0 | 2019-01-01 |

我没有在 dask 文档中看到任何描述的内容。

有人知道如何在对 parquet 数据进行分区时保存索引吗？

Answer 1

问题出在 pyarrow 的后端。我在他们的 JIRA 网页上提交了错误报告： https://issues.apache.org/jira/browse/ARROW-7782

Answer 2

我似乎试图回避这个问题，但我的建议是沿着索引进行分区。这也将确保分区中的索引不重叠。

这类似于dd.from_pandas(data_to_export, npartitions=3)，然后在to_parquet 中跳过partition_on 和write_index。必须对索引进行排序。

这会保留索引并正确设置分区。

请注意，您不能保证获得您使用partitions 请求的确切分区数，尤其是对于小型数据集。