使用带有分区的 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 文档中任何地方的描述。

有谁知道如何在分区镶木地板数据时保存索引？

Answer 1

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

Answer 2

我似乎试图回避这个问题，但我的建议是沿着索引进行分区。这也将确保分区中的非重叠索引。这就像dd.from_pandas(data_to_export, npartitions=3)然后跳过partition_on和write_index进入to_parquet。必须对索引进行排序。这会保留索引并正确设置分区。请注意，您不能保证获得您请求的确切分区数partitions，尤其是对于小型数据集。