python - pyarrow.ParquetDataset > 分区列的架构

Question

1. 我有一个 pandas 数据框:

import pandas as pd

df = pd.DataFrame(data={"col1": [1, 2], "col2": [3.0, 4.0], "col3": ["foo", "bar"]})

1. 使用s3fs :

from s3fs import S3FileSystem

s3fs = S3FileSystem(**kwargs)

1. 我可以把它写成 Parquet 数据集

import pyarrow as pa
import pyarrow.parquet as pq

tbl = pa.Table.from_pandas(df)
root_path = "../parquet_dataset/foo"

pq.write_to_dataset(
    table=tbl,
    root_path=root_path,
    filesystem=s3fs,
    partition_cols=["col3"],
    partition_filename_cb=lambda _: "data.parquet",
)

1. 稍后，我需要用于转储 DataFrame 的 pq.ParquetSchema。

import pyarrow as pa
import pyarrow.parquet as pq


dataset = pq.ParquetDataset(root_path, filesystem=s3fs)
schema = dataset.schema

但是 Parquet 数据集 -> “模式”不包括分区 cols 模式。

如何获取分区列的架构？

Answer 1

我认为您需要为 ParquetDataset 提供分区键模式的提示。

partition_schema = pa.schema([pa.field('col3', pa.string())])
partitioning = pa.dataset.partitioning(schema=partition_schema)

partitionaldataset = pq.ParquetDataset(
    root_path, 
    partitioning=partitioning,
)

这给了你这个模式:

col1: int64
col2: double
col3: string

PS:我无法完全重现您的示例(我无权访问 S3)并且我不得不在写入和读取数据集时添加 use_legacy_dataset=False。